diff --git a/data_plane/src/drivers/ingest/otel.rs b/data_plane/src/drivers/ingest/otel.rs index 305556fc..14b627f2 100644 --- a/data_plane/src/drivers/ingest/otel.rs +++ b/data_plane/src/drivers/ingest/otel.rs @@ -2153,12 +2153,20 @@ fn decode_modified_otlp_sketch_bytes( // frame applies its matrix delta + heap onto a heap // accumulator. Fall back to the plain CountSketch decode for // heap-less msgpack frames (byte-parity path, PR I). - use asap_sketchlib::CountMinSketchWithHeap; - if let Ok(heap) = CountMinSketchWithHeap::from_msgpack(bytes) { + // + // Uses the real `CountSketchWithHeap` (median-of-signed-rows), + // NOT `CountMinSketchWithHeap` — the two share the same wire + // envelope shape (structural peek only), but decoding a real + // CountSketch's matrix through the CMS wrapper would silently + // apply CMS's min-of-rows math to CountSketch data forever + // after (the same conflation bug fixed on the write side in + // `accumulator_factory.rs`). + use asap_sketchlib::CountSketchWithHeap; + if let Ok(heap) = CountSketchWithHeap::from_msgpack(bytes) { if !heap.topk_heap_items().is_empty() { - use crate::precompute_engine::operators::CountMinSketchWithHeapAccumulator; + use crate::precompute_engine::operators::CountSketchWithHeapAccumulator; return Ok(Box::new( - CountMinSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(bytes)?, + CountSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(bytes)?, )); } } @@ -2216,7 +2224,7 @@ fn empty_accumulator_for_delta_bootstrap( encoding: i32, ) -> Option> { use crate::precompute_engine::operators::{ - CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, CountSketchAccumulator, + CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, HllSketchAccumulator, }; use crate::storage_engines::sketch_db::index::SketchConfig; @@ -2237,14 +2245,14 @@ fn empty_accumulator_for_delta_bootstrap( (SketchKind::CountSketch, SketchConfig::CountSketch { rows, cols }) => { // A heap-bearing DELTA-HEAP frame must reconstruct onto a heap // accumulator (the apply path downcasts to - // `CountMinSketchWithHeapAccumulator`); a plain matrix delta + // `CountSketchWithHeapAccumulator`); a plain matrix delta // reconstructs onto a vanilla CountSketch. Pick the base shape // from the encoding so the subsequent // `apply_modified_otlp_delta_bytes` downcast succeeds. if encoding == ENCODING_MSGPACK_DELTA { // heap_size 0 is fine — the DELTA-HEAP apply REPLACES the // heap wholesale from the frame's full heap. - Some(Box::new(CountMinSketchWithHeapAccumulator::new( + Some(Box::new(CountSketchWithHeapAccumulator::new( *rows as usize, *cols as usize, 0, @@ -2290,7 +2298,7 @@ pub(crate) fn apply_modified_otlp_delta_bytes( bytes: &[u8], ) -> Result<(), Box> { use crate::precompute_engine::operators::{ - CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, CountSketchAccumulator, + CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, DDSketchAccumulator, HllSketchAccumulator, }; @@ -2349,13 +2357,14 @@ pub(crate) fn apply_modified_otlp_delta_bytes( // window boundary, so applying the delta reconstructs the // window's own matrix and replaces the heap. Decoded generically // in `apply_msgpack_heap_delta_bytes` (rmp_serde, no - // `asap_sketchlib` delta API). + // `asap_sketchlib` delta API). Real `CountSketchWithHeapAccumulator` + // (median-of-signed-rows), not the CMS-family wrapper. let heap = existing .as_any_mut() - .downcast_mut::() + .downcast_mut::() .ok_or( "apply_modified_otlp_delta_bytes: existing accumulator is \ - not a CountMinSketchWithHeapAccumulator (heap-bearing \ + not a CountSketchWithHeapAccumulator (heap-bearing \ CountSketch delta requires a heap base — the window-1 \ full frame must have promoted the sid)", )?; diff --git a/data_plane/src/precompute_engine/accumulator_factory.rs b/data_plane/src/precompute_engine/accumulator_factory.rs index 43115aa4..9da5bed5 100644 --- a/data_plane/src/precompute_engine/accumulator_factory.rs +++ b/data_plane/src/precompute_engine/accumulator_factory.rs @@ -1,7 +1,8 @@ use crate::precompute_engine::operators::{ - CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, DDSketchAccumulator, - DatasketchesKLLAccumulator, HydraKllSketchAccumulator, IncreaseAccumulator, MinMaxAccumulator, - MultipleIncreaseAccumulator, MultipleMinMaxAccumulator, MultipleSumAccumulator, SumAccumulator, + CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, CountSketchAccumulator, + CountSketchWithHeapAccumulator, DDSketchAccumulator, DatasketchesKLLAccumulator, + HydraKllSketchAccumulator, IncreaseAccumulator, MinMaxAccumulator, MultipleIncreaseAccumulator, + MultipleMinMaxAccumulator, MultipleSumAccumulator, SumAccumulator, }; use crate::storage_engines::types::{ AggregateCore, AggregationType, KeyByLabelValues, Measurement, @@ -674,6 +675,120 @@ impl AccumulatorUpdater for CmsHeapAccumulatorUpdater { } } +// --------------------------------------------------------------------------- +// CountSketchAccumulatorUpdater (real median-of-signed-rows CountSketch) +// --------------------------------------------------------------------------- + +/// Keyed point-frequency updater backed by a real `asap_sketchlib::CountSketch` +/// (signed rows, median-of-rows estimator) — distinct math from +/// `CmsAccumulatorUpdater`'s CMS (min-of-rows). Closes, on the raw-metric +/// ingest path, the conflation bug where `SummaryKind::CountSketch` silently +/// shared `CmsAccumulatorUpdater` with bare CMS. +pub struct CountSketchAccumulatorUpdater { + acc: CountSketchAccumulator, + row_num: usize, + col_num: usize, +} + +impl CountSketchAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + acc: CountSketchAccumulator::new(row_num, col_num), + row_num, + col_num, + } + } +} + +impl AccumulatorUpdater for CountSketchAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.inner.update(&key.to_semicolon_str(), value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = CountSketchAccumulator::new(self.row_num, self.col_num); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() + } +} + +// --------------------------------------------------------------------------- +// CountSketchWithHeapAccumulatorUpdater (real CountSketch + top-k heap) +// --------------------------------------------------------------------------- + +/// Keyed top-k updater backed by a real `CountSketchWithHeap` (signed-row +/// CountSketch matrix PLUS a size-`heap_size` top-k heap). Distinct math from +/// `CmsHeapAccumulatorUpdater`'s CMS-with-heap (min-of-rows); shares the same +/// [`TopkWeight`] semantics and heap payload shape. +pub struct CountSketchWithHeapAccumulatorUpdater { + acc: CountSketchWithHeapAccumulator, + row_num: usize, + col_num: usize, + heap_size: usize, + weight: TopkWeight, +} + +impl CountSketchWithHeapAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize, heap_size: usize, weight: TopkWeight) -> Self { + Self { + acc: CountSketchWithHeapAccumulator::new(row_num, col_num, heap_size), + row_num, + col_num, + heap_size, + weight, + } + } +} + +impl AccumulatorUpdater for CountSketchWithHeapAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + let weighted = match self.weight { + TopkWeight::Value => value, + TopkWeight::Count => 1.0, + }; + self.acc.inner.update(&key.to_semicolon_str(), weighted); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = CountSketchWithHeapAccumulator::new(self.row_num, self.col_num, self.heap_size); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() + + self.heap_size * (std::mem::size_of::() + 32) + } +} + // --------------------------------------------------------------------------- // HydraKllAccumulatorUpdater // --------------------------------------------------------------------------- @@ -914,29 +1029,33 @@ pub fn create_accumulator_updater(config: &AggregationConfig) -> Box { + // Bare CMS: point-frequency only, min-of-rows estimator. `keyed=false` + // can't actually arise here today (no `AggregationType` resolves to + // bare Cms unkeyed — see accumulator_spec.rs), matched anyway as a + // safe default. + (SummaryKind::Cms, _) => { let (row_num, col_num) = cms_dims(&spec.params); Box::new(CmsAccumulatorUpdater::new(row_num, col_num)) } - // Heap-bearing top-k variants (raw-input ingest path): route to - // the real `CmsHeapAccumulatorUpdater` so the per-policy top-k - // heap is BUILT (heap-less CMS could not answer `topk(...)` — - // recall 0). Keyed by the configured group-by `aggregated_labels` - // (e.g. `host`), ranked by Σ value per key by default - // (`weight_mode: value`), or Σ count for genuine frequency-top-k - // (`weight_mode: count`). `CountSketchWithHeap` shares the wire - // shape (heap is the distinguishing payload), so it routes here - // too. The OTLP modified-sketch path builds the heap agent-side - // and uses `SketchEnvelope` ingest, not this raw arm. - (SummaryKind::CmsWithHeap, _) | (SummaryKind::CountSketchWithHeap, _) => { + // Bare CountSketch: real median-of-signed-rows estimator, via the + // dedicated `CountSketchAccumulatorUpdater` (previously conflated + // with `CmsAccumulatorUpdater`'s CMS min-math — see that struct's + // doc). + (SummaryKind::CountSketch, _) => { + let (row_num, col_num) = cms_dims(&spec.params); + Box::new(CountSketchAccumulatorUpdater::new(row_num, col_num)) + } + + // Heap-bearing top-k variant (raw-input ingest path): route to the + // real `CmsHeapAccumulatorUpdater` so the per-policy top-k heap is + // BUILT (heap-less CMS could not answer `topk(...)` — recall 0). + // Keyed by the configured group-by `aggregated_labels` (e.g. `host`), + // ranked by Σ value per key by default (`weight_mode: value`), or Σ + // count for genuine frequency-top-k (`weight_mode: count`). The OTLP + // modified-sketch path builds the heap agent-side and uses + // `SketchEnvelope` ingest, not this raw arm. + (SummaryKind::CmsWithHeap, _) => { let (row_num, col_num, heap_size) = cms_heap_dims(&spec.params); Box::new(CmsHeapAccumulatorUpdater::new( row_num, @@ -946,6 +1065,19 @@ pub fn create_accumulator_updater(config: &AggregationConfig) -> Box { + let (row_num, col_num, heap_size) = cms_heap_dims(&spec.params); + Box::new(CountSketchWithHeapAccumulatorUpdater::new( + row_num, + col_num, + heap_size, + topk_weight_param(config), + )) + } + (SummaryKind::DDSketch, _) => Box::new(DDSketchAccumulatorUpdater::new(ddsketch_alpha( &spec.params, ))), @@ -1275,6 +1407,23 @@ mod tests { items.into_iter().map(|i| (i.key, i.value)).collect() } + /// Same as `ranked_topk`, but for the real `CountSketchWithHeapAccumulator` + /// (median-of-signed-rows) built by `SummaryKind::CountSketchWithHeap` — + /// no longer conflated with the CMS-family accumulator above. + fn ranked_topk_cs(acc: &dyn AggregateCore) -> Vec<(String, f64)> { + let heap = acc + .as_any() + .downcast_ref::() + .expect("CountSketchWithHeap config must build a CountSketchWithHeapAccumulator"); + let mut items = heap.inner.topk_heap_items(); + items.sort_by(|a, b| { + b.value + .partial_cmp(&a.value) + .unwrap_or(std::cmp::Ordering::Equal) + }); + items.into_iter().map(|i| (i.key, i.value)).collect() + } + fn host_key(h: &str) -> KeyByLabelValues { KeyByLabelValues::new_with_labels(vec![h.to_string()]) } @@ -1353,13 +1502,15 @@ mod tests { #[test] fn countsketch_with_heap_also_routes_to_value_weighted_heap() { - // CountSketchWithHeap shares the heap path — same value-weighted default. + // CountSketchWithHeap gets its OWN dedicated updater/accumulator + // (real median-of-signed-rows math) — same value-weighted default + // as the CMS-family heap path, but no longer conflated with it. let config = topk_config(AggregationType::CountSketchWithHeap, None); let mut updater = create_accumulator_updater(&config); feed_stream(&mut *updater); let acc = updater.take_accumulator(); - assert_eq!(acc.type_name(), "CountMinSketchWithHeapAccumulator"); - let ranked = ranked_topk(&*acc); + assert_eq!(acc.type_name(), "CountSketchWithHeapAccumulator"); + let ranked = ranked_topk_cs(&*acc); assert_eq!(ranked[0].0, "host-a"); assert_eq!(ranked[0].1, 100.0); } diff --git a/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs b/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs index bd0b63fc..885c0114 100644 --- a/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs +++ b/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs @@ -13,19 +13,24 @@ //! `asap_sketchlib::proto::sketchlib::CountSketchState`). Mirrors //! the CountMin decoder in PR B. //! -//! Query semantics (median-of-estimators heavy-hitter tracking, -//! `TopKState` integration) are intentionally deferred — queries -//! against stored CountSketch data return a placeholder error today. -//! The wire format carries the matrix losslessly, so the merge + store -//! round-trip works end-to-end without that richer query surface. +//! Per-key point queries go through `query_key`/`MultipleSubpopulationAggregate`, +//! which delegate to `asap_sketchlib::CountSketch::estimate` (the real, +//! hash-spec-compatible median-of-signed-rows estimator) — this used to +//! be a hand-rolled, non-compatible hash (fixed alongside the raw-metric +//! ingest dispatch bug, see `accumulator_factory.rs`). Top-k heap +//! tracking (a distinct capability from a bare median estimate) is a +//! separate concern — see `count_sketch_with_heap_accumulator.rs`. use crate::storage_engines::types::{ - AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink, + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, }; use asap_sketchlib::{CountSketch, CountSketchDelta, MessagePackCodec}; use serde_json::Value; use std::collections::HashMap; +use asap_types::Statistic; + /// Count Sketch accumulator — inner matrix of signed counts. #[derive(Debug, Clone)] pub struct CountSketchAccumulator { @@ -39,6 +44,14 @@ impl CountSketchAccumulator { } } + /// Median-of-signed-rows point estimate for `key`, via the real + /// `asap_sketchlib::CountSketch::estimate` — the canonical, hash-spec- + /// compatible estimator (see `AggregateCore::query_statistic`'s doc for + /// why this replaced a hand-rolled, non-compatible hash). + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + self.inner.estimate(&key.to_semicolon_str()) + } + /// Decode from the modified OTLP wire format's /// `CountSketchDataPoint.sketch` bytes when /// `encoding = COUNT_SKETCH_ENCODING_MSGPACK`. The bytes are the @@ -282,28 +295,34 @@ impl AggregateCore for CountSketchAccumulator { fn query_statistic( &self, statistic: asap_types::Statistic, - _key: &Option, + key: &Option, query_kwargs: &HashMap, ) -> Result> { use asap_types::Statistic; - // Use median-of-row estimator for a specific key when the - // caller provides one in `query_kwargs["key"]`. Without a - // key, fall back to summing the absolute counter values - // (rough total-volume signal — useful for sanity checks - // but not a heavy-hitter answer). Hash compatibility note: - // this relies on the agent and backend using the - // sketchlib HashSpec; sketchlib-go's `portableHashSpec` - // is the canonical seed list, and `asap_sketchlib::CountSketch` - // hashes against the same spec. + // Key-provided path: route to MultipleSubpopulationAggregate::query + // (the canonical "what's the count of this key?" lookup), same + // pattern as CountMinSketchAccumulator. Fixed from a hand-rolled + // `DefaultHasher`-based estimator that did NOT use the sketchlib + // hash spec (its own doc admitted this — "not the sketchlib hash + // spec... the canonical compatibility path requires plumbing the + // sketchlib seeds through") — `asap_sketchlib::CountSketch::estimate` + // already hashes against the correct portable spec, so this is a + // genuine correctness fix, not just a refactor. + if let Some(key_val) = key.as_ref() { + return self.query(statistic, key_val, Some(query_kwargs)); + } + if let Some(k) = query_kwargs.get("key") { + let key_val = KeyByLabelValues::new_with_labels(vec![k.clone()]); + return self.query(statistic, &key_val, Some(query_kwargs)); + } + // No-key path: unchanged from before this fix -- CountSketch's + // signed rows have no CMS-style "min-row-sum = true total" + // property, so these are documented approximations, not a + // heavy-hitter answer. Not touched by this fix (only the + // key-provided path above had the hash-compatibility bug). match statistic { Statistic::Topk | Statistic::Count => { let matrix = self.inner.sketch(); - if let Some(key) = query_kwargs.get("key") { - return Ok(count_sketch_query_key(matrix, key)); - } - // No key → return total absolute volume across the - // sketch as a rough activity proxy. Better than - // erroring out; documented limitation. let total: f64 = matrix.iter().flatten().map(|v| v.abs()).sum(); let rows = matrix.len() as f64; Ok(if rows > 0.0 { total / rows } else { 0.0 }) @@ -323,50 +342,94 @@ impl AggregateCore for CountSketchAccumulator { } } -/// Median-of-row count estimator for CountSketch. Computes one -/// signed estimate per row at `key`'s hash position and returns -/// the median (canonical CountSketch query). -/// -/// Hash compatibility with the agent is via the sketchlib hash -/// spec; the agent's `sketchlib-go::CountSketch` and the -/// backend's `asap_sketchlib::CountSketch` must use -/// the same seed list (sketchlib's `portableHashSpec` / -/// `default_hash_spec`). -fn count_sketch_query_key(matrix: &Vec>, key: &str) -> f64 { - use std::collections::hash_map::DefaultHasher; - use std::hash::{Hash, Hasher}; - if matrix.is_empty() { - return 0.0; - } - let cols = matrix[0].len(); - if cols == 0 { - return 0.0; - } - let mut estimates: Vec = Vec::with_capacity(matrix.len()); - for (i, row) in matrix.iter().enumerate() { - let mut hasher = DefaultHasher::new(); - // Salt with the row index so each row uses a distinct - // hash. Note: this is *not* the sketchlib hash spec — the - // canonical compatibility path requires plumbing the - // sketchlib seeds through to the backend (tracked as a - // follow-up; the wire format already carries the seed - // list, but the accumulator drops it on decode today). - i.hash(&mut hasher); - key.hash(&mut hasher); - let h = hasher.finish() as usize; - let col = h % cols; - // Sign hash: +1 / -1 alternating by a second hash bit. - let sign = if (h >> 32) & 1 == 0 { 1.0 } else { -1.0 }; - estimates.push(sign * row[col]); - } - estimates.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); - estimates[estimates.len() / 2] +impl MultipleSubpopulationAggregate for CountSketchAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountSketchAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } } #[cfg(test)] mod tests { use super::*; + #[test] + fn test_query_key_uses_real_sketchlib_estimator() { + // Regression: query_key/MultipleSubpopulationAggregate used to go + // through a hand-rolled DefaultHasher-based estimator that did NOT + // use the sketchlib hash spec (its own doc admitted this). Fixed + // to delegate to `asap_sketchlib::CountSketch::estimate` directly + // -- prove `query_key` and `.inner.estimate(..)` now agree exactly. + let mut cs = CountSketchAccumulator::new(4, 1000); + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + cs.inner.update(&key.to_semicolon_str(), 10.0); + assert_eq!( + cs.query_key(&key), + cs.inner.estimate(&key.to_semicolon_str()) + ); + } + + #[test] + fn test_multiple_subpopulation_aggregate_query() { + let mut cs = CountSketchAccumulator::new(4, 1000); + let key = KeyByLabelValues::new_with_labels(vec!["checkout".to_string()]); + cs.inner.update(&key.to_semicolon_str(), 25.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cs; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert_eq!(result, cs.query_key(&key)); + + // query_statistic (the AggregateCore entry point) must route a + // provided key through the same path. + let core: &dyn AggregateCore = &cs; + let via_core = core + .query_statistic(Statistic::Sum, &Some(key.clone()), &HashMap::new()) + .unwrap(); + assert_eq!(via_core, cs.query_key(&key)); + } + + #[test] + fn test_mergeable_accumulator_merge_accumulators() { + let cs1 = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![1.0, -2.0], vec![3.0, -4.0]], 2, 2), + }; + let cs2 = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![-1.0, 2.0], vec![-3.0, 4.0]], 2, 2), + }; + let merged = CountSketchAccumulator::merge_accumulators(vec![cs1, cs2]).unwrap(); + assert_eq!(merged.inner.sketch(), &vec![vec![0.0, 0.0], vec![0.0, 0.0]]); + } + + #[test] + fn test_mergeable_accumulator_rejects_empty() { + let result = CountSketchAccumulator::merge_accumulators(vec![]); + assert!(result.is_err()); + } + fn encode_state( rows: u32, cols: u32, diff --git a/data_plane/src/precompute_engine/operators/count_sketch_with_heap_accumulator.rs b/data_plane/src/precompute_engine/operators/count_sketch_with_heap_accumulator.rs new file mode 100644 index 00000000..e4cc527e --- /dev/null +++ b/data_plane/src/precompute_engine/operators/count_sketch_with_heap_accumulator.rs @@ -0,0 +1,575 @@ +//! Count Sketch with Heap accumulator — wraps +//! `asap_sketchlib::CountSketchWithHeap`. +//! +//! Port of `count_min_sketch_with_heap_accumulator.rs` for the distinct +//! `CountSketchWithHeap` (median-of-signed-rows estimator) rather than +//! `CountMinSketchWithHeap` (min-over-rows estimator). The two are +//! different sketch algorithms that happen to share a storage shape and +//! wire layout -- see `asap_sketchlib::CountSketchWithHeap`'s own doc and +//! this session's `delta_apply.rs`/`decoders.rs` fix on the read side. +//! Before this file existed, `accumulator_factory.rs`'s raw-metric +//! ingest dispatch built a `CountMinSketchWithHeapAccumulator` (CMS math) +//! for `SummaryKind::CountSketchWithHeap` sids -- the same conflation bug +//! already fixed on the read side, now closed on the write side too. + +use crate::storage_engines::types::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CountSketchWithHeap, CsHeapItem, MessagePackCodec}; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashMap; + +use asap_types::Statistic; + +/// Local serde view of the DELTA-HEAP wire frame (encoding `MSGPACK_DELTA`). +/// Identical shape to `count_min_sketch_with_heap_accumulator.rs`'s +/// `HeapDeltaWire`/`MatrixDeltaWire` -- the wire frame is generic (sparse +/// cell deltas + a full heap), not CMS-specific. See that file's doc for +/// the exact rmp_serde positional layout. +#[derive(Debug, Deserialize)] +struct HeapDeltaWire { + is_delta: bool, + matrix_delta: MatrixDeltaWire, + topk_heap: Vec<(String, f64)>, + #[allow(dead_code)] + heap_size: u64, +} + +#[derive(Debug, Deserialize)] +struct MatrixDeltaWire { + rows: u32, + cols: u32, + cells: Vec<(u32, u32, i64)>, +} + +/// Validated/flattened view of a decoded DELTA-HEAP frame. +struct HeapDeltaFrame { + rows: u32, + cols: u32, + heap_size: u64, + cells: Vec<(u32, u32, i64)>, + heap: Vec<(String, f64)>, +} + +impl HeapDeltaFrame { + fn from_msgpack(buffer: &[u8]) -> Result> { + let wire: HeapDeltaWire = rmp_serde::from_slice(buffer) + .map_err(|e| format!("decode CountSketchWithHeap delta msgpack: {e}"))?; + if !wire.is_delta { + return Err("CountSketchWithHeap delta frame has is_delta=false".into()); + } + Ok(Self { + rows: wire.matrix_delta.rows, + cols: wire.matrix_delta.cols, + heap_size: wire.heap_size, + cells: wire.matrix_delta.cells, + heap: wire.topk_heap, + }) + } +} + +/// Count Sketch with Heap accumulator — wraps `asap_sketchlib::CountSketchWithHeap`. +/// Core struct, update/merge/serde logic live in +/// `asap_sketchlib::message_pack_format::portable::countsketch_topk`. This +/// file retains QE-specific trait impls, legacy deserializers, and JSON +/// output -- same split as `CountMinSketchWithHeapAccumulator`. +#[derive(Debug, Clone)] +pub struct CountSketchWithHeapAccumulator { + pub inner: CountSketchWithHeap, +} + +impl CountSketchWithHeapAccumulator { + pub fn new(row_num: usize, col_num: usize, heap_size: usize) -> Self { + Self { + inner: CountSketchWithHeap::new(row_num, col_num, heap_size), + } + } + + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + let key_string = key.labels.join(";"); + self.inner.estimate(&key_string) + } + + /// Decode a heap-bearing CountSketch FULL msgpack frame into a heap + /// accumulator -- the window-1 / full-frame base for the DELTA-HEAP + /// delta path. Mirrors `CountMinSketchWithHeapAccumulator::from_msgpack_with_heap_bytes`. + pub fn from_msgpack_with_heap_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketchWithHeap msgpack: {e}"))?, + }) + } + + /// Apply a DELTA-HEAP msgpack frame (encoding `MSGPACK_DELTA`) onto this + /// accumulator IN PLACE. Mirrors + /// `CountMinSketchWithHeapAccumulator::apply_msgpack_heap_delta_bytes` + /// exactly -- the frame decode/apply logic is generic, not tied to + /// which estimator the rebuilt sketch uses. + pub fn apply_msgpack_heap_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + let frame = HeapDeltaFrame::from_msgpack(buffer)?; + + let rows = self.inner.rows(); + let cols = self.inner.cols(); + let heap_size = self.inner.heap_size; + + let mut matrix = self.inner.sketch_matrix(); + for (r, c, dc) in &frame.cells { + let (r, c) = (*r as usize, *c as usize); + if r >= rows || c >= cols { + continue; + } + matrix[r][c] += *dc as f64; + } + + let heap: Vec = frame + .heap + .into_iter() + .map(|(key, value)| CsHeapItem { key, value }) + .collect(); + + self.inner = CountSketchWithHeap::from_legacy_matrix(matrix, heap, rows, cols, heap_size); + Ok(()) + } + + /// Reconstruct a heap accumulator STANDALONE from a single DELTA-HEAP + /// msgpack frame, with no cached per-series base. Mirrors + /// `CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes`. + pub fn from_msgpack_heap_delta_bytes( + buffer: &[u8], + ) -> Result> { + let frame = HeapDeltaFrame::from_msgpack(buffer)?; + if frame.rows == 0 || frame.cols == 0 { + return Err(format!( + "CountSketchWithHeap delta frame has zero dims (rows={}, cols={})", + frame.rows, frame.cols + ) + .into()); + } + let mut acc = Self::new( + frame.rows as usize, + frame.cols as usize, + frame.heap_size as usize, + ); + acc.apply_msgpack_heap_delta_bytes(buffer)?; + Ok(acc) + } + + /// Value-weighted heavy-hitter update -- see + /// `CountMinSketchWithHeapAccumulator::insert_value`'s doc for why + /// this (not a `+1`-per-occurrence update) is the correct semantics + /// for `topk(k, sum by (label) (metric))`-shaped queries. + pub fn insert_value(&mut self, group_label: &str, value: f64) { + self.inner.update(group_label, value); + } + + /// Read the top-`k` groups ranked by summed value (descending, tie-broken + /// by key for determinism). Mirrors `CountMinSketchWithHeapAccumulator::topk_by_value`. + pub fn topk_by_value(&self, k: usize) -> Vec<(String, f64)> { + let mut items: Vec<(String, f64)> = self + .inner + .topk_heap_items() + .into_iter() + .map(|it| (it.key, it.value)) + .collect(); + items.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.0.cmp(&b.0)) + }); + items.truncate(k); + items + } + + /// Get all keys from the top-k heap. + pub fn get_topk_keys(&self) -> Vec { + self.inner + .topk_heap_items() + .iter() + .map(|item| { + let labels: Vec = item.key.split(';').map(|s| s.to_string()).collect(); + KeyByLabelValues { labels } + }) + .collect() + } +} + +impl SerializableToSink for CountSketchWithHeapAccumulator { + fn serialize_to_json(&self) -> Value { + let heap_items: Vec = self + .inner + .topk_heap_items() + .iter() + .map(|item| { + serde_json::json!({ + "key": item.key, + "value": item.value + }) + }) + .collect(); + + serde_json::json!({ + "row_num": self.inner.rows(), + "col_num": self.inner.cols(), + "heap_size": self.inner.heap_size, + "sketch": self.inner.sketch_matrix(), + "topk_heap": heap_items + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountSketchWithHeapAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountSketchWithHeapAccumulator" + } + + /// Per-window base rotation -- mirrors + /// `CountMinSketchWithHeapAccumulator::reset_to_empty`. + fn reset_to_empty(&mut self) { + self.inner = + CountSketchWithHeap::new(self.inner.rows(), self.inner.cols(), self.inner.heap_size); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountSketchWithHeapAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_cs = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountSketchWithHeapAccumulator")?; + + let merged = Self::merge_accumulators(vec![self.clone(), other_cs.clone()])?; + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountSketchWithHeap + } + + fn get_keys(&self) -> Option> { + Some(self.get_topk_keys()) + } + + fn query_statistic( + &self, + statistic: asap_types::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::storage_engines::types::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for CountSketchWithHeapAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for CountSketchWithHeapAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountSketchWithHeapAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_count_sketch_with_heap_creation() { + let cs = CountSketchWithHeapAccumulator::new(4, 1000, 20); + assert_eq!(cs.inner.rows(), 4); + assert_eq!(cs.inner.cols(), 1000); + assert_eq!(cs.inner.heap_size, 20); + assert_eq!(cs.inner.topk_heap_items().len(), 0); + } + + #[test] + fn test_count_sketch_with_heap_query() { + let cs = CountSketchWithHeapAccumulator::new(2, 10, 5); + let key = KeyByLabelValues::new(); + assert_eq!(cs.query_key(&key), 0.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cs; + assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); + } + + #[test] + fn test_count_sketch_with_heap_merge() { + let sketch1 = vec![ + vec![10.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap1 = vec![ + CsHeapItem { + key: "key1".to_string(), + value: 100.0, + }, + CsHeapItem { + key: "key2".to_string(), + value: 50.0, + }, + ]; + let sketch2 = vec![ + vec![5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 15.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap2 = vec![ + CsHeapItem { + key: "key3".to_string(), + value: 75.0, + }, + CsHeapItem { + key: "key1".to_string(), + value: 80.0, + }, + ]; + + let cs1 = CountSketchWithHeapAccumulator { + inner: CountSketchWithHeap::from_legacy_matrix(sketch1, heap1, 2, 10, 5), + }; + let cs2 = CountSketchWithHeapAccumulator { + inner: CountSketchWithHeap::from_legacy_matrix(sketch2, heap2, 2, 10, 3), + }; + + let result = CountSketchWithHeapAccumulator::merge_accumulators(vec![cs1, cs2]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.sketch_matrix()[0][0], 15.0); + assert_eq!(merged.inner.sketch_matrix()[1][1], 35.0); + assert_eq!(merged.inner.heap_size, 3); + assert!(merged.inner.topk_heap_items().len() <= 3); + } + + #[test] + fn test_count_sketch_with_heap_merge_single() { + let cs = CountSketchWithHeapAccumulator::new(2, 3, 5); + let result = CountSketchWithHeapAccumulator::merge_accumulators(vec![cs.clone()]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.rows(), cs.inner.rows()); + assert_eq!(merged.inner.cols(), cs.inner.cols()); + assert_eq!(merged.inner.heap_size, cs.inner.heap_size); + } + + #[test] + fn test_count_sketch_with_heap_merge_dimension_mismatch() { + let cs1 = CountSketchWithHeapAccumulator::new(2, 10, 5); + let cs2 = CountSketchWithHeapAccumulator::new(3, 10, 5); + let result = CountSketchWithHeapAccumulator::merge_accumulators(vec![cs1, cs2]); + assert!(result.is_err()); + } + + #[test] + fn test_count_sketch_with_heap_as_aggregate_core() { + let cs = CountSketchWithHeapAccumulator::new(2, 3, 5); + assert_eq!(cs.type_name(), "CountSketchWithHeapAccumulator"); + } + + #[test] + fn test_get_topk_keys() { + let mut cs = CountSketchWithHeapAccumulator::new(2, 3, 5); + cs.inner.update("label1;label2", 100.0); + cs.inner.update("label3;label4", 50.0); + + let keys = cs.get_topk_keys(); + assert_eq!(keys.len(), 2); + let label_sets: std::collections::HashSet<_> = + keys.iter().map(|k| k.labels.clone()).collect(); + assert!(label_sets.contains(&vec!["label1".to_string(), "label2".to_string()])); + assert!(label_sets.contains(&vec!["label3".to_string(), "label4".to_string()])); + } + + #[test] + fn test_multiple_subpopulation_aggregate() { + let cs = CountSketchWithHeapAccumulator::new(3, 50, 10); + let key = KeyByLabelValues::new(); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cs; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert_eq!(result, 0.0); + + let keys = multi_trait.get_keys(); + assert!(keys.is_some()); + assert_eq!(keys.unwrap().len(), 0); + } + + #[test] + fn test_pwr_full_then_delta_then_delta_reconstructs_per_window() { + use asap_sketchlib::MessagePackCodec; + + let w1 = CountSketchWithHeap::from_legacy_matrix( + vec![vec![300.0; 4]; 5], + vec![CsHeapItem { + key: "k".into(), + value: 300.0, + }], + 5, + 4, + 20, + ); + let w1_bytes = w1.to_msgpack().expect("w1 full msgpack"); + let mut base = CountSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(&w1_bytes) + .expect("decode w1 full frame as heap accumulator"); + assert_eq!(base.inner.sketch_matrix()[0][0], 300.0); + + let w2_frame = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); + base.reset_to_empty(); + assert_eq!( + base.inner.sketch_matrix()[0][0], + 0.0, + "reset_to_empty cleared matrix" + ); + base.apply_msgpack_heap_delta_bytes(&w2_frame) + .expect("apply w2 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 50.0, "window-2 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 50.0); + assert_eq!(base.inner.sketch_matrix()[2][2], 0.0); + let h2: Vec<_> = base.inner.topk_heap_items(); + assert_eq!(h2.len(), 1); + assert_eq!(h2[0].key, "k"); + assert_eq!(h2[0].value, 50.0); + + let w3_frame = encode_delta_heap(5, 4, &[(0, 0, 80)], &[("k", 80.0)], 20); + base.reset_to_empty(); + base.apply_msgpack_heap_delta_bytes(&w3_frame) + .expect("apply w3 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 80.0, "window-3 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 0.0, "no window-2 leakage"); + let h3 = base.inner.topk_heap_items(); + assert_eq!(h3.len(), 1); + assert_eq!(h3[0].value, 80.0); + } + + #[test] + fn test_apply_delta_rejects_full_frame_and_garbage() { + use asap_sketchlib::MessagePackCodec; + let mut acc = CountSketchWithHeapAccumulator::new(2, 4, 5); + let full = CountSketchWithHeap::from_legacy_matrix( + vec![vec![1.0; 4]; 2], + vec![CsHeapItem { + key: "a".into(), + value: 1.0, + }], + 2, + 4, + 5, + ) + .to_msgpack() + .unwrap(); + assert!(acc.apply_msgpack_heap_delta_bytes(&full).is_err()); + assert!(acc.apply_msgpack_heap_delta_bytes(b"not msgpack").is_err()); + } + + fn encode_delta_heap( + rows: u32, + cols: u32, + cells: &[(u32, u32, i64)], + heap: &[(&str, f64)], + heap_size: u64, + ) -> Vec { + #[derive(serde::Serialize)] + struct W<'a>( + bool, + (u32, u32, &'a [(u32, u32, i64)]), + Vec<(String, f64)>, + u64, + ); + let heap_owned: Vec<(String, f64)> = + heap.iter().map(|(k, v)| (k.to_string(), *v)).collect(); + let w = W(true, (rows, cols, cells), heap_owned, heap_size); + rmp_serde::to_vec(&w).expect("encode delta-heap") + } + + #[test] + fn insert_value_accumulates_summed_value_in_heap() { + let mut acc = CountSketchWithHeapAccumulator::new(4, 1024, 8); + acc.insert_value("g", 10.0); + acc.insert_value("g", 25.0); + let top = acc.topk_by_value(1); + assert_eq!(top.len(), 1); + assert_eq!(top[0].0, "g"); + assert!( + (top[0].1 - 35.0).abs() < 1e-6, + "summed value should be 35 (10+25), got {}", + top[0].1 + ); + } + + /// The core proof this file exists at all: `CountSketchWithHeapAccumulator` + /// wraps the real, distinct `asap_sketchlib::CountSketchWithHeap` -- + /// not the CMS-family `CountMinSketchWithHeap` a collapsed dispatch + /// used to substitute (the exact bug this file fixes on the ingest + /// side, mirroring the already-fixed read side). Two different Rust + /// types means `merge_with` rejects mixing them at the type-check + /// level, same as any other mismatched-family merge attempt -- + /// verified directly rather than via a numeric estimate comparison + /// (asap_sketchlib's own test suite already proves the median vs + /// min-over-rows divergence at the sketch-math level). + #[test] + fn test_rejects_merge_with_cms_family_accumulator() { + use crate::precompute_engine::operators::count_min_sketch_with_heap_accumulator::CountMinSketchWithHeapAccumulator; + + let cs = CountSketchWithHeapAccumulator::new(4, 64, 10); + let cms = CountMinSketchWithHeapAccumulator::new(4, 64, 10); + let result = cs.merge_with(&cms); + assert!( + result.is_err(), + "CountSketchWithHeapAccumulator must not merge with CountMinSketchWithHeapAccumulator \ + -- different algorithms sharing only a storage shape" + ); + } +} diff --git a/data_plane/src/precompute_engine/operators/mod.rs b/data_plane/src/precompute_engine/operators/mod.rs index 84c98947..45881271 100644 --- a/data_plane/src/precompute_engine/operators/mod.rs +++ b/data_plane/src/precompute_engine/operators/mod.rs @@ -1,6 +1,7 @@ pub mod count_min_sketch_accumulator; pub mod count_min_sketch_with_heap_accumulator; pub mod count_sketch_accumulator; +pub mod count_sketch_with_heap_accumulator; pub mod datasketches_kll_accumulator; pub mod dd_sketch_accumulator; pub mod edge_runtime_adapter; @@ -17,6 +18,7 @@ pub mod sum_accumulator; pub use count_min_sketch_accumulator::*; pub use count_min_sketch_with_heap_accumulator::*; pub use count_sketch_accumulator::*; +pub use count_sketch_with_heap_accumulator::*; pub use datasketches_kll_accumulator::*; pub use dd_sketch_accumulator::*; pub use hll_sketch_accumulator::*;