From 7701feb9f06c5ace2cf3e553010d3566fd9f9ab9 Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Mon, 6 Jul 2026 11:39:13 +0200 Subject: [PATCH 1/2] introduce a fast masked key for sorting --- crates/geo_filters/Cargo.toml | 5 ++ crates/geo_filters/README.md | 45 ++++++++++ crates/geo_filters/evaluation/masked_sort.rs | 89 +++++++++++++++++++ crates/geo_filters/src/diff_count.rs | 90 ++++++++++++++++++++ 4 files changed, 229 insertions(+) create mode 100644 crates/geo_filters/evaluation/masked_sort.rs diff --git a/crates/geo_filters/Cargo.toml b/crates/geo_filters/Cargo.toml index 62670385..1d837747 100644 --- a/crates/geo_filters/Cargo.toml +++ b/crates/geo_filters/Cargo.toml @@ -49,6 +49,11 @@ path = "evaluation/performance.rs" harness = false required-features = ["evaluation"] +[[bench]] +name = "masked_sort" +path = "evaluation/masked_sort.rs" +harness = false + [[bin]] name = "accuracy" path = "evaluation/accuracy.rs" diff --git a/crates/geo_filters/README.md b/crates/geo_filters/README.md index 8b5ebf2f..6f9c2e87 100644 --- a/crates/geo_filters/README.md +++ b/crates/geo_filters/README.md @@ -61,6 +61,51 @@ Our implementation is at least as fast as the HLL++ port for all sizes. The `GeoDistinctCount` solution only needs to count the number of occupied buckets and the offset of the first bucket. Both numbers are tracked during incremental updates, such that the distinct count can be computed in constant time. +## Sorting filters by masked similarity + +`GeoDiffCount::cmp_masked` compares two filters after applying a repeated bit mask, a locality-sensitive projection that induces a total order in which similar filters end up close together. +Sorting or bucketing a large collection this way is dominated by the cost of the pairwise `cmp_masked` calls. + +`GeoDiffCount::masked_sort_key` builds a compact `u64` sort key from the most significant masked bits of a filter. +Comparing two keys numerically reproduces the `cmp_masked` order whenever the keys differ, so a sort only falls back to the (much slower) `cmp_masked` on the rare key ties: + +```rust +use geo_filters::diff_count::GeoDiffCount7; +use geo_filters::Count; + +// A repeated 20-bit mask with 3 bits set acts as a locality-sensitive projection. +const MASK: u64 = 0b0000_0100_0000_1000_0001; +const MASK_SIZE: usize = 20; + +let filters: Vec<_> = (0..5u64) + .map(|s| { + let mut f = GeoDiffCount7::default(); + (0..1000).for_each(|i| f.push(s * 1000 + i)); + f + }) + .collect(); + +// Precompute one key per filter, then sort by comparing keys, only using the exact +// `cmp_masked` to break ties. +let keys: Vec = filters.iter().map(|f| f.masked_sort_key(MASK, MASK_SIZE)).collect(); +let mut order: Vec = (0..filters.len()).collect(); +order.sort_by(|&i, &j| { + keys[i] + .cmp(&keys[j]) + .then_with(|| filters[i].cmp_masked(&filters[j], MASK, MASK_SIZE)) +}); +assert_eq!(order.len(), filters.len()); +``` + +Sorting 1,000 filters of 1,000 items each with the 20-bit mask above, measured with `cargo bench --bench masked_sort`: + +| operation | `GeoDiffConfig7` | `GeoDiffConfig13` | +| --- | --- | --- | +| sort via `cmp_masked` | 675 µs | 862 µs | +| sort via `masked_sort_key` (incl. key construction) | 25 µs | 18 µs | +| speed-up | ~27× | ~47× | +| key construction only | 13 µs | 7 µs | + ## Evaluation Accuracy and performance evaluations for the predefined filter configurations are included in the repository: diff --git a/crates/geo_filters/evaluation/masked_sort.rs b/crates/geo_filters/evaluation/masked_sort.rs new file mode 100644 index 00000000..ae46055b --- /dev/null +++ b/crates/geo_filters/evaluation/masked_sort.rs @@ -0,0 +1,89 @@ +use std::hint::black_box; + +use criterion::{criterion_group, criterion_main, Criterion}; +use geo_filters::config::GeoConfig; +use geo_filters::diff_count::{GeoDiffConfig13, GeoDiffConfig7, GeoDiffCount}; +use geo_filters::{Count, Diff}; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +/// Repeated 20-bit mask with 3 bits set, applied to the filters before comparison. +const MASK: u64 = 0b0000_0100_0000_1000_0001; +const MASK_SIZE: usize = 20; + +/// Number of filters that are compared / sorted. +const N: usize = 1000; +/// Number of items inserted into each filter. +const ITEMS: usize = 1000; + +fn random_filters + Default>( + n: usize, + items: usize, + seed: u64, +) -> Vec> { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + (0..n) + .map(|_| { + let mut f = GeoDiffCount::::new(C::default()); + for _ in 0..items { + f.push_hash(rng.next_u64()); + } + f + }) + .collect() +} + +fn bench_config + Default>(c: &mut Criterion, name: &str) { + let filters = random_filters::(N, ITEMS, 42); + + let mut group = c.benchmark_group(format!("masked_sort/{name}")); + // Baseline: sort the filters by comparing them directly with the exact masked comparison. + group.bench_function("cmp_masked", |b| { + b.iter(|| { + let mut idx: Vec = (0..N as u32).collect(); + idx.sort_unstable_by(|&i, &j| { + filters[i as usize].cmp_masked(&filters[j as usize], MASK, MASK_SIZE) + }); + black_box(idx) + }) + }); + // Sort by building the sort keys first (construction is included in the measurement), then + // comparing them numerically, falling back to `cmp_masked` only on ties. + group.bench_function("sort_key", |b| { + b.iter(|| { + let keys: Vec = filters + .iter() + .map(|f| f.masked_sort_key(MASK, MASK_SIZE)) + .collect(); + let mut idx: Vec = (0..N as u32).collect(); + idx.sort_unstable_by(|&i, &j| { + let (i, j) = (i as usize, j as usize); + keys[i] + .cmp(&keys[j]) + .then_with(|| filters[i].cmp_masked(&filters[j], MASK, MASK_SIZE)) + }); + black_box(idx) + }) + }); + group.finish(); + + // Key construction on its own, i.e. the extra work the `sort_key` variant pays up front. + c.benchmark_group(format!("build_keys/{name}")) + .bench_function("masked_sort_key", |b| { + b.iter(|| { + let keys: Vec = filters + .iter() + .map(|f| f.masked_sort_key(MASK, MASK_SIZE)) + .collect(); + black_box(keys) + }) + }); +} + +fn criterion_benchmark(c: &mut Criterion) { + bench_config::(c, "geo_diff_count_7"); + bench_config::(c, "geo_diff_count_13"); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/crates/geo_filters/src/diff_count.rs b/crates/geo_filters/src/diff_count.rs index 728cd1f3..a401b38e 100644 --- a/crates/geo_filters/src/diff_count.rs +++ b/crates/geo_filters/src/diff_count.rs @@ -145,6 +145,49 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { } } + /// Builds a sort key from the most significant bits of the masked filter. + /// + /// The key packs the largest bucket positions of the masked filter into a single `u64`, + /// most significant position first. Because masking distributes over the xor used by + /// [`Self::cmp_masked`], comparing two keys numerically yields the same ordering as + /// [`Self::cmp_masked`] whenever the keys differ. When two keys are equal the ordering is + /// undetermined and the caller must fall back to [`Self::cmp_masked`], e.g. + /// `a_key.cmp(&b_key).then_with(|| a.cmp_masked(b, mask, mask_size))`. + /// + /// Each position occupies `C::BucketType::BITS` bits, so the key holds + /// `64 / C::BucketType::BITS` positions (4 for `u16`, 2 for `u32`). + pub fn masked_sort_key(&self, mask: u64, mask_size: usize) -> u64 { + let bits = C::BucketType::BITS; + debug_assert!( + (1..=32).contains(&bits) && u64::BITS % bits == 0, + "sort key packing requires a bucket type of at most 32 bits" + ); + let per_word = (u64::BITS / bits) as usize; + + // The most significant bits are stored sparsely and sorted from largest to smallest, so we + // can test each of them against the periodic mask directly, avoiding the more expensive + // merge performed by `bit_chunks`. The least significant bits are dense and always below + // every most significant bit, so we mask whole 64-bit blocks at once. Chaining the two + // therefore yields the masked one-bit positions from largest to smallest. + let msb = self + .msb + .iter() + .map(|&p| p.into_usize()) + .filter(move |&pos| (mask >> (pos % mask_size)) & 1 != 0) + .map(|pos| pos as u64); + let lsb = iter_ones::( + mask_bit_chunks(self.lsb.bit_chunks(), mask, mask_size).peekable(), + ) + .map(|b| b.into_usize() as u64); + let mut positions = msb.chain(lsb); + + let mut key = 0u64; + for _ in 0..per_word { + key = (key << bits) | positions.next().unwrap_or(0); + } + key + } + fn test_bit(&self, index: C::BucketType) -> bool { if index.into_usize() < self.lsb.num_bits() { self.lsb.test_bit(index.into_usize()) @@ -650,6 +693,53 @@ mod tests { }); } + #[test] + fn test_masked_sort_key() { + let masks: &[(u64, usize)] = &[ + (0b1, 1), // keeps every bit, i.e. a full comparison + (0b10, 2), // keeps every other bit + (0b110, 3), // keeps two out of every three bits + (0b110100, 6), + (0b100001100000, 12), + ]; + + fn check + Default>(rnd: &mut ChaCha12Rng, masks: &[(u64, usize)]) { + let mut build = || { + let mut f = GeoDiffCount::::new(C::default()); + for _ in 0..1000 { + f.push_hash(rnd.next_u64()); + } + f + }; + let a = build(); + let b = build(); + for &(mask, mask_size) in masks { + let ka = a.masked_sort_key(mask, mask_size); + let kb = b.masked_sort_key(mask, mask_size); + let expected = a.cmp_masked(&b, mask, mask_size); + // The key comparison plus fall back must always agree with the exact comparison. + assert_eq!( + ka.cmp(&kb).then_with(|| a.cmp_masked(&b, mask, mask_size)), + expected, + "keyed comparison mismatch for mask {mask:b}/{mask_size}", + ); + // Whenever the keys already differ, they alone must yield the exact order. + if ka != kb { + assert_eq!( + ka.cmp(&kb), + expected, + "key ordering mismatch for mask {mask:b}/{mask_size}", + ); + } + } + } + + prng_test_harness(20, |rnd| { + check::(rnd, masks); + check::(rnd, masks); + }); + } + #[test] fn test_bit_chunks() { prng_test_harness(100, |rnd| { From 708f39aa41d5d6af61c05151f63ff68f544c6219 Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Mon, 6 Jul 2026 14:12:05 +0200 Subject: [PATCH 2/2] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- crates/geo_filters/src/diff_count.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/geo_filters/src/diff_count.rs b/crates/geo_filters/src/diff_count.rs index a401b38e..0cb794d4 100644 --- a/crates/geo_filters/src/diff_count.rs +++ b/crates/geo_filters/src/diff_count.rs @@ -156,7 +156,11 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { /// /// Each position occupies `C::BucketType::BITS` bits, so the key holds /// `64 / C::BucketType::BITS` positions (4 for `u16`, 2 for `u32`). - pub fn masked_sort_key(&self, mask: u64, mask_size: usize) -> u64 { +pub fn masked_sort_key(&self, mask: u64, mask_size: usize) -> u64 { + assert!( + (1..u64::BITS as usize).contains(&mask_size), + "mask_size must be in 1..=63 (got {mask_size})" + ); let bits = C::BucketType::BITS; debug_assert!( (1..=32).contains(&bits) && u64::BITS % bits == 0,