feat(cms,countsketch): checked f64 updates and heap rescoring

CHANGELOG.md

@@ -10,6 +10,18 @@ signals a backwards-compatible change.
 
 ## [Unreleased]
 
+### Added
+
+- **Checked f64 updates and heap rescoring for the Count-Min and Count Sketch
+  heap sketches.** Over `Vector2D<f64>` counters, `check_weight` and
+  `try_update_weighted` on `CountMinSketch` and `CountSketch` reject a
+  non-finite weight (and a negative one for Count-Min) and a counter
+  overflow, leaving the sketch unchanged; `CountMinSketchWithHeap` and
+  `CountSketchWithHeap` gain the same `try_update_weighted`, which then
+  updates the heap. `rescore_heap` on both heap sketches re-reads every
+  resident's count from the counters, since an update rescores only the key
+  it touches.
+
 ## [0.4.0] - 2026-10-20
 
 Breaking release. ASAPv1 is the one serialized form: protobuf, the legacy

docs/api/api_count_min_sketch.md

@@ -35,6 +35,11 @@ fn update_weighted_batch(&mut self, values: &[(DataInput, S::Counter)])
 fn update_with_hash(&mut self, hashed_val: &H::HashType)
 fn update_weighted_with_hash(&mut self, hashed_val: &H::HashType, weight: S::Counter)
 fn clear(&mut self)
+
+// Vector2D<f64>, RegularPath: weights finite and nonnegative; an error (a bad
+// weight or a counter overflow) leaves the sketch unchanged
+fn check_weight(weight: f64) -> Result<(), String>
+fn try_update_weighted(&mut self, value: &DataInput, weight: f64) -> Result<(), String>
 ```
 
 `clear` zeroes every counter and returns the Nitro sampler to its start, keeping

docs/api/api_count_min_sketch_with_heap.md

@@ -34,6 +34,13 @@ fn update_weighted(&mut self, key: &DataInput, weight: S::Counter)
 fn update_batch(&mut self, values: &[DataInput])
 fn clear_heap(&mut self)
 fn clear(&mut self)
+// Vector2D<f64>, RegularPath: CountMinSketch::try_update_weighted, then the heap;
+// an error leaves the counters and the heap unchanged
+fn check_weight(weight: f64) -> Result<(), String>  // the inner sketch's
+fn try_update_weighted(&mut self, key: &DataInput, weight: f64) -> Result<(), String>
+// RegularPath: re-reads every resident's count from the counters (an update
+// rescores only its own key); call before reading the heap for current counts
+fn rescore_heap(&mut self)
 ```
 
 `clear_heap` empties the heap; `clear` also zeroes the counters and restarts the

docs/api/api_count_sketch.md

@@ -39,6 +39,11 @@ fn update_weighted(&mut self, value: &DataInput, weight: S::Counter)
 fn update_with_hash(&mut self, hashed_val: &H::HashType)
 fn clear(&mut self)
 
+// Vector2D<f64>, RegularPath: weights finite (signed); an error (a bad weight
+// or a counter overflow) leaves the sketch unchanged
+fn check_weight(weight: f64) -> Result<(), String>
+fn try_update_weighted(&mut self, value: &DataInput, weight: f64) -> Result<(), String>
+
 // CountSketchWithL2Norm
 fn update_weighted(&mut self, val: &DataInput, weight: i64)
 fn update_weighted_with_hash(&mut self, hashed_val: u128, weight: i64)

docs/api/api_count_sketch_with_heap.md

@@ -35,6 +35,13 @@ fn update_weighted(&mut self, key: &DataInput, weight: S::Counter)
 fn update_batch(&mut self, values: &[DataInput])
 fn clear_heap(&mut self)
 fn clear(&mut self)
