Implement a metric space variant of geofilter

crates/geo_filters/Cargo.toml

@@ -49,6 +49,11 @@ path = "evaluation/performance.rs"
 harness = false
 required-features = ["evaluation"]
 
+[[bench]]
+name = "nearest_neighbor"
+path = "evaluation/nearest_neighbor.rs"
+harness = false
+
 [[bin]]
 name = "accuracy"
 path = "evaluation/accuracy.rs"

crates/geo_filters/README.md

@@ -61,6 +61,61 @@ 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.
 
+## Nearest-neighbor similarity metric
+
+`GeoDiffCount` also doubles as a compact similarity sketch. The number of differing one-bits between
+two filters — the Hamming distance of their bit representations — is a simple, uncalibrated distance
+that grows with the true symmetric-difference size. This is useful for nearest-neighbor search, e.g.
+to find the most similar repository or document.
+
+[`GeoDiffMetric`](https://docs.rs/geo_filters/latest/geo_filters/diff_count/struct.GeoDiffMetric.html)
+wraps a `GeoDiffCount`, caches its one-bit count, and implements the `MetricSpace` trait:
+
+- `size()` returns the filter's one-bit count as a `Metric` value;
+- `symmetric_diff_size(other, bound)` returns the exact one-bit distance to another filter, but
+  abandons the computation once it reaches `bound` (returning an "infinite" value), so far-away
+  candidates are rejected cheaply while scanning a candidate list;
+- `size().abs_diff(&other.size())` is an `O(1)` reverse-triangle lower bound on that distance.
+
+```rust
+use geo_filters::diff_count::{GeoDiffCount7, GeoDiffMetric, OnesMetric};
+use geo_filters::{Count, Metric, MetricSpace};
+
+let mut a = GeoDiffCount7::default();
+(0..1000u64).for_each(|i| a.push(i));
+let mut b = GeoDiffCount7::default();
+(500..1500u64).for_each(|i| b.push(i));
+
+let a = GeoDiffMetric::new(a);
+let b = GeoDiffMetric::new(b);
+
+// O(1) lower bound from the cached sizes, then the exact distance (unbounded).
+let lower_bound = a.size().abs_diff(&b.size());
+let distance = a.symmetric_diff_size(&b, OnesMetric::infinite());
+assert!(lower_bound <= distance);
+```
+
+Time to compare a query filter (1M items) with a candidate, given the distance to a known near
+neighbor as the pruning bound. A `far` candidate is disjoint (farther than the bound, so it is
+rejected), while a `near` candidate is a closer near-duplicate (below the bound, so it becomes the
+new best and is scanned in full). Release build; absolute numbers are hardware-dependent:
+
+| configuration     | candidate | `estimate` | `symmetric_diff_size` | `symmetric_diff_size` (capped) |
+| ----------------- | --------- | ---------- | --------------------- | ------------------------------ |
+| `GeoDiffConfig7`  | far       | 115 ns     | 60 ns                 | 12 ns                          |
+| `GeoDiffConfig7`  | near      | 149 ns     | 50 ns                 | 51 ns                          |
+| `GeoDiffConfig10` | far       | 633 ns     | 374 ns                | 74 ns                          |
+| `GeoDiffConfig10` | near      | 918 ns     | 298 ns                | 300 ns                         |
+| `GeoDiffConfig13` | far       | 4.69 µs    | 2.44 µs               | 337 ns                         |
+| `GeoDiffConfig13` | near      | 6.92 µs    | 1.90 µs               | 1.92 µs                        |
+
+`estimate` is the calibrated `size_with_sketch` estimate and `symmetric_diff_size` the exact bit
+distance; the capped variant abandons once the running distance reaches the bound. For a `far`
+candidate it abandons almost immediately (~7× faster than the exact distance, ~14× faster than the
+estimate), whereas a `near` candidate is below the bound and therefore scanned in full — the capped
+and exact costs match. Since a nearest-neighbor scan rejects far more candidates than it keeps, the
+early abandon dominates the search cost. Reproduce with `cargo bench --bench nearest_neighbor`.
+
 ## Evaluation
 
 Accuracy and performance evaluations for the predefined filter configurations are included in the repository:

crates/geo_filters/evaluation/nearest_neighbor.rs

@@ -0,0 +1,128 @@
+//! Benchmarks the different ways to measure (dis)similarity of `GeoDiffCount` filters, in a
+//! nearest-neighbor setting where a `query` is compared against many candidates.
+//!
+//! Two families are compared:
+//!
+//! * The calibrated *size estimate* ([`Count::size`] / [`Count::size_with_sketch`]), which scans
+//!   only a bounded window and upscales.
+//! * The exact *bit count* via [`GeoDiffMetric`], a simple, uncalibrated metric based on the number
+//!   of differing one-bits (Hamming distance). It caches each filter's one-bit count (`size`) and
+//!   exposes `symmetric_diff_size`, the exact distance, abandoned once a given bound is reached.
+//!
+//! Both are measured for a `far` candidate (disjoint, rejected by the bound) and a `kept` candidate
+//! (a near-duplicate, roughly at the bound and scanned in full), to contrast pruning vs. keeping.
+//!
+//! Groups:
+//! * `size:*` - single filter: `estimate` (calibrated) vs. `metric` (exact one-bit count).
+//! * `diff:*` - pairwise: `estimate`/`symmetric_diff_size`/`symmetric_diff_size_capped` for the
+//!   `far` and `kept` candidates.
+
+use std::hint::black_box;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use geo_filters::config::GeoConfig;
+use geo_filters::diff_count::{
+    GeoDiffConfig10, GeoDiffConfig13, GeoDiffConfig7, GeoDiffCount, GeoDiffMetric, OnesMetric,
+};
+use geo_filters::{Count, Diff, Metric, MetricSpace};
+use rand::{Rng, SeedableRng};
+use rand_chacha::ChaCha12Rng;
+
+fn build<C: GeoConfig<Diff> + Default>(
+    rng: &mut ChaCha12Rng,
+    n: usize,
+) -> GeoDiffCount<'static, C> {
+    let mut f = GeoDiffCount::<C>::default();
+    for _ in 0..n {
+        f.push_hash(rng.next_u64());
+    }
+    f
+}
+
+/// Returns a near-duplicate of `base` that differs from it by roughly `diff` items.
+fn near<C: GeoConfig<Diff> + Default>(
+    rng: &mut ChaCha12Rng,
+    base: &GeoDiffCount<'static, C>,
+    diff: usize,
+) -> GeoDiffCount<'static, C> {
+    let mut f = base.clone();
+    for _ in 0..diff {
+        f.push_hash(rng.next_u64());
+    }
+    f
+}
+
+fn bench_config<C: GeoConfig<Diff> + Default>(c: &mut Criterion, name: &str, items: usize) {
+    let mut rng = ChaCha12Rng::seed_from_u64(42);
+
+    // Each filter is owned by its metric, which caches the one-bit count, as a real index would.
+    let query_m = GeoDiffMetric::new(build::<C>(&mut rng, items));
+    // A far-away candidate (an independent random set), and a near neighbor differing by ~10% of the base filter
+    // whose distance provides the prior bound.
+    let far_m = GeoDiffMetric::new(build::<C>(&mut rng, items));
+    let bound_m = GeoDiffMetric::new(near(&mut rng, query_m.filter(), (items / 10).max(1)));
+    // A kept candidate: a closer ~5% near-duplicate, i.e. below the bound and hence scanned in full.
+    let kept_m = GeoDiffMetric::new(near(&mut rng, query_m.filter(), (items / 20).max(1)));
+    let ones_bound = query_m.symmetric_diff_size(&bound_m, OnesMetric::infinite());
+
+    // Single-filter size metric: the estimate vs. the exact bit count computed when constructing a
+    // `GeoDiffMetric` (the filter is cloned in the untimed setup so only the count is measured).
+    let mut group = c.benchmark_group(format!("size:{name}/{items}"));
+    group.bench_function("estimate", |b| {
+        b.iter(|| black_box(query_m.filter().size()))
+    });
+    group.bench_function("metric", |b| {
+        b.iter_batched(
+            || query_m.filter().clone(),
+            |f| black_box(GeoDiffMetric::new(f).size()),
+            BatchSize::SmallInput,
+        )
+    });
+    group.finish();
+
+    // Diff between two filters, for a `far` candidate (disjoint, pruned by the bound) and a `kept`
+    // candidate (a ~10% near-duplicate, roughly at the bound and scanned in full):
+    // * `estimate`: calibrated size of the symmetric difference (`Count::size_with_sketch`);
+    // * `symmetric_diff_size`: exact one-bit distance (`MetricSpace::symmetric_diff_size` with an
+    //   infinite bound), scanning both filters in full;
+    // * `symmetric_diff_size_capped`: same, but abandons once `ones_bound` differing bits are reached.
+    let mut group = c.benchmark_group(format!("diff:{name}/{items}"));
+    group.bench_function("estimate_far", |b| {
+        b.iter(|| black_box(query_m.filter().size_with_sketch(black_box(far_m.filter()))))
+    });
+    group.bench_function("symmetric_diff_size_far", |b| {
+        b.iter(|| black_box(query_m.symmetric_diff_size(black_box(&far_m), OnesMetric::infinite())))
+    });
+    group.bench_function("symmetric_diff_size_capped_far", |b| {
+        b.iter(|| black_box(query_m.symmetric_diff_size(black_box(&far_m), ones_bound)))
+    });
+    group.bench_function("estimate_kept", |b| {
+        b.iter(|| {
+            black_box(
+                query_m
+                    .filter()
+                    .size_with_sketch(black_box(kept_m.filter())),
+            )
+        })
+    });
+    group.bench_function("symmetric_diff_size_kept", |b| {
+        b.iter(|| {
+            black_box(query_m.symmetric_diff_size(black_box(&kept_m), OnesMetric::infinite()))
+        })
+    });
+    group.bench_function("symmetric_diff_size_capped_kept", |b| {
+        b.iter(|| black_box(query_m.symmetric_diff_size(black_box(&kept_m), ones_bound)))
+    });
+    group.finish();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    for items in [1_000, 10_000, 100_000, 1_000_000] {
+        bench_config::<GeoDiffConfig7>(c, "config7", items);
+        bench_config::<GeoDiffConfig10>(c, "config10", items);
+        bench_config::<GeoDiffConfig13>(c, "config13", items);
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

crates/geo_filters/src/diff_count.rs

@@ -14,10 +14,12 @@ use crate::{Count, Diff};
 
 mod bitvec;
 mod config;
+mod metric;
 mod sim_hash;
 
 use bitvec::*;
 pub use config::{GeoDiffConfig10, GeoDiffConfig13, GeoDiffConfig7};
+pub use metric::{GeoDiffMetric, OnesMetric};
 pub use sim_hash::SimHash;
 
 /// Diff count filter with a relative error standard deviation of ~0.125.

crates/geo_filters/src/diff_count/bitvec.rs

@@ -72,6 +72,19 @@ impl BitVec<'_> {
         self.num_bits
     }
 
+    /// The total number of one-bits set in the vector.
+    pub fn count_ones(&self) -> usize {
+        self.blocks
+            .iter()
+            .map(|block| block.count_ones() as usize)
+            .sum()
+    }
+
+    /// The raw 64-bit blocks backing the vector, from least to most significant.
+    pub fn blocks(&self) -> &[u64] {
+        self.blocks.deref()
+    }
+
     pub fn is_empty(&self) -> bool {
         self.num_bits() == 0
     }

crates/geo_filters/src/diff_count/config.rs

@@ -63,7 +63,7 @@ impl Lookups for Diff {
 // and the loop be aborted.
 // Note: Increasing the constant to switch earlier into approximation mode will mostly increase the
 // approximation error, but not reduce the runtime of the function by much.
-fn expected_diff_buckets(phi: f64, items: f64) -> (f64, f64) {
+pub(super) fn expected_diff_buckets(phi: f64, items: f64) -> (f64, f64) {
     let mut sum = 0.0;
     let mut derivative = 0.0;
     for k in 0.. {