fix(frequencies): define emptiness by stream weight

CHANGELOG.md

@@ -6,6 +6,7 @@ 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.
 * Move `SearchCriteria` from `req` to `common` and remove its `Default` implementation. Import `datasketches::common::SearchCriteria` and explicitly choose `Inclusive` or `Exclusive` for each query.
 
@@ -23,6 +24,7 @@ All significant changes to this project will be documented in this file.
 
 ### Bug fixes
 
+* `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.
 

datasketches/src/frequencies/serialization.rs

@@ -25,9 +25,9 @@ use crate::error::Error;
 /// Serialization version.
 pub const SERIAL_VERSION: u8 = 1;
 
-/// Preamble longs for empty sketch.
+/// Preamble longs for a sketch with zero stream weight.
 pub const PREAMBLE_LONGS_EMPTY: u8 = 1;
-/// Preamble longs for non-empty sketch.
+/// Preamble longs for a sketch with positive stream weight, even if it retains no items.
 pub const PREAMBLE_LONGS_NONEMPTY: u8 = 4;
 
 /// Empty flag mask (both bits for compatibility).

datasketches/src/frequencies/sketch.rs

@@ -178,17 +178,13 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
         Ok(Self::with_lg_map_sizes(lg_max_map_size, LG_MIN_MAP_SIZE))
     }
 
-    /// Returns `true` if the sketch has no active items.
+    /// Returns `true` if the total stream weight is zero.
     ///
     /// A purge can remove all active items while retaining a non-zero total weight and
-    /// maximum error. Use [`Self::total_weight`] to distinguish that state from a newly created
-    /// or reset sketch.
+    /// maximum error. Such a sketch is not empty. Use [`Self::num_active_items`] to check
+    /// whether any items are retained.
     pub fn is_empty(&self) -> bool {
-        self.hash_map.num_active() == 0
-    }
-
-    fn is_initial_state(&self) -> bool {
-        self.stream_weight == 0 && self.offset == 0 && self.hash_map.num_active() == 0
+        self.stream_weight == 0
     }
 
     /// Returns the number of active items being tracked.
@@ -335,6 +331,10 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
 
     /// Updates the sketch with a count of one.
     ///
+    /// # Panics
+    ///
+    /// Panics without modifying the sketch if the total stream weight would exceed `u64::MAX`.
+    ///
     /// # Examples
     ///
     /// ```
@@ -352,6 +352,10 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
     ///
     /// A count of zero is a no-op.
     ///
+    /// # Panics
+    ///
+    /// Panics without modifying the sketch if the total stream weight would exceed `u64::MAX`.
+    ///
     /// # Examples
     ///
     /// ```
@@ -365,7 +369,10 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
         if count == 0 {
             return;
         }
-        self.stream_weight += count;
+        self.stream_weight = self
+            .stream_weight
+            .checked_add(count)
+            .expect("total stream weight overflow");
         self.hash_map.adjust_or_put_value(item, count);
         self.maybe_resize_or_purge();
     }
@@ -376,6 +383,10 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
     /// only allocates an owned item when it is newly inserted, so updating an
     /// already-tracked item is allocation free.
     ///
+    /// # Panics
+    ///
+    /// Panics without modifying the sketch if the total stream weight would exceed `u64::MAX`.
+    ///
     /// # Examples
     ///
     /// ```
@@ -400,6 +411,10 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
     /// item by reference and only allocates an owned item when it is newly
     /// inserted. A count of zero is a no-op.
     ///
+    /// # Panics
+    ///
+    /// Panics without modifying the sketch if the total stream weight would exceed `u64::MAX`.
+    ///
     /// # Examples
     ///
     /// ```
@@ -417,7 +432,10 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
         if count == 0 {
             return;
         }
-        self.stream_weight += count;
+        self.stream_weight = self
+            .stream_weight
+            .checked_add(count)
+            .expect("total stream weight overflow");
         self.hash_map.adjust_or_put_value_ref(item, count);
         self.maybe_resize_or_purge();
     }
