add a fast constructor if all items are available in one go

Author: aneubeck

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,50 @@ 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 (`from_hashes`). 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_from_hashes", |b| {
+            b.iter(|| {
+                black_box(GeoDiffCount7::from_hashes(
+                    Default::default(),
+                    hashes.iter().copied(),
+                ));
+            })
+        });
+        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_from_hashes", |b| {
+            b.iter(|| {
+                black_box(GeoDiffCount13::from_hashes(
+                    Default::default(),
+                    hashes.iter().copied(),
+                ));
+            })
+        });
+    }
 }
 
 criterion_group!(benches, criterion_benchmark);

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
     }
 }
 

crates/geo_filters/src/diff_count.rs

@@ -120,6 +120,55 @@ impl<'a, C: GeoConfig<Diff>> GeoDiffCount<'a, C> {
         result
     }
 
+    /// Constructs a [`GeoDiffCount`] directly from an iterator of item hashes.
+    ///
+    /// Inserting hashes one by one via [`Count::push_hash`] repeatedly toggles individual bits
+    /// and rebalances the sparse/dense split, which is wasteful when the hashes are known up
+    /// front. Instead, this function uses the (exact) number of hashes to estimate the split
+    /// point between the sparsely stored most-significant buckets (the "numbers") and the densely
+    /// stored least-significant buckets (the "bits"):
+    ///
+    ///   * Buckets at or above the estimated split point are collected, sorted, and reduced
+    ///     modulo two so that buckets occurring an even number of times cancel out. This matches
+    ///     the symmetric-difference semantics of [`Diff`] and only sorts the comparatively few
+    ///     most-significant buckets.
+    ///   * Buckets below the split point are folded directly into a dense bit vector, where
+    ///     repeated buckets cancel out as well.
+    ///
+    /// Both pieces are then combined into a single descending stream of one bits and handed to
+    /// [`Self::from_bit_chunks`], which re-establishes the filter invariants. The resulting filter
+    /// is identical to the one produced by pushing the same hashes one by one.
+    pub fn from_hashes(config: C, hashes: impl ExactSizeIterator<Item = u64>) -> Self {
+        let split = estimate_split_bucket(&config, hashes.len());
+
+        // Buckets >= split are sparse and collected for sorting; buckets < split are dense and
+        // folded directly into the bit vector, where repeated buckets cancel via xor. The toggler
+        // resolves the bit vector's owned storage once, keeping the per-bit work out of the loop.
+        let mut numbers = Vec::new();
+        let mut bits = BitVec::default();
+        bits.resize(split);
+        {
+            let mut toggler = bits.toggler();
+            for hash in hashes {
+                let bucket = config.hash_to_bucket(hash).into_usize();
+                if bucket >= split {
+                    numbers.push(bucket);
+                } else {
+                    toggler.toggle(bucket);
+                }
+            }
+        }
+
+        // Sort the most-significant buckets and drop those occurring an even number of times.
+        numbers.sort_unstable_by(|a, b| b.cmp(a));
+        erase_even_occurrences(&mut numbers);
+
+        Self::from_bit_chunks(
+            config,
+            iter_bit_chunks(numbers.into_iter(), bits.bit_chunks()),
+        )
+    }
+
     /// Compare two geometric filters after applying the specified mask.
     ///
     /// To reduce the number of operations, the implementation first xors the bit chunks together,
@@ -429,6 +478,47 @@ pub(crate) fn xor<C: GeoConfig<Diff>>(
     )
 }
 
