Skip to content

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

Open
zzylol wants to merge 2 commits into
mainfrom
feat/f64-heap-sketch-checks
Open

zzylol wants to merge 2 commits into
mainfrom
feat/f64-heap-sketch-checks

Conversation

@zzylol

@zzylol zzylol commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

Why

ASAPPlanner's executor runs Float64-weighted top-k on main's CountMinSketchWithHeap<Vector2D<f64>> and CountSketchWithHeap<Vector2D<f64>>, keyed by TupleKey bytes. main already has check_shape, clear and the memory estimates (#242). Two gaps remain:

  • Weight checks. update_weighted over f64 counters accepts NaN and ±inf, and Count-Min also accepts negative weights. Any of these silently corrupts every later estimate. Count-Min's estimate is an upper bound only while every update adds. A counter can also overflow to inf partway through an update, after some rows have already changed.
  • Current heap counts. An update rescores only the key it touches. Another resident that shares a cell keeps its earlier count, so a top-k read can report stale scores.

This is what remains of the closed WeightedFrequency stack (#160, #225, #226, #227), redone on main's sketches.

What

  • On CountMinSketch<Vector2D<f64>, RegularPath> and CountSketch<Vector2D<f64>, RegularPath>:
    • check_weight(weight). Count-Min accepts finite, nonnegative weights. Count Sketch accepts any finite weight, including negative ones.
    • try_update_weighted(key, weight). A weight check_weight rejects, or an update that would carry a probed counter past f64::MAX, returns Err and leaves the sketch unchanged.
  • On CountMinSketchWithHeap and CountSketchWithHeap:
    • check_weight (over Vector2D<f64>, RegularPath): forwards to the inner sketch's, the way check_shape is forwarded.
    • try_update_weighted (over Vector2D<f64>, RegularPath): the checked update, then the heap.
    • rescore_heap() (any RegularPath storage): re-reads every resident's count from the counters and keeps the same residents. It is opt-in: call it before reading the heap.

Existing methods are unchanged.

Not changed: main's check_shape accepts more than MATRIX_MAX_ROWS (20) rows. Past that, a row reuses another row's seed, and to_asapv1 refuses the sketch. A caller that needs independent rows and an encodable state must also check rows <= MATRIX_MAX_ROWS. ASAPPlanner's executor will.

How

  • try_update_weighted computes each row's column once, plus the signed delta for Count Sketch, in a SmallVec sized MATRIX_MAX_ROWS. It checks that every probed counter stays finite, then applies the update. So it hashes once per row, as update_weighted does.
  • rescore_heap estimates each resident and writes the count back through TopKHeap::update_owned. That re-sifts the resident and never evicts it. It costs one estimate per resident.
  • Possible follow-up: rescore_heap clones each resident's key and looks it up in the index. A TopKHeap method that rewrites each slot's count by position and re-heapifies once in O(k), rebuilding the index, would avoid that.
  • WeightedFrequency rescored every candidate on each update, which cost O(k) estimates per update. Here the caller pays that once per read. One difference remains: admission still compares a new key against the root's possibly stale count, as main's heap always has.

Key code interfaces added / modified

src/sketches/countminsketch.rs

impl<H: SketchHasher> CountMinSketch<Vector2D<f64>, RegularPath, H> {
    /// Ok for a finite, nonnegative weight; lets a caller reject an input row up front.
    pub fn check_weight(weight: f64) -> Result<(), String>;   // new
    /// `update_weighted` that returns Err, sketch unchanged, for a weight `check_weight`
    /// rejects or a counter overflow (bad input).
    pub fn try_update_weighted(&mut self, value: &DataInput, weight: f64) -> Result<(), String>;   // new
}

src/sketches/countsketch.rs

impl<H: SketchHasher> CountSketch<Vector2D<f64>, RegularPath, H> {
    /// Ok for a finite weight; a negative weight is a signed update.
    pub fn check_weight(weight: f64) -> Result<(), String>;   // new
    /// `update_weighted` that returns Err, sketch unchanged, for a non-finite weight or a
    /// counter overflow in either direction (bad input).
    pub fn try_update_weighted(&mut self, value: &DataInput, weight: f64) -> Result<(), String>;   // new
}

src/sketches/countminsketch_heap.rs, src/sketches/countsketch_heap.rs (the same on CountSketchWithHeap)

impl<H: SketchHasher> CountMinSketchWithHeap<Vector2D<f64>, RegularPath, H> {
    /// The inner sketch's `check_weight`, so a caller checks a weight on the type it updates.
    pub fn check_weight(weight: f64) -> Result<(), String>;   // new
    /// The inner sketch's `try_update_weighted`, then the heap; on Err neither changes.
    pub fn try_update_weighted(&mut self, key: &DataInput, weight: f64) -> Result<(), String>;   // new
}

impl<S: MatrixStorage, H: SketchHasher> CountMinSketchWithHeap<S, RegularPath, H> {
    /// Re-reads every resident's count from the counters, keeping the same residents;
    /// the executor calls it before reading top-k so every score is a current estimate.
    pub fn rescore_heap(&mut self);   // new
}

Usage, end to end: an executor runs one window of topk(2, sum by (region, port) (bytes_rate)) on a Count-Min plan. This was run as a test at the head commit; the output is shown in the comments.

use asap_sketchlib::{
    CountMinSketchWithHeap, DataInput, KeyColumn, MATRIX_MAX_ROWS, RegularPath, TupleKey, Vector2D,
};

type Cms = CountMinSketchWithHeap<Vector2D<f64>, RegularPath>;
// `rows`, `cols`: the plan's shape. `main`'s check_shape refuses a shape that cannot build;
// the executor also caps rows so every row has its own seed and the state encodes.
let (rows, cols) = (5, 2048);
CountMinSketchWithHeap::check_shape(rows, cols)?;
assert!(rows <= MATRIX_MAX_ROWS);
// `state`: one window's sketch with a heap of 2, reused across windows by `clear`.
let mut state = Cms::new(rows, cols, 2);
// `input`: the window's (region, port, rate) rows; rates are fractional weights.
let input = [
    ("us-east", 443, 0.125),
    ("us-east", 443, 0.25),
    ("eu-west", 80, 1.5),
    ("ap-south", 22, 0.0625),
    ("eu-west", 80, f64::NAN),
];
for (region, port, rate) in input {
    // `key`: the ASAPv1 tuple key of the group-by columns.
    let key = TupleKey::new(vec![KeyColumn::String(region.into()), KeyColumn::I64(port)]).to_bytes()?;
    // Ok(()) for the first four rows; the NaN row returns
    // Err("CMS weight NaN must be finite and nonnegative") and leaves `state` unchanged.
    let outcome = state.try_update_weighted(&DataInput::Bytes(&key), rate);
}
state.rescore_heap(); // every resident's count is now its current estimate
// The heap holds [String("us-east"), I64(443)] at 0.375 and
// [String("eu-west"), I64(80)] at 1.5, decoded with `TupleKey::from_bytes`.
// `bytes`: the window's ASAPv1 encoding (counter_type "f64"), 92,620 bytes, sent to the merger.
let bytes = state.to_asapv1()?;
state.clear(); // main's clear: the next window reuses the allocation

Before this PR

A NaN or negative rate in one input row silently corrupted the window's Count-Min estimates. An overflowing update could leave some rows changed and others not. A top-k read could report a resident's count from before later updates moved its cells.

After this PR

The executor checks each weight on update and gets bad input back as an Err with the state intact. It calls rescore_heap before reading top-k when it needs current scores. The state is main's f64 heap sketch, serialized as ASAPv1.

Verification

  • New unit tests (they failed to compile before the methods existed):
    • f64_checked_updates_are_all_or_nothing (both heap sketches): checks fractional and signed sums and the heap count. Bad weights and an f64::MAX overflow return Err, with counters and heap identical before and after. The old unchecked update_weighted fails these checks: a NaN weight changes the probed cells, and a second f64::MAX drives the estimate to inf (to -inf for Count Sketch). This was checked with a throwaway test at the head commit.
    • rescore_heap_rereads_every_resident (both heap sketches): 60 colliding updates in 8 columns leave some residents stale, and after rescore_heap every resident's count equals its estimate.
  • cargo fmt --all -- --check: clean.
  • cargo clippy --workspace --all-targets --locked -- -D warnings, with and without --all-features: clean.
  • cargo test --locked and cargo test --all-features --locked: all pass.

Rebase note: rebased onto main eb04e7f. #242 already added check_shape, clear and the memory estimates, so this PR's versions of those were dropped in favor of main's. Only check_weight, try_update_weighted and rescore_heap remain.

@zzylol zzylol added the enhancement New feature or request label Oct 20, 2026
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 try_update_weighted,
which updates the heap after the checked update, and rescore_heap, which
re-reads every resident's count from the counters.
@zzylol zzylol changed the title feat(cms,countsketch): checked f64 updates, shape checks, clear and heap rescoring feat(cms,countsketch): checked f64 updates and heap rescoring Oct 20, 2026
@zzylol
zzylol force-pushed the feat/f64-heap-sketch-checks branch from 9085765 to c9c7e7a Compare October 20, 2026 23:41
CountMinSketchWithHeap and CountSketchWithHeap over Vector2D<f64> forward
check_weight to the inner sketch, the way check_shape is forwarded, so a
caller checks a weight against the type it updates.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant