diff --git a/CHANGELOG.md b/CHANGELOG.md index 904006b5..fa299761 100644 --- a/CHANGELOG.md +++ b/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` 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 diff --git a/docs/api/api_count_min_sketch.md b/docs/api/api_count_min_sketch.md index d20526dd..ca987ec7 100644 --- a/docs/api/api_count_min_sketch.md +++ b/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, 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 diff --git a/docs/api/api_count_min_sketch_with_heap.md b/docs/api/api_count_min_sketch_with_heap.md index 2742d5e5..7a2faaed 100644 --- a/docs/api/api_count_min_sketch_with_heap.md +++ b/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, 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 diff --git a/docs/api/api_count_sketch.md b/docs/api/api_count_sketch.md index 8a314e61..967c0956 100644 --- a/docs/api/api_count_sketch.md +++ b/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, 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) diff --git a/docs/api/api_count_sketch_with_heap.md b/docs/api/api_count_sketch_with_heap.md index 6393c9e4..45b8c258 100644 --- a/docs/api/api_count_sketch_with_heap.md +++ b/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, 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 diff --git a/src/sketches/countminsketch.rs b/src/sketches/countminsketch.rs index 6e217287..cebe1dd3 100644 --- a/src/sketches/countminsketch.rs +++ b/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 = CountMinSketch, RegularPath, H>; +// Checked updates for f64 counters. +impl CountMinSketch, 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::>(); + 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, H: SketchHasher> CountMinSketch { /// Inserts an observation, emitting a delta at every promotion of the /// default threshold `OCTO_COUNT_MIN_DEFAULT_THRESHOLD`. diff --git a/src/sketches/countminsketch_heap.rs b/src/sketches/countminsketch_heap.rs index 04fbdeaf..230f462f 100644 --- a/src/sketches/countminsketch_heap.rs +++ b/src/sketches/countminsketch_heap.rs @@ -251,6 +251,26 @@ where } } +// -- Checked f64 updates ------------------------------------------------------ + +impl CountMinSketchWithHeap, 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::, 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 CountMinSketchWithHeap @@ -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::>(); + 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, RegularPath>, + ) -> Vec<(crate::OwnedDataInput, f64)> { + let mut entries = sketch + .heap() + .heap() + .iter() + .map(|item| (item.key.clone(), item.count)) + .collect::>(); + 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::, 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::, 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::, 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"); + } } diff --git a/src/sketches/countsketch.rs b/src/sketches/countsketch.rs index 20e8027d..a389a5e4 100644 --- a/src/sketches/countsketch.rs +++ b/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 CountSketch, 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::>(); + 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 CountSketch where diff --git a/src/sketches/countsketch_heap.rs b/src/sketches/countsketch_heap.rs index 1b2cb032..68f639e3 100644 --- a/src/sketches/countsketch_heap.rs +++ b/src/sketches/countsketch_heap.rs @@ -307,6 +307,26 @@ where } } +// -- Checked f64 updates ------------------------------------------------------ + +impl CountSketchWithHeap, RegularPath, H> { + /// [`CountSketch::check_weight`]: Ok for a finite weight, the weights + /// [`Self::try_update_weighted`] accepts. + pub fn check_weight(weight: f64) -> Result<(), String> { + CountSketch::, RegularPath, H>::check_weight(weight) + } + + /// [`Self::update_weighted`] through [`CountSketch::try_update_weighted`]: + /// rejects a 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_sketch.try_update_weighted(key, weight)?; + let est = self.count_sketch.estimate(key); + self.heap.update(key, est); + Ok(()) + } +} + // -- RegularPath insert / estimate / merge ----------------------------------- impl CountSketchWithHeap @@ -362,6 +382,26 @@ where self.heap.update(&key_ref, S::Counter::heap_count(est)); } } + + /// 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_sketch.estimate(&item.key.as_data_input()); + (item.key.clone(), S::Counter::heap_count(est)) + }) + .collect::>(); + for (key, count) in rescored { + self.heap.update_owned(&key, count); + } + } } // -- FastPath insert / estimate / merge -------------------------------------- @@ -1070,4 +1110,81 @@ mod tests { "heap overlap too low: overlap={overlap}, regular={regular_heap_keys:?}, fast={fast_heap_keys:?}" ); } + + fn f64_heap_entries( + sketch: &CountSketchWithHeap, RegularPath>, + ) -> Vec<(crate::OwnedDataInput, f64)> { + let mut entries = sketch + .heap() + .heap() + .iter() + .map(|item| (item.key.clone(), item.count)) + .collect::>(); + entries.sort_by(|a, b| a.1.total_cmp(&b.1)); + entries + } + + // A checked f64 update adds signed fractional weights, rejects a + // non-finite weight and a counter overflow in either direction, 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 = CountSketchWithHeap::, RegularPath>::new(5, 4096, 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); + let resident = sketch + .heap() + .heap() + .iter() + .find(|item| item.key.as_data_input() == key) + .map(|item| item.count); + assert_eq!(resident, Some(-0.375)); + sketch + .try_update_weighted(&DataInput::I64(8), -f64::MAX) + .unwrap(); + + let cells = sketch.count_sketch().as_storage().as_slice().to_vec(); + let heap = f64_heap_entries(&sketch); + for weight in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + assert!( + CountSketchWithHeap::, 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_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 = CountSketchWithHeap::, 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"); + } }