Pack richer compact masked sort keys

crates/geo_filters/evaluation/masked_sort.rs

@@ -2,7 +2,7 @@ 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::diff_count::{GeoDiffConfig10, GeoDiffConfig13, GeoDiffConfig7, GeoDiffCount};
 use geo_filters::{Count, Diff};
 use rand::{Rng, SeedableRng};
 use rand_chacha::ChaCha8Rng;
@@ -82,6 +82,7 @@ fn bench_config<C: GeoConfig<Diff> + Default>(c: &mut Criterion, name: &str) {
 
 fn criterion_benchmark(c: &mut Criterion) {
     bench_config::<GeoDiffConfig7>(c, "geo_diff_count_7");
+    bench_config::<GeoDiffConfig10>(c, "geo_diff_count_10");
     bench_config::<GeoDiffConfig13>(c, "geo_diff_count_13");
 }
 

crates/geo_filters/src/config.rs

@@ -42,6 +42,14 @@ pub trait GeoConfig<M: Method>: Clone + Eq + Sized {
 
     fn bits_per_level(&self) -> usize;
 
+    /// The number of bits required to store positions returned by [`Self::hash_to_bucket`].
+    ///
+    /// The default uses the full bucket type width. Configurations can override this with a tighter
+    /// proven width to enable more compact representations of bucket positions.
+    fn bucket_position_bits(&self) -> u32 {
+        Self::BucketType::BITS
+    }
+
     /// The granularity of the geometric buckets.
     /// The size of the i-th bucket is determined by the formula:
     ///    (1 - phi) * phi^i
@@ -119,6 +127,11 @@ impl<
         1 << B
     }
 
+    #[inline]
+    fn bucket_position_bits(&self) -> u32 {
+        bucket_position_bits(B)
+    }
+
     #[inline]
     fn phi(&self) -> f32 {
         phi(B)
@@ -270,6 +283,11 @@ impl<M: Method, T: IsBucketType + 'static, H: ReproducibleBuildHasher> GeoConfig
         1 << self.b
     }
 
+    #[inline]
+    fn bucket_position_bits(&self) -> u32 {
+        bucket_position_bits(self.b)
+    }
+
     #[inline]
     fn phi(&self) -> f32 {
         phi(self.b)

crates/geo_filters/src/config/buckets.rs

@@ -145,23 +145,25 @@ impl IsBucketType for usize {
     }
 }
 
-/// Computes the largest bucket index for 64-bit hashes given that (1 << B) bits cover half
-/// the hash space.
+/// Computes the bits required to store bucket positions for 64-bit hashes given that (1 << B)
+/// buckets cover half the hash space.
 ///
 /// (1 << B) buckets cover half the hash space, i.e., buckets [k * (1<<B), (k+1) * (1<<B) cover
-/// the hashes with k leading zeros. For a 64-bit hash, this gives us 64 * (1<<B) buckets.
+/// the hashes with k leading zeros. The zero hash has 64 leading zeros and maps to the last bucket
+/// of the following level, so the inclusive maximum position is 65 * (1<<B) - 1.
 #[inline]
-pub(crate) fn largest_bucket(b: usize) -> usize {
-    64 * (1 << b)
+pub(crate) fn bucket_position_bits(b: usize) -> u32 {
+    u32::try_from(b).expect("B must fit in u32") + 7
 }
 
 #[inline]
 pub(crate) fn assert_bucket_type_large_enough<T: IsBucketType>(b: usize) {
+    let required_bits = bucket_position_bits(b);
     assert!(
-        largest_bucket(b).ilog2() < T::BITS,
-        "bucket type has {} bits, which is too small for B = {}, requires bits > {}",
+        required_bits <= T::BITS,
+        "bucket type has {} bits, which is too small for B = {}, requires {} bits",
         T::BITS,
         b,
-        largest_bucket(b).ilog2()
+        required_bits
     );
 }

crates/geo_filters/src/diff_count.rs

@@ -149,26 +149,28 @@ impl<'a, C: GeoConfig<Diff>> 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.
+    /// The key packs as many complete bucket positions as fit into a `u64`, most significant
+    /// position first. Positions use the smallest width that can hold the configuration's
+    /// bucket positions, as reported by [`GeoConfig::bucket_position_bits`]. Any remaining key bits
+    /// contain the most-significant portion of the next position. Because masking distributes over
+    /// the xor used by
+    /// [`Self::cmp_masked`] and the partial position is an order-preserving prefix, 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 {
         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,
-            "sort key packing requires a bucket type of at most 32 bits"
+        let position_bits = self.config.bucket_position_bits();
+        assert!(
+            (1..=u64::BITS).contains(&position_bits),
+            "bucket position width must be in 1..=64 (got {position_bits})"
         );
-        let per_word = (u64::BITS / bits) as usize;
+        let complete_positions = (u64::BITS / position_bits) as usize;
+        let remaining_bits = u64::BITS % position_bits;
 
         // 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
@@ -188,8 +190,21 @@ impl<'a, C: GeoConfig<Diff>> GeoDiffCount<'a, C> {
         let mut positions = msb.chain(lsb);
 
         let mut key = 0u64;
-        for _ in 0..per_word {
-            key = (key << bits) | positions.next().unwrap_or(0);
+        for _ in 0..complete_positions {
+            let position = positions.next().unwrap_or(0);
+            debug_assert!(
+                position_bits == u64::BITS || position < 1u64 << position_bits,
+                "bucket position {position} exceeds configured width {position_bits}"
+            );
+            key = (key << position_bits) | position;
+        }
+        if remaining_bits > 0 {
+            let position = positions.next().unwrap_or(0);
+            debug_assert!(
+                position_bits == u64::BITS || position < 1u64 << position_bits,
+                "bucket position {position} exceeds configured width {position_bits}"
+            );
+            key = (key << remaining_bits) | (position >> (position_bits - remaining_bits));
         }
         key
     }
@@ -1062,7 +1077,7 @@ mod tests {
 
     #[test]
     fn test_masked_sort_key() {
-        let masks: &[(u64, usize)] = &[
+        let fixed_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
@@ -1071,42 +1086,117 @@ mod tests {
         ];
 
         fn check<C: GeoConfig<Diff> + Default>(rnd: &mut ChaCha12Rng, masks: &[(u64, usize)]) {
-            let mut build = || {
-                let mut f = GeoDiffCount::<C>::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}",
-                    );
+            let filters = (0..8)
+                .map(|_| {
+                    let mut f = GeoDiffCount::<C>::new(C::default());
+                    let items = 250 + rnd.next_u64() as usize % 1500;
+                    for _ in 0..items {
+                        f.push_hash(rnd.next_u64());
+                    }
+                    f
+                })
+                .collect_vec();
+
+            for a_index in 0..filters.len() {
+                for b_index in (a_index + 1)..filters.len() {
+                    let a = &filters[a_index];
+                    let b = &filters[b_index];
+                    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::<GeoDiffConfig7>(rnd, masks);
-            check::<GeoDiffConfig13>(rnd, masks);
+            let mut masks = fixed_masks.to_vec();
+            for _ in 0..12 {
+                let mask_size = 1 + rnd.next_u64() as usize % 63;
+                let mask = (rnd.next_u64() & (u64::MAX >> (64 - mask_size))) | 1;
+                masks.push((mask, mask_size));
+            }
+            check::<GeoDiffConfig7>(rnd, &masks);
+            check::<GeoDiffConfig10>(rnd, &masks);
+            check::<GeoDiffConfig13>(rnd, &masks);
         });
     }
 
+    #[test]
+    fn test_masked_sort_key_packing_boundaries() {
+        fn check<C: GeoConfig<Diff> + Default>(
+            expected_bits: u32,
+            expected_complete: usize,
+            expected_remaining: u32,
+        ) {
+            let config = C::default();
+            let max_position = 65 * config.bits_per_level() - 1;
+            let position_bits = config.bucket_position_bits();
+            assert_eq!(position_bits, expected_bits);
+            assert_eq!((u64::BITS / position_bits) as usize, expected_complete);
+            assert_eq!(u64::BITS % position_bits, expected_remaining);
+
+            let positions = (0..expected_complete + 2)
+                .map(|offset| C::BucketType::from_usize(max_position - offset))
+                .collect_vec();
+            let filter = GeoDiffCount::<C>::from_ones(positions.iter().copied());
+            let actual = filter.masked_sort_key(1, 1);
+
+            let mut expected = 0;
+            for &position in positions.iter().take(expected_complete) {
+                expected = (expected << position_bits) | position.into_usize() as u64;
+            }
+            if expected_remaining > 0 {
+                expected = (expected << expected_remaining)
+                    | (positions[expected_complete].into_usize() as u64
+                        >> (position_bits - expected_remaining));
+            }
+            assert_eq!(actual, expected);
+
+            let common = positions[..expected_complete].iter().copied();
+            let prefix_step = 1usize << (position_bits - expected_remaining);
+            let lower = C::BucketType::from_usize(prefix_step - 1);
+            let higher = C::BucketType::from_usize(prefix_step);
+            let lower_filter = GeoDiffCount::<C>::from_ones(common.clone().chain([lower]));
+            let higher_filter = GeoDiffCount::<C>::from_ones(common.chain([higher]));
+            let lower_key = lower_filter.masked_sort_key(1, 1);
+            let higher_key = higher_filter.masked_sort_key(1, 1);
+            assert_eq!(lower_key.cmp(&higher_key), Ordering::Less);
+            assert_eq!(
+                lower_key.cmp(&higher_key),
+                lower_filter.cmp_masked(&higher_filter, 1, 1)
+            );
+        }
+
+        check::<GeoDiffConfig7>(14, 4, 8);
+        check::<GeoDiffConfig10>(17, 3, 13);
+        check::<GeoDiffConfig13>(20, 3, 4);
+
+        let empty = GeoDiffCount10::default();
+        let bucket_zero = GeoDiffCount10::from_ones([0]);
+        assert_eq!(
+            empty.masked_sort_key(1, 1),
+            bucket_zero.masked_sort_key(1, 1)
+        );
+        assert_eq!(empty.cmp_masked(&bucket_zero, 1, 1), Ordering::Less);
+    }
+
     #[test]
     fn test_bit_chunks() {
         prng_test_harness(100, |rnd| {