[zzylol]

Why

A planner prices a top-k sketch before it builds one, and today it has to guess the heap's bytes. A TopKHeap resident takes more than its count: a 48-byte HeapEntry<f64>, an 8-byte digest, a digest-index entry and the key's own bytes. None of the heap types say what they hold, so every caller re-derives the layout and drifts from it. ASAPPlanner charges 24 bytes per heap entry.

This redoes #233 (closed with the WeightedFrequency stack) on main's TopKHeap and its f64 heap sketches.

What

• TopKHeap::<C>::estimated_memory_bytes(heap_size, key): the bytes a heap of capacity heap_size holds once heap_size keys shaped like key are resident.
• CountMinSketchWithHeap::<Vector2D<T>, M, H>::estimated_memory_bytes(rows, cols, heap_size, key) and the same on CountSketchWithHeap: the heap's bytes plus the rows x cols counters and the struct.

Behavior is otherwise unchanged.

How

The heap's count adds up these parts, using the types' own size_of:

• the struct;
• the HeapEntry<C> array and the parallel u64 digest array, which reserve min(heap_size, 1024) up front and then grow by std's doubling;
• the digest index: the std HashMap buckets for that reservation, grown as a fresh fill grows them, plus control bytes and the trailing group, using the build target's group width;
• each resident's owned key bytes: a string's or byte slice's length, 0 for a number.

The heap sketches add rows * cols * size_of::<T>() for their counters. Vector2D allocates exactly that.

Left out, and stated in the docs: evictions leave index tombstones that can make the index rehash into twice its buckets. Measured on a 5 x 4096 Count-Min with 9-byte keys after 20 x k distinct updates:

k	Measured / estimate
1, 7, 64, 1,000	exact
100	1.048
5,000	1.358 (the index doubled)

Whether the rehash happens depends on where the tombstones land.

Key code interfaces added / modified

src/common/heap.rs

impl<C: HeapCount> TopKHeap<C> {
    /// Bytes `TopKHeap::new(heap_size)` holds once `heap_size` distinct keys shaped like
    /// `key` are resident: struct, entry array, digest array, digest index, owned key
    /// bytes. A memory model calls it to price a top-k heap before building one.
    pub fn estimated_memory_bytes(heap_size: usize, key: &DataInput) -> usize;   // new
}

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

impl<T, M, H: SketchHasher> CountMinSketchWithHeap<Vector2D<T>, M, H> {
    /// Bytes `new(rows, cols, heap_size)` holds once its heap is full of keys shaped like
    /// `key`: struct, counters and the heap (`TopKHeap::estimated_memory_bytes`).
    pub fn estimated_memory_bytes(rows: usize, cols: usize, heap_size: usize, key: &DataInput) -> usize;   // new
}

Usage, end to end: a planner prices one group's state for topk(100, sum by (region, port) (bytes_rate)) before choosing a plan. The numbers come from running these calls at the head commit.

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

// `key_shape`: a representative group key. ("us-east", 443) encodes as a 21-byte
// ASAPv1 tuple key; only its length matters to the estimate.
let key_bytes = TupleKey::new(vec![KeyColumn::String("us-east".into()), KeyColumn::I64(443)]).to_bytes()?;
let key_shape = DataInput::Bytes(&key_bytes);
// `cms_bytes`: a 5 x 2048 f64 Count-Min with a heap of 100. The planner multiplies it by
// the expected group count and compares that with the memory budget.
let cms_bytes = CountMinSketchWithHeap::<Vector2D<f64>, RegularPath>::estimated_memory_bytes(5, 2048, 100, &key_shape);
assert_eq!(cms_bytes, 95_224); // 240 struct + 81,920 counters + 13,064 heap
// `cs_bytes`: the Count Sketch candidate of the same shape, priced the same way.
let cs_bytes = CountSketchWithHeap::<Vector2D<f64>, RegularPath>::estimated_memory_bytes(5, 2048, 100, &key_shape);
assert_eq!(cs_bytes, 95_224);
// A UnivMon-style layer heap of 256 u64 keys, priced next to that layer's counters.
assert_eq!(TopKHeap::<f64>::estimated_memory_bytes(256, &DataInput::U64(0)), 35_440);

Before this PR

A caller had no way to ask what a heap or a heap sketch holds, and guessed. ASAPPlanner priced 24 bytes per heap entry. The heap above holds about 130 bytes per resident.

After this PR

estimated_memory_bytes returns the bytes the state holds once its heap is full, before the state exists. ASAPPlanner's plan selection can call it instead of hardcoding a per-entry size once its pin includes this PR.

Verification

• tests/heap_memory_bytes.rs (new; it failed to compile before the functions existed) checks the estimates against size_of and against the bytes actually allocated, counted by a global allocator:
  • TopKHeap<f64> filled with k u64 keys and TopKHeap<i64> filled with k 9-byte keys, for k ∈ {1, 3, 7, 8, 100, 1024, 2000, 5000}. These cover both sides of the 1,024-slot reservation.
  • CountMinSketchWithHeap (3 x 64) and CountSketchWithHeap (5 x 64) over f64, filled with k 9-byte keys, for k ∈ {1, 7, 64, 100, 1000}.
  • All are equal.
• 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.
• Limitation: the index size follows std HashMap's table layout on the build target (16-byte control groups with SSE2 on x86, 8 elsewhere). A std change to that layout would show up as a failure of this test.

[zzylol]

Closing as superseded by #242, merged after this PR was opened. main now has TopKHeap::estimated_memory_bytes and estimated_memory_bytes on both heap sketches as a steady-state upper bound that includes the index rehash under churn. It also has live memory_bytes(&self) on TopKHeap and both heap sketches.

Checked on main (eb04e7f) against a counting allocator after 20 x k distinct updates, for k in {1, 3, 7, 64, 100, 1000, 3000, 5000}, with u64 keys and 21-byte tuple keys, plus a 5 x 2048 f64 Count-Min with a heap: measured bytes never exceeded the estimate. They matched it exactly where the index had rehashed, and were 0.63 to 0.73 of it where it had not.

Not carried over from the review of this PR, and left as possible follow-ups on main:

• renaming the shape-based estimate (for example to memory_bytes_when_full);
• a doc note that the std HashMap/RawVec sizing it mirrors may drift on a toolchain bump.

The branch is kept.