@@ -427,6 +445,10 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
     /// The other sketch may have a different map size. The merged sketch respects the
     /// larger error tolerance of the inputs.
     ///
+    /// # Panics
+    ///
+    /// Panics without modifying the sketch if the combined stream weight would exceed `u64::MAX`.
+    ///
     /// # Examples
     ///
     /// ```
@@ -443,10 +465,13 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
     where
         T: Clone,
     {
-        if other.is_initial_state() {
+        if other.is_empty() {
             return;
         }
-        let merged_total = self.stream_weight + other.stream_weight;
+        let merged_total = self
+            .stream_weight
+            .checked_add(other.stream_weight)
+            .expect("total stream weight overflow");
         for (item, count) in other.hash_map.iter() {
             self.update_with_count_ref(item, count);
         }
@@ -576,7 +601,7 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
         count_serialize_size: CountSerializeSize<T>,
         serialize_item: SerializeItem<T>,
     ) -> Vec<u8> {
-        if self.is_initial_state() {
+        if self.is_empty() {
             let mut bytes = SketchBytes::with_capacity(PREAMBLE_LONGS_EMPTY as usize * 8);
             bytes.write_u8(PREAMBLE_LONGS_EMPTY);
             bytes.write_u8(SERIAL_VERSION);
@@ -649,15 +674,17 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
         ensure_serial_version_is(SERIAL_VERSION, serial_version)?;
         validate_lg_map_sizes(lg_max, lg_cur)?;
 
-        let is_empty = (flags & EMPTY_FLAG_MASK) != 0;
+        ensure_preamble_longs_in(&[PREAMBLE_LONGS_EMPTY, PREAMBLE_LONGS_NONEMPTY], pre_longs)?;
+        let is_empty = pre_longs == PREAMBLE_LONGS_EMPTY;
+        if ((flags & EMPTY_FLAG_MASK) != 0) != is_empty {
+            return Err(Error::deserial("empty flag does not match preamble longs"));
+        }
         if is_empty {
-            ensure_preamble_longs_in(&[PREAMBLE_LONGS_EMPTY], pre_longs)?;
             // Java also restores empty images at the minimum size. `lg_cur` does
             // not carry item state here and must not control an eager allocation.
             return Ok(Self::with_lg_map_sizes(lg_max, LG_MIN_MAP_SIZE));
         }
 
-        ensure_preamble_longs_in(&[PREAMBLE_LONGS_NONEMPTY], pre_longs)?;
         let active_items = cursor
             .read_u32_le()
             .map_err(insufficient_data("active_items"))?;
@@ -675,6 +702,11 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
         let stream_weight = cursor
             .read_u64_le()
             .map_err(insufficient_data("stream_weight"))?;
+        if stream_weight == 0 {
+            return Err(Error::deserial(
+                "non-empty sketch must have a positive stream weight",
+            ));
+        }
         let offset_val = cursor.read_u64_le().map_err(insufficient_data("offset"))?;
 
         // Each active item has an eight-byte weight before its encoded key. Check
@@ -691,10 +723,16 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
         }
 
         let mut values = Vec::with_capacity(active_items);
+        let mut remaining_weight = stream_weight;
         for i in 0..active_items {
-            values.push(cursor.read_u64_le().map_err(|error| {
+            let value = cursor.read_u64_le().map_err(|error| {
                 Error::insufficient_data_of("frequent item weight", error).with_context("index", i)
-            })?);
+            })?;
+            // Reconstructing the map must not overflow or exceed the declared stream weight.
+            remaining_weight = remaining_weight
+                .checked_sub(value)
+                .ok_or_else(|| Error::deserial("item weights exceed total stream weight"))?;
+            values.push(value);
         }
 
         let items = deserialize_items(cursor, active_items)?;

tests-integration/tests/frequencies_test/update.rs

@@ -15,6 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::hash::Hash;
+use std::panic::AssertUnwindSafe;
+use std::panic::catch_unwind;
+
 use datasketches::error::ErrorKind;
 use datasketches::frequencies::ErrorType;
 use datasketches::frequencies::FrequentItemsSketch;
@@ -506,24 +510,80 @@ fn test_items_merge_empty_is_noop() {
     assert_eq!(sketch.estimate(&1), 1);
 }
 
-#[test]
-fn test_merge_preserves_purged_empty_state() {
-    let mut purged: FrequentItemsSketch<i64> = FrequentItemsSketch::new(32).unwrap();
-    for item in 0..=(32 * 3 / 4) {
-        purged.update(item);
+fn check_purged_state<T: Eq + Hash + Clone>(make_item: impl Fn(i64) -> T) {
+    let mut purged = FrequentItemsSketch::new(256).unwrap();
+    for item in 0..193 {
+        purged.update(make_item(item));
     }
-    assert!(purged.is_empty());
-    assert_eq!(purged.total_weight(), 25);
+    assert!(!purged.is_empty());
+    assert_eq!(purged.num_active_items(), 0);
+    assert_eq!(purged.total_weight(), 193);
     assert_eq!(purged.maximum_error(), 1);
 
-    let mut merged: FrequentItemsSketch<i64> = FrequentItemsSketch::new(32).unwrap();
+    let mut merged = FrequentItemsSketch::new(256).unwrap();
     merged.merge(&purged);
 
-    assert!(merged.is_empty());
+    assert!(!merged.is_empty());
     assert_eq!(merged.num_active_items(), 0);
     assert_eq!(merged.total_weight(), purged.total_weight());
     assert_eq!(merged.maximum_error(), purged.maximum_error());
-    assert_eq!(merged.upper_bound(&1000), purged.upper_bound(&1000));
+    assert_eq!(merged.upper_bound(&make_item(0)), 1);
+
+    let mut nonempty = FrequentItemsSketch::new(256).unwrap();
+    nonempty.update(make_item(1000));
+    nonempty.merge(&purged);
+    assert_eq!(nonempty.total_weight(), 194);
+    assert_eq!(nonempty.maximum_error(), 1);
+    assert_eq!(nonempty.upper_bound(&make_item(0)), 1);
+
+    purged.reset();
+    assert!(purged.is_empty());
+    assert_eq!(purged.total_weight(), 0);
+    assert_eq!(purged.maximum_error(), 0);
+}
+
+#[test]
+fn test_longs_purged_state() {
+    check_purged_state(|item| item);
+}
+
+#[test]
+fn test_items_purged_state() {
+    check_purged_state(|item| item.to_string());
+}
+
+#[test]
+fn test_update_weight_overflow_preserves_state() {
+    let mut sketch = FrequentItemsSketch::<String>::new(8).unwrap();
+    sketch.update_with_count_ref("a", u64::MAX - 1);
+    sketch.update("b".to_string());
+    assert_eq!(sketch.total_weight(), u64::MAX);
+    let before = sketch.serialize();
+
+    assert!(catch_unwind(AssertUnwindSafe(|| sketch.update("c".to_string()))).is_err());
+    assert_eq!(sketch.serialize(), before);
+    assert!(catch_unwind(AssertUnwindSafe(|| sketch.update_ref("a"))).is_err());
+    assert_eq!(sketch.serialize(), before);
+    assert!(!sketch.is_empty());
+}
+
+#[test]
+fn test_merge_weight_overflow_preserves_state() {
+    let mut sketch = FrequentItemsSketch::<i64>::new(8).unwrap();
+    sketch.update_with_count(1, u64::MAX - 7);
+    let mut purged = FrequentItemsSketch::<i64>::new(8).unwrap();
+    for item in 0..7 {
+        purged.update(item);
+    }
+    assert_eq!(purged.num_active_items(), 0);
+
+    sketch.merge(&purged);
+    assert_eq!(sketch.total_weight(), u64::MAX);
+    assert_eq!(sketch.maximum_error(), 1);
+    let before = sketch.serialize();
+    assert!(catch_unwind(AssertUnwindSafe(|| sketch.merge(&purged))).is_err());
+    assert_eq!(sketch.serialize(), before);
+    assert!(!sketch.is_empty());
 }
 
 #[test]