From 9b57222bce86b5ea43744e8927131a9a0655623f Mon Sep 17 00:00:00 2001 From: zzylol Date: Tue, 26 May 2026 13:02:45 -0600 Subject: [PATCH 1/2] KLL: structural clone + direct bit-exact reconstruction from portable state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit KllSketch::clone serialized to msgpack and deserialized back on every clone. The backing KLL derives Clone, so clone it structurally instead — identical sketch, no rmp encode/decode and far fewer allocations. Add KLL::from_portable_state to rebuild a sketch directly from its portable wire form (k + level-ordered items + the levels[] boundary array): place the items straight into the internal buffer at the correct offset and fix up the level boundaries, then recompute the capacity cache. This replaces reconstructing a sketch by replaying every retained item through update() — it is bit-exact (identical quantiles, vs a lossy statistical reconstruction) and several times faster (no per-item compaction/sort/RNG). Exposed via KllSketch::from_portable_state for proto/envelope decoders; the empty-sketch case is handled. Add a test asserting from_portable_state reproduces the source sketch's quantiles exactly across the quantile range, plus the empty-sketch case. Co-Authored-By: Claude Opus 4.7 --- src/message_pack_format/portable/kll.rs | 24 ++++++- src/sketches/kll.rs | 93 +++++++++++++++++++++++++ 2 files changed, 115 insertions(+), 2 deletions(-) diff --git a/src/message_pack_format/portable/kll.rs b/src/message_pack_format/portable/kll.rs index fd3d5e7..55f5cf9 100644 --- a/src/message_pack_format/portable/kll.rs +++ b/src/message_pack_format/portable/kll.rs @@ -135,14 +135,34 @@ impl KllSketch { impl Clone for KllSketch { fn clone(&self) -> Self { - let bytes = bytes_from_sketchlib_kll(&self.backend); + // Structural clone: the backing `KLL` derives `Clone` (its fields + // are `Box<[..]>`/`Vec`/scalars), so copy it directly instead of a full + // msgpack serialize -> deserialize -> rebuild round-trip. Produces an + // identical sketch far more cheaply (no rmp encode/decode, far fewer + // allocations). Self { k: self.k, - backend: sketchlib_kll_from_bytes(&bytes).unwrap(), + backend: self.backend.clone(), } } } +impl KllSketch { + /// Reconstruct directly from portable wire state (k + level-ordered items + + /// level boundaries) without replaying items through `update()`. Bit-exact. + pub fn from_portable_state( + k: u16, + items: &[f64], + levels: &[usize], + num_levels: usize, + ) -> Result { + Ok(Self { + k, + backend: SketchlibKll::from_portable_state(k as usize, items, levels, num_levels)?, + }) + } +} + impl std::fmt::Debug for KllSketch { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("KllSketch") diff --git a/src/sketches/kll.rs b/src/sketches/kll.rs index 1d9b9b5..ee11b3e 100644 --- a/src/sketches/kll.rs +++ b/src/sketches/kll.rs @@ -268,6 +268,66 @@ impl KLL { s } + /// Reconstruct a KLL directly from portable wire state — the retained + /// `items` in level order plus the `levels` boundary array (proto contract: + /// `levels[0] == 0`, `levels[num_levels] == items.len()`). Avoids replaying + /// every item through `update()`: it places the items straight into the + /// internal buffer and fixes up the level boundaries, so the result is + /// bit-identical to the source sketch (identical quantiles) at a fraction + /// of the cost. Errors if the supplied state is inconsistent. + pub fn from_portable_state( + k: usize, + items: &[T], + levels: &[usize], + num_levels: usize, + ) -> Result { + let mut s = Self::init_kll(k as i32); + if items.is_empty() { + return Ok(s); + } + if num_levels == 0 || num_levels >= s.levels.len() { + return Err(format!( + "from_portable_state: invalid num_levels {num_levels}" + )); + } + if levels.len() != num_levels + 1 { + return Err(format!( + "from_portable_state: levels.len()={} != num_levels+1={}", + levels.len(), + num_levels + 1 + )); + } + if levels[0] != 0 || levels[num_levels] != items.len() { + return Err(format!( + "from_portable_state: level bounds [{}..{}] inconsistent with items.len()={}", + levels[0], + levels[num_levels], + items.len() + )); + } + if items.len() > s.max_capacity { + return Err(format!( + "from_portable_state: {} items exceed max_capacity {} for k={}", + items.len(), + s.max_capacity, + k + )); + } + // Live data occupies the high end of the buffer; free space at the front. + let offset = s.max_capacity - items.len(); + let max_cap = s.max_capacity; + s.items[offset..].clone_from_slice(items); + for (dst, &src) in s.levels.iter_mut().zip(levels[..=num_levels].iter()) { + *dst = src + offset; + } + for dst in s.levels.iter_mut().skip(num_levels + 1) { + *dst = max_cap; + } + s.num_levels = num_levels; + s.rebuild_capacity_cache(); // recomputes top_height + level0_capacity for num_levels + Ok(s) + } + /// Hot-path insert: decrement `levels[0]`, write item, check capacity. #[inline] fn push_value(&mut self, value: T) { @@ -813,6 +873,39 @@ mod tests { use super::*; use crate::test_utils::{sample_uniform_f64, sample_zipf_f64}; + // Direct reconstruction from portable wire state must be BIT-EXACT: + // identical quantiles to the source sketch (unlike a lossy + // replay-through-`update()` reconstruction). + #[test] + fn from_portable_state_reproduces_source_exactly() { + let k = 200usize; + let mut src = KLL::::init_kll(k as i32); + for i in 0..200_000u32 { + src.update(&((i as f64) * 0.0007 + 3.0)); + } + // Extract the compacted portable form (proto contract: levels[0]==0, + // levels[num_levels]==items.len()). + let off = src.levels[0]; + let items: Vec = src.items[off..src.max_capacity].to_vec(); + let levels: Vec = (0..=src.num_levels).map(|h| src.levels[h] - off).collect(); + let n = src.num_levels; + + // (A) bit-exact: direct reconstruction reproduces the source exactly. + let rebuilt = KLL::::from_portable_state(k, &items, &levels, n).unwrap(); + for &q in &[0.0, 0.01, 0.1, 0.25, 0.5, 0.75, 0.9, 0.99, 1.0] { + assert_eq!( + rebuilt.quantile(q), + src.quantile(q), + "direct reconstruction quantile mismatch at q={q}" + ); + } + + // An empty sketch round-trips to an empty sketch. + let empty = KLL::::from_portable_state(k, &[], &[], 0).unwrap(); + assert_eq!(empty.num_levels, 1); + assert_eq!(empty.levels[0], empty.max_capacity); + } + // Ensure each 64-bit chunk is consumed bit-by-bit before refilling. #[test] fn coin_bit_cache_behavior() { From 5528b17778d9abb29de284ef997559ca14bf01e5 Mon Sep 17 00:00:00 2001 From: zzylol Date: Tue, 26 May 2026 13:24:29 -0600 Subject: [PATCH 2/2] cms-heap: structural Clone for CountMinSketchWithHeap CMSHeap's fields (CountMin and HHHeap) both derive Clone, so derive Clone on CMSHeap and clone CountMinSketchWithHeap structurally instead of extracting the count matrix + top-k heap to wire form and rebuilding the backend on every clone. Same sketch, far less work per clone. Co-Authored-By: Claude Opus 4.7 --- .../portable/countminsketch_topk.rs | 13 ++++--------- src/sketches/countminsketch_topk.rs | 1 + 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/src/message_pack_format/portable/countminsketch_topk.rs b/src/message_pack_format/portable/countminsketch_topk.rs index 552ae35..fa850f9 100644 --- a/src/message_pack_format/portable/countminsketch_topk.rs +++ b/src/message_pack_format/portable/countminsketch_topk.rs @@ -138,19 +138,14 @@ impl std::fmt::Debug for CountMinSketchWithHeap { impl Clone for CountMinSketchWithHeap { fn clone(&self) -> Self { - let sketch = matrix_from_sketchlib_cms_heap(&self.backend); - let wire_heap: Vec = heap_to_wire(&self.backend); + // Structural clone: `CMSHeap` derives `Clone` (its `CountMin` + `HHHeap` + // fields are `Clone`), so copy the backend directly instead of + // extracting the matrix + heap to wire form and rebuilding it. Self { rows: self.rows, cols: self.cols, heap_size: self.heap_size, - backend: sketchlib_cms_heap_from_matrix_and_heap( - self.rows, - self.cols, - self.heap_size, - &sketch, - &wire_heap, - ), + backend: self.backend.clone(), } } } diff --git a/src/sketches/countminsketch_topk.rs b/src/sketches/countminsketch_topk.rs index 3dc6172..85d55af 100644 --- a/src/sketches/countminsketch_topk.rs +++ b/src/sketches/countminsketch_topk.rs @@ -15,6 +15,7 @@ const DEFAULT_TOP_K: usize = 32; /// A Count-Min Sketch paired with a top-k heavy-hitter heap. /// /// Generic over the same type parameters as [`CountMin`]. +#[derive(Clone)] pub struct CMSHeap< S: MatrixStorage = Vector2D, Mode = RegularPath,