feat(heap): estimated_memory_bytes for TopKHeap and the heap sketches

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

docs/api/api_common_heap.md

@@ -44,6 +44,10 @@ fn heap(&self) -> &[HeapEntry<C>]
 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

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

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

src/common/heap.rs

@@ -191,6 +191,30 @@ impl<C: HeapCount> TopKHeap<C> {
         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::<HeapEntry<C>>());
+        let digest_slots = vec_slots_after_pushes(reserved, heap_size, size_of::<u64>());
+        size_of::<Self>()
+            + entry_slots * size_of::<HeapEntry<C>>()
+            + digest_slots * size_of::<u64>()
+            + 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<C>) -> bool) -> Option<usize> {
@@ -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<T>(reserved: usize, len: usize) -> usize {
+    let entry_bytes = size_of::<T>();
+    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::<T>().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)]

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::<Self>()
+            + rows.saturating_mul(cols).saturating_mul(size_of::<T>())
+            + TopKHeap::<HeapCountOf<Vector2D<T>>>::estimated_memory_bytes(heap_size, key)
+            - size_of::<TopKHeap<HeapCountOf<Vector2D<T>>>>()
+    }
 }
 
 // -- Construction from any MatrixStorage -------------------------------------

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::<Self>()
+            + rows.saturating_mul(cols).saturating_mul(size_of::<T>())
+            + TopKHeap::<HeapCountOf<Vector2D<T>>>::estimated_memory_bytes(heap_size, key)
+            - size_of::<TopKHeap<HeapCountOf<Vector2D<T>>>>()
+    }
 }
 
 // -- Construction from any MatrixStorage -------------------------------------

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<T>(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::<T>())
+}
+
+/// A 9-byte key, the length of a one-column `TupleKey` holding an `i64`.
+fn tuple_like_key(i: u64) -> Vec<u8> {
+    [&[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::<i64>::estimated_memory_bytes(0, &DataInput::U64(0)),
+        std::mem::size_of::<TopKHeap<i64>>()
+    );
+    assert!(
+        TopKHeap::<f64>::estimated_memory_bytes(1024, &DataInput::U64(0))
+            > std::mem::size_of::<TopKHeap<f64>>()
+                + 1024 * (std::mem::size_of::<HeapEntry<f64>>() + 8)
+    );
+    assert_eq!(
+        CountMinSketchWithHeap::<Vector2D<f64>, RegularPath>::estimated_memory_bytes(
+            3,
+            64,
+            0,
+            &DataInput::U64(0)
+        ),
+        std::mem::size_of::<CountMinSketchWithHeap<Vector2D<f64>, 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::<f64>::new(k);
+            for i in 0..k as u64 {
+                heap.update(&DataInput::U64(i), i as f64 + 0.5);
+            }
+            heap
+        });
+        assert_eq!(
+            bytes,
+            TopKHeap::<f64>::estimated_memory_bytes(k, &DataInput::U64(0)),
+            "u64 keys, k = {k}"
+        );
+
+        let (_heap, bytes) = measured_bytes(|| {
+            let mut heap = TopKHeap::<i64>::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::<i64>::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::<Vector2D<f64>, 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::<Vector2D<f64>, RegularPath>::estimated_memory_bytes(
+                3, 64, k, &key_shape
+            ),
+            "CountMinSketchWithHeap, k = {k}"
+        );
+
+        let (_sketch, bytes) = measured_bytes(|| {
+            let mut sketch = CountSketchWithHeap::<Vector2D<f64>, 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::<Vector2D<f64>, RegularPath>::estimated_memory_bytes(
+                5, 64, k, &key_shape
+            ),
+            "CountSketchWithHeap, k = {k}"
+        );
+    }
+}