From acefc5354b3cf05005fb72a2d0ba9bd9d04e9c07 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 26 Sep 2026 16:54:52 +0800 Subject: [PATCH 1/4] fix: reject sketch weight overflow before mutation --- CHANGELOG.md | 2 + datasketches/src/countmin/sketch.rs | 44 ++++++++++++--- datasketches/src/countmin/value.rs | 21 ++++++-- datasketches/src/kll/sketch.rs | 12 ++--- datasketches/src/req/sketch.rs | 14 +++-- datasketches/src/tdigest/sketch.rs | 12 +++++ .../tests/countmin_test/sketch.rs | 54 +++++++++++++++++++ tests-integration/tests/kll_test/main.rs | 1 + tests-integration/tests/kll_test/overflow.rs | 44 +++++++++++++++ tests-integration/tests/req_test/main.rs | 1 + tests-integration/tests/req_test/overflow.rs | 44 +++++++++++++++ .../tests/serde_tests/countmin.rs | 18 +++++++ tests-integration/tests/tdigest_test/main.rs | 1 + .../tests/tdigest_test/overflow.rs | 47 ++++++++++++++++ 14 files changed, 293 insertions(+), 22 deletions(-) create mode 100644 tests-integration/tests/kll_test/overflow.rs create mode 100644 tests-integration/tests/req_test/overflow.rs create mode 100644 tests-integration/tests/tdigest_test/overflow.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 7144ea5f..0e5a6615 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,8 @@ All significant changes to this project will be documented in this file. ### Bug fixes +* Count-Min, REQ, and T-Digest now reject total-weight overflow before modifying the sketch, including in release builds. Updates and T-Digest merges panic; Count-Min and REQ merges return an error. KLL updates now detect overflow before changing extrema or compacting retained items. +* Count-Min upper bounds now clamp to the counter type's maximum instead of overflowing. Deserialization rejects negative total absolute weights and counters whose magnitudes exceed that weight. * `FrequentItemsSketch` updates and merges now panic before modifying the sketch if the total stream weight would overflow, including in release builds. Deserialization rejects non-empty images with zero stream weight or item weights whose sum exceeds the declared stream weight. * Fix T-Digest `merge` so it preserves `min`/`max` from the other digest instead of re-deriving them from centroid means after compression. * T-Digest deserialization now rejects unknown or conflicting flags, reversed extrema, out-of-range values, unsorted centroids, and non-empty images without stored values. diff --git a/datasketches/src/countmin/sketch.rs b/datasketches/src/countmin/sketch.rs index 500f4396..3b9f7b71 100644 --- a/datasketches/src/countmin/sketch.rs +++ b/datasketches/src/countmin/sketch.rs @@ -119,7 +119,7 @@ impl CountMinSketch { self.seed } - /// Returns the total weight inserted into the sketch. + /// Returns the sum of absolute update weights, scaled by any halving or decay. pub fn total_weight(&self) -> T { self.total_weight } @@ -177,6 +177,10 @@ impl CountMinSketch { /// Updates the sketch with a single occurrence of the item. /// + /// # Panics + /// + /// Panics without modifying the sketch if the total absolute weight would exceed `T::MAX`. + /// /// # Examples /// /// ``` @@ -192,6 +196,11 @@ impl CountMinSketch { /// Updates the sketch with the given item and weight. /// + /// # Panics + /// + /// Panics without modifying the sketch if the absolute weight or the total absolute weight + /// cannot be represented by `T`. + /// /// # Examples /// /// ``` @@ -205,8 +214,12 @@ impl CountMinSketch { if weight == T::ZERO { return; } - let abs_weight = weight.abs(); - self.total_weight = self.total_weight + abs_weight; + // Each counter's magnitude is bounded by the total absolute weight, so this also + // protects the per-bucket additions. Deserialization validates the same invariant. + self.total_weight = weight + .checked_abs() + .and_then(|weight| self.total_weight.checked_add(weight)) + .expect("total absolute weight overflow"); let num_buckets = self.num_buckets as usize; for (row, seed) in self.hash_seeds.iter().enumerate() { let bucket = self.bucket_index(&item, *seed); @@ -246,17 +259,20 @@ impl CountMinSketch { } /// Returns the upper bound on the true frequency of the given item. + /// + /// Clamps the bound to `T::MAX` if adding the error would overflow. pub fn upper_bound(&self, item: I) -> T { let estimate = self.estimate(item); let error = self.total_weight.scale(self.relative_error()); - estimate + error + estimate.checked_add(error).unwrap_or(T::MAX) } /// Merges another sketch into this one. /// /// # Errors /// - /// Returns an error if the sketches have different numbers of hashes, bucket counts, or seeds. + /// 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`. /// /// # Examples /// @@ -281,10 +297,13 @@ impl CountMinSketch { "Count-Min sketches must have matching numbers of hashes, bucket counts, and seeds", )); } + self.total_weight = self + .total_weight + .checked_add(other.total_weight) + .ok_or_else(|| Error::invalid_argument("total absolute weight overflow"))?; for (count, other_count) in self.counts.iter_mut().zip(&other.counts) { *count = *count + *other_count; } - self.total_weight = self.total_weight + other.total_weight; Ok(()) } @@ -445,8 +464,21 @@ impl CountMinSketch { } sketch.total_weight = read_value(&mut cursor, "total_weight")?; + if sketch.total_weight < T::ZERO { + return Err(Error::deserial( + "total absolute weight must be non-negative", + )); + } for count in &mut sketch.counts { *count = read_value(&mut cursor, "counts")?; + if count + .checked_abs() + .is_none_or(|weight| weight > sketch.total_weight) + { + return Err(Error::deserial( + "counter magnitude exceeds total absolute weight", + )); + } } Ok(sketch) } diff --git a/datasketches/src/countmin/value.rs b/datasketches/src/countmin/value.rs index 70fd5427..6e9ec93e 100644 --- a/datasketches/src/countmin/value.rs +++ b/datasketches/src/countmin/value.rs @@ -37,7 +37,8 @@ mod private { const ONE: Self; const MAX: Self; - fn abs(self) -> Self; + fn checked_abs(self) -> Option; + fn checked_add(self, other: Self) -> Option; fn scale(self, factor: f64) -> Self; fn to_bytes(self) -> [u8; 8]; fn try_from_bytes(bytes: [u8; 8]) -> Result; @@ -56,8 +57,13 @@ macro_rules! impl_signed { const MAX: Self = $max; #[inline(always)] - fn abs(self) -> Self { - if self >= 0 { self } else { -self } + fn checked_abs(self) -> Option { + self.checked_abs() + } + + #[inline(always)] + fn checked_add(self, other: Self) -> Option { + self.checked_add(other) } #[inline(always)] @@ -102,8 +108,13 @@ macro_rules! impl_unsigned { const MAX: Self = $max; #[inline(always)] - fn abs(self) -> Self { - self + fn checked_abs(self) -> Option { + Some(self) + } + + #[inline(always)] + fn checked_add(self, other: Self) -> Option { + self.checked_add(other) } #[inline(always)] diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index 38b9245f..76e76fb5 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -136,8 +136,9 @@ impl KllSketch { /// /// # Panics /// - /// Panics if the stream weight would exceed [`u64::MAX`]. + /// Panics without modifying the sketch if the stream weight would exceed [`u64::MAX`]. pub fn update(&mut self, item: T) { + assert!(self.n < u64::MAX, "total stream weight overflow"); self.update_min_max(&item); self.internal_update(item); } @@ -708,13 +709,8 @@ impl KllSketch { if self.num_retained >= self.capacity { self.compress_while_updating(); } - self.n = self.n.checked_add(1).unwrap_or_else(|| { - panic!( - "cannot update KLL sketch: stream weight is {}, maximum is {}", - self.n, - u64::MAX - ) - }); + // Both update and merge check the final stream weight before modifying the sketch. + self.n += 1; self.num_retained += 1; self.is_level_zero_sorted = false; self.levels[0].push(item); diff --git a/datasketches/src/req/sketch.rs b/datasketches/src/req/sketch.rs index 89a3a70c..2373b1f2 100644 --- a/datasketches/src/req/sketch.rs +++ b/datasketches/src/req/sketch.rs @@ -133,7 +133,12 @@ where } /// Updates the sketch with a new item. + /// + /// # Panics + /// + /// Panics without modifying the sketch if the stream weight would exceed `u64::MAX`. pub fn update(&mut self, item: T) { + self.n = self.n.checked_add(1).expect("total stream weight overflow"); match &mut self.min_item { None => self.min_item = Some(item.clone()), Some(cur) if item.cmp(cur).is_lt() => *cur = item.clone(), @@ -146,7 +151,6 @@ where } self.compactors[0].append(item); - self.n += 1; self.num_retained += 1; if self.num_retained >= self.max_nom_size { @@ -285,7 +289,8 @@ where /// /// # Errors /// - /// Returns an error if the two sketches have different `rank_accuracy`. + /// Returns an error without modifying the sketch if the two sketches have different + /// `rank_accuracy` or their combined stream weight exceeds `u64::MAX`. /// /// # Examples /// @@ -317,7 +322,10 @@ where return Ok(()); } - self.n += other.n; + self.n = self + .n + .checked_add(other.n) + .ok_or_else(|| Error::invalid_argument("total stream weight overflow"))?; if let Some(m) = &other.min_item { match &self.min_item { diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 1c23867a..34a4209c 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -264,6 +264,10 @@ impl TDigestMut { /// /// [f64::NAN], [f64::INFINITY], and [f64::NEG_INFINITY] values are ignored. /// + /// # Panics + /// + /// Panics without modifying the digest if the total weight would exceed `u64::MAX`. + /// /// # Examples /// /// ``` @@ -277,6 +281,7 @@ impl TDigestMut { if !value.is_finite() { return; } + assert!(self.total_weight() < u64::MAX, "total weight overflow"); let max_unmerged = self.max_unmerged(); if self.buffer.unmerged_len() >= max_unmerged { @@ -322,6 +327,10 @@ impl TDigestMut { /// Merges the given t-digest into this one. /// + /// # Panics + /// + /// Panics without modifying the digest if the combined total weight would exceed `u64::MAX`. + /// /// # Examples /// /// ``` @@ -338,6 +347,9 @@ impl TDigestMut { if other.is_empty() { return; } + self.total_weight() + .checked_add(other.total_weight()) + .expect("total weight overflow"); // Preserve true extrema from `other`. Compression only sees centroid means, which can // differ from `min`/`max` after ordinary compression or deserialization. diff --git a/tests-integration/tests/countmin_test/sketch.rs b/tests-integration/tests/countmin_test/sketch.rs index f35eb78a..785b86e1 100644 --- a/tests-integration/tests/countmin_test/sketch.rs +++ b/tests-integration/tests/countmin_test/sketch.rs @@ -15,12 +15,66 @@ // specific language governing permissions and limitations // under the License. +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; + use datasketches::countmin::CountMinSketch; use datasketches::error::ErrorKind; use googletest::assert_that; use googletest::prelude::ge; use googletest::prelude::le; +#[test] +fn weight_overflow_preserves_unsigned_state() { + let mut sketch = CountMinSketch::::new(2, 8).unwrap(); + sketch.update_with_weight("x", u8::MAX - 1); + let mut one = CountMinSketch::::new(2, 8).unwrap(); + one.update("x"); + sketch.merge(&one).unwrap(); + assert_eq!(sketch.total_weight(), u8::MAX); + let before = sketch.clone(); + + sketch.update_with_weight("x", 0); + assert!(catch_unwind(AssertUnwindSafe(|| sketch.update("x"))).is_err()); + assert_eq!(sketch, before); + assert_eq!( + sketch.merge(&one).unwrap_err().kind(), + ErrorKind::InvalidArgument + ); + assert_eq!(sketch, before); +} + +#[test] +fn weight_overflow_preserves_signed_state() { + let mut sketch = CountMinSketch::::new(2, 8).unwrap(); + let empty = sketch.clone(); + assert!(catch_unwind(AssertUnwindSafe(|| sketch.update_with_weight("x", i8::MIN))).is_err()); + assert_eq!(sketch, empty); + + sketch.update_with_weight("x", -i8::MAX); + assert_eq!(sketch.total_weight(), i8::MAX); + assert_eq!(sketch.estimate("x"), -i8::MAX); + let before = sketch.clone(); + // Cancellation reduces the counter, but still increases the absolute stream weight. + assert!(catch_unwind(AssertUnwindSafe(|| sketch.update("x"))).is_err()); + assert_eq!(sketch, before); + assert_eq!( + CountMinSketch::::deserialize(&sketch.serialize()).unwrap(), + sketch + ); +} + +#[test] +fn upper_bound_clamps_on_overflow() { + let mut unsigned = CountMinSketch::::new(2, 8).unwrap(); + unsigned.update_with_weight("x", u8::MAX); + assert_eq!(unsigned.upper_bound("x"), u8::MAX); + + let mut signed = CountMinSketch::::new(2, 8).unwrap(); + signed.update_with_weight("x", i8::MAX); + assert_eq!(signed.upper_bound("x"), i8::MAX); +} + #[test] fn test_init_defaults() { let sketch = CountMinSketch::::new(3, 5).unwrap(); diff --git a/tests-integration/tests/kll_test/main.rs b/tests-integration/tests/kll_test/main.rs index 26441bc2..ec24d90f 100644 --- a/tests-integration/tests/kll_test/main.rs +++ b/tests-integration/tests/kll_test/main.rs @@ -18,4 +18,5 @@ mod core; mod generic; mod merge; +mod overflow; mod query; diff --git a/tests-integration/tests/kll_test/overflow.rs b/tests-integration/tests/kll_test/overflow.rs new file mode 100644 index 00000000..b2ae8382 --- /dev/null +++ b/tests-integration/tests/kll_test/overflow.rs @@ -0,0 +1,44 @@ +// 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 std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; + +use datasketches::error::ErrorKind; +use datasketches::kll::KllSketch; + +#[test] +fn weight_overflow_preserves_state() { + let mut one = KllSketch::::new(8).unwrap(); + one.update(0); + let mut sketch = one.clone(); + // Doubling and adding one reaches the exact limit through valid public operations. + for _ in 0..63 { + sketch.merge(&sketch.clone()).unwrap(); + sketch.update(0); + } + assert_eq!(sketch.n(), u64::MAX); + let before = sketch.serialize(); + + assert!(catch_unwind(AssertUnwindSafe(|| sketch.update(1))).is_err()); + assert!(sketch.serialize() == before, "overflow changed the sketch"); + assert_eq!( + sketch.merge(&one).unwrap_err().kind(), + ErrorKind::InvalidArgument + ); + assert!(sketch.serialize() == before, "overflow changed the sketch"); +} diff --git a/tests-integration/tests/req_test/main.rs b/tests-integration/tests/req_test/main.rs index a76589bd..88ad9afe 100644 --- a/tests-integration/tests/req_test/main.rs +++ b/tests-integration/tests/req_test/main.rs @@ -33,6 +33,7 @@ mod bounds; mod core; mod generic; mod merge; +mod overflow; mod property; mod query; mod sorted_view_api; diff --git a/tests-integration/tests/req_test/overflow.rs b/tests-integration/tests/req_test/overflow.rs new file mode 100644 index 00000000..633379ae --- /dev/null +++ b/tests-integration/tests/req_test/overflow.rs @@ -0,0 +1,44 @@ +// 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 std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; + +use datasketches::error::ErrorKind; +use datasketches::req::ReqSketch; + +#[test] +fn weight_overflow_preserves_state() { + let mut one = ReqSketch::::default(); + one.update(0); + let mut sketch = one.clone(); + // Doubling and adding one reaches the exact limit through valid public operations. + for _ in 0..63 { + sketch.merge(&sketch.clone()).unwrap(); + sketch.update(0); + } + assert_eq!(sketch.n(), u64::MAX); + let before = sketch.serialize(); + + assert!(catch_unwind(AssertUnwindSafe(|| sketch.update(1))).is_err()); + assert!(sketch.serialize() == before, "overflow changed the sketch"); + assert_eq!( + sketch.merge(&one).unwrap_err().kind(), + ErrorKind::InvalidArgument + ); + assert!(sketch.serialize() == before, "overflow changed the sketch"); +} diff --git a/tests-integration/tests/serde_tests/countmin.rs b/tests-integration/tests/serde_tests/countmin.rs index be490804..07786b42 100644 --- a/tests-integration/tests/serde_tests/countmin.rs +++ b/tests-integration/tests/serde_tests/countmin.rs @@ -25,6 +25,24 @@ use tests_integration::ZERO_HASH_SEED; use crate::serialization_test_data; +#[test] +fn rejects_weights_that_cannot_bound_counter_arithmetic() { + let mut sketch = CountMinSketch::::new(2, 8).unwrap(); + sketch.update("x"); + let bytes = sketch.serialize(); + for (total, count) in [(-1i64, 0i64), (1, 2), (1, -2), (127, -128)] { + let mut corrupt = bytes.clone(); + corrupt[16..24].copy_from_slice(&total.to_le_bytes()); + corrupt[24..32].copy_from_slice(&count.to_le_bytes()); + assert_eq!( + CountMinSketch::::deserialize(&corrupt) + .unwrap_err() + .kind(), + ErrorKind::InvalidData + ); + } +} + // This test validates binary format compatibility (deserialize + byte round-trip) for // C++ CountMin snapshots. It intentionally does not assert estimate equivalence against // original input keys because per-row hash seed derivation differs across implementations. diff --git a/tests-integration/tests/tdigest_test/main.rs b/tests-integration/tests/tdigest_test/main.rs index 25ecc796..ae0ab1c4 100644 --- a/tests-integration/tests/tdigest_test/main.rs +++ b/tests-integration/tests/tdigest_test/main.rs @@ -15,5 +15,6 @@ // specific language governing permissions and limitations // under the License. +mod overflow; mod property; mod sketch; diff --git a/tests-integration/tests/tdigest_test/overflow.rs b/tests-integration/tests/tdigest_test/overflow.rs new file mode 100644 index 00000000..4f923863 --- /dev/null +++ b/tests-integration/tests/tdigest_test/overflow.rs @@ -0,0 +1,47 @@ +// 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 std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; + +use datasketches::tdigest::TDigestMut; + +#[test] +fn weight_overflow_preserves_buffer_and_extrema() { + let mut one = TDigestMut::new(20).unwrap(); + one.update(0.0); + let mut sketch = one.clone(); + for _ in 0..63 { + sketch.merge(&sketch.clone()); + sketch.update(0.0); + } + assert_eq!(sketch.total_weight(), u64::MAX); + let before = sketch.clone().serialize(); + + sketch.update(f64::NAN); + sketch.merge(&TDigestMut::default()); + assert!(catch_unwind(AssertUnwindSafe(|| sketch.update(1.0))).is_err()); + assert_eq!(sketch.clone().serialize(), before); + assert!(catch_unwind(AssertUnwindSafe(|| sketch.merge(&one))).is_err()); + assert_eq!(sketch.clone().serialize(), before); + + let mut target = TDigestMut::default(); + target.update(1.0); + let before = target.clone().serialize(); + assert!(catch_unwind(AssertUnwindSafe(|| target.merge(&sketch))).is_err()); + assert_eq!(target.serialize(), before); +} From 06c1c881f3bdda1c28bee24a378c80d8c8cf7495 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 26 Sep 2026 17:13:42 +0800 Subject: [PATCH 2/4] refactor: simplify sketch weight handling --- datasketches/src/countmin/sketch.rs | 5 +- datasketches/src/tdigest/sketch.rs | 37 +++++---------- tests-integration/tests/kll_test/core.rs | 25 ++++++++++ tests-integration/tests/kll_test/main.rs | 1 - tests-integration/tests/kll_test/overflow.rs | 44 ----------------- tests-integration/tests/req_test/core.rs | 25 ++++++++++ tests-integration/tests/req_test/main.rs | 1 - tests-integration/tests/req_test/overflow.rs | 44 ----------------- tests-integration/tests/tdigest_test/main.rs | 1 - .../tests/tdigest_test/overflow.rs | 47 ------------------- .../tests/tdigest_test/sketch.rs | 28 +++++++++++ 11 files changed, 94 insertions(+), 164 deletions(-) delete mode 100644 tests-integration/tests/kll_test/overflow.rs delete mode 100644 tests-integration/tests/req_test/overflow.rs delete mode 100644 tests-integration/tests/tdigest_test/overflow.rs diff --git a/datasketches/src/countmin/sketch.rs b/datasketches/src/countmin/sketch.rs index 3b9f7b71..bffdca15 100644 --- a/datasketches/src/countmin/sketch.rs +++ b/datasketches/src/countmin/sketch.rs @@ -49,6 +49,9 @@ pub struct CountMinSketch { num_buckets: u32, seed: u64, seed_hash: u16, + // 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. total_weight: T, counts: Vec, hash_seeds: Vec, @@ -214,8 +217,6 @@ impl CountMinSketch { if weight == T::ZERO { return; } - // Each counter's magnitude is bounded by the total absolute weight, so this also - // protects the per-bucket additions. Deserialization validates the same invariant. self.total_weight = weight .checked_abs() .and_then(|weight| self.total_weight.checked_add(weight)) diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 34a4209c..e0d59535 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -347,7 +347,8 @@ impl TDigestMut { if other.is_empty() { return; } - self.total_weight() + let total_weight = self + .total_weight() .checked_add(other.total_weight()) .expect("total weight overflow"); @@ -356,9 +357,8 @@ impl TDigestMut { self.min = self.min.min(other.min); self.max = self.max.max(other.max); - let self_unmerged_weight = self.buffer.unmerged_len() as u64; let centroids = std::mem::take(&mut self.buffer).into_merged_centroids(&other.buffer); - self.compress_sorted_centroids(centroids, self_unmerged_weight + other.total_weight()) + self.compress_sorted_centroids(centroids, total_weight); } /// Converts this mutable t-digest into an immutable one. @@ -894,38 +894,27 @@ impl TDigestMut { /// Processes unmerged values and merges centroids if needed. fn compress(&mut self) { - let additional_weight = self.buffer.unmerged_len() as u64; - if additional_weight == 0 { + if self.buffer.unmerged_len() == 0 { // Also preserves fully compressed deserialized images verbatim. return; } - let centroids = std::mem::take(&mut self.buffer).into_centroids_for_compression(); - self.compress_centroids(centroids, additional_weight); - } - - /// Compresses the given centroids into this t-digest. - /// - /// # Contract - /// - /// * `centroids` must contain at least one centroid. - /// * `centroids` contains every centroid to be merged, including all centroids previously - /// stored in `self`. - /// * `additional_weight` is the total weight not yet included in `self.compressed_weight`. - /// * Every centroid mean in `centroids` is finite. - /// * `self.buffer` has no unmerged values before returning. - fn compress_centroids(&mut self, mut centroids: Vec, additional_weight: u64) { - debug_assert!(!centroids.is_empty()); + let total_weight = self.total_weight(); + let mut centroids = std::mem::take(&mut self.buffer).into_centroids_for_compression(); centroids.sort_by(centroid_cmp); - self.compress_sorted_centroids(centroids, additional_weight); + self.compress_sorted_centroids(centroids, total_weight); } - fn compress_sorted_centroids(&mut self, mut centroids: Vec, additional_weight: u64) { + /// Compresses nonempty, sorted centroids whose combined weight is `total_weight`. + /// + /// Includes all retained and incoming values, with finite means and nonzero weights. + /// Callers ensure the total fits in `u64` before taking the buffer. + fn compress_sorted_centroids(&mut self, mut centroids: Vec, total_weight: u64) { debug_assert!(!centroids.is_empty()); debug_assert!(centroids_are_sorted(¢roids)); if self.reverse_merge { centroids.reverse(); } - self.compressed_weight += additional_weight; + self.compressed_weight = total_weight; let mut num_centroids = 1; let len = centroids.len(); diff --git a/tests-integration/tests/kll_test/core.rs b/tests-integration/tests/kll_test/core.rs index c0bcaec0..0505649f 100644 --- a/tests-integration/tests/kll_test/core.rs +++ b/tests-integration/tests/kll_test/core.rs @@ -15,6 +15,9 @@ // specific language governing permissions and limitations // under the License. +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; + use datasketches::error::ErrorKind; use datasketches::kll::KllFloat; use datasketches::kll::KllSketch; @@ -85,3 +88,25 @@ fn retained_count_stays_consistent_through_compaction_and_roundtrip() { assert_eq!(decoded.min_item(), Some(&0)); assert_eq!(decoded.max_item(), Some(&99_999)); } + +#[test] +fn weight_overflow_preserves_state() { + let mut one = KllSketch::::new(8).unwrap(); + one.update(0); + let mut sketch = one.clone(); + // Doubling and adding one reaches the exact limit through valid public operations. + for _ in 0..63 { + sketch.merge(&sketch.clone()).unwrap(); + sketch.update(0); + } + assert_eq!(sketch.n(), u64::MAX); + let before = sketch.serialize(); + + assert!(catch_unwind(AssertUnwindSafe(|| sketch.update(1))).is_err()); + assert!(sketch.serialize() == before, "overflow changed the sketch"); + assert_eq!( + sketch.merge(&one).unwrap_err().kind(), + ErrorKind::InvalidArgument + ); + assert!(sketch.serialize() == before, "overflow changed the sketch"); +} diff --git a/tests-integration/tests/kll_test/main.rs b/tests-integration/tests/kll_test/main.rs index ec24d90f..26441bc2 100644 --- a/tests-integration/tests/kll_test/main.rs +++ b/tests-integration/tests/kll_test/main.rs @@ -18,5 +18,4 @@ mod core; mod generic; mod merge; -mod overflow; mod query; diff --git a/tests-integration/tests/kll_test/overflow.rs b/tests-integration/tests/kll_test/overflow.rs deleted file mode 100644 index b2ae8382..00000000 --- a/tests-integration/tests/kll_test/overflow.rs +++ /dev/null @@ -1,44 +0,0 @@ -// 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 std::panic::AssertUnwindSafe; -use std::panic::catch_unwind; - -use datasketches::error::ErrorKind; -use datasketches::kll::KllSketch; - -#[test] -fn weight_overflow_preserves_state() { - let mut one = KllSketch::::new(8).unwrap(); - one.update(0); - let mut sketch = one.clone(); - // Doubling and adding one reaches the exact limit through valid public operations. - for _ in 0..63 { - sketch.merge(&sketch.clone()).unwrap(); - sketch.update(0); - } - assert_eq!(sketch.n(), u64::MAX); - let before = sketch.serialize(); - - assert!(catch_unwind(AssertUnwindSafe(|| sketch.update(1))).is_err()); - assert!(sketch.serialize() == before, "overflow changed the sketch"); - assert_eq!( - sketch.merge(&one).unwrap_err().kind(), - ErrorKind::InvalidArgument - ); - assert!(sketch.serialize() == before, "overflow changed the sketch"); -} diff --git a/tests-integration/tests/req_test/core.rs b/tests-integration/tests/req_test/core.rs index 737b4892..626632db 100644 --- a/tests-integration/tests/req_test/core.rs +++ b/tests-integration/tests/req_test/core.rs @@ -17,6 +17,9 @@ //! Core ReqSketch construction and update behavior. +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; + use datasketches::common::SearchCriteria; use datasketches::error::Error; use datasketches::error::ErrorKind; @@ -254,3 +257,25 @@ fn new_validates_k() { ); assert!(ReqSketch::::new(12, RankAccuracy::HighRank).is_ok()); } + +#[test] +fn weight_overflow_preserves_state() { + let mut one = ReqSketch::::default(); + one.update(0); + let mut sketch = one.clone(); + // Doubling and adding one reaches the exact limit through valid public operations. + for _ in 0..63 { + sketch.merge(&sketch.clone()).unwrap(); + sketch.update(0); + } + assert_eq!(sketch.n(), u64::MAX); + let before = sketch.serialize(); + + assert!(catch_unwind(AssertUnwindSafe(|| sketch.update(1))).is_err()); + assert!(sketch.serialize() == before, "overflow changed the sketch"); + assert_eq!( + sketch.merge(&one).unwrap_err().kind(), + ErrorKind::InvalidArgument + ); + assert!(sketch.serialize() == before, "overflow changed the sketch"); +} diff --git a/tests-integration/tests/req_test/main.rs b/tests-integration/tests/req_test/main.rs index 88ad9afe..a76589bd 100644 --- a/tests-integration/tests/req_test/main.rs +++ b/tests-integration/tests/req_test/main.rs @@ -33,7 +33,6 @@ mod bounds; mod core; mod generic; mod merge; -mod overflow; mod property; mod query; mod sorted_view_api; diff --git a/tests-integration/tests/req_test/overflow.rs b/tests-integration/tests/req_test/overflow.rs deleted file mode 100644 index 633379ae..00000000 --- a/tests-integration/tests/req_test/overflow.rs +++ /dev/null @@ -1,44 +0,0 @@ -// 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 std::panic::AssertUnwindSafe; -use std::panic::catch_unwind; - -use datasketches::error::ErrorKind; -use datasketches::req::ReqSketch; - -#[test] -fn weight_overflow_preserves_state() { - let mut one = ReqSketch::::default(); - one.update(0); - let mut sketch = one.clone(); - // Doubling and adding one reaches the exact limit through valid public operations. - for _ in 0..63 { - sketch.merge(&sketch.clone()).unwrap(); - sketch.update(0); - } - assert_eq!(sketch.n(), u64::MAX); - let before = sketch.serialize(); - - assert!(catch_unwind(AssertUnwindSafe(|| sketch.update(1))).is_err()); - assert!(sketch.serialize() == before, "overflow changed the sketch"); - assert_eq!( - sketch.merge(&one).unwrap_err().kind(), - ErrorKind::InvalidArgument - ); - assert!(sketch.serialize() == before, "overflow changed the sketch"); -} diff --git a/tests-integration/tests/tdigest_test/main.rs b/tests-integration/tests/tdigest_test/main.rs index ae0ab1c4..25ecc796 100644 --- a/tests-integration/tests/tdigest_test/main.rs +++ b/tests-integration/tests/tdigest_test/main.rs @@ -15,6 +15,5 @@ // specific language governing permissions and limitations // under the License. -mod overflow; mod property; mod sketch; diff --git a/tests-integration/tests/tdigest_test/overflow.rs b/tests-integration/tests/tdigest_test/overflow.rs deleted file mode 100644 index 4f923863..00000000 --- a/tests-integration/tests/tdigest_test/overflow.rs +++ /dev/null @@ -1,47 +0,0 @@ -// 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 std::panic::AssertUnwindSafe; -use std::panic::catch_unwind; - -use datasketches::tdigest::TDigestMut; - -#[test] -fn weight_overflow_preserves_buffer_and_extrema() { - let mut one = TDigestMut::new(20).unwrap(); - one.update(0.0); - let mut sketch = one.clone(); - for _ in 0..63 { - sketch.merge(&sketch.clone()); - sketch.update(0.0); - } - assert_eq!(sketch.total_weight(), u64::MAX); - let before = sketch.clone().serialize(); - - sketch.update(f64::NAN); - sketch.merge(&TDigestMut::default()); - assert!(catch_unwind(AssertUnwindSafe(|| sketch.update(1.0))).is_err()); - assert_eq!(sketch.clone().serialize(), before); - assert!(catch_unwind(AssertUnwindSafe(|| sketch.merge(&one))).is_err()); - assert_eq!(sketch.clone().serialize(), before); - - let mut target = TDigestMut::default(); - target.update(1.0); - let before = target.clone().serialize(); - assert!(catch_unwind(AssertUnwindSafe(|| target.merge(&sketch))).is_err()); - assert_eq!(target.serialize(), before); -} diff --git a/tests-integration/tests/tdigest_test/sketch.rs b/tests-integration/tests/tdigest_test/sketch.rs index d13c2bac..38a94b58 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -16,6 +16,8 @@ // under the License. use std::mem::size_of; +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; use datasketches::tdigest::TDigestMut; use googletest::assert_that; @@ -437,3 +439,29 @@ fn test_merge_preserves_min_max_from_other() { assert_eq!(left.quantile(0.0), Some(-10.0)); assert_eq!(left.quantile(1.0), Some(80.0)); } + +#[test] +fn weight_overflow_preserves_buffer_and_extrema() { + let mut one = TDigestMut::new(20).unwrap(); + one.update(0.0); + let mut sketch = one.clone(); + for _ in 0..63 { + sketch.merge(&sketch.clone()); + sketch.update(0.0); + } + assert_eq!(sketch.total_weight(), u64::MAX); + let before = sketch.clone().serialize(); + + sketch.update(f64::NAN); + sketch.merge(&TDigestMut::default()); + assert!(catch_unwind(AssertUnwindSafe(|| sketch.update(1.0))).is_err()); + assert_eq!(sketch.clone().serialize(), before); + assert!(catch_unwind(AssertUnwindSafe(|| sketch.merge(&one))).is_err()); + assert_eq!(sketch.clone().serialize(), before); + + let mut target = TDigestMut::default(); + target.update(1.0); + let before = target.clone().serialize(); + assert!(catch_unwind(AssertUnwindSafe(|| target.merge(&sketch))).is_err()); + assert_eq!(target.serialize(), before); +} From 95e1fbdb9e2adfd20814e4727a940205b0ab97ff Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 26 Sep 2026 17:40:56 +0800 Subject: [PATCH 3/4] docs: consolidate unreleased changes by user impact --- CHANGELOG.md | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e5a6615..8ee26cc7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,29 +6,26 @@ All significant changes to this project will be documented in this file. ### Breaking changes -* `FrequentItemsSketch::is_empty` now checks whether the total stream weight is zero. A sketch whose counters were all removed by a purge remains non-empty; use `num_active_items() == 0` to check whether any items are retained. -* Remove `BloomFilter::invert`. Bit inversion has no sound set-membership interpretation; use the new `BloomFilter::difference` for the approximate set-difference (A NOT B) use case it was meant to serve. +* `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. +* Remove `BloomFilter::invert`; 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. * Move `SearchCriteria` from `req` to `common` and remove its `Default` implementation. Import `datasketches::common::SearchCriteria` and explicitly choose `Inclusive` or `Exclusive` for each query. ### New features -* Add `BloomFilter::difference` for approximate set difference: the result excludes the other filter's items exactly, while items unique to the left filter are kept unless their hash positions collide with the right filter. * Add KLL sketches behind the `kll` feature, with rank, quantile, PMF, and CDF queries, merging, totally ordered custom item types, a `KllFloat` adapter for non-NaN floating-point values, and serialization. ### Improvements -* The crate no longer has any runtime dependencies. The `kll` and `req` features previously pulled in `rand`; compaction now draws its coin from an in-tree generator. -* Improve truncated-input diagnostics across sketch deserializers. -* Improve hash-backed sketch update performance for integer and raw-byte inputs. -* Improve Bloom filter membership-and-insert performance and simplify Theta-family hash table thresholds. +* Remove the `rand` dependency; the crate now has no runtime dependencies. +* Improve error messages for truncated serialized sketches. +* Speed up hash-based sketch updates for integer and raw-byte inputs, and `BloomFilter::contains_and_insert`. ### Bug fixes -* Count-Min, REQ, and T-Digest now reject total-weight overflow before modifying the sketch, including in release builds. Updates and T-Digest merges panic; Count-Min and REQ merges return an error. KLL updates now detect overflow before changing extrema or compacting retained items. -* Count-Min upper bounds now clamp to the counter type's maximum instead of overflowing. Deserialization rejects negative total absolute weights and counters whose magnitudes exceed that weight. -* `FrequentItemsSketch` updates and merges now panic before modifying the sketch if the total stream weight would overflow, including in release builds. Deserialization rejects non-empty images with zero stream weight or item weights whose sum exceeds the declared stream weight. -* Fix T-Digest `merge` so it preserves `min`/`max` from the other digest instead of re-deriving them from centroid means after compression. -* T-Digest deserialization now rejects unknown or conflicting flags, reversed extrema, out-of-range values, unsorted centroids, and non-empty images without stored values. +* Count-Min, REQ, T-Digest, and `FrequentItemsSketch` updates and merges now reject total-weight overflow without changing the sketch, including in release builds. Count-Min and REQ merges return an error; other affected operations panic. +* Count-Min `upper_bound` now clamps to the counter type's maximum instead of overflowing. +* T-Digest merges now preserve the true minimum and maximum from both inputs, including compressed digests. +* Deserialization now returns `InvalidData` for inconsistent weights in Count-Min and `FrequentItemsSketch`, and invalid flags, extrema, or stored values in T-Digest. ## v0.5.0 (2026-09-04) From 7c38502df87f760fb30f7ae36c8a307db7020bcb Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 26 Sep 2026 17:58:52 +0800 Subject: [PATCH 4/4] docs: organize unreleased entries around sketch APIs --- CHANGELOG.md | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ee26cc7..ba590aec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,26 +6,37 @@ All significant changes to this project will be documented in this file. ### Breaking changes +* `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. -* Remove `BloomFilter::invert`; 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. -* Move `SearchCriteria` from `req` to `common` and remove its `Default` implementation. Import `datasketches::common::SearchCriteria` and explicitly choose `Inclusive` or `Exclusive` for each query. +* `ReqSketch` queries now use `datasketches::common::SearchCriteria` instead of `datasketches::req::SearchCriteria`. `SearchCriteria` no longer implements `Default`; explicitly choose `Inclusive` or `Exclusive`. ### New features -* Add KLL sketches behind the `kll` feature, with rank, quantile, PMF, and CDF queries, merging, totally ordered custom item types, a `KllFloat` adapter for non-NaN floating-point values, and serialization. +* `KllSketch` is now available behind the `kll` feature, with rank, quantile, PMF, and CDF queries, merging, serialization, custom ordered item types, and a `KllFloat` adapter for non-NaN floating-point values. ### Improvements -* Remove the `rand` dependency; the crate now has no runtime dependencies. -* Improve error messages for truncated serialized sketches. -* Speed up hash-based sketch updates for integer and raw-byte inputs, and `BloomFilter::contains_and_insert`. +* `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. +* `CpcSketch` updates are faster for integer and raw-byte inputs. +* `FrequentItemsSketch` updates are faster for integer and raw-byte keys. +* `HllSketch` updates are faster for integer and raw-byte inputs. +* `ThetaSketch` updates are faster for integer and raw-byte inputs. +* `TupleSketch` updates are faster for integer and raw-byte inputs. +* Library-wide: the crate no longer depends on `rand` and has no runtime dependencies. +* Library-wide: sketch deserializers report clearer errors for truncated input. ### Bug fixes -* Count-Min, REQ, T-Digest, and `FrequentItemsSketch` updates and merges now reject total-weight overflow without changing the sketch, including in release builds. Count-Min and REQ merges return an error; other affected operations panic. -* Count-Min `upper_bound` now clamps to the counter type's maximum instead of overflowing. -* T-Digest merges now preserve the true minimum and maximum from both inputs, including compressed digests. -* Deserialization now returns `InvalidData` for inconsistent weights in Count-Min and `FrequentItemsSketch`, and invalid flags, extrema, or stored values in T-Digest. +* `CountMinSketch` updates now panic and merges return `InvalidArgument` if the total absolute weight would exceed the counter type's maximum. Both leave the sketch unchanged, including in release builds. +* `CountMinSketch::upper_bound` now clamps to the counter type's maximum instead of overflowing. +* `CountMinSketch` deserialization now returns `InvalidData` if the total absolute weight is negative or any counter's magnitude exceeds it. +* `FrequentItemsSketch` updates and merges now panic without changing the sketch if the total stream weight would overflow, including in release builds. +* `FrequentItemsSketch` deserialization now returns `InvalidData` if a non-empty image declares zero stream weight or the item weights sum to more than the declared stream weight. +* `ReqSketch` updates now panic and merges return `InvalidArgument` if the stream weight would exceed `u64::MAX`. Both leave the sketch unchanged, including in release builds. +* `TDigestMut` updates and merges now panic without changing the digest if the total weight would exceed `u64::MAX`, including in release builds. +* `TDigestMut::merge` now preserves the true minimum and maximum from both inputs, including compressed digests. +* `TDigest` and `TDigestMut` deserialization now returns `InvalidData` for invalid flags or extrema, out-of-range values, unsorted centroids, or non-empty images with no stored values. ## v0.5.0 (2026-09-04)