1 Billion Row Challenge

The Problem

The One Billion Rows Challenge (1BRC) is a challenge from a few years ago published by Gunnar Morling for Java. The problem is pretty straight forward: you have a text file with rows of weather stations, a semi-colon, and a temperature reading.

Wareham;2.0
Alston;12.1
Maesteg;1.2
Kendal;-7.2
Onchan;-4.4
Dinton;17.2
Loughton;8.8
Ware;-2.4
LlannerchYMedd;2.8
Hoylake;6.0

You need to, as fast as possible, iterate through the file and calculate for each unique station the lowest, highest and median temperature before printing all the stations (and the data) to the terminal in alphabetical order.

There are some restrictions on the data, but that's pretty much the gist of it. It's an approximately 14Gb file on a disk.

Baseline

To begin with let's take the naive approach and build something that uses maps, floating point arithmetic and no parallelism.

fn main() {
    let file = fs::File::open("../weather_stations/data.csv").expect("Unable to open file");
    let file = BufReader::new(file);
    // B-Tree for ordering on insert
    let mut stations: BTreeMap<String, Station> = BTreeMap::new();

    // This part consumes the file and stores all the values for each entry
    for line in file.lines() {
        let line = line.unwrap().clone();
        let mut line = line.split(';');
        let station_val = line.next().unwrap().to_string();
        let temp = line.next().unwrap();

        let mut station = stations.get_mut(&station_val);

        if station.is_none() {
            stations.insert(
                station_val.clone(),
                Station {
                    measurements: vec![],
                },
            );

            station = stations.get_mut(&station_val);
        }

        match station {
            Some(val) => {
                val.measurements.push(temp.parse().unwrap());
            }
            None => {
                println!("Error!");
                return;
            }
        }
    }

    for (station, temps) in stations.iter() {
        let mut min = f64::MAX;
        let mut max = f64::MIN;
        let mut sum: f64 = 0.0;

        print!("{{"}});
        for temp in &temps.measurements {
            let temp = *temp;

            if temp > max { max = temp; }
            if temp < min { min = temp; }

            sum = sum + temp;
        }

        let mean = sum / temps.measurements.len() as f64;
        print!("{station}={min}/{mean:.2}/{max}, ");
    }

    println!("}}");
}

We take this pretty simple: we use an ordered B-Tree which will maintain the ordering on inserts (shouldn't happen too often because the number of rows is vastly exceeding the number of possible unique stations) and we read each line one by one.

When I run this on my M1 MBP (5 years old now, wow…) it takes: approximately 170 seconds.

Basic Optimisations

Looking at the benchmarks we can see a few obvious improvements, especially around the string handling and operations since these are accounting for a huge amount of the workload.

Further down, but still impactful, we can see that the B-Tree may be a bit slower than a regular HashMap and merging; we can experiment to find out. Additionally we might be taking some pain from the repeated loads which we can optimise.

    let mut file = BufReader::with_capacity(8 * 8 * 1024, file);
    let mut parser = Parser::new();

    loop {
        let read_data = parser.parse_line(&mut file);

        if read_data.is_none() {
            // Move onto the sorting stage
            break;
        }
    }

    let mut stations: Vec<Station> = parser.stations_vec;
    stations.sort();

Here we are using the BufReader to consume multiple rows of data at once; minimising the OS interaction time.

#[derive(PartialEq, PartialOrd, Eq)]
pub struct Station {
    pub name: String,
    pub measurements: Vec<usize>,
}

impl Ord for Station {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.name.cmp(&other.name)
    }
}

pub struct Parser {
    station_buffer: String,
    temperature_buffer: String,
    stations_vec: Vec<Station>,
    stations_index_cache: HashMap<String, usize>,
    char_buffer: [u8; 1],
}

impl Parser {
    pub fn new() -> Self {
        Self {
            station_buffer: String::with_capacity(100),
            temperature_buffer: String::with_capacity(100),
            stations_vec: Vec::with_capacity(10_000),
            stations_index_cache: HashMap::with_capacity(10_000),
            char_buffer: [0],
        }
    }