+/// Estimates the bucket id that separates the sparse most-significant buckets ("numbers") from
+/// the dense least-significant buckets ("bits") for a filter built from `n` hashes.
+///
+/// The expected number of hashes falling into buckets `>= s` is `n * phi^s`. We pick `s` such
+/// that this is a small multiple of the most-significant bucket budget, so that after parity
+/// reduction the most-significant buckets are almost always supplied by the collected numbers.
+/// Choosing a slightly lower split only makes the collected set a bit larger; correctness does
+/// not depend on the estimate, since [`GeoDiffCount::from_bit_chunks`] re-splits the combined
+/// stream regardless.
+fn estimate_split_bucket<C: GeoConfig<Diff>>(config: &C, n: usize) -> usize {
+    let target = 2 * config.max_msb_len();
+    if n <= target {
+        return 0;
+    }
+    let ratio = target as f64 / n as f64;
+    let split = (ratio.ln() / config.phi_f64().ln()).floor() as usize;
+    // No bucket can ever exceed this bound, so never allocate a larger bit vector.
+    split.min(64 * config.bits_per_level())
+}
+
+/// Given a slice sorted in descending order, keeps a single copy of every value occurring an odd
+/// number of times and drops every value occurring an even number of times. The retained values
+/// stay sorted in descending order.
+fn erase_even_occurrences(values: &mut Vec<usize>) {
+    let mut write = 0;
+    let mut read = 0;
+    while read < values.len() {
+        let value = values[read];
+        let mut count = 0;
+        while read < values.len() && values[read] == value {
+            read += 1;
+            count += 1;
+        }
+        if count % 2 == 1 {
+            values[write] = value;
+            write += 1;
+        }
+    }
+    values.truncate(write);
+}
+
 impl<C: GeoConfig<Diff>> Count<Diff> for GeoDiffCount<'_, C> {
     fn push_hash(&mut self, hash: u64) {
         self.xor_bit(self.config.hash_to_bucket(hash));
@@ -542,6 +632,40 @@ mod tests {
         assert_eq!(m.iter_ones().count(), 101);
     }
 
+    /// Building a filter from an iterator of hashes must produce exactly the same filter as
+    /// pushing those hashes one by one.
+    #[test]
+    fn test_from_hashes() {
+        fn assert_from_hashes_matches<C: GeoConfig<Diff> + Default>(hashes: &[u64], n: usize) {
+            let mut expected: GeoDiffCount<'_, C> = GeoDiffCount::new(C::default());
+            for &hash in hashes {
+                expected.push_hash(hash);
+            }
+            let actual: GeoDiffCount<'_, C> =
+                GeoDiffCount::from_hashes(C::default(), hashes.iter().copied());
+            assert_eq!(expected, actual, "filter mismatch for n = {n}");
+            assert_eq!(
+                expected.iter_ones().collect_vec(),
+                actual.iter_ones().collect_vec(),
+                "ones mismatch for n = {n}",
+            );
+        }
+
+        prng_test_harness(4, |rnd| {
+            for n in [0usize, 1, 5, 50, 500, 5000, 50000] {
+                // Draw from a smaller pool so that buckets are hit multiple times, exercising the
+                // even-occurrence cancellation on both the dense bits and the sparse numbers path.
+                let pool: Vec<u64> = (0..n.div_ceil(2).max(1)).map(|_| rnd.next_u64()).collect();
+                let hashes: Vec<u64> = (0..n)
+                    .map(|_| *pool.iter().choose(rnd).expect("pool is non-empty"))
+                    .collect();
+
+                assert_from_hashes_matches::<GeoDiffConfig7>(&hashes, n);
+                assert_from_hashes_matches::<GeoDiffConfig13>(&hashes, n);
+            }
+        });
+    }
+
     #[test]
     fn test_estimate_fast() {
         prng_test_harness(1, |rnd| {

crates/geo_filters/src/diff_count/bitvec.rs

@@ -90,6 +90,17 @@ impl BitVec<'_> {
         self.blocks.to_mut()[block_idx] ^= bit_idx.into_block();
     }
 
+    /// Returns a [`BitToggler`] that toggles many bits without re-resolving the `Cow` on every
+    /// access. [`Self::toggle`] resolves `self.blocks.to_mut()` on every call, which keeps a
+    /// branch in the caller's hot loop even when the storage is already owned; resolving it once
+    /// up front avoids that overhead when toggling a large number of bits.
+    pub fn toggler(&mut self) -> BitToggler<'_> {
+        BitToggler {
+            num_bits: self.num_bits,
+            blocks: self.blocks.to_mut(),
+        }
+    }
+
     /// Returns an iterator over all blocks in reverse order.
     /// The blocks are represented as `BitChunk`s.
     pub fn bit_chunks(&self) -> impl Iterator<Item = BitChunk> + '_ {
@@ -200,6 +211,23 @@ impl Index<usize> for BitVec<'_> {
     }
 }
 
+/// Toggles bits in an already-owned [`BitVec`] without re-resolving the `Cow` on every call.
+/// Obtained via [`BitVec::toggler`].
+pub(crate) struct BitToggler<'a> {
+    num_bits: usize,
+    blocks: &'a mut [u64],
+}
+
+impl BitToggler<'_> {
+    /// Toggles the bit at the given zero-based position. The position must be `< num_bits`.
+    #[inline]
+    pub fn toggle(&mut self, index: usize) {
+        debug_assert!(index < self.num_bits);
+        let (block_idx, bit_idx) = index.into_index_and_bit();
+        self.blocks[block_idx] ^= bit_idx.into_block();
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;