Introduce an optimized extend_by_hashes function

crates/geo_filters/evaluation/performance.rs

@@ -1,9 +1,10 @@
+use std::hash::BuildHasher;
 use std::hint::black_box;
 
 use criterion::{criterion_group, criterion_main, Criterion};
 use geo_filters::build_hasher::UnstableDefaultBuildHasher;
 use geo_filters::config::VariableConfig;
-use geo_filters::diff_count::{GeoDiffCount, GeoDiffCount13};
+use geo_filters::diff_count::{GeoDiffCount, GeoDiffCount13, GeoDiffCount7};
 use geo_filters::distinct_count::GeoDistinctCount13;
 use geo_filters::evaluation::hll::Hll14;
 use geo_filters::Count;
@@ -130,6 +131,94 @@ fn criterion_benchmark(c: &mut Criterion) {
             })
         });
     }
+
+    // Compare building a diff filter from a precomputed slice of hashes one by one
+    // (`push_hash`) versus in a single batch (`extend_by_hashes` on an empty filter). The hashes
+    // are precomputed so that only the construction cost is measured.
+    for size in [1000usize, 10000, 100000, 1000000] {
+        let mut group = c.benchmark_group(format!("construct:{size}"));
+        let build_hasher = UnstableDefaultBuildHasher::default();
+        let hashes: Vec<u64> = (0..size).map(|i| build_hasher.hash_one(i)).collect();
+
+        group.bench_function("geo_diff_count_7_push", |b| {
+            b.iter(|| {
+                let mut gc = GeoDiffCount7::default();
+                for &hash in &hashes {
+                    gc.push_hash(hash);
+                }
+                black_box(&gc);
+            })
+        });
+        group.bench_function("geo_diff_count_7_extend", |b| {
+            b.iter(|| {
+                let mut gc = GeoDiffCount7::default();
+                gc.extend_by_hashes(hashes.iter().copied());
+                black_box(&gc);
+            })
+        });
+        group.bench_function("geo_diff_count_13_push", |b| {
+            b.iter(|| {
+                let mut gc = GeoDiffCount13::default();
+                for &hash in &hashes {
+                    gc.push_hash(hash);
+                }
+                black_box(&gc);
+            })
+        });
+        group.bench_function("geo_diff_count_13_extend", |b| {
+            b.iter(|| {
+                let mut gc = GeoDiffCount13::default();
+                gc.extend_by_hashes(hashes.iter().copied());
+                black_box(&gc);
+            })
+        });
+    }
+
+    // Insert a batch of hashes into an existing filter: one-by-one `push_hash` vs the batched
+    // `extend_by_hashes`. The existing filter is cloned in the (untimed) batched setup so that only
+    // the insertion cost is measured, not the clone.
+    for (existing_n, batch_n) in [
+        (10000usize, 10000usize),
+        (100000, 10000),
+        (100000, 100000),
+        (1000, 100000),
+    ] {
+        let mut group = c.benchmark_group(format!("insert_batch:{existing_n}+{batch_n}"));
+        let build_hasher = UnstableDefaultBuildHasher::default();
+        let batch: Vec<u64> = (existing_n..existing_n + batch_n)
+            .map(|i| build_hasher.hash_one(i))
+            .collect();
+        let base = {
+            let mut gc = GeoDiffCount13::default();
+            for i in 0..existing_n {
+                gc.push_hash(build_hasher.hash_one(i));
+            }
+            gc
+        };
+
+        group.bench_function("push", |b| {
+            b.iter_batched(
+                || base.clone(),
+                |mut gc| {
+                    for &hash in &batch {
+                        gc.push_hash(hash);
+                    }
+                    gc
+                },
+                criterion::BatchSize::SmallInput,
+            )
+        });
+        group.bench_function("extend_by_hashes", |b| {
+            b.iter_batched(
+                || base.clone(),
+                |mut gc| {
+                    gc.extend_by_hashes(batch.iter().copied());
+                    gc
+                },
+                criterion::BatchSize::SmallInput,
+            )
+        });
+    }
 }
 
 criterion_group!(benches, criterion_benchmark);

crates/geo_filters/src/config/bitchunks.rs

@@ -20,6 +20,11 @@ impl BitChunk {
     }
 }
 
+/// Merges a descending stream of distinct one-bit positions (`leading`) with a descending stream
+/// of `BitChunk`s (`trailing`) into a single descending `BitChunk` stream. All leading positions
+/// must be more significant than all trailing bits, except that the least-significant leading block
+/// may overlap the most-significant trailing block (the two are or-ed). Leading positions must be
+/// distinct.
 pub(crate) fn iter_bit_chunks(
     leading: impl Iterator<Item = usize>,
     trailing: impl Iterator<Item = BitChunk>,
@@ -55,8 +60,7 @@ impl<I: Iterator<Item = BitChunk>, J: Iterator<Item = usize>> Iterator for BitCh
                     _ => break,
                 }
             }