    fn put(&mut self, temperature: usize) -> &Station {
        let station_name: &String = &self.station_buffer;

        match self.stations_index_cache.get(station_name) {
            Some(station_index) => {
                let station = self.stations_vec.index_mut(*station_index);
                station.measurements.push(temperature);
            }
            None => {
                let station_index = self.stations_vec.len();
                self.stations_index_cache
                    .insert(station_name.clone(), station_index);
                let mut new_station = Station {
                    name: station_name.clone(),
                    measurements: Vec::with_capacity(100),
                };
                new_station.measurements.push(temperature);
                self.stations_vec.push(new_station);
            }
        }

        &self.stations_vec[*self.stations_index_cache.get(station_name).expect("Error?")]
    }

    // Returns the station and the temp
    fn parse_line(&mut self, file: &mut BufReader<File>) -> Option<&Station> {
        self.station_buffer.clear();
        self.temperature_buffer.clear();

        loop {
            self.char_buffer[0] = 0;
            let result = file.read_exact(&mut self.char_buffer);
            if result.is_err() { break; }
            let char = self.char_buffer[0] as char;
            if char == 0 as char || char == ';' { break; }
            self.station_buffer.push(char);
        }

        if self.station_buffer.is_empty() { return None; }

        loop {
            self.char_buffer[0] = 0;
            let result = file.read_exact(&mut self.char_buffer);
            if result.is_err() { break; }
            let char = self.char_buffer[0] as char;
            if char == 0 as char || char == '\n' { break; }
            self.temperature_buffer.push(char as char);
        }

        let temp = self
            .temperature_buffer
            .parse::<f64>()
            .expect("Unable to parse temp")
            * 100.0;

        Some(self.put(temp as usize))
    }
}

Reading this top to bottom: we add a struct to maintain the running list of measurements for each station and we'll calculate the results from there at the end. The massive and clear pain point here is that we need to store something for each row in the file, but we'll deal with that later.

We make our first attempts to reduce the number of String operations with the Parser struct.

    pub fn new() -> Self {
        Self {
            station_buffer: String::with_capacity(100),
            temperature_buffer: String::with_capacity(100),
            stations_vec: Vec::with_capacity(10_000),
            stations_index_cache: HashMap::with_capacity(10_000),
            char_buffer: [0],
        }
    }

Since we don't have to worry about any threading issues right now, we know we are processing one row (i.e one station; measurement pair) that we can just copy each row into memory and use that as the look-up in our HashMap rather than producing a new String each time.

We construct the data in the main loop:

    // Returns the station and the temp
    fn parse_line(&mut self, file: &mut BufReader<File>) -> Option<&Station> {
        self.station_buffer.clear();
        self.temperature_buffer.clear();

        loop {
            self.char_buffer[0] = 0;
            let result = file.read_exact(&mut self.char_buffer);
            if result.is_err() { break; }
            let char = self.char_buffer[0] as char;
            if char == 0 as char || char == ';' { break; }
            self.station_buffer.push(char);
        }

        if self.station_buffer.is_empty() { return None; }

        loop {
            self.char_buffer[0] = 0;
            let result = file.read_exact(&mut self.char_buffer);
            if result.is_err() { break; }
            let char = self.char_buffer[0] as char;
            if char == 0 as char || char == '\n' { break; }
            self.temperature_buffer.push(char as char);
        }

        let temp = self
            .temperature_buffer
            .parse::<f64>()
            .expect("Unable to parse temp")
            * 100.0;

        Some(self.put(temp as usize))
    }

This obviously isn't going to be optimal, but does yield good improvements. On my M1 this comes in at: 81 seconds.

I think this highlights something important: this is a good enough time for so many use cases and we haven't really done anything smart yet. But this is a challenge to be the fastest, so let's keep going.

SIMD and OS Optimisations

Looking at this benchmark we can still see a few obvious optimisations. Our put method is expensive because of the hashing algorithm; our String operations (e.g: is_empty, parse, etc.) aren't cheap; and our reading still involves copies.

Let's tackle those one by one starting with the reading of the file.

let mut mmap: Mmap = unsafe { Mmap::map(&file).unwrap() };

