perf: Countmin performance improvements

CHANGELOG.md

@@ -26,7 +26,7 @@ All significant changes to this project will be documented in this file.
 * Improve the readability of `Debug` output for HLL and CPC sketches and unions.
 * `KllSketch::quantile` and `ReqSketch::quantile` are faster at ranks `0.0` and `1.0`, with query work independent of the number of retained samples.
 * `BloomFilter::insert` is faster for integer and raw-byte inputs. `BloomFilter::contains_and_insert` is also faster when checking already-present integer values.
-* `CountMinSketch` updates are faster for integer and raw-byte inputs.
+* `CountMinSketch` updates and estimates are faster for all item types, with the largest gains for integer and raw-byte inputs. Counter placement and serialized bytes are unchanged.
 * `CountMinSketch::decay(1.0)` now takes constant time.
 * `CpcSketch` updates are faster for integer and raw-byte inputs.
 * `FrequentItemsSketch` updates are faster for integer and raw-byte keys.

benchmarks/countmin/estimate.rs

@@ -0,0 +1,58 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use datasketches::countmin::CountMinSketch;
+use datasketches::hash::value::raw_bytes;
+use divan::Bencher;
+use divan::black_box;
+use divan::counter::ItemsCount;
+
+use crate::hash_inputs::ITEMS;
+use crate::hash_inputs::bytes_32_values;
+
+#[divan::bench]
+fn u64(bencher: Bencher) {
+    let mut sketch = CountMinSketch::<u64>::new(4, 16_384).unwrap();
+    for value in 0..ITEMS as u64 {
+        sketch.update(value);
+    }
+
+    bencher.counter(ItemsCount::new(ITEMS)).bench_local(|| {
+        let mut total = 0u64;
+        for value in 0..ITEMS as u64 {
+            total = total.wrapping_add(sketch.estimate(black_box(value)));
+        }
+        black_box(total)
+    });
+}
+
+#[divan::bench]
+fn bytes_32(bencher: Bencher) {
+    let values = bytes_32_values();
+    let mut sketch = CountMinSketch::<u64>::new(4, 16_384).unwrap();
+    for value in &values {
+        sketch.update(raw_bytes::from_slice(value));
+    }
+
+    bencher.counter(ItemsCount::new(ITEMS)).bench_local(|| {
+        let mut total = 0u64;
+        for value in &values {
+            total = total.wrapping_add(sketch.estimate(raw_bytes::from_slice(black_box(value))));
+        }
+        black_box(total)
+    });
+}

benchmarks/countmin/mod.rs

@@ -16,4 +16,5 @@
 // under the License.
 
 mod decay;
+mod estimate;
 mod update;

datasketches/src/countmin/sketch.rs

@@ -55,6 +55,8 @@ pub struct CountMinSketch<T: CountMinValue> {
     total_weight: T,
     counts: Vec<T>,
     hash_seeds: Vec<u64>,
+    /// Multiplier for [`fast_remainder`], derived from `num_buckets`.
+    bucket_multiplier: u128,
 }
 
 impl<T: CountMinValue> CountMinSketch<T> {
@@ -511,14 +513,20 @@ impl<T: CountMinValue> CountMinSketch<T> {
             total_weight: T::ZERO,
             counts,
             hash_seeds,
+            bucket_multiplier: u128::MAX / u128::from(num_buckets) + 1,
         }
     }
 
     fn bucket_index<I: Hash>(&self, item: &I, seed: u64) -> usize {
         let mut hasher = MurmurHash3X64128::with_seed(seed);
         item.hash(&mut hasher);
         let (h1, _) = hasher.finish128();
-        (h1 % self.num_buckets as u64) as usize
+        if self.num_buckets.is_power_of_two() {
+            // Same result as `%` for a power-of-two divisor.
+            (h1 & (u64::from(self.num_buckets) - 1)) as usize
+        } else {
+            fast_remainder(h1, self.bucket_multiplier, self.num_buckets) as usize
+        }
     }
 }
 
@@ -626,3 +634,57 @@ fn make_hash_seeds(seed: u64, num_hashes: u8) -> Vec<u64> {
     }
     seeds
 }
+
+/// Returns `a % d` without a hardware division, using Lemire's direct remainder
+/// ("Faster Remainder by Direct Computation", 2019) with `multiplier` =
+/// `u128::MAX / d + 1`. Exact for every `u64` `a` and `u32` `d >= 2`.
+#[inline]
+fn fast_remainder(a: u64, multiplier: u128, d: u32) -> u64 {
+    let low_bits = multiplier.wrapping_mul(u128::from(a));
+    // High 64 bits of the 192-bit product `low_bits * d`.
+    let d = u128::from(d);
+    let lo = u128::from(low_bits as u64) * d;
+    let hi = (low_bits >> 64) * d;
+    ((hi + (lo >> 64)) >> 64) as u64
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_fast_remainder_matches_remainder() {
+        let mut x = 0x9E37_79B9_7F4A_7C15u64;
+        let check = |d: u32, a: u64| {
+            let multiplier = u128::MAX / u128::from(d) + 1;
+            assert_eq!(
+                fast_remainder(a, multiplier, d),
+                a % u64::from(d),
+                "a={a} d={d}"
+            );
+        };
+        for d in [3, 5, 7, 1000, 20_000, 65_537, u32::MAX - 1, u32::MAX] {
+            for a in [
+                0,
+                1,
+                u64::MAX,
+                u64::MAX - 1,
+                1 << 63,
+                (1 << 32) - 1,
+                1 << 32,
+            ] {
+                check(d, a);
+            }
+            for _ in 0..1_000_000 {
+                x ^= x << 13;
+                x ^= x >> 7;
+                x ^= x << 17;
+                check(d, x);
+            }
+        }
+        for d in (3..u32::MAX).step_by(1_000_003) {
+            check(d, u64::MAX);
+            check(d, u64::from(d) * 12_345 + 7);
+        }
+    }
+}

tests-integration/tests/countmin_test/sketch.rs

@@ -445,11 +445,11 @@ fn test_increment_single_key_like_rust_count_min_sketch() {
 #[test]
 fn test_estimated_size() {
     let mut sketch = CountMinSketch::<i64>::new(4, 128).unwrap();
-    assert_eq!(sketch.estimated_size(), 4200);
+    assert_eq!(sketch.estimated_size(), 4224);
 
     // The backing tables are allocated up front; updates do not grow the sketch.
     sketch.update("apple");
-    assert_eq!(sketch.estimated_size(), 4200);
+    assert_eq!(sketch.estimated_size(), 4224);
 }
 
 #[test]