-            // All leading bits were consumed, test whether it can be merged with
-            // trailing bits.
+            // All leading bits were consumed, test whether it can be merged with trailing bits.
             match self.trailing.peek() {
                 Some(BitChunk {
                     index: other_index,
@@ -78,6 +82,39 @@ impl<I: Iterator<Item = BitChunk>, J: Iterator<Item = usize>> Iterator for BitCh
     }
 }
 
+/// Turns a descending stream of one-bit positions into a descending `BitChunk` stream containing
+/// each position's parity. Unlike [`iter_bit_chunks`], positions may repeat: a position occurring
+/// an even number of times cancels out (xor), and a block that cancels out entirely is skipped.
+pub(crate) fn parity_bit_positions(
+    positions: impl Iterator<Item = usize>,
+) -> impl Iterator<Item = BitChunk> {
+    ParityBitPositions(positions.peekable())
+}
+
+struct ParityBitPositions<I: Iterator<Item = usize>>(Peekable<I>);
+
+impl<I: Iterator<Item = usize>> Iterator for ParityBitPositions<I> {
+    type Item = BitChunk;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            let (index, bit) = self.0.next()?.into_index_and_bit();
+            let mut block: u64 = bit.into_block();
+            while let Some(&other) = self.0.peek() {
+                if other.into_index() != index {
+                    break;
+                }
+                block ^= other.into_bit().into_block();
+                self.0.next();
+            }
+            // The block is empty only if its positions cancelled out entirely; skip it.
+            if block != 0 {
+                return Some(BitChunk::new(index, block));
+            }
+        }
+    }
+}
+
 /// Combine-s two `BitChunk` iterators using the given operator. Both iterators have to output
 /// `BitChunk`s from most to least significant and must not report duplicates!
 struct BinOpBitChunksIterator<
@@ -314,7 +351,7 @@ impl<T: IsBucketType, I: Iterator<Item = BitChunk>> Iterator for BitChunksOnes<T
 mod tests {
     use itertools::Itertools;
 
-    use super::{iter_ones, BitChunk};
+    use super::{iter_bit_chunks, iter_ones, parity_bit_positions, BitChunk};
 
     #[test]
     fn test_iter_ones() {
@@ -338,4 +375,40 @@ mod tests {
             iter_ones::<usize, _>(chunks.into_iter().peekable()).collect_vec()
         );
     }
+
+    #[test]
+    fn test_iter_bit_chunks() {
+        // Distinct leading bits merge within a block (via or) and merge with the trailing block at
+        // the boundary index.
+        let chunks = iter_bit_chunks(
+            vec![70, 67, 5].into_iter(),
+            vec![BitChunk::new(0, 1 << 2)].into_iter(),
+        )
+        .collect_vec();
+        assert_eq!(
+            chunks,
+            vec![
+                BitChunk::new(1, (1 << 6) | (1 << 3)), // 70, 67
+                BitChunk::new(0, (1 << 5) | (1 << 2)), // 5 (leading) and bit 2 (trailing)
+            ]
+        );
+    }
+
+    #[test]
+    fn test_parity_bit_positions() {
+        // Equal positions cancel (xor): 70 twice cancels, 67 stays; 5 three times stays, 3 once.
+        let chunks = parity_bit_positions(vec![70, 70, 67, 5, 5, 5, 3].into_iter()).collect_vec();
+        assert_eq!(
+            chunks,
+            vec![
+                BitChunk::new(1, 1 << 3),              // 67
+                BitChunk::new(0, (1 << 5) | (1 << 3)), // 5, 3
+            ]
+        );
+
+        // A block whose positions cancel out entirely is skipped.
+        assert!(parity_bit_positions(vec![5, 5].into_iter())
+            .collect_vec()
+            .is_empty());
+    }
 }

crates/geo_filters/src/config/lookup.rs

@@ -2,22 +2,24 @@ use crate::config::phi_f64;
 
 pub(crate) struct HashToBucketLookup {
     b: usize,
-    buckets: Vec<(usize, usize)>,
+    buckets: Vec<(u32, u32)>,
 }
 
 impl HashToBucketLookup {
     pub(crate) fn new(b: usize) -> Self {
-        let mut buckets = vec![(0, 0); 2 << b];
+        let mut buckets = vec![(0u32, 0u32); 2 << b];
         let mut last_filled_bucket = buckets.len();
         let phi = phi_f64(b);
         for bucket in 0..(1 << b) {
             let lower_bucket_limit = phi.powf((bucket + 1) as f64);
+            // `lower_hash_limit` is a 32-bit hash threshold: `lower_bucket_limit` lies in
+            // `[0.5, 1)`, so this value is always in `[0, 2^32)` and fits losslessly into a `u32`.
             let lower_hash_limit = ((lower_bucket_limit - 0.5) * 2.0f64.powf(33.0)) as usize;
             let lower_hash_bucket = lower_hash_limit >> (32 - b - 1);
             assert!(lower_hash_bucket < last_filled_bucket);
             while last_filled_bucket > lower_hash_bucket {
                 last_filled_bucket -= 1;
-                buckets[last_filled_bucket] = (bucket, lower_hash_limit);
+                buckets[last_filled_bucket] = (bucket as u32, lower_hash_limit as u32);
             }
         }
         assert_eq!(last_filled_bucket, 0);
@@ -38,8 +40,12 @@ impl HashToBucketLookup {
         } & 0xFFFFFFFF) as usize;
         // From those, the first B bits determine the bucket index in our lookup table.
         let idx = hash >> (32 - self.b - 1);
-        let offset = (hash < self.buckets[idx].1) as usize;
-        offset + self.buckets[idx].0 + (1 << self.b) * levels
+        // SAFETY: `hash` was masked to 32 bits, so `idx = hash >> (31 - b)` holds at most `b + 1`
+        // significant bits and is therefore always `< 2^(b+1) == 2 << b == self.buckets.len()`.
+        debug_assert!(idx < self.buckets.len());
+        let (base, threshold) = *unsafe { self.buckets.get_unchecked(idx) };
+        let offset = (hash < threshold as usize) as usize;
+        offset + base as usize + (1 << self.b) * levels
     }
 }