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}" + ); + } +}