This allows us to index into the file on disk directly without having to copy it into some representation. Far faster, and will have high cache hit rates as we are reading the file sequentially.

    // Returns the station and the temp
    #[inline(always)]
    fn parse_file(&mut self, file: &mut Mmap) {
        let mut start_index = 0;
        for end_index in memchr::memchr_iter(NEW_LINE, &file) {
            let line: &[u8] = &file[start_index..end_index];
            self.parse_line(&line);
            start_index = end_index + 1;
        }
    }

This allows us to use memchr (man page) from libc in Rust. It is a very OS specific (and highly optimised) search for a char in an array of chars, and it can utilise SIMD to work on those simultaneously. Between these two optimisations we can remove the read_exact and parse, is_empty from our code.

We can then also remove the need for an intermediate String representation of the temperature with:

    #[inline(always)]
    pub fn add_u8(&mut self, bytes: &[u8]) {
        // We do have some invariants in the data
        // but these are not represented at compile time
        // Starts with - or a number
        // At most two digits before the decimal and one after
        // 0 is ASCII 48
        // eg: -10.23 = [45, 49, 48, 46, 50, 51]
        // eg:  10.23 = [49, 48, 46, 50, 51]
        let mut value: i32 = 0;
        let mut factor = 1;
        let mut offset = 0;
        if bytes[0] == NEGATIVE {
            factor = -1;
            offset = 1;
        }

        // now we look for the .
        if bytes[offset + 1] == DOT as u8 {
            value += (bytes[offset] as i32 - 48) * 100;
            offset = offset + 2;
        } else {
            value += (bytes[offset] as i32 - 48) * 1000;
            value += (bytes[offset + 1] as i32 - 48) * 100;
            offset = offset + 3;
        }

        // the data always has one value after the decimal
        value += bytes[offset] as i32 - 48;

        self.add(value * factor);
    }

I did experiment with a few variants of this, but given the known format this works pretty well.

And finally we can use a far faster hash function (which is not cryptographically secure):

stations: FxHashMap::with_capacity_and_hasher(40_000, FxBuildHasher::default()),

That then gets us down to 25 seconds on a single core on a five year old M1 MBP.

Parallelism

Before we begin to parallelise there is still one optimisation we can make:

pub struct Station<'a> {
    pub name: &'a [u8],
    pub min: i32,
    pub max: i32,
    pub sum: i32,
    pub count: i32,
}

And:

let station_name = &file[start_index..semi_colon_index];
let temperature = &file[semi_colon_index + 1..end_index];
let key = Key::new(&station_name);

stations
    .entry(key)
    .and_modify(|station| station.add_u8(temperature))
    .or_insert_with(|| Station::new(station_name, temperature));

We can borrow the raw bytes from the file to use as an index since we will do a comparison, and we can modify the map in place.

let num_threads: usize = 4;
let mut offset = 0;
let read_size = 2_usize.pow(32); // 4Gb each thread
let mut maps = Vec::with_capacity(num_threads);

for _ in 0..num_threads {
    maps.push(FxHashMap::with_capacity_and_hasher(
        20_000,
        FxBuildHasher::default(),
    ));
}

let mut mmaps = Vec::with_capacity(num_threads);
for _ in 0..num_threads - 1 {
    let mmap: Mmap = unsafe {
        MmapOptions::new()
            .offset(offset)
            .len(read_size)
            .map(&file)
            .unwrap()
    };

    match memrchr(NEW_LINE, &mmap) {
        None => panic!("badly set read size"),
        Some(val) => {
            offset += val as u64 + 1;
        }
    }

    mmaps.push(mmap);
}

mmaps.push(unsafe { MmapOptions::new().offset(offset).map(&file).unwrap() });

thread::scope(|scope| {
    for (mut map, mmap) in maps.iter_mut().zip(&mmaps) {
        scope.spawn(move || {
            parse_file(&mut map, mmap);
        });
    }
});

The key with the parallelism here is that we can split the file into chunks and read it simultaneously. This means we do not have to do any locking!

A really interesting point here is how fast sorting algorithms are. We can merge and sort without using any substantial amounts of compute.

This gets us the final code which is on GitHub!

Now to get the time for this we'll need to move to a dedicated box because my MBP is exhausting its ability to read from disk.

Results

Sections 01–04 are measured on an M1 MBP; the final result is on the same reference hardware as the challenge.

The fastest Rust implementation I could find was around 8 seconds, meaning we come in just a touch faster.