From 3b6cac9b54702ca5e8da2c1b99776b1a5ee4b655 Mon Sep 17 00:00:00 2001 From: zz_y Date: Thu, 15 Oct 2026 20:36:38 +0000 Subject: [PATCH] fix(kmv): read and merge by shared reference, report memory, skip the duplicate scan for hashes past the k-th estimate took &mut self and merge took the other sketch as &mut KMV without mutating either, so a caller holding a shared state had to clone it (k hashes) to read an estimate or to merge it in. Both now borrow shared; &mut call sites still compile. approx_memory_bytes reports the struct plus the hash slots the heap has allocated, so callers stop guessing the footprint. (The heap reserves up to 1024 slots and grows toward k, so neither the retained count nor k alone is the allocation.) insert_by_hash scanned all k retained hashes for a duplicate before offering a hash to the heap, so every steady-state insert cost O(k). A full sketch now turns away a hash at or above its largest retained one first; the heap rejected those anyway, so the retained set and every estimate are unchanged. --- CHANGELOG.md | 8 +++++ docs/api/api_kmv.md | 9 +++-- docs/tests.md | 2 ++ src/common/structures/heap.rs | 7 ++++ src/sketches/kmv.rs | 63 +++++++++++++++++++++++++++++++++-- src/sketches/kmv/wire.rs | 6 ++-- tests/e2e/experimental.rs | 9 +++-- tests/e2e/wire.rs | 2 +- tests/pbt/kmv.rs | 28 ++++++++-------- 9 files changed, 106 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 32b069ed..e635df99 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -92,6 +92,14 @@ signals a backwards-compatible change. ### Fixed +- **`KMV` reads, merges and reports memory without a mutable borrow or a + copy.** `estimate` takes `&self` and `merge` takes the other sketch as + `&KMV`, so a caller holding a shared state no longer clones it to read an + estimate or to merge it in; existing `&mut` call sites still compile. + `approx_memory_bytes` reports the struct plus the hash slots the heap has + allocated. A full sketch turns away a hash at or above its largest + retained one before scanning the `k` retained hashes for a duplicate, so + steady-state inserts no longer cost `O(k)`; the retained set is unchanged. - `KLLDynamic` panicked on every quantile query once a NaN had entered the stream: the CDF was sorted with `partial_cmp(..).unwrap()`. It sorts with `total_cmp`, as the three sibling sorts beside it already did. diff --git a/docs/api/api_kmv.md b/docs/api/api_kmv.md index 73d2038b..89ee2f65 100644 --- a/docs/api/api_kmv.md +++ b/docs/api/api_kmv.md @@ -29,13 +29,18 @@ fn insert_by_hash(&mut self, hash_value: u64) ## Query ```rust -fn estimate(&mut self) -> f64 +fn estimate(&self) -> f64 +fn approx_memory_bytes(&self) -> usize ``` +`approx_memory_bytes` is the struct plus the hash slots the heap has +allocated: up to 1024 at construction, growing toward `k` as distinct hashes +arrive. + ## Merge ```rust -fn merge(&mut self, other: &mut KMV) +fn merge(&mut self, other: &KMV) ``` ## Serialization diff --git a/docs/tests.md b/docs/tests.md index ccb1571c..7a08535d 100644 --- a/docs/tests.md +++ b/docs/tests.md @@ -707,6 +707,8 @@ Wire tests: [`src/sketches/kmv/wire.rs`](../src/sketches/kmv/wire.rs) | test_name | test_description | what_is_tested | | --- | --- | --- | +| `estimate_and_merge_borrow_inputs_shared` | Reading an estimate and merging another sketch need no mutable binding or copy. | Merges an immutable `k = 64` sketch over `400..1000` into a copy of one over `0..600`, verifies the merged estimate equals the build over `0..1000`, and that the left input's estimate is unchanged. | +| `memory_report_covers_the_retained_hashes` | The memory report covers the retained hashes and does not move on a rejected insert. | For `k = 3000`, verifies the empty and the filled (20,000 keys) sketches each report at least the struct plus `8` bytes per retained hash, that the filled one reports more, and that a duplicate key and a hash of `u64::MAX` leave the report and the estimate unchanged. | | `assert_serialization_round_trip` | Assert serialization round trip. | Serialization/deserialization preserves component state and behavior after round trip. | | `kmv_round_trip_serialization` | The envelope frames a sketch under kind_id `0x0e 0x00`, and the bytes are stable across a round trip. | For `k = 64` filled by 5,000 distinct keys, verifies the bytes open with the ASAPv1 magic, a `kind_id_len` of `2`, and `0x0e 0x00`, that the metadata carries `k` `64` and the canonical seed index, that the 64 emitted hashes are strictly ascending, and that the decode matches on `k`, capacity, retained hashes, heap root, and `estimate`, re-serializing identically. | | `kmv_emitted_order_is_independent_of_insertion_order` | The emitted order follows the hash value, not the arrival order. | Fills two `k = 32` sketches with the same 2,000 keys forward and backward, verifies their heap arrays differ, and that both serialize to the same bytes, which a decode re-encodes exactly. | diff --git a/src/common/structures/heap.rs b/src/common/structures/heap.rs index 78c202b3..51b8e513 100644 --- a/src/common/structures/heap.rs +++ b/src/common/structures/heap.rs @@ -93,6 +93,13 @@ impl> CommonHeap { self.size } + /// Element slots the heap has allocated: at least [`Self::len`], and + /// grown on demand toward [`Self::capacity`]. For memory accounting. + #[inline] + pub fn allocated_slots(&self) -> usize { + self.data.capacity() + } + /// Returns true if the heap is at full capacity. #[inline] pub fn is_full(&self) -> bool { diff --git a/src/sketches/kmv.rs b/src/sketches/kmv.rs index af17a8c3..c33dbee2 100644 --- a/src/sketches/kmv.rs +++ b/src/sketches/kmv.rs @@ -46,6 +46,11 @@ impl KMV { } pub fn insert_by_hash(&mut self, hash_value: u64) { + // A full sketch keeps only hashes below its largest retained one, so + // most inserts are turned away here without scanning the k hashes. + if self.k_vals.is_full() && self.k_vals.peek().is_some_and(|max| hash_value >= *max) { + return; + } if self.k_vals.iter().any(|value| *value == hash_value) { return; } @@ -54,7 +59,7 @@ impl KMV { /// `(k - 1) / U(k)`, where `U(k)` is the k-th smallest hash's position in /// the 64-bit range, mapped into (0, 1]. - pub fn estimate(&mut self) -> f64 { + pub fn estimate(&self) -> f64 { if self.k_vals.len() < self.k { return self.k_vals.len() as f64; } @@ -67,7 +72,7 @@ impl KMV { (self.k - 1) as f64 / mapped } - pub fn merge(&mut self, other: &mut KMV) { + pub fn merge(&mut self, other: &KMV) { assert_eq!( self.k, other.k, "Two KMV sketch have different k size, not mergeable" @@ -76,6 +81,13 @@ impl KMV { self.insert_by_hash(value); } } + + /// Approximate in-memory footprint in bytes: the struct plus the hash + /// slots its heap has allocated. The heap reserves up to 1024 slots at + /// construction and grows toward `k` as distinct hashes arrive. + pub fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::().saturating_add(self.k_vals.allocated_slots().saturating_mul(8)) + } } #[cfg(test)] @@ -87,6 +99,51 @@ mod tests { const ERROR_TOLERANCE: f64 = 0.02; const SERDE_SAMPLE: usize = 100_000; + fn build(k: usize, items: std::ops::Range) -> KMV { + let mut sketch = KMV::new(k); + for item in items { + sketch.insert(&DataInput::U64(item)); + } + sketch + } + + /// Reading an estimate and merging in another sketch borrow their inputs + /// shared, so neither needs a mutable binding or a copy, and the merged + /// estimate equals the union build's. + #[test] + fn estimate_and_merge_borrow_inputs_shared() { + let left = build(64, 0..600); + let right = build(64, 400..1_000); + let union = build(64, 0..1_000); + let mut merged = left.clone(); + merged.merge(&right); + assert_eq!(merged.estimate(), union.estimate()); + assert_eq!(left.estimate(), build(64, 0..600).estimate()); + } + + /// The memory report covers the hashes the heap holds, grows as it fills, + /// and is the same after a duplicate or a hash past the retained k. + #[test] + fn memory_report_covers_the_retained_hashes() { + let k = 3_000; + let empty = KMV::new(k); + let full = build(k, 0..20_000); + assert_eq!(full.k_vals.len(), k); + for sketch in [&empty, &full] { + assert!( + sketch.approx_memory_bytes() + >= std::mem::size_of::() + sketch.k_vals.len() * 8 + ); + } + assert!(full.approx_memory_bytes() > empty.approx_memory_bytes()); + let mut again = full.clone(); + let before = again.approx_memory_bytes(); + again.insert(&DataInput::U64(7)); + again.insert_by_hash(u64::MAX); + assert_eq!(again.approx_memory_bytes(), before); + assert_eq!(again.estimate(), full.estimate()); + } + #[test] fn assert_serialization_round_trip() { let mut sketch: KMV = KMV::default(); @@ -103,7 +160,7 @@ mod tests { "KMV serialization output should not be empty" ); - let mut decoded: KMV = KMV::deserialize_from_bytes(&encoded) + let decoded: KMV = KMV::deserialize_from_bytes(&encoded) .unwrap_or_else(|err| panic!("KMV deserialize_from_bytes failed: {err}")); let reencoded = decoded diff --git a/src/sketches/kmv/wire.rs b/src/sketches/kmv/wire.rs index 30d6f969..90ef11f5 100644 --- a/src/sketches/kmv/wire.rs +++ b/src/sketches/kmv/wire.rs @@ -242,7 +242,7 @@ mod tests { #[test] fn kmv_round_trip_serialization() { - let mut sketch = populated(64, 5_000); + let sketch = populated(64, 5_000); assert_eq!(sketch.k_vals.len(), 64, "the bound was never reached"); let encoded = sketch.serialize_to_bytes().expect("serialize KMV"); @@ -260,7 +260,7 @@ mod tests { assert!(pair[0] < pair[1], "the emitted hashes are not ascending"); } - let mut decoded = KMV::::deserialize_from_bytes(&encoded).expect("decode"); + let decoded = KMV::::deserialize_from_bytes(&encoded).expect("decode"); assert_eq!(decoded.k, sketch.k); assert_eq!(decoded.k_vals.capacity(), sketch.k_vals.capacity()); assert_eq!(decoded.wire_hashes(), sketch.wire_hashes()); @@ -317,7 +317,7 @@ mod tests { assert_eq!(metadata_of(&bytes).k, 16); assert!(payload_of(&bytes).hashes.is_empty()); - let mut decoded = KMV::::deserialize_from_bytes(&bytes).expect("decode"); + let decoded = KMV::::deserialize_from_bytes(&bytes).expect("decode"); assert_eq!(decoded.k, 16); assert_eq!(decoded.k_vals.len(), 0); assert_eq!(decoded.k_vals.capacity(), 16); diff --git a/tests/e2e/experimental.rs b/tests/e2e/experimental.rs index 512ecd89..18aa78a0 100644 --- a/tests/e2e/experimental.rs +++ b/tests/e2e/experimental.rs @@ -122,7 +122,7 @@ fn kmv_estimates_stay_inside_their_relative_standard_error_band_over_independent for regime in KMV_REGIMES { let n = regime.n(k); namespace = namespace.wrapping_add(NAMESPACE_STRIDE); - let mut sketch = kmv_trial(k, n, seed_idx, namespace); + let sketch = kmv_trial(k, n, seed_idx, namespace); let estimate = sketch.estimate(); let outcome = spec.check(estimate, n); tally.record(outcome.is_ok(), || { @@ -175,7 +175,7 @@ fn kmv_is_exact_below_k_and_estimates_at_k() { ); for n in [1usize, 2, K / 2, K - 2, K - 1] { - let mut sketch: KMV = kmv_trial(K, n, 0, NAMESPACE_STRIDE * 900 + n as u64 * 4096); + let sketch: KMV = kmv_trial(K, n, 0, NAMESPACE_STRIDE * 900 + n as u64 * 4096); assert_eq!( sketch.estimate(), n as f64, @@ -188,7 +188,7 @@ fn kmv_is_exact_below_k_and_estimates_at_k() { // to k, so the exact-regime branch must not claim it. let mut estimates = Vec::new(); for (i, &seed_idx) in KMV_HASH_SEEDS.iter().enumerate() { - let mut sketch = kmv_trial(K, K, seed_idx, NAMESPACE_STRIDE * (1_000 + i as u64)); + let sketch = kmv_trial(K, K, seed_idx, NAMESPACE_STRIDE * (1_000 + i as u64)); let est = sketch.estimate(); assert!( spec.check(est, K).is_ok(), @@ -255,8 +255,7 @@ fn kmv_duplicates_are_inert_and_a_shard_merge_reproduces_the_single_pass_exactly ); let mut merged = a.clone(); - let mut rhs = b.clone(); - merged.merge(&mut rhs); + merged.merge(&b); assert_eq!( merged.estimate(), single, diff --git a/tests/e2e/wire.rs b/tests/e2e/wire.rs index 6ab4277d..9eb8e435 100644 --- a/tests/e2e/wire.rs +++ b/tests/e2e/wire.rs @@ -461,7 +461,7 @@ fn the_experimental_envelopes_round_trip_and_name_their_kinds() { } let kmv_bytes = kmv.serialize_to_bytes().expect("encode KMV"); assert_kind(&kmv_bytes, 0x0e, 0x00, "KMV"); - let mut decoded_kmv = + let decoded_kmv = KMV::::deserialize_from_bytes(&kmv_bytes).expect("decode KMV"); assert_eq!(decoded_kmv.estimate(), kmv.estimate(), "KMV estimate"); diff --git a/tests/pbt/kmv.rs b/tests/pbt/kmv.rs index d7dba430..804ed978 100644 --- a/tests/pbt/kmv.rs +++ b/tests/pbt/kmv.rs @@ -143,8 +143,8 @@ fn from_hashes(k: usize, stream: &[u64]) -> Kmv { fn merged(k: usize, left: &[u64], right: &[u64]) -> Kmv { let mut a = from_hashes(k, left); - let mut b = from_hashes(k, right); - a.merge(&mut b); + let b = from_hashes(k, right); + a.merge(&b); a } @@ -241,7 +241,7 @@ proptest! { fn a_full_sketch_estimates_k_minus_one_over_its_mapped_k_th_smallest_hash( (k, stream, raise) in full_stream(), ) { - let mut sketch = from_hashes(k, &stream); + let sketch = from_hashes(k, &stream); prop_assert_eq!(sketch.k_vals.len(), k); let hashes = retained(&sketch); let kth = *hashes.last().expect("a full sketch retains k hashes"); @@ -333,7 +333,7 @@ proptest! { fn a_sketch_under_its_bound_estimates_its_distinct_count_exactly( (k, distinct, stream) in under_the_bound(48), ) { - let mut sketch = from_hashes(k, &stream); + let sketch = from_hashes(k, &stream); prop_assert_eq!(sketch.k_vals.len(), distinct.len()); prop_assert_eq!( @@ -353,8 +353,8 @@ proptest! { fn a_full_sketch_estimates_above_one_holding_a_larger_k_th_hash( (k, lower, higher) in two_full_streams(), ) { - let mut lower_sketch = from_hashes(k, &lower); - let mut higher_sketch = from_hashes(k, &higher); + let lower_sketch = from_hashes(k, &lower); + let higher_sketch = from_hashes(k, &higher); prop_assert_eq!(lower_sketch.k_vals.len(), k); prop_assert_eq!(higher_sketch.k_vals.len(), k); @@ -379,9 +379,9 @@ proptest! { right in hash_stream(200), ) { let mut a = from_hashes(k, &left); - let mut b = from_hashes(k, &right); + let b = from_hashes(k, &right); let read_before = retained(&b); - a.merge(&mut b); + a.merge(&b); let both: Vec = left.iter().chain(right.iter()).copied().collect(); prop_assert_eq!( @@ -411,12 +411,12 @@ proptest! { ); let mut left_first = merged(k, &first, &second); - let mut third_sketch = from_hashes(k, &third); - left_first.merge(&mut third_sketch); + let third_sketch = from_hashes(k, &third); + left_first.merge(&third_sketch); let mut right_first = from_hashes(k, &first); - let mut tail = merged(k, &second, &third); - right_first.merge(&mut tail); + let tail = merged(k, &second, &third); + right_first.merge(&tail); prop_assert_eq!( retained(&left_first), @@ -433,8 +433,8 @@ proptest! { ) { let mut sketch = from_hashes(k, &stream); let before = retained(&sketch); - let mut copy = sketch.clone(); - sketch.merge(&mut copy); + let copy = sketch.clone(); + sketch.merge(©); prop_assert_eq!( retained(&sketch),