Merge #76

76: feat: implement optimised insert_range r=Kerollmops a=domodwyer Implements an optimised `insert_range()` function for efficiently adding consecutive elements to the set. Previously the easiest way to do this was with `extends()` or calling `insert()` in a loop. To insert 100,000 elements this took about `1,000us` - with the `insert_range()` this completes in `~2us`. <details><summary> Bechmarks </summary> `insert_range()` vs. `insert()` in a loop inserting 0..N: ```text add_range/roaring/10 time: [191.80 ns 194.15 ns 197.16 ns] thrpt: [50.720 Melem/s 51.506 Melem/s 52.137 Melem/s] change: time: [-58.031% -57.320% -56.475%] (p = 0.00 < 0.05) thrpt: [+129.75% +134.30% +138.27%] Performance has improved. add_range/pre_populated_roaring/10 time: [177.77 ns 180.03 ns 182.21 ns] thrpt: [54.881 Melem/s 55.547 Melem/s 56.252 Melem/s] change: time: [-36.465% -34.812% -33.125%] (p = 0.00 < 0.05) thrpt: [+49.532% +53.404% +57.393%] Performance has improved. add_range/roaring/100 time: [342.94 ns 346.12 ns 351.19 ns] thrpt: [284.75 Melem/s 288.92 Melem/s 291.60 Melem/s] change: time: [-85.032% -84.862% -84.654%] (p = 0.00 < 0.05) thrpt: [+551.63% +560.58% +568.10%] Performance has improved. add_range/pre_populated_roaring/100 time: [366.68 ns 371.57 ns 376.29 ns] thrpt: [265.76 Melem/s 269.13 Melem/s 272.72 Melem/s] change: time: [-83.662% -83.417% -83.156%] (p = 0.00 < 0.05) thrpt: [+493.70% +503.04% +512.08%] Performance has improved. add_range/roaring/1000 time: [1.6535 us 1.6579 us 1.6630 us] thrpt: [601.33 Melem/s 603.16 Melem/s 604.78 Melem/s] change: time: [-91.569% -91.513% -91.457%] (p = 0.00 < 0.05) thrpt: [+1070.6% +1078.3% +1086.1%] Performance has improved. add_range/pre_populated_roaring/1000 time: [1.7225 us 1.7491 us 1.7824 us] thrpt: [561.03 Melem/s 571.72 Melem/s 580.56 Melem/s] change: time: [-94.873% -94.807% -94.737%] (p = 0.00 < 0.05) thrpt: [+1800.0% +1825.8% +1850.5%] Performance has improved. add_range/roaring/5000 time: [317.68 ns 319.43 ns 321.60 ns] thrpt: [15.547 Gelem/s 15.653 Gelem/s 15.739 Gelem/s] change: time: [-99.669% -99.665% -99.661%] (p = 0.00 < 0.05) thrpt: [+29373% +29744% +30135%] Performance has improved. add_range/pre_populated_roaring/5000 time: [844.07 ns 907.00 ns 966.49 ns] thrpt: [5.1733 Gelem/s 5.5127 Gelem/s 5.9237 Gelem/s] change: time: [-98.691% -98.541% -98.365%] (p = 0.00 < 0.05) thrpt: [+6015.9% +6753.8% +7541.4%] Performance has improved. add_range/roaring/10000 time: [373.12 ns 373.91 ns 374.77 ns] thrpt: [26.683 Gelem/s 26.745 Gelem/s 26.801 Gelem/s] change: time: [-99.736% -99.734% -99.732%] (p = 0.00 < 0.05) thrpt: [+37214% +37481% +37710%] Performance has improved. add_range/pre_populated_roaring/10000 time: [936.65 ns 1.0032 us 1.0666 us] thrpt: [9.3760 Gelem/s 9.9681 Gelem/s 10.676 Gelem/s] change: time: [-99.237% -99.132% -99.035%] (p = 0.00 < 0.05) thrpt: [+10258% +11415% +13008%] Performance has improved. add_range/roaring/100000 time: [1.4464 us 1.4534 us 1.4613 us] thrpt: [68.431 Gelem/s 68.806 Gelem/s 69.135 Gelem/s] change: time: [-99.866% -99.864% -99.861%] (p = 0.00 < 0.05) thrpt: [+71984% +73250% +74639%] Performance has improved. add_range/pre_populated_roaring/100000 time: [1.8819 us 1.9888 us 2.1003 us] thrpt: [47.612 Gelem/s 50.283 Gelem/s 53.137 Gelem/s] change: time: [-99.864% -99.849% -99.831%] (p = 0.00 < 0.05) thrpt: [+59136% +66238% +73485%] Performance has improved. ``` </details> The API exposed copies the `remove_range()` function, accepting a `Range<u64>` (a half-open range) in order to allow including `u32::MAX`. Co-authored-by: Dom <[email protected]>
RoaringBitmap · Dec 6, 2020 · 5e591ef · 5e591ef
2 parents ac237a9 + 03023f7
commit 5e591ef
Show file tree

Hide file tree

Showing 5 changed files with 404 additions and 3 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -22,6 +22,8 @@ byteorder = "1.0"
 
 [dev-dependencies]
 criterion = "0.3"
+quickcheck = "0.9"
+quickcheck_macros = "0.9"
 
 [[bench]]
 name = "lib"

diff --git a/benches/lib.rs b/benches/lib.rs
@@ -206,6 +206,30 @@ fn remove_range_bitmap(c: &mut Criterion) {
     });
 }
 
+fn insert_range_bitmap(c: &mut Criterion) {
+    for &size in &[10, 100, 1_000, 5_000, 10_000, 20_000] {
+        let mut group = c.benchmark_group("insert_range");
+        group.throughput(criterion::Throughput::Elements(size));
+        group.bench_function(format!("from_empty_{}", size), |b| {
+            let bm = RoaringBitmap::new();
+            b.iter_batched(
+                || bm.clone(),
+                |mut bm| black_box(bm.insert_range(0..size)),
+                criterion::BatchSize::SmallInput,
+            )
+        });
+        group.bench_function(format!("pre_populated_{}", size), |b| {
+            let mut bm = RoaringBitmap::new();
+            bm.insert_range(0..size);
+            b.iter_batched(
+                || bm.clone(),
+                |mut bm| black_box(bm.insert_range(0..size)),
+                criterion::BatchSize::SmallInput,
+            )
+        });
+    }
+}
+
 fn iter(c: &mut Criterion) {
     c.bench_function("iter", |b| {
         let bitmap: RoaringBitmap = (1..10_000).collect();
@@ -300,6 +324,7 @@ criterion_group!(
     is_subset,
     remove,
     remove_range_bitmap,
+    insert_range_bitmap,
     iter,
     is_empty,
     serialize,

diff --git a/src/bitmap/container.rs b/src/bitmap/container.rs
@@ -1,4 +1,4 @@
-use std::fmt;
+use std::{fmt, ops::Range};
 
 use super::store::{self, Store};
 use super::util;
@@ -38,6 +38,19 @@ impl Container {
         }
     }
 
+    pub fn insert_range(&mut self, range: Range<u16>) -> u64 {
+        // If the range is larger than the array limit, skip populating the
+        // array to then have to convert it to a bitmap anyway.
+        if matches!(self.store, Store::Array(_)) && range.end - range.start > ARRAY_LIMIT as u16 {
+            self.store = self.store.to_bitmap()
+        }
+
+        let inserted = self.store.insert_range(range);
+        self.len += inserted;
+        self.ensure_correct_store();
+        inserted
+    }
+
     pub fn push(&mut self, index: u16) {
         if self.store.push(index) {
             self.len += 1;

diff --git a/src/bitmap/inherent.rs b/src/bitmap/inherent.rs
@@ -43,6 +43,102 @@ impl RoaringBitmap {
         container.insert(index)
     }
 
+    /// Inserts a range of values from the set specific as [start..end). Returns
+    /// the number of inserted values.
+    ///
+    /// Note that due to the exclusive end this functions take indexes as u64
+    /// but you still can't index past 2**32 (u32::MAX + 1).
+    ///
+    /// # Safety
+    ///
+    /// This function panics if the range upper bound exceeds `u32::MAX`.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use roaring::RoaringBitmap;
+    ///
+    /// let mut rb = RoaringBitmap::new();
+    /// rb.insert_range(2..4);
+    /// assert!(rb.contains(2));
+    /// assert!(rb.contains(3));
+    /// assert!(!rb.contains(4));
+    /// ```
+    pub fn insert_range(&mut self, range: Range<u64>) -> u64 {
+        assert!(
+            range.end <= u64::from(u32::max_value()) + 1,
+            "can't index past 2**32"
+        );
+        if range.is_empty() {
+            return 0;
+        }
+
+        let (start_container_key, start_index) = util::split(range.start as u32);
+        let (end_container_key, end_index) = util::split((range.end) as u32);
+
+        // Find the container index for start_container_key
+        let start_i = match self
+            .containers
+            .binary_search_by_key(&start_container_key, |c| c.key)
+        {
+            Ok(loc) => loc,
+            Err(loc) => {
+                self.containers
+                    .insert(loc, Container::new(start_container_key));
+                loc
+            }
+        };
+
+        // If the end range value is in the same container, just call into
+        // the one container.
+        if start_container_key == end_container_key {
+            return self.containers[start_i].insert_range(start_index..end_index);
+        }
+
+        // For the first container, insert start_index..u16::MAX, with
+        // subsequent containers inserting 0..MAX.
+        //
+        // The last container (end_container_key) is handled explicitly outside
+        // the loop.
+        let mut low = start_index;
+        let mut inserted = 0;
+
+        // Walk through the containers until the container for end_container_key
+        let end_i = usize::from(end_container_key - start_container_key);
+        for i in start_i..end_i {
+            // Fetch (or upsert) the container for i
+            let c = match self.containers.get_mut(i) {
+                Some(c) => c,
+                None => {
+                    // For each i, the container key is start_container + i in
+                    // the upper u8 of the u16.
+                    let key = start_container_key + ((1 << 8) * i) as u16;
+                    self.containers.insert(i, Container::new(key));
+                    &mut self.containers[i]
+                }
+            };
+
+            // Insert the range subset for this container
+            inserted += c.insert_range(low..u16::MAX);
+
+            // After the first container, always fill the containers.
+            low = 0;
+        }
+
+        // Handle the last container
+        let c = match self.containers.get_mut(end_i) {
+            Some(c) => c,
+            None => {
+                let (key, _) = util::split(range.start as u32);
+                self.containers.insert(end_i, Container::new(key));
+                &mut self.containers[end_i]
+            }
+        };
+        c.insert_range(0..end_index);
+
+        inserted
+    }
+
     /// Adds a value to the set.
     /// The value **must** be greater or equal to the maximum value in the set.
     ///
@@ -131,7 +227,7 @@ impl RoaringBitmap {
             range.end <= u64::from(u32::max_value()) + 1,
             "can't index past 2**32"
         );
-        if range.start == range.end {
+        if range.is_empty() {
             return 0;
         }
         // inclusive bounds for start and end
@@ -292,3 +388,61 @@ impl Default for RoaringBitmap {
         RoaringBitmap::new()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use quickcheck_macros::quickcheck;
+
+    #[quickcheck]
+    fn insert_range(r: Range<u32>, checks: Vec<u32>) {
+        let r: Range<u64> = u64::from(r.start)..u64::from(r.end);
+
+        let mut b = RoaringBitmap::new();
+        let inserted = b.insert_range(r.clone());
+        if r.end > r.start {
+            assert_eq!(inserted, r.end - r.start);
+        } else {
+            assert_eq!(inserted, 0);
+        }
+
+        // Assert all values in the range are present
+        for i in r.clone() {
+            assert!(b.contains(i as u32), format!("does not contain {}", i));
+        }
+
+        // Run the check values looking for any false positives
+        for i in checks {
+            let bitmap_has = b.contains(i);
+            let range_has = r.contains(&u64::from(i));
+            assert!(
+                bitmap_has == range_has,
+                format!(
+                    "value {} in bitmap={} and range={}",
+                    i, bitmap_has, range_has
+                )
+            );
+        }
+    }
+
+    #[test]
+    fn test_insert_range_same_container() {
+        let mut b = RoaringBitmap::new();
+        let inserted = b.insert_range(1..5);
+        assert_eq!(inserted, 4);
+
+        for i in 1..5 {
+            assert!(b.contains(i));
+        }
+    }
+
+    #[test]
+    fn test_insert_range_pre_populated() {
+        let mut b = RoaringBitmap::new();
+        let inserted = b.insert_range(1..20_000);
+        assert_eq!(inserted, 19_999);
+
+        let inserted = b.insert_range(1..20_000);
+        assert_eq!(inserted, 0);
+    }
+}