+// Vector2D<f64>, RegularPath: CountSketch::try_update_weighted, then the heap;
+// an error leaves the counters and the heap unchanged
+fn check_weight(weight: f64) -> Result<(), String>  // the inner sketch's
+fn try_update_weighted(&mut self, key: &DataInput, weight: f64) -> Result<(), String>
+// RegularPath: re-reads every resident's count from the counters (an update
+// rescores only its own key); call before reading the heap for current counts
+fn rescore_heap(&mut self)
 ```
 
 `clear_heap` empties the heap; `clear` also zeroes the counters and restarts the

src/sketches/countminsketch.rs

@@ -19,8 +19,9 @@ use crate::octo_delta::{
 };
 use crate::{
     DataInput, DefaultMatrixI32, DefaultMatrixI64, DefaultMatrixI128, DefaultXxHasher, FastPath,
-    FastPathHasher, FixedMatrix, MatrixFastHash, MatrixStorage, NitroTarget, QuickMatrixI64,
-    QuickMatrixI128, RegularPath, SketchHasher, Vector2D, hash64_seeded, nitro_delta_saturated_i32,
+    FastPathHasher, FixedMatrix, MATRIX_MAX_ROWS, MatrixFastHash, MatrixStorage, NitroTarget,
+    QuickMatrixI64, QuickMatrixI128, RegularPath, SketchHasher, Vector2D, hash64_seeded,
+    nitro_delta_saturated_i32,
 };
 
 mod delta;
@@ -410,6 +411,42 @@ where
 /// Count-Min sketch with floating-point counters (no integer rounding).
 pub type CountMinF64<H = DefaultXxHasher> = CountMinSketch<Vector2D<f64>, RegularPath, H>;
 
+// Checked updates for f64 counters.
+impl<H: SketchHasher> CountMinSketch<Vector2D<f64>, RegularPath, H> {
+    /// Checks a weight [`Self::try_update_weighted`] accepts: finite and
+    /// nonnegative, since a Count-Min estimate is a minimum over rows and is an
+    /// upper bound only while every update adds.
+    pub fn check_weight(weight: f64) -> Result<(), String> {
+        if weight.is_finite() && weight >= 0.0 {
+            Ok(())
+        } else {
+            Err(format!(
+                "CMS weight {weight} must be finite and nonnegative"
+            ))
+        }
+    }
+
+    /// [`Self::update_weighted`] that rejects a weight [`Self::check_weight`]
+    /// rejects, and an update that would carry a probed counter past
+    /// `f64::MAX`. On an error the sketch is unchanged.
+    pub fn try_update_weighted(&mut self, value: &DataInput, weight: f64) -> Result<(), String> {
+        Self::check_weight(weight)?;
+        let cols = self.counts.cols();
+        let probed_cols = (0..self.counts.rows())
+            .map(|r| ((H::hash64_seeded(r, value) & LOWER_32_MASK) as usize) % cols)
+            .collect::<smallvec::SmallVec<[usize; MATRIX_MAX_ROWS]>>();
+        for (r, &col) in probed_cols.iter().enumerate() {
+            if !(self.counts.query_one_counter(r, col) + weight).is_finite() {
+                return Err(format!("CMS counter overflow adding weight {weight}"));
+            }
+        }
+        for (r, col) in probed_cols.into_iter().enumerate() {
+            self.counts.increment_by_row(r, col, weight);
+        }
+        Ok(())
+    }
+}
+
 impl<S: MatrixStorage<Counter = i32>, H: SketchHasher> CountMinSketch<S, RegularPath, H> {
     /// Inserts an observation, emitting a delta at every promotion of the
     /// default threshold `OCTO_COUNT_MIN_DEFAULT_THRESHOLD`.

src/sketches/countminsketch_heap.rs

@@ -251,6 +251,26 @@ where
     }
 }
 
+// -- Checked f64 updates ------------------------------------------------------
+
+impl<H: SketchHasher> CountMinSketchWithHeap<Vector2D<f64>, RegularPath, H> {
+    /// [`CountMinSketch::check_weight`]: Ok for a finite, nonnegative weight,
+    /// the weights [`Self::try_update_weighted`] accepts.
+    pub fn check_weight(weight: f64) -> Result<(), String> {
+        CountMinSketch::<Vector2D<f64>, RegularPath, H>::check_weight(weight)
+    }
+
+    /// [`Self::update_weighted`] through [`CountMinSketch::try_update_weighted`]:
+    /// rejects a negative or non-finite weight and a counter overflow, leaving
+    /// the counters and the heap unchanged on an error.
+    pub fn try_update_weighted(&mut self, key: &DataInput, weight: f64) -> Result<(), String> {
+        self.count_min_sketch.try_update_weighted(key, weight)?;
+        let est = self.count_min_sketch.estimate(key);
+        self.heap.update(key, est);
+        Ok(())
+    }
+}
+
 // -- RegularPath insert / estimate / merge -----------------------------------
 
 impl<S: MatrixStorage, H: SketchHasher> CountMinSketchWithHeap<S, RegularPath, H>
@@ -307,6 +327,26 @@ where
             self.heap.update(&key_ref, est.into());
         }
     }
+
+    /// Re-reads every heap resident's count from the counters, keeping the
+    /// same residents. An update rescores only the key it touches, so a
+    /// resident sharing a cell with it keeps its earlier count; a caller that
+    /// needs every resident's count to equal its current estimate calls this
+    /// before reading the heap. Costs one estimate per resident.
+    pub fn rescore_heap(&mut self) {
+        let rescored = self
+            .heap
+            .heap()
+            .iter()
+            .map(|item| {
+                let est = self.count_min_sketch.estimate(&item.key.as_data_input());
+                (item.key.clone(), est.into())
+            })
+            .collect::<Vec<_>>();
+        for (key, count) in rescored {
+            self.heap.update_owned(&key, count);
+        }
+    }
 }
 
 // -- FastPath insert / estimate / merge --------------------------------------
@@ -996,4 +1036,75 @@ mod tests {
             "heap overlap too low: overlap={overlap}, regular={regular_heap_keys:?}, fast={fast_heap_keys:?}"
         );
     }
+
+    fn f64_heap_entries(
+        sketch: &CountMinSketchWithHeap<Vector2D<f64>, RegularPath>,
+    ) -> Vec<(crate::OwnedDataInput, f64)> {
+        let mut entries = sketch
+            .heap()
+            .heap()
+            .iter()
+            .map(|item| (item.key.clone(), item.count))
+            .collect::<Vec<_>>();
+        entries.sort_by(|a, b| a.1.total_cmp(&b.1));
+        entries
+    }
+
+    // A checked f64 update adds fractional weights, rejects a negative or
+    // non-finite weight and a counter overflow, and on an error leaves the
+    // counters and the heap as they were.
+    #[test]
+    fn f64_checked_updates_are_all_or_nothing() {
+        let mut sketch = CountMinSketchWithHeap::<Vector2D<f64>, RegularPath>::new(3, 64, 4);
+        let key = DataInput::I64(7);
+        sketch.try_update_weighted(&key, 0.125).unwrap();
+        sketch.try_update_weighted(&key, 0.25).unwrap();
+        assert_eq!(sketch.estimate(&key), 0.375);
+        assert_eq!(heap_count_for_key(sketch.heap(), &key), Some(0.375));
+        sketch
+            .try_update_weighted(&DataInput::I64(8), f64::MAX)
+            .unwrap();
+
+        let cells = sketch.count_min_sketch().as_storage().as_slice().to_vec();
+        let heap = f64_heap_entries(&sketch);
+        for weight in [-0.5, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
+            assert!(
+                CountMinSketchWithHeap::<Vector2D<f64>, RegularPath>::check_weight(weight).is_err()
+            );
+            assert!(sketch.try_update_weighted(&key, weight).is_err());
+        }
+        assert!(
+            sketch
+                .try_update_weighted(&DataInput::I64(8), f64::MAX)
+                .is_err()
+        );
+        assert_eq!(sketch.count_min_sketch().as_storage().as_slice(), cells);
+        assert_eq!(f64_heap_entries(&sketch), heap);
+    }
+
+    // Under heavy collisions an update moves cells other residents share, so
+    // their heap counts go stale; `rescore_heap` re-reads every resident and
+    // leaves each count equal to its current estimate.
+    #[test]
+    fn rescore_heap_rereads_every_resident() {
+        let mut sketch = CountMinSketchWithHeap::<Vector2D<f64>, RegularPath>::new(3, 8, 4);
+        let mut stale = 0;
+        for step in 0..60i64 {
+            sketch
+                .try_update_weighted(&DataInput::I64(step % 7), 0.5 + (step % 3) as f64)
+                .unwrap();
+            stale += sketch
+                .heap()
+                .heap()
+                .iter()
+                .filter(|item| item.count != sketch.estimate(&item.key.as_data_input()))
+                .count();
+            sketch.rescore_heap();
+            assert!(sketch.heap().len() <= 4);
+            for item in sketch.heap().heap() {
+                assert_eq!(item.count, sketch.estimate(&item.key.as_data_input()));
+            }
+        }
+        assert!(stale > 0, "8 columns must leave some resident stale");
+    }
 }

src/sketches/countsketch.rs

@@ -12,8 +12,9 @@ use crate::common::hash::check_matrix_shape;
 use crate::common::structure_utils::AdmittedRows;
 use crate::{
     DataInput, DefaultMatrixI32, DefaultMatrixI64, DefaultMatrixI128, DefaultXxHasher, FastPath,
-    FastPathHasher, FixedMatrix, MatrixFastHash, MatrixStorage, NitroTarget, QuickMatrixI64,
-    QuickMatrixI128, RegularPath, SketchHasher, Vector2D, hash64_seeded, nitro_delta_saturated_i32,
+    FastPathHasher, FixedMatrix, MATRIX_MAX_ROWS, MatrixFastHash, MatrixStorage, NitroTarget,
+    QuickMatrixI64, QuickMatrixI128, RegularPath, SketchHasher, Vector2D, hash64_seeded,
+    nitro_delta_saturated_i32,
 };
 use std::marker::PhantomData;
 use std::ops::Neg;
@@ -407,6 +408,51 @@ where
     }
 }
 
+// Checked updates for f64 counters.
+impl<H: SketchHasher> CountSketch<Vector2D<f64>, RegularPath, H> {
+    /// Checks a weight [`Self::try_update_weighted`] accepts: finite. A
+    /// negative weight is a signed update.
+    pub fn check_weight(weight: f64) -> Result<(), String> {
+        if weight.is_finite() {
+            Ok(())
+        } else {
+            Err(format!("Count Sketch weight {weight} must be finite"))
+        }
+    }
+
+    /// [`Self::update_weighted`] that rejects a weight [`Self::check_weight`]
+    /// rejects, and an update that would carry a probed counter past
+    /// `f64::MAX` in either direction. On an error the sketch is unchanged.
+    pub fn try_update_weighted(&mut self, value: &DataInput, weight: f64) -> Result<(), String> {
+        Self::check_weight(weight)?;
+        let cols = self.counts.cols();
+        let probed = (0..self.counts.rows())
+            .map(|r| {
+                let hashed = H::hash64_seeded(r, value);
+                let col = ((hashed & LOWER_32_MASK) as usize) % cols;
+                let delta = if (hashed >> 63) & 1 == 1 {
+                    weight
+                } else {
+                    -weight
+                };
+                (col, delta)
+            })
+            .collect::<smallvec::SmallVec<[(usize, f64); MATRIX_MAX_ROWS]>>();
+        for (r, &(col, delta)) in probed.iter().enumerate() {
+            if !(self.counts.query_one_counter(r, col) + delta).is_finite() {
+                return Err(format!(
+                    "Count Sketch counter overflow adding weight {weight}"
+                ));
+            }
+        }
+        for (r, (col, delta)) in probed.into_iter().enumerate() {
+            self.counts
+                .update_one_counter(r, col, |a, b| *a += b, delta);
+        }
+        Ok(())
+    }
+}
+
 // Fast-path Count operations using precomputed hashes.
 impl<S, H: SketchHasher> CountSketch<S, FastPath, H>
 where