From 6a7ba6d80001c24838fce774bb0cbfed62b96363 Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 20 Oct 2026 19:06:16 +0000 Subject: [PATCH] feat(heap): estimated_memory_bytes for TopKHeap and the heap sketches TopKHeap::estimated_memory_bytes(heap_size, key) gives the bytes a heap holds once heap_size keys shaped like key are resident: the struct, the HeapEntry and digest arrays, the digest index (std HashMap layout on the build target) and each resident's owned key bytes. CountMinSketchWithHeap and CountSketchWithHeap over Vector2D add their rows x cols counters, so a memory model prices a top-k sketch before building one. Evictions can make the index rehash into twice its buckets, which the count leaves out. --- CHANGELOG.md | 7 ++ docs/api/api_common_heap.md | 4 + docs/api/api_count_min_sketch_with_heap.md | 3 + docs/api/api_count_sketch_with_heap.md | 3 + src/common/heap.rs | 118 +++++++++++++++++ src/sketches/countminsketch_heap.rs | 17 +++ src/sketches/countsketch_heap.rs | 17 +++ tests/heap_memory_bytes.rs | 139 +++++++++++++++++++++ 8 files changed, 308 insertions(+) create mode 100644 tests/heap_memory_bytes.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index b6c88904..f439a67c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,13 @@ signals a backwards-compatible change. ### Added +- **`estimated_memory_bytes` for `TopKHeap`, `CountMinSketchWithHeap` and + `CountSketchWithHeap`**, so a memory model prices a top-k sketch before + building one. `TopKHeap::estimated_memory_bytes(heap_size, key)` counts the + struct, the `HeapEntry` and digest arrays, the digest index and each + resident's owned key bytes once `heap_size` keys shaped like `key` are + resident; the heap sketches add their `rows x cols` counters. Evictions can + make the index rehash into twice its buckets, which the count leaves out. - **`TupleKey` and `KeyColumn`**, a heap key of typed columns (null, bool, `i64`, `f64`, string) encoded into one canonical byte string, so a GROUP BY over several columns keeps one CountMinSketchWithHeap or CountSketchWithHeap candidate per tuple under diff --git a/docs/api/api_common_heap.md b/docs/api/api_common_heap.md index 40717c2c..6a5f2fb9 100644 --- a/docs/api/api_common_heap.md +++ b/docs/api/api_common_heap.md @@ -44,6 +44,10 @@ fn heap(&self) -> &[HeapEntry] fn len(&self) -> usize fn is_empty(&self) -> bool fn capacity(&self) -> usize +// Bytes a heap of `heap_size` holds once full of keys shaped like `key` (a +// string's or byte slice's length counts; a number's does not). Eviction +// tombstones can make the index rehash into twice its buckets, not counted. +fn estimated_memory_bytes(heap_size: usize, key: &DataInput) -> usize ``` ## Merge diff --git a/docs/api/api_count_min_sketch_with_heap.md b/docs/api/api_count_min_sketch_with_heap.md index 79c108be..2e03d025 100644 --- a/docs/api/api_count_min_sketch_with_heap.md +++ b/docs/api/api_count_min_sketch_with_heap.md @@ -20,6 +20,9 @@ and `i128`. Each entry's count is the sketch's estimate, unrounded. fn new(rows: usize, cols: usize, heap_size: usize) -> Self fn from_storage(storage: S, heap_size: usize) -> Self fn default() -> Self +// Vector2D storage: bytes `new(rows, cols, heap_size)` holds once its heap is +// full of keys shaped like `key`, priced before building +fn estimated_memory_bytes(rows: usize, cols: usize, heap_size: usize, key: &DataInput) -> usize ``` ## Insert/Update diff --git a/docs/api/api_count_sketch_with_heap.md b/docs/api/api_count_sketch_with_heap.md index 08e0ad0e..09dd2eee 100644 --- a/docs/api/api_count_sketch_with_heap.md +++ b/docs/api/api_count_sketch_with_heap.md @@ -21,6 +21,9 @@ and `i128`. Each entry's count is the sketch's median estimate, through fn new(rows: usize, cols: usize, heap_size: usize) -> Self fn from_storage(storage: S, heap_size: usize) -> Self fn default() -> Self +// Vector2D storage: bytes `new(rows, cols, heap_size)` holds once its heap is +// full of keys shaped like `key`, priced before building +fn estimated_memory_bytes(rows: usize, cols: usize, heap_size: usize, key: &DataInput) -> usize ``` ## Insert/Update diff --git a/src/common/heap.rs b/src/common/heap.rs index 60ff8036..c2185101 100644 --- a/src/common/heap.rs +++ b/src/common/heap.rs @@ -191,6 +191,30 @@ impl TopKHeap { self.heap_size } + /// Bytes a heap built by [`Self::new`]`(heap_size)` holds once + /// `heap_size` distinct keys shaped like `key` are resident: the struct, + /// the [`HeapEntry`] array, the parallel digest array, the digest index and + /// each resident's owned key bytes (a string's or byte slice's length, 0 + /// for a number). + /// + /// The arrays and the index reserve `min(heap_size, 1024)` up front and + /// double from there, so above 1,024 the arrays hold the next doubling at + /// or above `heap_size`. Evictions leave index tombstones that can make the + /// index rehash into twice its buckets, which this count leaves out. The + /// index size follows std's `HashMap` table layout on the build target. + /// + /// A memory model calls this to price a top-k heap before building one. + pub fn estimated_memory_bytes(heap_size: usize, key: &DataInput) -> usize { + let reserved = heap_size.min(PREALLOCATED_SLOTS); + let entry_slots = vec_slots_after_pushes(reserved, heap_size, size_of::>()); + let digest_slots = vec_slots_after_pushes(reserved, heap_size, size_of::()); + size_of::() + + entry_slots * size_of::>() + + digest_slots * size_of::() + + hash_index_bytes::<(u64, Slot)>(reserved, heap_size) + + heap_size * owned_key_bytes(key) + } + // -- index maintenance --------------------------------------------------- fn lookup(&self, slot: u64, matches: impl Fn(&HeapEntry) -> bool) -> Option { @@ -307,6 +331,100 @@ fn drop_entry(positions: &mut Index, digest: u64, idx: usize) { } } +/// Bytes a key shaped like `key` owns once [`DataInput::to_owned_input`] +/// copies it into a heap entry: a string's or a byte slice's length, 0 for a +/// number. +fn owned_key_bytes(key: &DataInput) -> usize { + match key { + DataInput::Str(s) => s.len(), + DataInput::String(s) => s.len(), + DataInput::Bytes(b) => b.len(), + _ => 0, + } +} + +/// Element slots a `Vec` holds after growing from `reserved` slots to `len` +/// elements one push at a time: std's amortized growth doubles the capacity +/// (at least 4 slots, 8 for 1-byte elements, 1 above 1 KiB). +fn vec_slots_after_pushes(reserved: usize, len: usize, element_bytes: usize) -> usize { + let minimum = match element_bytes { + 1 => 8, + 2..=1024 => 4, + _ => 1, + }; + let mut slots = reserved; + while slots < len { + slots = (slots * 2).max(slots + 1).max(minimum); + } + slots +} + +/// Width of std's `HashMap` control group on this target: 16 with SSE2 on +/// x86, 8 elsewhere. +const HASH_GROUP_WIDTH: usize = if cfg!(all( + any(target_arch = "x86", target_arch = "x86_64"), + target_feature = "sse2" +)) { + 16 +} else { + 8 +}; + +/// Buckets std's `HashMap` allocates for `capacity` entries of `entry_bytes`. +fn hash_buckets_for(capacity: usize, entry_bytes: usize) -> usize { + if capacity == 0 { + return 0; + } + if capacity < 15 { + // Small tables round up to a minimum capacity so the control-group + // padding is not wasted. + let minimum = match (HASH_GROUP_WIDTH, entry_bytes) { + (16, 0..=1) => 14, + (16, 2..=3) => 7, + (8, 0..=1) => 7, + _ => 3, + }; + return match minimum.max(capacity) { + 0..=3 => 4, + 4..=7 => 8, + _ => 16, + }; + } + (capacity * 8 / 7).next_power_of_two() +} + +/// Entries a table of `buckets` holds before it grows (7/8 of its buckets). +fn hash_capacity_of(buckets: usize) -> usize { + if buckets <= 8 { + buckets - 1 + } else { + buckets / 8 * 7 + } +} + +/// Bytes std's `HashMap` allocates for entries of type `T` when created with +/// capacity `reserved` and then filled with `len` distinct keys: the entry +/// slots padded to the control alignment, one control byte per bucket and one +/// trailing control group. A full table grows to the buckets for one more +/// entry than it holds. +fn hash_index_bytes(reserved: usize, len: usize) -> usize { + let entry_bytes = size_of::(); + let mut buckets = hash_buckets_for(reserved, entry_bytes); + while (buckets == 0 && len > 0) || (buckets > 0 && hash_capacity_of(buckets) < len) { + let held = if buckets == 0 { + 0 + } else { + hash_capacity_of(buckets) + }; + buckets = hash_buckets_for(held + 1, entry_bytes); + } + if buckets == 0 { + return 0; + } + let control_align = align_of::().max(HASH_GROUP_WIDTH); + (buckets * entry_bytes).next_multiple_of(control_align) + buckets + HASH_GROUP_WIDTH +} + /// Every element sits in heap order beneath the parent the array layout gives /// it, with the parent index computed here rather than taken from the heap. #[cfg(test)] diff --git a/src/sketches/countminsketch_heap.rs b/src/sketches/countminsketch_heap.rs index 7de2a51d..b5d773b5 100644 --- a/src/sketches/countminsketch_heap.rs +++ b/src/sketches/countminsketch_heap.rs @@ -50,6 +50,23 @@ where heap: TopKHeap::new(heap_size), } } + + /// Bytes a sketch built by [`Self::new`]`(rows, cols, heap_size)` holds + /// once its heap is full of keys shaped like `key`: the struct, the + /// `rows x cols` counters and the heap's allocations + /// ([`TopKHeap::estimated_memory_bytes`], whose caveats apply). A memory + /// model calls this to price a top-k sketch before building one. + pub fn estimated_memory_bytes( + rows: usize, + cols: usize, + heap_size: usize, + key: &DataInput, + ) -> usize { + size_of::() + + rows.saturating_mul(cols).saturating_mul(size_of::()) + + TopKHeap::>>::estimated_memory_bytes(heap_size, key) + - size_of::>>>() + } } // -- Construction from any MatrixStorage ------------------------------------- diff --git a/src/sketches/countsketch_heap.rs b/src/sketches/countsketch_heap.rs index a5574a81..480ab901 100644 --- a/src/sketches/countsketch_heap.rs +++ b/src/sketches/countsketch_heap.rs @@ -111,6 +111,23 @@ where heap: TopKHeap::new(heap_size), } } + + /// Bytes a sketch built by [`Self::new`]`(rows, cols, heap_size)` holds + /// once its heap is full of keys shaped like `key`: the struct, the + /// `rows x cols` counters and the heap's allocations + /// ([`TopKHeap::estimated_memory_bytes`], whose caveats apply). A memory + /// model calls this to price a top-k sketch before building one. + pub fn estimated_memory_bytes( + rows: usize, + cols: usize, + heap_size: usize, + key: &DataInput, + ) -> usize { + size_of::() + + rows.saturating_mul(cols).saturating_mul(size_of::()) + + TopKHeap::>>::estimated_memory_bytes(heap_size, key) + - size_of::>>>() + } } // -- Construction from any MatrixStorage ------------------------------------- diff --git a/tests/heap_memory_bytes.rs b/tests/heap_memory_bytes.rs new file mode 100644 index 00000000..7d333f66 --- /dev/null +++ b/tests/heap_memory_bytes.rs @@ -0,0 +1,139 @@ +//! `TopKHeap::estimated_memory_bytes` and the heap sketches' +//! `estimated_memory_bytes`, checked against `size_of` and against the bytes +//! actually allocated, counted by a global allocator. One `#[test]` so no other +//! test allocates while a count is taken. + +use asap_sketchlib::{ + CountMinSketchWithHeap, CountSketchWithHeap, DataInput, HeapEntry, RegularPath, TopKHeap, + Vector2D, +}; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicIsize, Ordering}; + +/// Counts the bytes currently allocated through it. +struct CountingAllocator; + +static LIVE_BYTES: AtomicIsize = AtomicIsize::new(0); + +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + LIVE_BYTES.fetch_add(layout.size() as isize, Ordering::SeqCst); + unsafe { System.alloc(layout) } + } + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + LIVE_BYTES.fetch_sub(layout.size() as isize, Ordering::SeqCst); + unsafe { System.dealloc(ptr, layout) } + } + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + LIVE_BYTES.fetch_add(new_size as isize - layout.size() as isize, Ordering::SeqCst); + unsafe { System.realloc(ptr, layout, new_size) } + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +fn live_bytes() -> isize { + LIVE_BYTES.load(Ordering::SeqCst) +} + +/// Heap bytes `build` leaves allocated, plus the size of the value it returns. +fn measured_bytes(build: impl FnOnce() -> T) -> (T, usize) { + let before = live_bytes(); + let value = build(); + let heap = (live_bytes() - before) as usize; + (value, heap + std::mem::size_of::()) +} + +/// A 9-byte key, the length of a one-column `TupleKey` holding an `i64`. +fn tuple_like_key(i: u64) -> Vec { + [&[3u8][..], &i.to_be_bytes()].concat() +} + +#[test] +fn estimated_memory_bytes_match_size_of_and_the_allocation() { + // size_of: a resident is one `HeapEntry` in the heap array, one `u64` + // digest and one index entry; a heap of capacity 0 allocates nothing. + assert_eq!( + TopKHeap::::estimated_memory_bytes(0, &DataInput::U64(0)), + std::mem::size_of::>() + ); + assert!( + TopKHeap::::estimated_memory_bytes(1024, &DataInput::U64(0)) + > std::mem::size_of::>() + + 1024 * (std::mem::size_of::>() + 8) + ); + assert_eq!( + CountMinSketchWithHeap::, RegularPath>::estimated_memory_bytes( + 3, + 64, + 0, + &DataInput::U64(0) + ), + std::mem::size_of::, RegularPath>>() + 3 * 64 * 8 + ); + + // TopKHeap filled with `k` distinct keys (no eviction), for capacities on + // both sides of the 1,024-slot up-front reservation. + for k in [1usize, 3, 7, 8, 100, 1024, 2000, 5000] { + let (_heap, bytes) = measured_bytes(|| { + let mut heap = TopKHeap::::new(k); + for i in 0..k as u64 { + heap.update(&DataInput::U64(i), i as f64 + 0.5); + } + heap + }); + assert_eq!( + bytes, + TopKHeap::::estimated_memory_bytes(k, &DataInput::U64(0)), + "u64 keys, k = {k}" + ); + + let (_heap, bytes) = measured_bytes(|| { + let mut heap = TopKHeap::::new(k); + for i in 0..k as u64 { + heap.update(&DataInput::Bytes(&tuple_like_key(i)), i as i64 + 1); + } + heap + }); + assert_eq!( + bytes, + TopKHeap::::estimated_memory_bytes(k, &DataInput::Bytes(&[0; 9])), + "9-byte keys, k = {k}" + ); + } + + // The f64 heap sketches filled with `k` distinct 9-byte keys. + for k in [1usize, 7, 64, 100, 1000] { + let key_shape = DataInput::Bytes(&[0; 9]); + let (_sketch, bytes) = measured_bytes(|| { + let mut sketch = CountMinSketchWithHeap::, RegularPath>::new(3, 64, k); + for i in 0..k as u64 { + sketch.update_weighted(&DataInput::Bytes(&tuple_like_key(i)), 0.5); + } + sketch + }); + assert_eq!( + bytes, + CountMinSketchWithHeap::, RegularPath>::estimated_memory_bytes( + 3, 64, k, &key_shape + ), + "CountMinSketchWithHeap, k = {k}" + ); + + let (_sketch, bytes) = measured_bytes(|| { + let mut sketch = CountSketchWithHeap::, RegularPath>::new(5, 64, k); + for i in 0..k as u64 { + sketch.update_weighted(&DataInput::Bytes(&tuple_like_key(i)), 0.5); + } + sketch + }); + assert_eq!( + bytes, + CountSketchWithHeap::, RegularPath>::estimated_memory_bytes( + 5, 64, k, &key_shape + ), + "CountSketchWithHeap, k = {k}" + ); + } +}