From e99211eab560133089442a5c8a495d09ec9888f0 Mon Sep 17 00:00:00 2001 From: tison Date: Thu, 22 Oct 2026 15:48:39 +0800 Subject: [PATCH] fix(countmin): add portable v2 hashing --- CHANGELOG.md | 1 + datasketches/src/countmin/serialization.rs | 3 +- datasketches/src/countmin/sketch.rs | 97 +++++++++++++------ datasketches/src/hash/mod.rs | 6 +- datasketches/src/hash/xxhash.rs | 10 ++ .../tests/countmin_test/sketch.rs | 49 ++++++++++ .../tests/serde_tests/countmin.rs | 59 ++++++++++- 7 files changed, 191 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2e4e3ca3..08133bac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ All significant changes to this project will be documented in this file. ### Breaking changes +* New `CountMinSketch` values use serial version 2 with XxHash64-derived row seeds and item hashes, plus unsigned bucket mapping. Deserialization still accepts legacy Rust version 1 images and preserves their version when reserializing, but version 1 images from other implementations remain incompatible for queries, updates, and merges. Merging now requires matching serial versions. * `BloomFilter::invert` is removed; use `BloomFilter::difference` for approximate A-not-B. The result excludes items in the right filter, but hash collisions can also remove items unique to the left filter. * `FrequentItemsSketch::is_empty` now returns `false` when the stream weight is nonzero, even if no items are retained. Use `num_active_items() == 0` to test for zero retained items. * `ReqSketch` queries now use `datasketches::common::SearchCriteria` instead of `datasketches::req::SearchCriteria`. `SearchCriteria` no longer implements `Default`; explicitly choose `Inclusive` or `Exclusive`. diff --git a/datasketches/src/countmin/serialization.rs b/datasketches/src/countmin/serialization.rs index 7a930e01..6ec67456 100644 --- a/datasketches/src/countmin/serialization.rs +++ b/datasketches/src/countmin/serialization.rs @@ -16,6 +16,7 @@ // under the License. pub const PREAMBLE_LONGS_SHORT: u8 = 2; -pub const SERIAL_VERSION: u8 = 1; +pub const SERIAL_VERSION_1: u8 = 1; +pub const SERIAL_VERSION_2: u8 = 2; pub const FLAGS_IS_EMPTY: u8 = 1 << 0; pub const LONG_SIZE_BYTES: usize = 8; diff --git a/datasketches/src/countmin/sketch.rs b/datasketches/src/countmin/sketch.rs index 24a60e44..781b3ea1 100644 --- a/datasketches/src/countmin/sketch.rs +++ b/datasketches/src/countmin/sketch.rs @@ -21,7 +21,6 @@ use std::hash::Hasher; use crate::codec::SketchBytes; use crate::codec::SketchSlice; use crate::codec::assert::ensure_preamble_longs_in; -use crate::codec::assert::ensure_serial_version_is; use crate::codec::assert::insufficient_data; use crate::codec::family::Family; use crate::countmin::CountMinValue; @@ -29,11 +28,13 @@ use crate::countmin::UnsignedCountMinValue; use crate::countmin::serialization::FLAGS_IS_EMPTY; use crate::countmin::serialization::LONG_SIZE_BYTES; use crate::countmin::serialization::PREAMBLE_LONGS_SHORT; -use crate::countmin::serialization::SERIAL_VERSION; +use crate::countmin::serialization::SERIAL_VERSION_1; +use crate::countmin::serialization::SERIAL_VERSION_2; use crate::error::Error; use crate::error::ErrorKind; use crate::hash::DEFAULT_UPDATE_SEED; use crate::hash::MurmurHash3X64128; +use crate::hash::XxHash64; use crate::hash::check_seed_hash; use crate::hash::compute_seed_hash; @@ -49,6 +50,7 @@ pub struct CountMinSketch { num_buckets: u32, seed: u64, seed_hash: u16, + serial_version: u8, // Every bucket satisfies |count| <= total_weight, so a checked total also bounds bucket // additions during update/merge: |a + b| <= |a| + |b|. Deserialization validates this // invariant; unsigned halving and decay preserve it. @@ -103,6 +105,7 @@ impl CountMinSketch { num_buckets, seed, seed_hash, + SERIAL_VERSION_2, entries, )) } @@ -271,7 +274,8 @@ impl CountMinSketch { /// # Errors /// /// Returns an error without modifying the sketch if the sketches have different numbers of - /// hashes, bucket counts, or seeds, or their combined total absolute weight exceeds `T::MAX`. + /// hashes, bucket counts, seeds, or serial versions, or their combined total absolute weight + /// exceeds `T::MAX`. /// /// # Examples /// @@ -291,9 +295,10 @@ impl CountMinSketch { if self.num_hashes != other.num_hashes || self.num_buckets != other.num_buckets || self.seed != other.seed + || self.serial_version != other.serial_version { return Err(Error::invalid_argument( - "Count-Min sketches must have matching numbers of hashes, bucket counts, and seeds", + "Count-Min sketches must have matching numbers of hashes, bucket counts, seeds, and serial versions", )); } self.total_weight = self @@ -308,8 +313,12 @@ impl CountMinSketch { /// Serializes this sketch into the Count-Min binary format. /// - /// DataSketches C++ can read the serialized data, but item queries and updates there are - /// unreliable because the two implementations map items to counters differently. + /// New sketches use serial version 2. Sketches deserialized from version 1 images retain the + /// legacy Rust item-to-counter mapping and are serialized as version 1 again. + /// + /// Version 2 derives each row seed by hashing the row's 8-byte little-endian index with + /// XxHash64 and the sketch seed. It uses XxHash64 for item hashing and shifts the hash right by + /// one before taking the bucket modulo. /// /// # Examples /// @@ -333,7 +342,7 @@ impl CountMinSketch { let mut bytes = SketchBytes::with_capacity(header_size + payload_size); bytes.write_u8(PREAMBLE_LONGS_SHORT); - bytes.write_u8(SERIAL_VERSION); + bytes.write_u8(self.serial_version); bytes.write_u8(Family::COUNTMIN.id); bytes.write_u8(if self.is_empty() { FLAGS_IS_EMPTY } else { 0 }); bytes.write_u32_le(0); // unused @@ -360,9 +369,9 @@ impl CountMinSketch { /// Deserializes a sketch from bytes using the default seed. /// - /// Can read sketches serialized by DataSketches C++. - /// Item queries and updates on those sketches are unreliable because the two - /// implementations map items to counters differently. + /// Accepts serial versions 1 and 2. Version 1 uses the legacy Rust item-to-counter mapping; + /// queries, updates, and merges on version 1 images from other implementations remain + /// unreliable. /// /// # Errors /// @@ -386,9 +395,9 @@ impl CountMinSketch { /// Deserializes a sketch from bytes using the provided seed. /// - /// Can read sketches serialized by DataSketches C++. - /// Item queries and updates on those sketches are unreliable because the two - /// implementations map items to counters differently. + /// Accepts serial versions 1 and 2. Version 1 uses the legacy Rust item-to-counter mapping; + /// queries, updates, and merges on version 1 images from other implementations remain + /// unreliable. /// /// # Errors /// @@ -430,7 +439,11 @@ impl CountMinSketch { .map_err(insufficient_data(""))?; Family::COUNTMIN.validate_id(family_id)?; - ensure_serial_version_is(SERIAL_VERSION, serial_version)?; + if serial_version != SERIAL_VERSION_1 && serial_version != SERIAL_VERSION_2 { + return Err(Error::deserial(format!( + "unsupported serial version: expected {SERIAL_VERSION_1} or {SERIAL_VERSION_2}, got {serial_version}" + ))); + } ensure_preamble_longs_in(&[PREAMBLE_LONGS_SHORT], preamble_longs)?; let num_buckets = cursor @@ -468,7 +481,14 @@ impl CountMinSketch { } } - let mut sketch = Self::make(num_hashes, num_buckets, seed, expected_seed_hash, entries); + let mut sketch = Self::make( + num_hashes, + num_buckets, + seed, + expected_seed_hash, + serial_version, + entries, + ); if is_empty { return Ok(sketch); } @@ -500,14 +520,22 @@ impl CountMinSketch { + self.hash_seeds.capacity() * size_of::() } - fn make(num_hashes: u8, num_buckets: u32, seed: u64, seed_hash: u16, entries: usize) -> Self { + fn make( + num_hashes: u8, + num_buckets: u32, + seed: u64, + seed_hash: u16, + serial_version: u8, + entries: usize, + ) -> Self { let counts = vec![T::ZERO; entries]; - let hash_seeds = make_hash_seeds(seed, num_hashes); + let hash_seeds = make_hash_seeds(seed, num_hashes, serial_version); CountMinSketch { num_hashes, num_buckets, seed, seed_hash, + serial_version, total_weight: T::ZERO, counts, hash_seeds, @@ -515,10 +543,20 @@ impl CountMinSketch { } fn bucket_index(&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 + let hash = match self.serial_version { + SERIAL_VERSION_1 => { + let mut hasher = MurmurHash3X64128::with_seed(seed); + item.hash(&mut hasher); + hasher.finish128().0 + } + SERIAL_VERSION_2 => { + let mut hasher = XxHash64::with_seed(seed); + item.hash(&mut hasher); + hasher.finish64() >> 1 + } + _ => unreachable!("unsupported Count-Min serial version"), + }; + (hash % self.num_buckets as u64) as usize } } @@ -615,14 +653,19 @@ fn entries_for_config_checked(num_hashes: u8, num_buckets: u32) -> Result Vec { +fn make_hash_seeds(seed: u64, num_hashes: u8, serial_version: u8) -> Vec { let mut seeds = Vec::with_capacity(num_hashes as usize); for i in 0..num_hashes { - // Derive per-row hash seeds deterministically from the sketch seed. - let mut hasher = MurmurHash3X64128::with_seed(seed); - hasher.write(&u64::from(i).to_le_bytes()); - let (h1, _) = hasher.finish128(); - seeds.push(h1); + let row_seed = match serial_version { + SERIAL_VERSION_1 => { + let mut hasher = MurmurHash3X64128::with_seed(seed); + hasher.write(&u64::from(i).to_le_bytes()); + hasher.finish128().0 + } + SERIAL_VERSION_2 => XxHash64::hash_u64(u64::from(i), seed), + _ => unreachable!("unsupported Count-Min serial version"), + }; + seeds.push(row_seed); } seeds } diff --git a/datasketches/src/hash/mod.rs b/datasketches/src/hash/mod.rs index 2084902b..d0860293 100644 --- a/datasketches/src/hash/mod.rs +++ b/datasketches/src/hash/mod.rs @@ -38,9 +38,9 @@ mod murmurhash; ))] pub(crate) use self::murmurhash::*; -#[cfg(feature = "bloom")] +#[cfg(any(feature = "bloom", feature = "countmin"))] mod xxhash; -#[cfg(feature = "bloom")] +#[cfg(any(feature = "bloom", feature = "countmin"))] pub(crate) use self::xxhash::*; #[cfg(any( @@ -82,7 +82,7 @@ pub(crate) use self::seed::*; ))] pub(crate) const DEFAULT_UPDATE_SEED: u64 = 9001; -#[cfg(feature = "bloom")] +#[cfg(any(feature = "bloom", feature = "countmin"))] #[inline(always)] fn read_u32_le(bytes: &[u8]) -> u32 { u32::from_le_bytes(bytes.try_into().expect("four-byte hash input")) diff --git a/datasketches/src/hash/xxhash.rs b/datasketches/src/hash/xxhash.rs index 1ea55035..e3886a33 100644 --- a/datasketches/src/hash/xxhash.rs +++ b/datasketches/src/hash/xxhash.rs @@ -141,6 +141,16 @@ impl Hasher for XxHash64 { self.finish64() } + #[inline(always)] + fn write_i64(&mut self, value: i64) { + self.write(&value.to_le_bytes()); + } + + #[inline(always)] + fn write_u64(&mut self, value: u64) { + self.write(&value.to_le_bytes()); + } + #[inline(always)] fn write(&mut self, bytes: &[u8]) { self.total_len = self.total_len.wrapping_add(bytes.len() as u64); diff --git a/tests-integration/tests/countmin_test/sketch.rs b/tests-integration/tests/countmin_test/sketch.rs index ba2ab679..3ddbcb6a 100644 --- a/tests-integration/tests/countmin_test/sketch.rs +++ b/tests-integration/tests/countmin_test/sketch.rs @@ -400,6 +400,55 @@ fn test_serialize_deserialize_non_empty_u64() { assert_eq!(decoded.estimate(42u64), sketch.estimate(42u64)); } +#[test] +fn test_v2_known_answer_vectors() { + // Independently computed with XXH64 over little-endian `u64` values. + const NUM_BUCKETS: usize = 1000; + + for (seed, one_buckets, nine_buckets) in [ + (9001, [185, 898, 4], [490, 217, 440]), + (42, [468, 459, 696], [841, 428, 322]), + ] { + let mut sketch = CountMinSketch::::with_seed(3, NUM_BUCKETS as u32, seed).unwrap(); + sketch.update_with_weight(1u64, 10); + sketch.update_with_weight(9u64, 810); + + let bytes = sketch.serialize(); + assert_eq!(bytes[1], 2); + let counts: Vec = bytes[24..] + .chunks_exact(8) + .map(|bytes| u64::from_le_bytes(bytes.try_into().unwrap())) + .collect(); + assert_eq!(counts.len(), 3 * NUM_BUCKETS); + + for row in 0..3 { + let actual: Vec<(usize, u64)> = counts[row * NUM_BUCKETS..(row + 1) * NUM_BUCKETS] + .iter() + .copied() + .enumerate() + .filter(|(_, count)| *count != 0) + .collect(); + let mut expected = vec![(one_buckets[row], 10), (nine_buckets[row], 810)]; + expected.sort_unstable_by_key(|(bucket, _)| *bucket); + assert_eq!(actual, expected, "seed {seed}, row {row}"); + } + + let mut decoded = CountMinSketch::::deserialize_with_seed(&bytes, seed).unwrap(); + assert_eq!(decoded.total_weight(), 820); + assert_eq!(decoded.estimate(1u64), 10); + assert_eq!(decoded.estimate(9u64), 810); + assert_eq!(decoded.serialize(), bytes); + + decoded.update_with_weight(1u64, 5); + assert_eq!(decoded.estimate(1u64), 15); + let mut merged = CountMinSketch::::with_seed(3, NUM_BUCKETS as u32, seed).unwrap(); + merged.merge(&decoded).unwrap(); + assert_eq!(merged.total_weight(), 825); + assert_eq!(merged.estimate(1u64), 15); + assert_eq!(merged.estimate(9u64), 810); + } +} + #[test] fn test_truncated_non_empty_payload_is_rejected_before_table_allocation() { let mut bytes = CountMinSketch::::new(1, 3).unwrap().serialize(); diff --git a/tests-integration/tests/serde_tests/countmin.rs b/tests-integration/tests/serde_tests/countmin.rs index 4a62faaf..eb282c5b 100644 --- a/tests-integration/tests/serde_tests/countmin.rs +++ b/tests-integration/tests/serde_tests/countmin.rs @@ -25,6 +25,59 @@ use tests_integration::ZERO_HASH_SEED; use crate::serialization_test_data; +// Generated by CountMinSketch v0.5.0 with one row, five buckets, seed 123, and key 1 at weight 10. +const RUST_V1_SERIALIZED: [u8; 64] = [ + 2, 1, 18, 0, 0, 0, 0, 0, 5, 0, 0, 0, 1, 182, 5, 0, 10, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 0, 0, 0, 0, + 0, 0, 0, +]; + +#[test] +fn legacy_rust_v1_image_preserves_queries_updates_and_merges() { + let mut sketch = + CountMinSketch::::deserialize_with_seed(&RUST_V1_SERIALIZED, 123).unwrap(); + assert_eq!(sketch.total_weight(), 10); + assert_eq!(sketch.estimate(1u64), 10); + assert_eq!(sketch.serialize(), RUST_V1_SERIALIZED); + + sketch.update_with_weight(1u64, 5); + assert_eq!(sketch.total_weight(), 15); + assert_eq!(sketch.estimate(1u64), 15); + assert_eq!(sketch.serialize()[1], 1); + + let mut merged = + CountMinSketch::::deserialize_with_seed(&RUST_V1_SERIALIZED, 123).unwrap(); + merged.merge(&sketch).unwrap(); + assert_eq!(merged.total_weight(), 25); + assert_eq!(merged.estimate(1u64), 25); +} + +#[test] +fn rejects_merge_across_serial_versions() { + let legacy = CountMinSketch::::deserialize_with_seed(&RUST_V1_SERIALIZED, 123).unwrap(); + let mut current = CountMinSketch::::with_seed(1, 5, 123).unwrap(); + current.update_with_weight(1u64, 10); + let before = current.clone(); + + assert_eq!( + current.merge(&legacy).unwrap_err().kind(), + ErrorKind::InvalidArgument + ); + assert_eq!(current, before); +} + +#[test] +fn rejects_unknown_serial_version() { + let mut bytes = CountMinSketch::::new(1, 5).unwrap().serialize(); + bytes[1] = 3; + let error = CountMinSketch::::deserialize(&bytes).unwrap_err(); + assert_eq!(error.kind(), ErrorKind::InvalidData); + assert_that!( + error.message(), + contains_substring("expected 1 or 2, got 3") + ); +} + #[test] fn rejects_weights_that_cannot_bound_counter_arithmetic() { let mut sketch = CountMinSketch::::new(2, 8).unwrap(); @@ -43,9 +96,9 @@ fn rejects_weights_that_cannot_bound_counter_arithmetic() { } } -// This test validates binary format compatibility (deserialize + byte round-trip) for -// datasketches-cpp CountMin snapshots. It intentionally does not assert estimate equivalence -// against original input keys because per-row hash seed derivation differs across implementations. +// This helper validates binary format compatibility (deserialize + byte round-trip) for +// version 1 datasketches-cpp CountMin snapshots. It intentionally does not assert estimate +// equivalence because version 1 row-seed derivation differs across implementations. fn assert_cpp_snapshot( filename: &str, seed: u64,