From fe3d56ef26c51f4605b5c0787e1ed2e2d2c9c965 Mon Sep 17 00:00:00 2001 From: zzylol Date: Sun, 18 Oct 2026 15:22:49 +0000 Subject: [PATCH] feat(heap): estimated_memory_bytes for HHHeap and WeightedFrequency A memory model can price a top-k heap before building one: the struct, the heap and digest arrays, the digest index and the owned key bytes for HHHeap; the cells, candidate slots, identity tuples and encoded keys for WeightedFrequency. Checked against size_of and a counting allocator. A retained candidate's identity tuple is now allocated at its exact length; it held four slots for a one-value tuple. --- CHANGELOG.md | 6 ++ src/common/heap.rs | 117 ++++++++++++++++++++++++++ src/sketches/weighted_frequency.rs | 62 +++++++++++--- tests/heap_memory_bytes.rs | 128 +++++++++++++++++++++++++++++ 4 files changed, 303 insertions(+), 10 deletions(-) create mode 100644 tests/heap_memory_bytes.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index a8de475e..194b8e14 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,12 @@ signals a backwards-compatible change. into the rescoring instead of copying them, and scores them with one buffer, about halving its cost at capacity 100; the retained candidates and their scores are unchanged. + `WeightedFrequency::estimated_memory_bytes` and + `HHHeap::estimated_memory_bytes` give the bytes a state of a given shape + holds once its top-k heap is full of keys of a given shape, so a memory + model can price one before building it. A retained candidate's identity + tuple is allocated at its exact length (it held four slots for a + one-value tuple, 72 bytes per candidate more than needed). - **`Vector3D`**, filled in from a stub. A `rows x cols` grid whose every cell is a contiguous `depth`-length record, addressed by the same diff --git a/src/common/heap.rs b/src/common/heap.rs index 74ad00cf..97318ddd 100644 --- a/src/common/heap.rs +++ b/src/common/heap.rs @@ -167,6 +167,30 @@ impl HHHeap { self.k } + /// Bytes a heap built by [`Self::new`]`(k)` holds once `k` distinct keys + /// shaped like `key` are resident: the struct, the heap array of + /// [`HHItem`]s, the parallel digest array, the digest index and each + /// resident's owned key bytes (a `String` or `Bytes` key's length; 0 for a + /// number). + /// + /// The arrays and the index reserve `min(k, 1024)` up front and double + /// from there, so for `k` above 1,024 the arrays hold the next doubling + /// at or above `k`. Evictions leave 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(k: usize, key: &DataInput) -> usize { + let reserved = k.min(PREALLOCATED_SLOTS); + let array_slots = vec_slots_after_pushes(reserved, k, size_of::()); + let digest_slots = vec_slots_after_pushes(reserved, k, size_of::()); + size_of::() + + array_slots * size_of::() + + digest_slots * size_of::() + + hash_index_bytes::<(u64, Slot)>(reserved, k) + + k * owned_key_bytes(key) + } + // -- index maintenance --------------------------------------------------- fn lookup(&self, slot: u64, matches: impl Fn(&HHItem) -> bool) -> Option { @@ -301,6 +325,99 @@ fn drop_entry(positions: &mut Index, digest: u64, idx: usize) { } } +/// Bytes a key shaped like `key` owns once [`input_to_owned`] copies it into +/// a heap entry: a string's or a byte slice's length, 0 for a number. +pub(crate) 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). +pub(crate) 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. +pub(crate) 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/weighted_frequency.rs b/src/sketches/weighted_frequency.rs index 51ea13ff..97990946 100644 --- a/src/sketches/weighted_frequency.rs +++ b/src/sketches/weighted_frequency.rs @@ -55,17 +55,21 @@ impl Ord for Candidate { } /// Canonical identity tuple and its hash key: floating identities must be /// finite, and signed zeros share one identity. +/// The identity is allocated at its exact length, so a retained candidate +/// holds no spare slots. fn normalize(values: &[FrequencyIdentity]) -> Result<(Vec, Vec), Error> { - let identity = values - .iter() - .map(|value| match value { - FrequencyIdentity::Float64(v) if !v.is_finite() => Err(Error::Invalid( - "non-finite weighted frequency identity".into(), - )), - FrequencyIdentity::Float64(v) if *v == 0.0 => Ok(FrequencyIdentity::Float64(0.0)), - value => Ok(value.clone()), - }) - .collect::, _>>()?; + let mut identity = Vec::with_capacity(values.len()); + for value in values { + identity.push(match value { + FrequencyIdentity::Float64(v) if !v.is_finite() => { + return Err(Error::Invalid( + "non-finite weighted frequency identity".into(), + )); + } + FrequencyIdentity::Float64(v) if *v == 0.0 => FrequencyIdentity::Float64(0.0), + value => value.clone(), + }); + } let key = bincode::serialize(&identity).map_err(|e| Error::Update(e.to_string()))?; Ok((identity, key)) } @@ -127,6 +131,44 @@ impl WeightedFrequency { .map(|_| ()) .ok_or_else(|| Error::Invalid("invalid weighted frequency dimensions".into())) } + /// Bytes a state of shape `width × depth` holds once it retains + /// `capacity` candidates whose identity tuples are shaped like `identity`: + /// the struct, the `width × depth` 8-byte cells, the candidate heap's + /// slots, and per retained candidate its identity tuple (with a `Utf8` + /// value's bytes) and its encoded hash key. Only the identity's shape + /// matters (its variants and string lengths), not its values. + /// + /// The candidate heap is rebuilt on every update by pushing up to + /// `capacity + 1` candidates, so it holds `max(4, (capacity + 1) + /// rounded up to a power of two)` slots, and during an update the + /// previous heap's slots are held as well. A memory model calls this to + /// price a top-k frequency state before building one. + pub fn estimated_memory_bytes( + width: usize, + depth: usize, + capacity: usize, + identity: &[FrequencyIdentity], + ) -> usize { + let cells = width.saturating_mul(depth).saturating_mul(size_of::()); + if capacity == 0 { + return size_of::() + cells; + } + let heap_slots = + crate::common::heap::vec_slots_after_pushes(0, capacity + 1, size_of::()); + let identity_bytes = size_of_val(identity) + + identity + .iter() + .map(|value| match value { + FrequencyIdentity::Utf8(text) => text.len(), + _ => 0, + }) + .sum::(); + let key_bytes = bincode::serialized_size(identity).map_or(0, |bytes| bytes as usize); + size_of::() + + cells + + heap_slots * size_of::() + + capacity * (identity_bytes + key_bytes) + } pub fn new( algorithm: FrequencyAlgorithm, width: usize, diff --git a/tests/heap_memory_bytes.rs b/tests/heap_memory_bytes.rs new file mode 100644 index 00000000..be482b15 --- /dev/null +++ b/tests/heap_memory_bytes.rs @@ -0,0 +1,128 @@ +//! `HHHeap::estimated_memory_bytes` and `WeightedFrequency::estimated_memory_bytes` +//! checked against `size_of` and against the bytes the heap actually holds, +//! counted by a global allocator. One `#[test]` so no other test allocates +//! while a count is taken. + +use asap_sketchlib::{ + DataInput, FrequencyAlgorithm, FrequencyIdentity, HHHeap, HHItem, WeightedFrequency, +}; +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::()) +} + +fn int_identity(i: usize) -> Vec { + vec![FrequencyIdentity::Int64(i as i64)] +} + +fn text_identity(i: usize) -> Vec { + vec![ + FrequencyIdentity::Utf8(format!("{i:016}")), + FrequencyIdentity::Bool(i % 2 == 0), + ] +} + +#[test] +fn estimated_memory_bytes_match_size_of_and_the_allocation() { + // size_of: a resident is one `HHItem` in the heap array, one `u64` digest + // and one index entry; a heap of capacity 0 allocates nothing. + assert_eq!( + HHHeap::estimated_memory_bytes(0, &DataInput::U64(0)), + std::mem::size_of::() + ); + assert!( + HHHeap::estimated_memory_bytes(1024, &DataInput::U64(0)) + > std::mem::size_of::() + 1024 * (std::mem::size_of::() + 8) + ); + assert_eq!( + WeightedFrequency::estimated_memory_bytes(64, 3, 0, &int_identity(0)), + std::mem::size_of::() + 64 * 3 * 8 + ); + + // HHHeap 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 = HHHeap::new(k); + for i in 0..k as u64 { + heap.update(&DataInput::U64(i), i as i64 + 1); + } + heap + }); + assert_eq!( + bytes, + HHHeap::estimated_memory_bytes(k, &DataInput::U64(0)), + "u64 keys, k = {k}" + ); + + let (_heap, bytes) = measured_bytes(|| { + let mut heap = HHHeap::new(k); + for i in 0..k as u64 { + let key = [&[3u8][..], &i.to_le_bytes()].concat(); + heap.update(&DataInput::Bytes(&key), i as i64 + 1); + } + heap + }); + assert_eq!( + bytes, + HHHeap::estimated_memory_bytes(k, &DataInput::Bytes(&[0; 9])), + "9-byte keys, k = {k}" + ); + } + + // WeightedFrequency after four times `capacity` distinct identities, so the + // candidate heap is full and has turned keys away. + for capacity in [1usize, 7, 64, 100, 1000] { + for (name, identity) in [ + ("Int64", int_identity as fn(usize) -> Vec), + ("Utf8 + Bool", text_identity), + ] { + let (_state, bytes) = measured_bytes(|| { + let mut state = + WeightedFrequency::new(FrequencyAlgorithm::Cms, 64, 3, capacity).unwrap(); + for i in 0..4 * capacity { + state.update(&identity(i), 1.0).unwrap(); + } + state + }); + assert_eq!( + bytes, + WeightedFrequency::estimated_memory_bytes(64, 3, capacity, &identity(0)), + "{name} identities, capacity = {capacity}" + ); + } + } +}