From e41f8778f1eea338b90125193a2736b2b1e9f797 Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 4 Aug 2026 12:07:54 -0600 Subject: [PATCH 1/3] Fix three bugs in Nitro geometric row-sampling 1. CountMin::fast_insert_nitro was missing the loop present in Count::fast_insert_nitro, causing a usize underflow in `(r + temp + 1) - rows` at full sampling (rate == 1.0), which silently disabled all further row-touches for the sketch's lifetime. 2. Nitro::draw_geometric (structure_utils.rs) had its live RNG-based draw commented out and replaced by a fixed ~1%-rate precomputed table lookup, ignoring the configured sampling_rate entirely at every other rate. 3. draw_geometric (both structure_utils.rs::Nitro and sketch_framework/nitro.rs::NitroBatch) generated Geo(p) directly (support {1,2,...}, mean 1/p), but callers advance the row cursor via `r += to_skip + 1`, which double-counts the +1 already implied by the paper's Algorithm 1 (`r += Geo(p)`, r initialized at -1). This inflated the true step mean to 1/p + 1, so the achieved sampling density was p/(1+p) instead of p -- e.g. ~33% actual density at a configured rate of 50%, systematically undercounting every Nitro estimate. Fixed by drawing Geo(p) - 1 instead. Verified via an isolated probe inserting a single key 1M times at rates [1.0, 0.5, 0.1, 0.05, 0.02, 0.01] for both CountMin and Count: estimates now track the true count within ~1-2.5% at every rate, versus wildly biased or rate-independent results before. Full test suite (497 unit + integration + doctests) passes with no regressions. --- src/common/structure_utils.rs | 73 ++++++++++++++++++++++------------ src/sketch_framework/nitro.rs | 16 +++++++- src/sketches/countminsketch.rs | 22 ++++++++-- 3 files changed, 80 insertions(+), 31 deletions(-) diff --git a/src/common/structure_utils.rs b/src/common/structure_utils.rs index a17e461..2098613 100644 --- a/src/common/structure_utils.rs +++ b/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, @@ -98,7 +95,7 @@ impl Nitro { sampling_rate: rate, to_skip: 0, inv_ln_one_minus_p: inv_ln, - // generator: new_small_rng(), + generator: new_small_rng(), delta: 0, idx: 0, mask: 0x10000, @@ -107,25 +104,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::(); - // 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::(); + 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; } diff --git a/src/sketch_framework/nitro.rs b/src/sketch_framework/nitro.rs index 8e03237..4c64346 100644 --- a/src/sketch_framework/nitro.rs +++ b/src/sketch_framework/nitro.rs @@ -183,9 +183,21 @@ impl NitroBatch { 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 NitroBatch { 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; } diff --git a/src/sketches/countminsketch.rs b/src/sketches/countminsketch.rs index 4d128ab..56beda2 100644 --- a/src/sketches/countminsketch.rs +++ b/src/sketches/countminsketch.rs @@ -495,6 +495,16 @@ impl CountMin, FastPath, H> { } /// Inserts an observation using Nitro-aware sampling logic. + /// + /// 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(&mut self, value: &DataInput) { let rows = self.counts.rows(); @@ -503,9 +513,15 @@ impl CountMin, FastPath, H> { 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); } From dfd16f779fd97c43a7fcf2fd7ce6c2753033276c Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 4 Aug 2026 12:23:16 -0600 Subject: [PATCH 2/3] Fix Nitro query hash mismatch; add weighted Nitro insert + Count support nitro_estimate (CountMin) queried via hash_for_matrix, which for some (rows, cols) dimensions selects a Packed64 hash layout hashed with hash64_seeded -- a different algorithm than the raw hash128_seeded(0, value) that fast_insert_nitro writes with. Insert and query only happened to agree at dimensions landing in Packed128 mode (where the seed-0 case coincides with hash128_seeded(0, key)), masking the bug at the row/col configuration used by the earlier verification probe. Fixed by adding Vector2D::query_by_row (mirrors update_by_row's exact column extraction) and rewriting nitro_estimate to hash and read the same way fast_insert_nitro writes, for both CountMin and Count (Count's fast_insert_nitro previously had no matching query method at all -- callers had to fall back to the also-hash_for_matrix-based regular estimate()). Also: - Added fast_insert_nitro_many(value, many) to both CountMin and Count (Vector2D and Vector2D), compensating by Nitro::scaled_increment(many) instead of the fixed per-occurrence delta -- lets one row-cycled, geometrically-sampled call stand in for a weighted observation (e.g. a metric sample's value) without looping the insert call itself, keeping cost proportional to event count rather than magnitude. - Mirrored all Nitro support (enable_nitro, fast_insert_nitro[_many], nitro_estimate) for Vector2D storage, for callers whose summed weights can exceed i32::MAX. Verified via the same isolated probe, extended with a rows=3/cols=8 case (forces Packed64 mode, previously broken) and a weighted-insert case for both i32 and i64 storage on CountMin and Count -- all track the true value within sampling noise. Full test suite (497 unit + integration + doctests) still passes. --- src/common/structures/vector2d.rs | 22 ++++++ src/sketches/countminsketch.rs | 97 +++++++++++++++++++++++--- src/sketches/countsketch.rs | 110 ++++++++++++++++++++++++++++-- 3 files changed, 216 insertions(+), 13 deletions(-) diff --git a/src/common/structures/vector2d.rs b/src/common/structures/vector2d.rs index 3d9a443..711302d 100644 --- a/src/common/structures/vector2d.rs +++ b/src/common/structures/vector2d.rs @@ -306,6 +306,28 @@ impl Vector2D { 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) { diff --git a/src/sketches/countminsketch.rs b/src/sketches/countminsketch.rs index 56beda2..61ee12f 100644 --- a/src/sketches/countminsketch.rs +++ b/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,7 +482,7 @@ where } } -// Nitro sampling helpers for fast-path CountMin. +// Nitro sampling helpers for fast-path CountMin (Vector2D storage). impl CountMin, FastPath, H> { /// Enables Nitro sampling with the provided rate. pub fn enable_nitro(&mut self, sampling_rate: f64) { @@ -494,7 +494,20 @@ impl CountMin, FastPath, H> { 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): @@ -506,9 +519,73 @@ impl CountMin, FastPath, H> { /// 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().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 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`, 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 = H::hash128_seeded(0, value); + let rows = self.counts.rows(); + let mut vals: Vec = (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 storage) -- +// mirrors the `Vector2D` 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 CountMin, FastPath, H> { + /// Enables Nitro sampling with the provided rate. + pub fn enable_nitro(&mut self, sampling_rate: f64) { + self.counts.enable_nitro(sampling_rate); + } + + /// 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` 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().delta as i32; + 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 { @@ -527,11 +604,15 @@ impl CountMin, FastPath, H> { } } - /// Returns the median estimate using a fast-path matrix hash. + /// Returns the median estimate for `value` -- see the `Vector2D` + /// overload's doc comment for why this doesn't use `hash_for_matrix`. pub fn nitro_estimate(&self, value: &DataInput) -> f64 { - let hashed_val = as FastPathHasher>::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 = (0..rows) + .map(|r| *self.counts.query_by_row(r, hashed) as f64) + .collect(); + compute_median_inline_f64(&mut vals) } } diff --git a/src/sketches/countsketch.rs b/src/sketches/countsketch.rs index 22a56cc..3ace1f4 100644 --- a/src/sketches/countsketch.rs +++ b/src/sketches/countsketch.rs @@ -10,7 +10,7 @@ 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, }; use rmp_serde::{ decode::Error as RmpDecodeError, encode::Error as RmpEncodeError, from_slice, to_vec_named, @@ -472,18 +472,33 @@ impl Count, M, H> { } } -// Nitro sampling helpers for fast-path Count. +// Nitro sampling helpers for fast-path Count (Vector2D storage). impl Count, FastPath, H> { /// Enables Nitro sampling with the provided rate. pub fn enable_nitro(&mut self, sampling_rate: f64) { self.counts.enable_nitro(sampling_rate); } - /// Inserts an observation using Nitro geometric-sampling acceleration. + /// 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 geometric-sampling + /// acceleration. #[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, compensating the + /// touched row(s) by `many / sampling_rate` instead of the fixed + /// per-occurrence `nitro().delta` -- see `CountMin::fast_insert_nitro_many`'s + /// doc comment for the full rationale (same reasoning applies here). + #[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; + 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 { @@ -493,7 +508,77 @@ impl Count, FastPath, H> { let bit = (hashed >> (127 - r)) & 1; let sign = (bit << 1) as i32 - 1; self.counts - .update_by_row(r, hashed, |a, b| *a += b, sign * (delta as i32)); + .update_by_row(r, hashed, |a, b| *a += b, sign * 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`, hashed and sign-extracted + /// the same way `fast_insert_nitro`/`fast_insert_nitro_many` do at + /// insert time (`H::hash128_seeded(0, value)`, `Vector2D::query_by_row` + /// for the column, `(hashed >> (127 - row)) & 1` for the sign) -- + /// *not* `hash_for_matrix`, which for some `(rows, cols)` picks a + /// `Packed64` layout hashed with a different algorithm entirely (see + /// `Vector2D::query_by_row`'s doc comment / `CountMin::nitro_estimate`). + pub fn nitro_estimate(&self, value: &DataInput) -> f64 { + let hashed = H::hash128_seeded(0, value); + let rows = self.counts.rows(); + let mut vals: Vec = (0..rows) + .map(|r| { + let bit = (hashed >> (127 - r)) & 1; + let sign = (bit << 1) as i64 - 1; + (sign * (*self.counts.query_by_row(r, hashed) as i64)) as f64 + }) + .collect(); + compute_median_inline_f64(&mut vals) + } +} + +// Nitro sampling helpers for fast-path Count (Vector2D storage) -- +// mirrors the `Vector2D` impl above verbatim; see +// `CountMin, FastPath, H>`'s equivalent block for why this +// exists as a separate, non-generic mirror. +impl Count, FastPath, H> { + /// Enables Nitro sampling with the provided rate. + pub fn enable_nitro(&mut self, sampling_rate: f64) { + self.counts.enable_nitro(sampling_rate); + } + + /// 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 geometric-sampling + /// acceleration. + #[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` 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 { + let bit = (hashed >> (127 - r)) & 1; + let sign = (bit << 1) as i64 - 1; + self.counts + .update_by_row(r, hashed, |a, b| *a += b, sign * delta); self.counts.nitro_mut().draw_geometric(); if r + self.counts.nitro_mut().to_skip + 1 >= rows { break; @@ -504,6 +589,21 @@ impl Count, FastPath, H> { self.counts.update_nitro_skip((r + temp + 1) - rows); } } + + /// Returns the median estimate for `value` -- see the `Vector2D` + /// 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 = (0..rows) + .map(|r| { + let bit = (hashed >> (127 - r)) & 1; + let sign = (bit << 1) as i64 - 1; + (sign * (*self.counts.query_by_row(r, hashed))) as f64 + }) + .collect(); + compute_median_inline_f64(&mut vals) + } } // NitroTarget integration for fast-path Count. From 861e969056699f0458bc12c7fad5ce3ee88458ae Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 4 Aug 2026 12:25:31 -0600 Subject: [PATCH 3/3] Add seeded Nitro sampling constructors for reproducible callers Nitro::init_nitro (and everything built on it: Vector2D::enable_nitro, CountMin/Count::enable_nitro) always seeds its geometric-skip RNG from OS entropy (new_small_rng()), with no way to inject a caller-supplied seed -- appropriate for production use, but it means a caller that needs deterministic, reproducible runs (e.g. a benchmark harness driven by its own --seed flag) can't get repeatable Nitro sampling decisions at all. Added Nitro::init_nitro_seeded(rate, seed), Vector2D::enable_nitro_seeded, and CountMin/Count::enable_nitro_seeded (both Vector2D and Vector2D) as siblings to the existing unseeded versions, seeding via SmallRng::seed_from_u64(seed) instead. --- src/common/structure_utils.rs | 19 ++++++++++++++++++- src/common/structures/vector2d.rs | 6 ++++++ src/sketches/countminsketch.rs | 12 ++++++++++++ src/sketches/countsketch.rs | 12 ++++++++++++ 4 files changed, 48 insertions(+), 1 deletion(-) diff --git a/src/common/structure_utils.rs b/src/common/structure_utils.rs index 2098613..f3248b3 100644 --- a/src/common/structure_utils.rs +++ b/src/common/structure_utils.rs @@ -81,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]" @@ -95,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, diff --git a/src/common/structures/vector2d.rs b/src/common/structures/vector2d.rs index 711302d..f4e7dbd 100644 --- a/src/common/structures/vector2d.rs +++ b/src/common/structures/vector2d.rs @@ -108,6 +108,12 @@ impl Vector2D { 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(); diff --git a/src/sketches/countminsketch.rs b/src/sketches/countminsketch.rs index 61ee12f..48c2ec7 100644 --- a/src/sketches/countminsketch.rs +++ b/src/sketches/countminsketch.rs @@ -489,6 +489,12 @@ impl CountMin, FastPath, H> { 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(); @@ -569,6 +575,12 @@ impl CountMin, FastPath, H> { 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(); diff --git a/src/sketches/countsketch.rs b/src/sketches/countsketch.rs index 3ace1f4..97870e9 100644 --- a/src/sketches/countsketch.rs +++ b/src/sketches/countsketch.rs @@ -479,6 +479,12 @@ impl Count, FastPath, H> { 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(); @@ -551,6 +557,12 @@ impl Count, FastPath, H> { 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();