Fix three bugs in Nitro geometric row-sampling

src/common/structure_utils.rs

@@ -3,12 +3,9 @@
 //! Vector2D:
 //! Vector3D:
 //! CommonHeap:
-// use rand::rngs::SmallRng;
-// use rand::{Rng, SeedableRng, rng};
+use rand::rngs::SmallRng;
+use rand::{Rng, SeedableRng, rng};
 use serde::{Deserialize, Serialize};
-
-// use crate::PRECOMPUTED_SAMPLE;
-use crate::PRECOMPUTED_SAMPLE_RATE_1PERCENT;
 /// Helper trait for converting sketch counter types to f64 for median calculation.
 pub trait ToF64 {
     /// Converts the value into `f64`.
@@ -52,19 +49,19 @@ pub struct Nitro {
     pub to_skip: usize,
     /// Precomputed: 1.0 / ln(1 - sampling_rate) for geometric sampling
     inv_ln_one_minus_p: f64,
-    // #[serde(skip)]
-    // #[serde(default = "new_small_rng")]
-    // // generator: SmallRng,
+    #[serde(skip)]
+    #[serde(default = "new_small_rng")]
+    generator: SmallRng,
     /// Weight applied to each sampled update.
     pub delta: u64,
     idx: usize,
     mask: usize,
 }
 
-// fn new_small_rng() -> SmallRng {
-//     let mut seed_rng = rng();
-//     SmallRng::from_rng(&mut seed_rng)
-// }
+fn new_small_rng() -> SmallRng {
+    let mut seed_rng = rng();
+    SmallRng::from_rng(&mut seed_rng)
+}
 
 impl Default for Nitro {
     fn default() -> Self {
@@ -73,7 +70,7 @@ impl Default for Nitro {
             sampling_rate: 0.0,
             to_skip: 0,
             inv_ln_one_minus_p: 0.0, // not used unless Nitro mode is enabled
-            // generator: new_small_rng(), // not used unless Nitro mode is enabled
+            generator: new_small_rng(), // not used unless Nitro mode is enabled
             delta: 0,
             idx: 0,
             mask: 0x10000,
@@ -84,6 +81,23 @@ impl Default for Nitro {
 impl Nitro {
     /// Creates a Nitro state with the given sampling rate.
     pub fn init_nitro(rate: f64) -> Self {
+        Self::init_nitro_with(rate, new_small_rng())
+    }
+
+    /// Creates a Nitro state with the given sampling rate, whose sampling
+    /// decisions are reproducible: the geometric-skip RNG is seeded from
+    /// `seed` (via `SmallRng::seed_from_u64`) instead of OS entropy.
+    /// `init_nitro`'s default is intentionally non-reproducible (a fresh
+    /// `new_small_rng()` draws from OS entropy every call, and the
+    /// generator field itself is `#[serde(skip)]`, so it never survives a
+    /// save/restore either) -- appropriate for production sampling, but
+    /// not for a caller that needs deterministic runs (e.g. a benchmark
+    /// harness driven by its own `--seed` flag).
+    pub fn init_nitro_seeded(rate: f64, seed: u64) -> Self {
+        Self::init_nitro_with(rate, SmallRng::seed_from_u64(seed))
+    }
+
+    fn init_nitro_with(rate: f64, generator: SmallRng) -> Self {
         assert!(
             !rate.is_nan() && rate > 0.0 && rate <= 1.0,
             "sample_rate must be within (0.0, 1.0]"
@@ -98,7 +112,7 @@ impl Nitro {
             sampling_rate: rate,
             to_skip: 0,
             inv_ln_one_minus_p: inv_ln,
-            // generator: new_small_rng(),
+            generator,
             delta: 0,
             idx: 0,
             mask: 0x10000,
@@ -107,25 +121,49 @@ impl Nitro {
         nitro
     }
 
-    // for profiling
     #[inline(always)]
     /// Draws the next geometric skip distance.
+    ///
+    /// Uses a live geometric draw from `self.generator`, matching
+    /// `NitroBatch::draw_geometric` in `sketch_framework/nitro.rs` (the
+    /// same formula, kept consistent between the two Nitro
+    /// implementations in this crate). Previously this used a fixed,
+    /// precomputed ~1%-rate lookup table regardless of the configured
+    /// `sampling_rate` (`idx` cycling through `PRECOMPUTED_SAMPLE_RATE_1PERCENT`),
+    /// which only coincidentally matched the requested rate at ~0.01 --
+    /// at every other configured rate, the actual skip distances (and
+    /// thus how often a row was genuinely touched) never tracked
+    /// `sampling_rate` at all, only the `delta` compensation did,
+    /// silently miscalibrating every estimate. Verified empirically: at
+    /// `sampling_rate=0.5` the true touch count is exactly the touch
+    /// count observed at `sampling_rate=0.01`, `0.02`, `0.05`, `0.1`, in
+    /// a Count-Min sketch with the table-driven version.
+    ///
+    /// The `- 1` matches NitroSketch paper Algorithm 1's `Update(p)`:
+    /// `r += Geo(p)` advances the row cursor by the raw geometric draw
+    /// (support `{1,2,...}`, mean `1/p`) with no extra offset. Callers
+    /// here (`fast_insert_nitro`, `NitroBatch::insert`) advance the
+    /// cursor via `r += to_skip + 1` -- treating `to_skip` as "positions
+    /// to skip *before* the touch" (support `{0,1,...}`, mean `1/p - 1`)
+    /// and adding the `+1` back for the touch itself. Without the `- 1`
+    /// here, that `+1` is double-counted: the true step mean becomes
+    /// `1/p + 1` instead of `1/p`, so the achieved sampling density is
+    /// `p/(1+p)` rather than `p` -- e.g. at `p=0.5` the sketch actually
+    /// samples at ~33%, not 50%, systematically undercounting every
+    /// estimate. Verified empirically via a standalone row-cycling
+    /// simulation matching this exact stepping logic.
     pub fn draw_geometric(&mut self) {
         if self.is_full_sampling() {
             self.to_skip = 0;
             return;
         }
-        // let k = loop {
-        //     let r = self.generator.random::<f64>();
-        //     if r != 0.0_f64 && r != 1.0_f64 {
-        //         break r;
-        //     }
-        // };
-        // self.to_skip = ((1.0 - k).ln() * self.inv_ln_one_minus_p).ceil() as usize;
-
-        // self.to_skip = (PRECOMPUTED_SAMPLE[self.idx] * self.inv_ln_one_minus_p).ceil() as usize;
-
-        self.to_skip = PRECOMPUTED_SAMPLE_RATE_1PERCENT[self.idx].ceil() as usize;
+        let k = loop {
+            let r = self.generator.random::<f64>();
+            if r != 0.0_f64 && r != 1.0_f64 {
+                break r;
+            }
+        };
+        self.to_skip = ((1.0 - k).ln() * self.inv_ln_one_minus_p).ceil() as usize - 1;
         self.idx = (self.idx + 1) & self.mask;
     }
 

src/common/structures/vector2d.rs

@@ -108,6 +108,12 @@ impl<T> Vector2D<T> {
         self.nitro = Nitro::init_nitro(sampling_rate);
     }
 
+    /// Enables Nitro sampling with the provided rate and a reproducible
+    /// RNG seed -- see `Nitro::init_nitro_seeded`'s doc comment.
+    pub fn enable_nitro_seeded(&mut self, sampling_rate: f64, seed: u64) {
+        self.nitro = Nitro::init_nitro_seeded(sampling_rate, seed);
+    }
+
     /// Disables Nitro sampling and resets the internal state.
     pub fn disable_nitro(&mut self) {
         self.nitro = Nitro::default();
@@ -306,6 +312,28 @@ impl<T> Vector2D<T> {
         self.update_one_counter(row, col, op, value);
     }
 
+    #[inline(always)]
+    /// Reads one row's counter using a packed hash value, with the exact
+    /// same column extraction as `update_by_row`.
+    ///
+    /// Nitro inserts (`fast_insert_nitro`) write via `update_by_row` using a
+    /// raw `H::hash128_seeded(0, value)`. Querying that data must extract
+    /// the same column from the same raw hash -- `hash_for_matrix`
+    /// (`common::hash::hash_for_matrix_seeded_with_mode_generic`) is *not*
+    /// equivalent: for dimensions where `hash_mode_for_matrix` selects
+    /// `Packed64`, it hashes with `hash64_seeded` instead of
+    /// `hash128_seeded` entirely (a different algorithm, not just a
+    /// truncation), so a query built on `hash_for_matrix` silently reads
+    /// the wrong cells for any Nitro sketch narrow/shallow enough to hit
+    /// that mode (only coincides for dimensions landing in `Packed128`
+    /// mode, where the seed-0 case matches `hash128_seeded(0, key)`
+    /// exactly). Use this for any query path that must round-trip with a
+    /// `fast_insert_nitro`/`update_by_row`-based insert.
+    pub fn query_by_row(&self, row: usize, hashed: u128) -> &T {
+        let col = (hashed >> (self.mask_bits as usize * row)) as usize & (self.mask as usize);
+        &self.data[row * self.cols + col]
+    }
+
     #[inline(always)]
     /// Decrements the Nitro skip counter by `c`.
     pub fn reduce_nitro_skip(&mut self, c: usize) {

src/sketch_framework/nitro.rs

@@ -183,9 +183,21 @@ impl<S: NitroTarget> NitroBatch<S> {
         nitro
     }
 
-    // for profiling
     #[inline(always)]
     /// Draws the next geometric skip distance.
+    ///
+    /// The `- 1` matches NitroSketch paper Algorithm 1's `Update(p)`:
+    /// `r += Geo(p)` advances the row cursor by the raw geometric draw
+    /// (support `{1,2,...}`, mean `1/p`) with no extra offset. `insert`
+    /// below advances the cursor via `position += self.to_skip + 1` --
+    /// treating `to_skip` as "positions to skip *before* the touch"
+    /// (support `{0,1,...}`, mean `1/p - 1`) and adding the `+1` back for
+    /// the touch itself. Without the `- 1` here, that `+1` is
+    /// double-counted: the true step mean becomes `1/p + 1` instead of
+    /// `1/p`, so the achieved sampling density is `p/(1+p)` rather than
+    /// `p`, systematically undercounting every estimate. See the
+    /// identical fix and derivation on `Nitro::draw_geometric` in
+    /// `common/structure_utils.rs`.
     pub fn draw_geometric(&mut self) {
         if self.is_full_sampling() {
             self.to_skip = 0;
@@ -197,7 +209,7 @@ impl<S: NitroTarget> NitroBatch<S> {
                 break r;
             }
         };
-        self.to_skip = ((1.0 - k).ln() * self.inv_ln_one_minus_p).ceil() as usize;
+        self.to_skip = ((1.0 - k).ln() * self.inv_ln_one_minus_p).ceil() as usize - 1;
         self.idx = (self.idx + 1) & self.mask;
     }
 

src/sketches/countminsketch.rs

@@ -15,7 +15,7 @@ use crate::octo_delta::{CM_PROMASK, CmDelta};
 use crate::{
     DataInput, DefaultMatrixI32, DefaultMatrixI64, DefaultMatrixI128, DefaultXxHasher, FastPath,
     FastPathHasher, FixedMatrix, MatrixFastHash, MatrixStorage, NitroTarget, QuickMatrixI64,
-    QuickMatrixI128, RegularPath, SketchHasher, Vector2D, hash64_seeded,
+    QuickMatrixI128, RegularPath, SketchHasher, Vector2D, compute_median_inline_f64, hash64_seeded,
 };
 
 mod wire;
@@ -482,40 +482,149 @@ where
     }
 }
 
-// Nitro sampling helpers for fast-path CountMin.
+// Nitro sampling helpers for fast-path CountMin (Vector2D<i32> storage).
 impl<H: SketchHasher> CountMin<Vector2D<i32>, FastPath, H> {
     /// Enables Nitro sampling with the provided rate.
     pub fn enable_nitro(&mut self, sampling_rate: f64) {
         self.counts.enable_nitro(sampling_rate);
     }
 
+    /// Enables Nitro sampling with a reproducible RNG seed -- see
+    /// `Nitro::init_nitro_seeded`'s doc comment.
+    pub fn enable_nitro_seeded(&mut self, sampling_rate: f64, seed: u64) {
+        self.counts.enable_nitro_seeded(sampling_rate, seed);
+    }
+
     /// Disables Nitro sampling and resets its internal state.
     pub fn disable_nitro(&mut self) {
         self.counts.disable_nitro();
     }
 
-    /// Inserts an observation using Nitro-aware sampling logic.
+    /// Inserts one occurrence of `value` using Nitro-aware sampling logic.
     #[inline(always)]
     pub fn fast_insert_nitro(&mut self, value: &DataInput) {
+        self.fast_insert_nitro_many(value, 1);
+    }
+
+    /// Inserts `many` occurrences of `value` at once using Nitro-aware
+    /// sampling logic, compensating the touched row(s) by `many /
+    /// sampling_rate` (via `Nitro::scaled_increment`) instead of the
+    /// fixed per-occurrence `nitro().delta` -- so a single call can stand
+    /// in for a weighted observation (e.g. a metric sample's value)
+    /// without looping the caller's own insert `many` times, which would
+    /// make cost scale with the observation's magnitude instead of with
+    /// event count.
+    ///
+    /// Loops until the drawn skip carries past the end of this item's
+    /// `rows` row-slots (matching `Count::fast_insert_nitro`'s pattern):
+    /// a single geometric draw can legitimately land more than one touch
+    /// within the same item when the sampling rate is high, and the final
+    /// `(r + temp + 1) - rows` is only guaranteed non-negative once that
+    /// invariant holds. Without the loop, full sampling (`rate == 1.0`)
+    /// always draws `r = 0, temp = 0`, so `(r + temp + 1) - rows`
+    /// underflows for `rows > 1` and silently disables all further
+    /// row-touches for the lifetime of the sketch.
+    #[inline(always)]
+    pub fn fast_insert_nitro_many(&mut self, value: &DataInput, many: i32) {
         let rows = self.counts.rows();
-        let delta = self.counts.nitro().delta as i32;
+        let delta = self.counts.nitro().scaled_increment(many as u64) as i32;
         if self.counts.nitro().to_skip >= rows {
             self.counts.reduce_nitro_skip(rows);
         } else {
             let hashed = H::hash128_seeded(0, value);
-            let r = self.counts.nitro().to_skip;
-            self.counts.update_by_row(r, hashed, |a, b| *a += b, delta);
-            self.counts.nitro_mut().draw_geometric();
+            let mut r = self.counts.nitro().to_skip;
+            loop {
+                self.counts.update_by_row(r, hashed, |a, b| *a += b, delta);
+                self.counts.nitro_mut().draw_geometric();
+                if r + self.counts.nitro_mut().to_skip + 1 >= rows {
+                    break;
+                }
+                r += self.counts.nitro_mut().to_skip + 1;
+            }
             let temp = self.counts.get_nitro_skip();
             self.counts.update_nitro_skip((r + temp + 1) - rows);
         }
     }
 
-    /// Returns the median estimate using a fast-path matrix hash.
+    /// Returns the median estimate for `value`, hashed the same way
+    /// `fast_insert_nitro`/`fast_insert_nitro_many` hash at insert time
+    /// (`H::hash128_seeded(0, value)` + `Vector2D::query_by_row`'s raw
+    /// shift-based column extraction) -- *not* `hash_for_matrix`, which
+    /// for some `(rows, cols)` picks a `Packed64` layout hashed with a
+    /// different algorithm (`hash64_seeded`) entirely, silently reading
+    /// the wrong cells for any Nitro-inserted data at those dimensions
+    /// (see `Vector2D::query_by_row`'s doc comment).
     pub fn nitro_estimate(&self, value: &DataInput) -> f64 {
-        let hashed_val = <Vector2D<i32> as FastPathHasher<H>>::hash_for_matrix(&self.counts, value);
-        self.counts
-            .fast_query_median(&hashed_val, |val, _, _| (*val) as f64)
+        let hashed = H::hash128_seeded(0, value);
+        let rows = self.counts.rows();
+        let mut vals: Vec<f64> = (0..rows)
+            .map(|r| *self.counts.query_by_row(r, hashed) as f64)
+            .collect();
+        compute_median_inline_f64(&mut vals)
+    }
+}
+
+// Nitro sampling helpers for fast-path CountMin (Vector2D<i64> storage) --
+// mirrors the `Vector2D<i32>` impl above verbatim (Nitro's row-cycling
+// mechanics don't depend on the counter width), for callers whose summed
+// weights can exceed `i32::MAX` (e.g. a heavy-tailed per-key weight
+// multiplied into every observation before insertion).
+impl<H: SketchHasher> CountMin<Vector2D<i64>, FastPath, H> {
+    /// Enables Nitro sampling with the provided rate.
+    pub fn enable_nitro(&mut self, sampling_rate: f64) {
+        self.counts.enable_nitro(sampling_rate);
+    }
+
+    /// Enables Nitro sampling with a reproducible RNG seed -- see
+    /// `Nitro::init_nitro_seeded`'s doc comment.
+    pub fn enable_nitro_seeded(&mut self, sampling_rate: f64, seed: u64) {
+        self.counts.enable_nitro_seeded(sampling_rate, seed);
+    }
+
+    /// Disables Nitro sampling and resets its internal state.
+    pub fn disable_nitro(&mut self) {
+        self.counts.disable_nitro();
+    }
+
+    /// Inserts one occurrence of `value` using Nitro-aware sampling logic.
+    #[inline(always)]
+    pub fn fast_insert_nitro(&mut self, value: &DataInput) {
+        self.fast_insert_nitro_many(value, 1);
+    }
+
+    /// Inserts `many` occurrences of `value` at once -- see the
+    /// `Vector2D<i32>` overload's doc comment for the full rationale.
+    #[inline(always)]
+    pub fn fast_insert_nitro_many(&mut self, value: &DataInput, many: i64) {
+        let rows = self.counts.rows();
+        let delta = self.counts.nitro().scaled_increment(many as u64) as i64;
+        if self.counts.nitro().to_skip >= rows {
+            self.counts.reduce_nitro_skip(rows);
+        } else {
+            let hashed = H::hash128_seeded(0, value);
+            let mut r = self.counts.nitro().to_skip;
+            loop {
+                self.counts.update_by_row(r, hashed, |a, b| *a += b, delta);
+                self.counts.nitro_mut().draw_geometric();
+                if r + self.counts.nitro_mut().to_skip + 1 >= rows {
+                    break;
+                }
+                r += self.counts.nitro_mut().to_skip + 1;
+            }
+            let temp = self.counts.get_nitro_skip();
+            self.counts.update_nitro_skip((r + temp + 1) - rows);
+        }
+    }
+
+    /// Returns the median estimate for `value` -- see the `Vector2D<i32>`
+    /// overload's doc comment for why this doesn't use `hash_for_matrix`.
+    pub fn nitro_estimate(&self, value: &DataInput) -> f64 {
+        let hashed = H::hash128_seeded(0, value);
+        let rows = self.counts.rows();
+        let mut vals: Vec<f64> = (0..rows)
+            .map(|r| *self.counts.query_by_row(r, hashed) as f64)
+            .collect();
+        compute_median_inline_f64(&mut vals)
     }
 }