diff --git a/datasketches/src/countmin/sketch.rs b/datasketches/src/countmin/sketch.rs index e69150c7..7e3c3edd 100644 --- a/datasketches/src/countmin/sketch.rs +++ b/datasketches/src/countmin/sketch.rs @@ -46,7 +46,7 @@ const MAX_TABLE_ENTRIES: usize = 1 << 30; #[derive(Debug, Clone, PartialEq)] pub struct CountMinSketch { num_hashes: u8, - num_buckets: u32, + buckets: BucketCount, seed: u64, seed_hash: u16, // Every bucket satisfies |count| <= total_weight, so a checked total also bounds bucket @@ -55,8 +55,45 @@ pub struct CountMinSketch { total_weight: T, counts: Vec, hash_seeds: Vec, - /// Multiplier for [`fast_remainder`], derived from `num_buckets`. - bucket_multiplier: u128, +} + +#[derive(Debug, Clone, PartialEq)] +struct BucketCount { + count: u32, + multiplier: u128, +} + +impl BucketCount { + fn new(count: u32) -> Self { + Self { + count, + multiplier: u128::MAX / u128::from(count) + 1, + } + } + + #[inline] + fn index(&self, hash: u64) -> usize { + if self.count.is_power_of_two() { + // Same result as `%` for a power-of-two divisor. + (hash & (u64::from(self.count) - 1)) as usize + } else { + self.remainder(hash) as usize + } + } + + /// Returns `a % self.count` without a hardware division, using Lemire's direct remainder + /// ("Faster Remainder by Direct Computation", 2019). + /// + /// Exact for every `u64` input and every `self.count >= 2`. + #[inline] + fn remainder(&self, a: u64) -> u64 { + let low_bits = self.multiplier.wrapping_mul(u128::from(a)); + // High 64 bits of the 192-bit product `low_bits * d`. + let d = u128::from(self.count); + let lo = u128::from(low_bits as u64) * d; + let hi = (low_bits >> 64) * d; + ((hi + (lo >> 64)) >> 64) as u64 + } } impl CountMinSketch { @@ -116,7 +153,7 @@ impl CountMinSketch { /// Returns the number of buckets per hash function. pub fn num_buckets(&self) -> u32 { - self.num_buckets + self.buckets.count } /// Returns the seed used by the sketch. @@ -131,7 +168,7 @@ impl CountMinSketch { /// Returns the relative error (epsilon) implied by the number of buckets. pub fn relative_error(&self) -> f64 { - std::f64::consts::E / self.num_buckets as f64 + std::f64::consts::E / self.buckets.count as f64 } /// Returns `true` if the total weight is zero. @@ -223,7 +260,7 @@ impl CountMinSketch { .checked_abs() .and_then(|weight| self.total_weight.checked_add(weight)) .expect("total absolute weight overflow"); - let num_buckets = self.num_buckets as usize; + let num_buckets = self.buckets.count as usize; for (row, seed) in self.hash_seeds.iter().enumerate() { let bucket = self.bucket_index(&item, *seed); let index = row * num_buckets + bucket; @@ -243,7 +280,7 @@ impl CountMinSketch { /// assert!(sketch.estimate("pear") >= 2); /// ``` pub fn estimate(&self, item: I) -> T { - let num_buckets = self.num_buckets as usize; + let num_buckets = self.buckets.count as usize; let mut min = T::MAX; for (row, seed) in self.hash_seeds.iter().enumerate() { let bucket = self.bucket_index(&item, *seed); @@ -291,7 +328,7 @@ impl CountMinSketch { /// ``` pub fn merge(&mut self, other: &CountMinSketch) -> Result<(), Error> { if self.num_hashes != other.num_hashes - || self.num_buckets != other.num_buckets + || self.buckets.count != other.buckets.count || self.seed != other.seed { return Err(Error::invalid_argument( @@ -340,7 +377,7 @@ impl CountMinSketch { bytes.write_u8(if self.is_empty() { FLAGS_IS_EMPTY } else { 0 }); bytes.write_u32_le(0); // unused - bytes.write_u32_le(self.num_buckets); + bytes.write_u32_le(self.buckets.count); bytes.write_u8(self.num_hashes); debug_assert_eq!( self.seed_hash, @@ -507,13 +544,12 @@ impl CountMinSketch { let hash_seeds = make_hash_seeds(seed, num_hashes); CountMinSketch { num_hashes, - num_buckets, + buckets: BucketCount::new(num_buckets), seed, seed_hash, total_weight: T::ZERO, counts, hash_seeds, - bucket_multiplier: u128::MAX / u128::from(num_buckets) + 1, } } @@ -521,12 +557,7 @@ impl CountMinSketch { let mut hasher = MurmurHash3X64128::with_seed(seed); item.hash(&mut hasher); let (h1, _) = hasher.finish128(); - 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 - } + self.buckets.index(h1) } } @@ -635,19 +666,6 @@ fn make_hash_seeds(seed: u64, num_hashes: u8) -> Vec { 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::*; @@ -656,12 +674,8 @@ mod tests { 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}" - ); + let buckets = BucketCount::new(d); + assert_eq!(buckets.remainder(a), 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 [ diff --git a/tests-integration/tests/countmin_test/sketch.rs b/tests-integration/tests/countmin_test/sketch.rs index a5d39189..9edae108 100644 --- a/tests-integration/tests/countmin_test/sketch.rs +++ b/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::::new(4, 128).unwrap(); - assert_eq!(sketch.estimated_size(), 4224); + assert_eq!(sketch.estimated_size(), 4240); // The backing tables are allocated up front; updates do not grow the sketch. sketch.update("apple"); - assert_eq!(sketch.estimated_size(), 4224); + assert_eq!(sketch.estimated_size(), 4240); } #[test]