From 3b880efc692ecad23a5a4efb29fc2da898891d05 Mon Sep 17 00:00:00 2001 From: zzylol Date: Wed, 30 Sep 2026 13:23:37 +0000 Subject: [PATCH 1/4] refactor: keep storage codecs over Planner kernel states The store now holds Planner's summary kernels directly. asap_summary_state keeps only their stored tags and byte encodings (stored_state::codec), edge wire decoding (stored_state::decoders), delta reconstruction and readout binding, and no second kernel implementation. - Delete the copied summary_kernels, the backend updater factory, physical.rs boundary conversions and the legacy per-statistic trait surface. - Current stored formats are unchanged; retired exact tags and the native SumAccumulatorV1 codec fail with a named rejection. - Edge-sampled frames are rejected: Planner kernels carry no sample_p. - OTLP SketchEnvelope attributes decode into Planner kernels. - Unkeyed Planner exact MIN/MAX panes feed the MIN/MAX rollup again. - UnivMon keeps a marked shim until Planner exposes its sketch. Co-Authored-By: Claude Opus 5.5 --- crates/asap_summary_state/src/lib.rs | 19 +- crates/asap_summary_state/src/physical.rs | 212 --- .../src/stored_state/codec.rs | 566 +++++++ .../src/stored_state/decoders.rs | 1045 +++++++++---- .../src/stored_state/delta_apply.rs | 858 ++++------- .../src/stored_state/mod.rs | 7 +- .../src/stored_state/native.rs | 196 ++- .../src/stored_state/readout.rs | 252 ++-- .../src/summary_kernels/count_min_sketch.rs | 1323 ----------------- .../count_min_sketch_with_heap.rs | 832 ----------- .../src/summary_kernels/count_sketch.rs | 678 --------- .../summary_kernels/count_sketch_with_heap.rs | 575 ------- .../src/summary_kernels/datasketches_kll.rs | 727 --------- .../src/summary_kernels/dd_sketch.rs | 726 --------- .../src/summary_kernels/exact.rs | 350 ----- .../src/summary_kernels/factory.rs | 1176 --------------- .../src/summary_kernels/hll_sketch.rs | 788 ---------- .../src/summary_kernels/hydra_kll.rs | 165 -- .../src/summary_kernels/increase.rs | 737 --------- .../summary_kernels/keyed_counter_state.rs | 530 ------- .../src/summary_kernels/keyed_max_state.rs | 336 ----- .../src/summary_kernels/keyed_min_state.rs | 336 ----- .../src/summary_kernels/keyed_sum_count.rs | 559 ------- .../src/summary_kernels/max.rs | 248 --- .../src/summary_kernels/min.rs | 253 ---- .../src/summary_kernels/mod.rs | 42 - .../src/summary_kernels/sketch_envelope.rs | 154 -- .../src/summary_kernels/sum.rs | 413 ----- .../src/summary_kernels/traits.rs | 357 ----- .../src/summary_kernels/weighted_frequency.rs | 75 - .../src/{summary_kernels => }/univmon.rs | 185 +-- crates/asap_types/src/lib.rs | 1 - crates/asap_types/src/traits.rs | 1 - data_plane/benches/sketch_db.rs | 13 +- data_plane/examples/univmon_erp_artifact.rs | 26 +- data_plane/src/drivers/ingest/otel.rs | 368 ++--- data_plane/src/lib.rs | 9 +- data_plane/src/precompute_engine/config.rs | 2 +- .../src/precompute_engine/ingest_handler.rs | 11 +- .../precompute_engine/maintenance_runtime.rs | 25 +- .../precompute_engine/native_precompute.rs | 18 +- .../src/precompute_engine/output_sink.rs | 45 +- data_plane/src/precompute_engine/raw_dag.rs | 37 +- data_plane/src/precompute_engine/revisions.rs | 12 +- .../src/precompute_engine/series_router.rs | 1 + data_plane/src/precompute_engine/worker.rs | 173 ++- .../accelerator.rs | 8 +- .../query_engines/asap_query_engine/engine.rs | 219 +-- .../asap_query_engine/exact_subqueries.rs | 9 +- .../asap_query_engine/live_serve.rs | 2 +- .../asap_query_engine/post_asap_readout.rs | 62 +- .../asap_query_engine/request_tests.rs | 2 +- .../asap_query_engine/summary_executor.rs | 46 +- .../sketch_db/backfill/processor.rs | 3 +- .../sketch_db/backfill/window_builder.rs | 14 +- .../src/storage_engines/sketch_db/data/mod.rs | 5 +- .../sketch_db/index/maintenance.rs | 57 +- .../storage_engines/sketch_db/index/mod.rs | 209 +-- .../storage_engines/sketch_db/index/native.rs | 151 +- .../sketch_db/lifecycle/eviction.rs | 3 +- .../sketch_db/query/window_merger.rs | 37 - data_plane/src/storage_engines/types/mod.rs | 2 +- data_plane/src/tests/accumulator_fixture.rs | 133 +- data_plane/src/tests/mod.rs | 1 - data_plane/src/tests/trait_design_tests.rs | 81 - data_plane/tests/edge_sketch_codec.rs | 18 +- .../tests/support/univmon_erp_process.rs | 16 +- 67 files changed, 2618 insertions(+), 13922 deletions(-) delete mode 100644 crates/asap_summary_state/src/physical.rs create mode 100644 crates/asap_summary_state/src/stored_state/codec.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/count_min_sketch.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/count_min_sketch_with_heap.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/count_sketch.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/count_sketch_with_heap.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/datasketches_kll.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/dd_sketch.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/exact.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/factory.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/hll_sketch.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/hydra_kll.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/increase.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/keyed_counter_state.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/keyed_max_state.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/keyed_min_state.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/keyed_sum_count.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/max.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/min.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/mod.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/sketch_envelope.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/sum.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/traits.rs delete mode 100644 crates/asap_summary_state/src/summary_kernels/weighted_frequency.rs rename crates/asap_summary_state/src/{summary_kernels => }/univmon.rs (51%) delete mode 100644 crates/asap_types/src/traits.rs delete mode 100644 data_plane/src/tests/trait_design_tests.rs diff --git a/crates/asap_summary_state/src/lib.rs b/crates/asap_summary_state/src/lib.rs index db5a1e9ac..60e394581 100644 --- a/crates/asap_summary_state/src/lib.rs +++ b/crates/asap_summary_state/src/lib.rs @@ -1,19 +1,16 @@ -//! Backend-owned summary state: the kernels that ingest and the sketch store -//! keep, their stored byte formats, and conversion to Planner physical states. +//! Backend-owned summary storage: stored byte formats of Planner kernel +//! states, edge wire decoding, delta reconstruction and readout binding. //! -//! Planner's `asap-physical-operators` keeps only in-memory computation state. -//! Storage formats, delta reconstruction and the legacy per-statistic kernels -//! are deployment concerns and live here. - -pub mod summary_kernels; -pub use summary_kernels::{factory, traits}; -pub use traits::*; +//! Planner's `asap-physical-operators` owns summary computation (update, +//! merge, estimate). The store keeps those kernel states directly; this crate +//! only encodes, decodes and reads them. mod aggregation_type; pub use aggregation_type::AggregationType; pub mod codec; -pub mod physical; pub mod stored_state; +pub mod univmon; -pub use asap_physical_operators::{KeyByLabelValues, Measurement, Statistic}; +pub use asap_physical_operators::{AggregateCore, KeyByLabelValues, Measurement, Statistic}; +pub use stored_state::codec::StoredState; diff --git a/crates/asap_summary_state/src/physical.rs b/crates/asap_summary_state/src/physical.rs deleted file mode 100644 index 80149e7dd..000000000 --- a/crates/asap_summary_state/src/physical.rs +++ /dev/null @@ -1,212 +0,0 @@ -//! Conversion between stored summary state and Planner physical state. -//! -//! Planner physical operators accept only their own in-memory kernels. Stored -//! kernels carry storage-only fields (for example edge `sample_p`), so only the -//! families Planner can bind as typed inputs convert, and only when those -//! fields are neutral. -use crate::summary_kernels as stored; -use crate::AggregateCore as StoredState; -use asap_physical_operators::summary_kernels as physical; -use asap_physical_operators::AggregateCore as PhysicalState; -use std::sync::Arc; - -type Error = Box; - -/// Bind a stored state as a Planner physical input. -/// -/// A stored `SumAccumulator` becomes Planner's unkeyed exact Sum state, which -/// is the only exact representation Planner operators accept. -pub fn to_physical(state: &dyn StoredState) -> Result, Error> { - let any = state.as_any(); - if let Some(s) = any.downcast_ref::() { - unsampled(s.sample_p)?; - return Ok(Arc::new(physical::DDSketchAccumulator { - inner: s.inner.clone(), - })); - } - if let Some(s) = any.downcast_ref::() { - unsampled(s.sample_p)?; - return Ok(Arc::new(physical::HllSketchAccumulator { - inner: s.inner.clone(), - })); - } - if let Some(s) = any.downcast_ref::() { - return Ok(Arc::new(physical::DatasketchesKLLAccumulator { - inner: s.inner.clone(), - })); - } - if let Some(s) = any.downcast_ref::() { - return Ok(Arc::new(exact_to_physical(s)?)); - } - if let Some(s) = any.downcast_ref::() { - return Ok(Arc::new(s.0.clone())); - } - if let Some(s) = any.downcast_ref::() { - use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; - let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); - let mut exact = physical::exact::ExactAccumulator::new(family, false)?; - exact.update(None, s.sum, 0); - return Ok(Arc::new(exact)); - } - Err(format!("{} has no Planner physical state", state.type_name()).into()) -} - -/// Keep a Planner physical output in the stored kernel family. -pub fn from_physical(state: &dyn PhysicalState) -> Result, Error> { - let any = state.as_any(); - if let Some(s) = any.downcast_ref::() { - return Ok(Box::new(stored::DDSketchAccumulator { - inner: s.inner.clone(), - sample_p: 1.0, - })); - } - if let Some(s) = any.downcast_ref::() { - return Ok(Box::new(stored::HllSketchAccumulator { - inner: s.inner.clone(), - sample_p: 1.0, - })); - } - if let Some(s) = any.downcast_ref::() { - return Ok(Box::new(stored::DatasketchesKLLAccumulator { - inner: s.inner.clone(), - })); - } - if let Some(s) = any.downcast_ref::() { - return Ok(Box::new(stored::weighted_frequency::WeightedFrequency( - s.clone(), - ))); - } - if let Some(s) = any.downcast_ref::() { - return Ok(Box::new( - stored::exact::ExactAccumulator::deserialize_from_bytes(&rmp_serde::to_vec_named(s)?)?, - )); - } - Err("Planner physical state has no stored kernel".into()) -} - -/// Stored kernel identity of a Planner physical state. -pub fn aggregation_type(state: &dyn PhysicalState) -> Result { - use crate::AggregationType as T; - let any = state.as_any(); - Ok(if any.is::() { - T::DDSketch - } else if any.is::() { - T::HLL - } else if any.is::() { - T::DatasketchesKLL - } else { - from_physical(state)?.get_accumulator_type() - }) -} - -/// Planner's weighted frequency state is serde-transparent over the sketchlib -/// kernel, which owns its persisted byte form. -pub(crate) fn frequency_kernel( - state: &physical::weighted_frequency::WeightedFrequency, -) -> Result { - Ok(rmp_serde::from_slice(&rmp_serde::to_vec(state)?)?) -} - -/// Both exact states share one serde shape; the stored decoder checks that the -/// payload matches its declared family. -fn exact_to_physical( - state: &stored::exact::ExactAccumulator, -) -> Result { - use crate::SerializableToSink; - Ok(rmp_serde::from_slice(&state.serialize_to_bytes())?) -} - -fn unsampled(sample_p: f64) -> Result<(), Error> { - if sample_p == 1.0 { - Ok(()) - } else { - Err("edge-sampled sketches have no Planner physical state".into()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::{KeyByLabelValues, SerializableToSink, Statistic}; - use planner_types::post_asap::{ExactKind, ExactParams, SketchQuery, SummaryFamilyType}; - - // Sketches keep their estimates in both directions. - #[test] - fn sketches_convert_both_ways() { - let mut dd = stored::DDSketchAccumulator::new(0.01); - let mut kll = stored::DatasketchesKLLAccumulator::new(200); - for v in 1..=100 { - dd.inner.update(f64::from(v)); - kll.update(f64::from(v)); - } - let q = SketchQuery::Quantile { q: 0.5 }; - for state in [&dd as &dyn StoredState, &kll] { - let physical = to_physical(state).unwrap(); - let median = physical.estimate(&q).unwrap(); - assert!((median - 50.0).abs() <= 2.0, "{median}"); - let back = from_physical(physical.as_ref()).unwrap(); - assert_eq!(back.serialize_to_bytes(), state.serialize_to_bytes()); - } - } - - // Edge-sampled sketches must not become unscaled Planner inputs. - #[test] - fn sampled_sketch_is_rejected() { - let mut dd = stored::DDSketchAccumulator::new(0.01); - dd.sample_p = 0.5; - assert!(to_physical(&dd).is_err()); - } - - // Exact state keeps family and value; a stored Sum binds as Planner exact Sum. - #[test] - fn exact_and_sum_bind_as_planner_exact_state() { - let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); - let mut exact = stored::exact::ExactAccumulator::new(family.clone(), false).unwrap(); - exact.update(None, 4.5, 10); - let physical = to_physical(&exact).unwrap(); - let read = |state: &Arc| { - state - .as_any() - .downcast_ref::() - .unwrap() - .readout(Statistic::Sum, None, None::<&KeyByLabelValues>) - .unwrap() - }; - assert_eq!(read(&physical), Some(4.5)); - let back = from_physical(physical.as_ref()).unwrap(); - assert_eq!(back.serialize_to_bytes(), exact.serialize_to_bytes()); - - let sum = stored::SumAccumulator::with_sum(7.0); - assert_eq!(read(&to_physical(&sum).unwrap()), Some(7.0)); - } - - // Heap state keeps its bytes and kernel identity through the stored form. - #[test] - fn weighted_frequency_converts_both_ways() { - use asap_physical_operators::values::Value; - use physical::weighted_frequency::{FrequencyAlgorithm, WeightedFrequency}; - let mut state = WeightedFrequency::new(FrequencyAlgorithm::Cms, 64, 5, 8).unwrap(); - state.update(&[Value::Utf8("a".into())], 0.5).unwrap(); - let stored = from_physical(&state).unwrap(); - assert_eq!( - stored.get_accumulator_type(), - crate::AggregationType::CountMinSketchWithHeap - ); - assert_eq!( - aggregation_type(&state).unwrap(), - crate::AggregationType::CountMinSketchWithHeap - ); - let back = to_physical(stored.as_ref()).unwrap(); - let back = back.as_any().downcast_ref::().unwrap(); - assert_eq!( - frequency_kernel(back).unwrap().to_bytes(), - frequency_kernel(&state).unwrap().to_bytes() - ); - } - - // Stored-only kernels have no Planner physical representation. - #[test] - fn stored_only_kernels_do_not_convert() { - assert!(to_physical(&stored::MinAccumulator::new()).is_err()); - } -} diff --git a/crates/asap_summary_state/src/stored_state/codec.rs b/crates/asap_summary_state/src/stored_state/codec.rs new file mode 100644 index 000000000..6f94f5932 --- /dev/null +++ b/crates/asap_summary_state/src/stored_state/codec.rs @@ -0,0 +1,566 @@ +//! Stored identity, byte form and statistic readout of Planner kernel states. +//! +//! The store keeps Planner's in-memory kernels directly. This module is the +//! only place that names their persisted type tags and byte encodings; Planner +//! owns their update, merge and estimate. +use super::native::NativeSummaryOutput; +use crate::univmon::UnivMonAccumulator; +use crate::{AggregationType, KeyByLabelValues}; +use asap_physical_operators::summary_kernels as k; +use asap_physical_operators::summary_kernels::weighted_frequency::{ + FrequencyAlgorithm, WeightedFrequency, +}; +pub use asap_physical_operators::AggregateCore; +use asap_sketchlib::MessagePackCodec; +use planner_types::post_asap::{ExactKind, SummaryFamilyType}; +use std::collections::HashMap; + +pub type Error = Box; + +/// Persisted type tag of Planner's exact state (its named msgpack form). +pub const EXACT_V1: &str = "PlannerExactAccumulatorV1"; + +/// Exact-state tags written before the store kept Planner kernels. Their byte +/// layouts are no longer decoded. +const RETIRED_EXACT_TAGS: &[&str] = &[ + "SumAccumulator", + "IncreaseAccumulator", + "MinAccumulator", + "MaxAccumulator", + "KeyedSumCountAccumulator", + "KeyedCounterState", + "KeyedMinState", + "KeyedMaxState", +]; + +/// Whether `type_name` is an exact-state tag whose layout is no longer decoded. +pub fn is_retired_exact(type_name: &str) -> bool { + RETIRED_EXACT_TAGS.contains(&type_name) +} + +/// Every state kind the store can persist, borrowed from a trait object. +enum View<'s> { + Exact(&'s k::exact::ExactAccumulator), + Dd(&'s k::DDSketchAccumulator), + Hll(&'s k::HllSketchAccumulator), + Kll(&'s k::DatasketchesKLLAccumulator), + Cms(&'s k::CountMinSketchAccumulator), + Cs(&'s k::CountSketchAccumulator), + CmsHeap(&'s k::CountMinSketchWithHeapAccumulator), + CsHeap(&'s k::CountSketchWithHeapAccumulator), + Hydra(&'s k::HydraKllSketchAccumulator), + Frequency(&'s WeightedFrequency), + UnivMon(&'s UnivMonAccumulator), + Native(&'s NativeSummaryOutput), +} + +fn view(state: &dyn AggregateCore) -> Option> { + let any = state.as_any(); + macro_rules! try_view { + ($($variant:ident => $ty:ty),* $(,)?) => { + $(if let Some(state) = any.downcast_ref::<$ty>() { + return Some(View::$variant(state)); + })* + }; + } + try_view!( + Exact => k::exact::ExactAccumulator, + Dd => k::DDSketchAccumulator, + Hll => k::HllSketchAccumulator, + Kll => k::DatasketchesKLLAccumulator, + Cms => k::CountMinSketchAccumulator, + Cs => k::CountSketchAccumulator, + CmsHeap => k::CountMinSketchWithHeapAccumulator, + CsHeap => k::CountSketchWithHeapAccumulator, + Hydra => k::HydraKllSketchAccumulator, + Frequency => WeightedFrequency, + UnivMon => UnivMonAccumulator, + Native => NativeSummaryOutput, + ); + None +} + +/// Planner's weighted frequency state is serde-transparent over the sketchlib +/// kernel, which owns the persisted `WeightedFrequencyV1` bytes. +// Planner keeps the kernel and its algorithm private; this round trip is the +// only public access until it exposes them. +pub(crate) fn frequency_kernel( + state: &WeightedFrequency, +) -> Result { + Ok(rmp_serde::from_slice(&rmp_serde::to_vec(state)?)?) +} + +pub(crate) fn frequency_state(bytes: &[u8]) -> Result { + let kernel = asap_sketchlib::WeightedFrequency::from_bytes(bytes) + .map_err(|e| format!("deserialize weighted frequency: {e:?}"))?; + Ok(rmp_serde::from_slice(&rmp_serde::to_vec(&kernel)?)?) +} + +fn exact_aggregation_type(state: &k::exact::ExactAccumulator) -> AggregationType { + match state.family() { + SummaryFamilyType::ExactAggregate(kind, _) => match kind { + ExactKind::Sum => AggregationType::Sum, + ExactKind::Count => AggregationType::Count, + ExactKind::Min => AggregationType::Min, + ExactKind::Max => AggregationType::Max, + ExactKind::Rate => AggregationType::Rate, + ExactKind::Increase => AggregationType::Increase, + other => unreachable!("Planner constructs no exact {other:?} state"), + }, + _ => unreachable!("Planner validates exact families at construction"), + } +} + +/// Backend storage view of a Planner kernel state held as a trait object. +pub trait StoredState { + /// Persisted type tag; [`decode`] selects the byte layout by it. + fn type_name(&self) -> &'static str; + fn get_accumulator_type(&self) -> AggregationType; + fn encode(&self) -> Result, Error>; + /// [`StoredState::encode`] for states admitted by [`check_storable`]. + fn serialize_to_bytes(&self) -> Vec { + self.encode().unwrap_or_default() + } +} + +impl<'a> StoredState for dyn AggregateCore + 'a { + fn type_name(&self) -> &'static str { + match view(self) { + Some(View::Exact(_)) => EXACT_V1, + Some(View::Dd(_)) => "DDSketchAccumulator", + Some(View::Hll(_)) => "HllSketchAccumulator", + Some(View::Kll(_)) => "DatasketchesKLLAccumulator", + Some(View::Cms(_)) => "CountMinSketchAccumulator", + Some(View::Cs(_)) => "CountSketchAccumulator", + Some(View::CmsHeap(_)) => "CountMinSketchWithHeapAccumulator", + Some(View::CsHeap(_)) => "CountSketchWithHeapAccumulator", + Some(View::Hydra(_)) => "HydraKllSketchAccumulator", + Some(View::Frequency(_)) => "WeightedFrequency", + Some(View::UnivMon(_)) => "UnivMonAccumulator", + Some(View::Native(_)) => super::native::NATIVE_OUTPUT_TYPE, + None => "UnstoredPlannerState", + } + } + + fn get_accumulator_type(&self) -> AggregationType { + use AggregationType as T; + match view(self) { + Some(View::Exact(state)) => exact_aggregation_type(state), + Some(View::Dd(_)) => T::DDSketch, + Some(View::Hll(_)) => T::HLL, + Some(View::Kll(_)) => T::DatasketchesKLL, + Some(View::Cms(_)) => T::CountMinSketch, + Some(View::Cs(_)) => T::CountSketch, + Some(View::CmsHeap(_)) => T::CountMinSketchWithHeap, + Some(View::CsHeap(_)) => T::CountSketchWithHeap, + Some(View::Hydra(_)) => T::HydraKLL, + Some(View::Frequency(state)) => match frequency_kernel(state).map(|k| k.algorithm()) { + Ok(FrequencyAlgorithm::CountSketch) => T::CountSketchWithHeap, + _ => T::CountMinSketchWithHeap, + }, + Some(View::UnivMon(_)) => T::UnivMon, + Some(View::Native(state)) => state.kind(), + None => T::MultipleSubpopulation, + } + } + + fn encode(&self) -> Result, Error> { + Ok( + match view(self).ok_or("Planner state has no stored codec")? { + View::Exact(state) => rmp_serde::to_vec_named(state)?, + View::Dd(state) => state.inner.to_msgpack()?, + View::Hll(state) => state.inner.to_msgpack()?, + View::Kll(state) => state.inner.to_msgpack()?, + View::Cms(state) => state.inner.to_msgpack()?, + View::Cs(state) => state.inner.to_msgpack()?, + View::CmsHeap(state) => state.inner.to_msgpack()?, + View::CsHeap(state) => state.inner.to_msgpack()?, + View::Hydra(state) => state.inner.to_msgpack()?, + View::Frequency(state) => frequency_kernel(state)?.to_bytes(), + View::UnivMon(state) => state.to_bytes()?, + View::Native(state) => state.bytes().to_vec(), + }, + ) + } +} + +/// Planner's unkeyed exact Sum, Count, Min or Max state after one observation. +pub fn exact_value(kind: ExactKind, value: f64) -> k::exact::ExactAccumulator { + use planner_types::post_asap::ExactParams as P; + let params = match kind { + ExactKind::Sum => P::Sum, + ExactKind::Count => P::Count, + ExactKind::Min => P::Min, + ExactKind::Max => P::Max, + other => panic!("{other:?} is a counter family, not a scalar value"), + }; + let mut state = + k::exact::ExactAccumulator::new(SummaryFamilyType::ExactAggregate(kind, params), false) + .expect("scalar exact family"); + state.update(None, value, 0); + state +} + +/// PromQL counter evaluation range, when the caller supplies one. +pub fn range_ms(parameters: &HashMap) -> Result, Error> { + match ( + parameters.get("range_start_ms"), + parameters.get("range_end_ms"), + ) { + (Some(start), Some(end)) => Ok(Some((start.parse()?, end.parse()?))), + (None, None) => Ok(None), + _ => Err("counter range requires both range_start_ms and range_end_ms".into()), + } +} + +/// Reject a Planner state before it enters the store when it has no stored +/// codec (for example Planner's UnivMon, whose sketch is private). +pub fn check_storable(state: &dyn AggregateCore) -> Result<(), Error> { + view(state) + .map(|_| ()) + .ok_or_else(|| "Planner state has no stored codec".into()) +} + +/// Decode a persisted `(type_name, bytes)` pair written by +/// [`StoredState::encode`]. +pub fn decode(type_name: &str, bytes: &[u8]) -> Result, Error> { + use super::decoders as d; + Ok(match type_name { + EXACT_V1 => Box::new(decode_exact(bytes)?), + "DDSketchAccumulator" => Box::new(k::DDSketchAccumulator { + inner: d::ddsketch_from_msgpack(bytes)?, + }), + "HllSketchAccumulator" => Box::new(k::HllSketchAccumulator { + inner: d::hll_from_msgpack(bytes)?, + }), + "DatasketchesKLLAccumulator" => Box::new(k::DatasketchesKLLAccumulator { + inner: d::kll_from_msgpack(bytes)?, + }), + "CountMinSketchAccumulator" => Box::new(k::CountMinSketchAccumulator { + inner: d::cms_from_msgpack(bytes)?, + }), + "CountSketchAccumulator" => Box::new(k::CountSketchAccumulator { + inner: d::cs_from_msgpack(bytes)?, + }), + "CountMinSketchWithHeapAccumulator" => Box::new(k::CountMinSketchWithHeapAccumulator { + inner: d::cms_with_heap_from_msgpack(bytes)?, + }), + "CountSketchWithHeapAccumulator" => Box::new(k::CountSketchWithHeapAccumulator { + inner: d::cs_with_heap_from_msgpack(bytes)?, + }), + "HydraKllSketchAccumulator" => Box::new(k::HydraKllSketchAccumulator { + inner: asap_sketchlib::HydraKllSketch::from_msgpack(bytes)?, + }), + "WeightedFrequency" => Box::new(frequency_state(bytes)?), + "UnivMonAccumulator" => Box::new(UnivMonAccumulator::from_bytes(bytes)?), + retired if is_retired_exact(retired) => { + return Err(format!( + "stored format {retired} is retired and no longer decoded; \ + exact state is stored as {EXACT_V1}" + ) + .into()) + } + other => return Err(format!("unknown stored state format {other}").into()), + }) +} + +/// Decode a bare OTLP `SketchEnvelope` attribute payload into the Planner +/// kernel of its sketch family. +pub fn decode_envelope(bytes: &[u8]) -> Result, Error> { + use super::decoders as d; + use asap_sketchlib::proto::sketchlib::{sketch_envelope::SketchState, SketchEnvelope}; + use prost::Message; + let envelope = + SketchEnvelope::decode(bytes).map_err(|e| format!("decode SketchEnvelope: {e}"))?; + Ok(match envelope.sketch_state { + Some(SketchState::Kll(_)) => Box::new(k::DatasketchesKLLAccumulator { + inner: d::kll_from_proto(bytes)?, + }), + Some(SketchState::Ddsketch(_)) => Box::new(k::DDSketchAccumulator { + inner: d::ddsketch_from_proto(bytes)?, + }), + Some(SketchState::Hll(_)) => Box::new(k::HllSketchAccumulator { + inner: d::hll_from_proto(bytes)?, + }), + Some(SketchState::CountMin(_)) => Box::new(k::CountMinSketchAccumulator { + inner: d::cms_from_proto(bytes)?, + }), + Some(SketchState::CountSketch(_)) => Box::new(k::CountSketchAccumulator { + inner: d::cs_from_proto(bytes)?, + }), + Some(other) => { + let family = match other { + SketchState::Univmon(_) => "UnivMon", + SketchState::Hydra(_) => "Hydra", + SketchState::Coco(_) => "CocoSketch", + SketchState::Elastic(_) => "Elastic", + _ => "this", + }; + return Err(format!("SketchEnvelope {family} family has no Planner kernel").into()); + } + None => return Err("SketchEnvelope carries no sketch state".into()), + }) +} + +/// Decode Planner's exact state, rejecting a payload whose population states +/// differ from its declared family. +pub fn decode_exact(bytes: &[u8]) -> Result { + // Planner's exact state is serde-derived without validation; this mirror + // of its persisted shape checks each population against the family. + #[derive(serde::Deserialize)] + struct Shape { + family: SummaryFamilyType, + scalar: Scalar, + keyed: Option>, + } + #[derive(serde::Deserialize)] + enum Scalar { + Sum(serde::de::IgnoredAny), + Count(serde::de::IgnoredAny), + Min(serde::de::IgnoredAny), + Max(serde::de::IgnoredAny), + Counter(serde::de::IgnoredAny), + } + let shape: Shape = rmp_serde::from_slice(bytes)?; + // Planner accepts only matching (kind, params) exact families. + k::exact::ExactAccumulator::new(shape.family.clone(), shape.keyed.is_some())?; + let SummaryFamilyType::ExactAggregate(expected, _) = &shape.family else { + return Err(format!("{:?} is not an exact family", shape.family).into()); + }; + let matches = |scalar: &Scalar| match scalar { + Scalar::Sum(_) => *expected == ExactKind::Sum, + Scalar::Count(_) => *expected == ExactKind::Count, + Scalar::Min(_) => *expected == ExactKind::Min, + Scalar::Max(_) => *expected == ExactKind::Max, + Scalar::Counter(_) => matches!(expected, ExactKind::Rate | ExactKind::Increase), + }; + if !matches(&shape.scalar) + || shape + .keyed + .iter() + .flat_map(HashMap::values) + .any(|s| !matches(s)) + { + return Err("exact payload differs from declared Planner family".into()); + } + Ok(rmp_serde::from_slice(bytes)?) +} + +/// An empty state of the same family and shape, so a window-reset base keeps +/// its configuration. +pub fn empty_like(state: &dyn AggregateCore) -> Result, Error> { + use asap_sketchlib::{ + CountMinSketch, CountMinSketchWithHeap, CountSketch, CountSketchWithHeap, DdSketch, + HllSketch, + }; + Ok( + match view(state).ok_or("Planner state has no stored codec")? { + View::Dd(s) => Box::new(k::DDSketchAccumulator { + inner: DdSketch::new(s.inner.alpha), + }), + View::Hll(s) => Box::new(k::HllSketchAccumulator { + inner: HllSketch::new(s.inner.variant, s.inner.precision), + }), + View::Cms(s) => Box::new(k::CountMinSketchAccumulator { + inner: CountMinSketch::new(s.inner.rows(), s.inner.cols()), + }), + View::Cs(s) => Box::new(k::CountSketchAccumulator { + inner: CountSketch::new(s.inner.rows, s.inner.cols), + }), + View::CmsHeap(s) => Box::new(k::CountMinSketchWithHeapAccumulator { + inner: CountMinSketchWithHeap::new( + s.inner.rows(), + s.inner.cols(), + s.inner.heap_size, + ), + }), + View::CsHeap(s) => Box::new(k::CountSketchWithHeapAccumulator { + inner: CountSketchWithHeap::new(s.inner.rows(), s.inner.cols(), s.inner.heap_size), + }), + View::UnivMon(s) => { + let mut empty = s.clone(); + empty.clear(); + Box::new(empty) + } + _ => return Err("only delta-capable sketch families reset to empty".into()), + }, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Statistic; + use asap_physical_operators::values::Value; + use planner_types::{post_asap::SketchQuery, pre_asap::ColumnRef}; + + /// Stored bytes written by the pre-Planner-kernel backend, one per family. + const GOLDEN: &[(&str, &str)] = &[ + ("DDSketchAccumulator", "93cb3f847ae147ae147bdc01040000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001000000000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001d0c0"), + ("HllSketchAccumulator", "415341507631010201010000013a0000001388b06d657461646174615f76657273696f6e01af686173685f70726f66696c655f6964bc70726f6a656374617361702e787868332e736565646c6973742e7631ae686173685f616c676f726974686dab787868335f36345f313238af736565645f64657269766174696f6eb4736565645f6c6973745f696e6465785f77726170ae696e7075745f656e636f64696e67b470726f6a656374617361702e696e7075742e7631a9736565645f6c697374dc0014cecafe3553cf000000ade3415118ce8cc70208ce2f024b2bce451a3df5ce6a09e667cebb67ae85ce3c6ef372cea54ff53ace510e527fce9b05688cce1f83d9abce5be0cd19cecbbb9d5dce629a292ace9159015ace152fecd8ce67332667ce8eb44a87cedb0c2e0db463616e6f6e6963616c5f736565645f696e64657805a9707265636973696f6e0491c41000000300010101000000010000000000"), + ("DatasketchesKLLAccumulator", "92ccc8dc006641534150763101020600000000280000002ccc84ccb06d657461646174615f76657273696f6e01cca16bccccccc8cca16d08cca96974656d5f74797065cca3663634cc93cc920003cc93cccb4008000000000000cccb3fccf0000000000000cccb4000000000000000cc93cccf560f2acc9b7e7e3ccca80000"), + ("CountMinSketchAccumulator", "939298cb0000000000000000cb4000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb000000000000000098cb0000000000000000cb4000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb00000000000000000208"), + ("CountSketchAccumulator", "9403089398cb0000000000000000cb4000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb000000000000000098cb0000000000000000cbc000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb000000000000000098cb0000000000000000cb4000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb000000000000000090"), + ("CountMinSketchWithHeapAccumulator", "93939298cb0000000000000000cb4008000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb000000000000000098cb4008000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb000000000000000002089192a161cb400800000000000002"), + ("CountSketchWithHeapAccumulator", "93939398cb0000000000000000cb4008000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb000000000000000098cbc008000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb000000000000000098cb0000000000000000cb0000000000000000cbc008000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb0000000000000000cb000000000000000003089192a161cb400800000000000002"), + (EXACT_V1, "83a666616d696c7981ae457861637441676772656761746592a353756da353756da67363616c617281a353756dcb4012000000000000a56b65796564c0"), + (EXACT_V1, "83a666616d696c7981ae457861637441676772656761746592a5436f756e74a5436f756e74a67363616c617281a5436f756e7400a56b657965648181a66c6162656c7391a16181a5436f756e7402"), + (EXACT_V1, "83a666616d696c7981ae457861637441676772656761746592a452617465a452617465a67363616c617281a7436f756e74657286b47374617274696e675f6d6561737572656d656e7481a576616c7565cb4024000000000000b27374617274696e675f74696d657374616d70cd03e8b56c6173745f7365656e5f6d6561737572656d656e7481a576616c7565cb4039000000000000b36c6173745f7365656e5f74696d657374616d70cd07d0ae746f74616c5f696e637265617365cb402e000000000000ac73616d706c655f636f756e7402a56b65796564c0"), + ("UnivMonAccumulator", "4153415076310102100000000155000000898bb06d657461646174615f76657273696f6e01af686173685f70726f66696c655f6964bc70726f6a656374617361702e787868332e736565646c6973742e7631ae686173685f616c676f726974686dab787868335f36345f313238af736565645f64657269766174696f6eb4736565645f6c6973745f696e6465785f77726170ae696e7075745f656e636f64696e67b470726f6a656374617361702e696e7075742e7631a9736565645f6c697374dc0014cecafe3553cf000000ade3415118ce8cc70208ce2f024b2bce451a3df5ce6a09e667cebb67ae85ce3c6ef372cea54ff53ace510e527fce9b05688cce1f83d9abce5be0cd19cecbbb9d5dce629a292ace9159015ace152fecd8ce67332667ce8eb44a87cedb0c2e0daa6c617965725f73697a6502aa736b657463685f726f7703aa736b657463685f636f6c10a9686561705f73697a6504a86b65795f74797065a375363498dc0060000000000000ff00000000000000010000000000000001000000000000ff00000000000000000000000000000001ff000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000009602020200000092020092cf3ff0000000000000cf400000000000000092010192c3c30201"), + ("WeightedFrequency", "415341502d57465245512d31000000000008000000000000000200000000000000020000000000000010000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000e03f0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000e03f01000000000000000100000000000000040000000100000000000000611500000000000000010000000000000004000000010000000000000061000000000000e03f"), + ]; + + fn golden(tag: &str, hex_bytes: &str) -> (Box, Vec) { + let bytes = hex::decode(hex_bytes).unwrap(); + (decode(tag, &bytes).unwrap(), bytes) + } + fn item() -> Option { + Some(KeyByLabelValues::new_with_labels(vec!["a".into()])) + } + /// Read one statistic through the Planner kernel's own readout. + fn read(state: &dyn AggregateCore, statistic: Statistic, key: Option) -> f64 { + let any = state.as_any(); + if let Some(exact) = any.downcast_ref::() { + return exact + .readout(statistic, None, key.as_ref()) + .unwrap() + .unwrap(); + } + if let Some(key) = key { + let keyed = [ + any.downcast_ref::() + .map(|s| s.query_key(&key)), + any.downcast_ref::() + .map(|s| s.query_key(&key)), + any.downcast_ref::() + .map(|s| s.query_key(&key)), + any.downcast_ref::() + .map(|s| s.query_key(&key)), + ]; + return keyed.into_iter().flatten().next().expect("keyed family"); + } + let query = match statistic { + Statistic::Quantile => SketchQuery::Quantile { q: 0.5 }, + Statistic::Cardinality => SketchQuery::Cardinality, + _ => SketchQuery::PointCount { + key: ColumnRef::SampleValue, + value: None, + }, + }; + state.estimate(&query).unwrap() + } + + // Every family's existing stored bytes decode under their tag and + // re-encode byte-identically. + #[test] + fn golden_bytes_decode_and_reencode_identically() { + for (tag, bytes) in GOLDEN { + let (state, bytes) = golden(tag, bytes); + assert_eq!(state.type_name(), *tag); + assert_eq!(state.encode().unwrap(), bytes, "{tag}"); + } + } + + // Decoded golden states answer the readouts they were written with. + #[test] + fn golden_states_keep_their_readouts() { + let state = |i: usize| golden(GOLDEN[i].0, GOLDEN[i].1).0; + assert_eq!(read(state(0).as_ref(), Statistic::Count, None), 4.0); + assert!((read(state(1).as_ref(), Statistic::Cardinality, None) - 5.0).abs() < 1.0); + assert_eq!(read(state(2).as_ref(), Statistic::Quantile, None), 2.0); + assert_eq!(read(state(3).as_ref(), Statistic::Count, item()), 2.0); + assert_eq!(read(state(4).as_ref(), Statistic::Count, item()), 2.0); + assert_eq!(read(state(5).as_ref(), Statistic::Count, item()), 3.0); + assert_eq!(read(state(6).as_ref(), Statistic::Count, item()), 3.0); + assert_eq!(read(state(7).as_ref(), Statistic::Sum, None), 4.5); + assert_eq!(read(state(8).as_ref(), Statistic::Count, item()), 2.0); + assert_eq!(read(state(9).as_ref(), Statistic::Rate, None), 15.0); + assert_eq!(read(state(10).as_ref(), Statistic::Count, None), 2.0); + let frequency = state(11); + let rows = frequency + .as_any() + .downcast_ref::() + .unwrap() + .rows(1); + assert!(matches!( + &rows[..], + [row] if matches!(&row[..], [Value::Utf8(item), Value::Float64(score)] + if item.as_ref() == "a" && *score == 0.5) + )); + assert_eq!( + frequency.get_accumulator_type(), + AggregationType::CountMinSketchWithHeap + ); + } + + // Retired exact layouts and unknown tags fail with a named rejection. + #[test] + fn retired_and_unknown_formats_are_rejected_by_name() { + let sum_v0 = 4.5f64.to_le_bytes(); + for tag in RETIRED_EXACT_TAGS { + let error = decode(tag, &sum_v0).err().unwrap().to_string(); + assert!(error.contains("retired") && error.contains(tag), "{error}"); + } + assert!(decode("SketchEnvelopeAccumulator", &[]) + .err() + .unwrap() + .to_string() + .contains("unknown stored state format")); + } + + // An exact payload whose population differs from its family is rejected. + #[test] + fn exact_payload_must_match_its_family() { + let (sum, mut bytes) = golden(GOLDEN[7].0, GOLDEN[7].1); + assert!(sum.as_any().is::()); + // Relabel the family as Min while keeping its Sum population. + let at = bytes.windows(3).position(|w| w == b"Sum").unwrap(); + bytes.splice(at..at + 3, b"Min".iter().copied()); + assert!(decode_exact(&bytes).is_err()); + } + + // A state without a stored codec cannot enter the store. + #[test] + fn planner_univmon_has_no_stored_codec() { + let planner = k::univmon::UnivMonAccumulator::new(4, 3, 16, 2).unwrap(); + assert!(check_storable(&planner).is_err()); + assert!((&planner as &dyn AggregateCore).encode().is_err()); + assert!(check_storable(&UnivMonAccumulator::new(4, 3, 16, 2).unwrap()).is_ok()); + } + + // An OTLP envelope attribute decodes into its family's Planner kernel; + // families without one are rejected. + #[test] + fn envelope_payload_decodes_into_its_planner_kernel() { + use asap_sketchlib::proto::sketchlib::{sketch_envelope::SketchState, SketchEnvelope}; + use prost::Message; + let mut dd = asap_sketchlib::DdSketch::new(0.01); + dd.update(3.0); + let state = decode_envelope(&asap_sketch_codec::encode_ddsketch(&dd)).unwrap(); + assert_eq!(state.get_accumulator_type(), AggregationType::DDSketch); + assert_eq!(read(state.as_ref(), Statistic::Count, None), 1.0); + let univmon = SketchEnvelope { + sketch_state: Some(SketchState::Univmon(Default::default())), + ..Default::default() + }; + assert!(decode_envelope(&univmon.encode_to_vec()).is_err()); + } + + // Window reset keeps each delta family's shape and drops its counts. + #[test] + fn empty_like_keeps_shape() { + let (dd, _) = golden(GOLDEN[0].0, GOLDEN[0].1); + let empty = empty_like(dd.as_ref()).unwrap(); + let empty = empty + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(empty.inner.total_count(), 0); + assert_eq!(empty.inner.alpha, 0.01); + let (cms, _) = golden(GOLDEN[3].0, GOLDEN[3].1); + let empty = empty_like(cms.as_ref()).unwrap(); + assert_eq!(read(empty.as_ref(), Statistic::Count, item()), 0.0); + let (kll, _) = golden(GOLDEN[2].0, GOLDEN[2].1); + assert!(empty_like(kll.as_ref()).is_err()); + } +} diff --git a/crates/asap_summary_state/src/stored_state/decoders.rs b/crates/asap_summary_state/src/stored_state/decoders.rs index 37928f8f4..53de7893c 100644 --- a/crates/asap_summary_state/src/stored_state/decoders.rs +++ b/crates/asap_summary_state/src/stored_state/decoders.rs @@ -1,237 +1,369 @@ -//! Shared sketch state reconstruction and decoding. -use asap_sketchlib::CountMinSketch; -use asap_sketchlib::CountMinSketchDelta; -use asap_sketchlib::CountMinSketchWithHeap; -use asap_sketchlib::CountSketch; -use asap_sketchlib::CountSketchDelta; -use asap_sketchlib::CountSketchWithHeap; -use asap_sketchlib::CsHeapItem; -use asap_sketchlib::MessagePackCodec; - -use crate::summary_kernels::count_min_sketch_with_heap::CountMinSketchWithHeapAccumulator; - -/// Decode a `CountMinSketch` from the modified-OTLP wire bytes. -/// MSGPACK path round-trips `CountMinSketch::deserialize_msgpack`; -/// PROTO path decodes a `SketchEnvelope{count_min: CountMinState}` -/// (or bare `CountMinState`) and re-projects to a flat matrix. Mirrors -/// `precompute_operators::count_min_sketch::from_sketchlib_proto_bytes`. -pub fn decode_cms_from_proto(buffer: &[u8]) -> Result { - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, CountMinState, CounterType, SketchEnvelope, - }; - use prost::Message; - - let state = match SketchEnvelope::decode(buffer) { - Ok(env) => match env.sketch_state { - Some(sketch_envelope::SketchState::CountMin(st)) => st, - Some(_) => return Err("SketchEnvelope contains non-CountMin sketch".to_string()), - None => { - CountMinState::decode(buffer).map_err(|e| format!("decode CountMinState: {e}"))? - } - }, - Err(_) => { - CountMinState::decode(buffer).map_err(|e| format!("decode CountMinState: {e}"))? - } - }; - let rows = state.rows as usize; - let cols = state.cols as usize; - if rows == 0 || cols == 0 { +//! Edge wire formats: sketchlib protobuf envelopes, msgpack frames and delta +//! frames, decoded into the sketchlib states that Planner kernels wrap. +//! +//! These are formats, not algorithms: every update, merge and estimate is done +//! by the sketchlib types (through Planner kernels) after decoding. +use asap_sketchlib::proto::sketchlib::{sketch_envelope, CounterType, SketchEnvelope}; +use asap_sketchlib::{ + CmsHeapItem, CountMinSketch, CountMinSketchDelta, CountMinSketchWithHeap, CountSketch, + CountSketchDelta, CountSketchWithHeap, CsHeapItem, DdSketch, DdSketchDelta, HllSketch, + HllVariant, KllSketch, MessagePackCodec, +}; +use prost::Message; + +/// Planner kernels carry no edge sampling probability, so a sampled frame +/// (`0 < sample_p < 1`) would silently read as unscaled counts. `0` (proto3 +/// default) and `1` both mean unsampled. +fn unsampled(what: &str, sample_p: f64) -> Result<(), String> { + if sample_p.is_finite() && sample_p > 0.0 && sample_p < 1.0 { return Err(format!( - "CountMinState has zero dims (rows={rows}, cols={cols})" + "{what} frame is edge-sampled (sample_p={sample_p}); sampled sketch \ + state is not supported by Planner kernels" )); } - let expected_len = rows * cols; - let counter_type = CounterType::try_from(state.counter_type) - .map_err(|_| format!("CountMinState unknown counter_type {}", state.counter_type))?; - let flat: Vec = match counter_type { - CounterType::Int32 | CounterType::Int64 => { - if state.counts_int.len() != expected_len { - return Err(format!( - "CountMinState counts_int has {} entries, expected {}", - state.counts_int.len(), - expected_len - )); - } - state.counts_int.iter().map(|&v| v as f64).collect() - } - CounterType::Float64 => { - if state.counts_float.len() != expected_len { - return Err(format!( - "CountMinState counts_float has {} entries, expected {}", - state.counts_float.len(), - expected_len - )); - } - state.counts_float.clone() - } - other => { - return Err(format!( - "CountMinState counter_type {other:?} not yet supported in reducer" - )); - } + Ok(()) +} + +/// Envelope-wrapped state, or the bare state for producers that omit the +/// envelope. Returns the envelope's `sample_p` (0 when absent). +fn enveloped( + buffer: &[u8], + what: &str, + pick: impl FnOnce(sketch_envelope::SketchState) -> Option, +) -> Result<(T, f64), String> { + let bare = || { + T::decode(buffer) + .map(|state| (state, 0.0)) + .map_err(|e| format!("decode {what}: {e}")) }; - let mut matrix = Vec::with_capacity(rows); - for r in 0..rows { - let start = r * cols; - matrix.push(flat[start..start + cols].to_vec()); + match SketchEnvelope::decode(buffer) { + Ok(SketchEnvelope { + sketch_state: Some(state), + sample_p, + .. + }) => pick(state) + .map(|state| (state, sample_p)) + .ok_or_else(|| format!("SketchEnvelope does not contain a {what}")), + _ => bare(), } - Ok(CountMinSketch::from_legacy_matrix(matrix, rows, cols)) } -/// Decode a `CountMinSketch` from msgpack bytes (sketch-core wire -/// format). Mirrors -/// `CountMinSketchAccumulator::from_msgpack_bytes`. -pub fn decode_cms_from_msgpack(buffer: &[u8]) -> Result { - CountMinSketch::from_msgpack(buffer) - .map_err(|e| format!("deserialize CountMinSketch msgpack: {e}")) +/// Whether `buffer` is a `SketchEnvelope` carrying a sketch state, as opposed +/// to a bare delta message. +pub fn carries_sketch_state(buffer: &[u8]) -> bool { + SketchEnvelope::decode(buffer).is_ok_and(|envelope| envelope.sketch_state.is_some()) } -/// Decode a `CountSketch` from the modified-OTLP proto wire bytes. -/// Mirrors -/// `precompute_operators::count_sketch::from_sketchlib_proto_bytes`. -pub fn decode_cs_from_proto(buffer: &[u8]) -> Result { - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, CountSketchState, CounterType, SketchEnvelope, - }; - use prost::Message; - - let state = match SketchEnvelope::decode(buffer) { - Ok(env) => match env.sketch_state { - Some(sketch_envelope::SketchState::CountSketch(st)) => st, - Some(_) => return Err("SketchEnvelope contains non-CountSketch sketch".to_string()), - None => CountSketchState::decode(buffer) - .map_err(|e| format!("decode CountSketchState: {e}"))?, - }, - Err(_) => { - CountSketchState::decode(buffer).map_err(|e| format!("decode CountSketchState: {e}"))? - } - }; - let rows = state.rows as usize; - let cols = state.cols as usize; - if rows == 0 || cols == 0 { +// --------------------------------------------------------------------------- +// DDSketch +// --------------------------------------------------------------------------- + +/// A `SketchEnvelope{DdSketchState}` frame. Bare states are rejected. +pub fn ddsketch_from_proto(buffer: &[u8]) -> Result { + let (state, sample_p) = asap_sketch_codec::ddsketch_state(buffer)?; + unsampled("DDSketch", sample_p)?; + if !(state.alpha > 0.0 && state.alpha < 1.0) { return Err(format!( - "CountSketchState has zero dims (rows={rows}, cols={cols})" + "DDSketchState alpha {} out of range (expected 0 < alpha < 1)", + state.alpha )); } - let expected_len = rows * cols; - let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { + Ok(DdSketch::from_proto(state)) +} + +pub fn ddsketch_from_msgpack(buffer: &[u8]) -> Result { + DdSketch::from_msgpack(buffer).map_err(|e| format!("deserialize DdSketch msgpack: {e}")) +} + +/// Add a `DDSketchDelta` bucket-count frame onto `sketch`. +pub fn apply_ddsketch_proto_delta(sketch: &mut DdSketch, buffer: &[u8]) -> Result<(), String> { + use asap_sketchlib::proto::sketchlib::DdSketchDelta as PbDelta; + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode DDSketchDelta: {e}"))?; + let delta = DdSketchDelta { + buckets: pb + .buckets + .into_iter() + .map(|b| (b.index, b.d_count)) + .collect(), + negative_buckets: pb + .negative_buckets + .into_iter() + .map(|b| (b.index, b.d_count)) + .collect(), + zero_count: pb.zero_count, + ..Default::default() + }; + sketch + .apply_delta(&delta) + .map_err(|error| format!("apply DDSketchDelta: {error}")) +} + +// --------------------------------------------------------------------------- +// KLL +// --------------------------------------------------------------------------- + +/// A `SketchEnvelope{KllState}` frame, reconstructed bit-exactly from its +/// level layout (or by replay when the producer sent no levels). +pub fn kll_from_proto(buffer: &[u8]) -> Result { + let state = asap_sketch_codec::kll_state(buffer)?; + if state.k < 8 { + return Err(format!("KllState.k must be >= 8 (got {})", state.k)); + } + let k = u16::try_from(state.k).map_err(|_| { format!( - "CountSketchState unknown counter_type {}", - state.counter_type + "KllState.k does not fit in u16 (got {}, max {})", + state.k, + u16::MAX ) })?; - let flat: Vec = match counter_type { - CounterType::Int32 | CounterType::Int64 => { - if state.counts_int.len() != expected_len { - return Err(format!( - "CountSketchState counts_int has {} entries, expected {}", - state.counts_int.len(), - expected_len - )); - } - state.counts_int.iter().map(|&v| v as f64).collect() + if state.levels.is_empty() { + let mut sketch = KllSketch::new(k); + for item in &state.items { + sketch.update(*item); } - CounterType::Float64 => { - if state.counts_float.len() != expected_len { - return Err(format!( - "CountSketchState counts_float has {} entries, expected {}", - state.counts_float.len(), - expected_len - )); - } - state.counts_float.clone() + return Ok(sketch); + } + if state.levels.len() as u32 != state.num_levels + 1 { + return Err(format!( + "KllState levels length = {}, expected num_levels+1 = {}", + state.levels.len(), + state.num_levels + 1 + )); + } + if state.levels[0] != 0 { + return Err(format!( + "KllState.levels[0] = {}, expected 0", + state.levels[0] + )); + } + if *state.levels.last().unwrap() as usize != state.items.len() { + return Err(format!( + "KllState.levels[{}] = {}, expected items.len() = {}", + state.num_levels, + state.levels.last().unwrap(), + state.items.len() + )); + } + if state + .levels + .windows(2) + .any(|bounds| bounds[0] > bounds[1] || bounds[1] as usize > state.items.len()) + { + return Err("KllState levels must be monotonic and within items".into()); + } + // KllState is highest level first; the in-memory constructor expects L0 + // first. Copying the wire order changes retained-item weights. + let mut items = Vec::with_capacity(state.items.len()); + let mut levels = vec![0]; + for bounds in state.levels.windows(2).rev() { + items.extend_from_slice(&state.items[bounds[0] as usize..bounds[1] as usize]); + levels.push(items.len()); + } + KllSketch::from_portable_state(k, &items, &levels, state.num_levels as usize) + .map_err(|e| e.to_string()) +} + +pub fn kll_from_msgpack(buffer: &[u8]) -> Result { + KllSketch::from_msgpack(buffer).map_err(|e| format!("deserialize KllSketch msgpack: {e}")) +} + +// --------------------------------------------------------------------------- +// HLL +// --------------------------------------------------------------------------- + +/// Decode one protobuf base-128 varint from the front of `buf`, returning +/// `(value, bytes_consumed)`. +fn read_uvarint(buf: &[u8]) -> Option<(u64, usize)> { + let mut result: u64 = 0; + for (i, &b) in buf.iter().enumerate() { + let shift = 7 * i as u32; + if shift >= 64 { + return None; } - other => { + result |= u64::from(b & 0x7f) << shift; + if b & 0x80 == 0 { + return Some((result, i + 1)); + } + } + None +} + +/// Expand sketchlib-go's sparse register form (varint `(index_delta, value)` +/// pairs in ascending index order) into the dense register array. +fn expand_sparse_hll_registers(packed: &[u8], num_registers: usize) -> Result, String> { + let mut regs = vec![0u8; num_registers]; + let (mut prev, mut pos) = (0u64, 0usize); + while pos < packed.len() { + let (delta, n1) = read_uvarint(&packed[pos..]) + .ok_or("HLLSparseRegisters.packed: truncated index_delta varint")?; + pos += n1; + let (value, n2) = read_uvarint(&packed[pos..]) + .ok_or("HLLSparseRegisters.packed: truncated value varint")?; + pos += n2; + let idx = prev + delta; + let i = usize::try_from(idx) + .map_err(|_| format!("HLLSparseRegisters: index {idx} overflows usize"))?; + if i >= num_registers { return Err(format!( - "CountSketchState counter_type {other:?} not yet supported in reducer" + "HLLSparseRegisters: register index {i} >= num_registers {num_registers}" )); } - }; - let mut matrix = Vec::with_capacity(rows); - for r in 0..rows { - let start = r * cols; - matrix.push(flat[start..start + cols].to_vec()); + regs[i] = u8::try_from(value) + .map_err(|_| format!("HLLSparseRegisters: register value {value} > 255"))?; + prev = idx; } - Ok(CountSketch::from_legacy_matrix(matrix, rows, cols)) + Ok(regs) } -/// Decode a `CountSketch` from msgpack bytes (sketch-core wire format). -pub fn decode_cs_from_msgpack(buffer: &[u8]) -> Result { - CountSketch::from_msgpack(buffer).map_err(|e| format!("deserialize CountSketch msgpack: {e}")) +/// A `SketchEnvelope{HyperLogLogState}` (or bare state) frame with dense, +/// sparse or absent (all-zero) registers. +pub fn hll_from_proto(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::{HllVariant as ProtoVariant, HyperLogLogState}; + let (state, sample_p) = + enveloped::(buffer, "HyperLogLogState", |s| match s { + sketch_envelope::SketchState::Hll(state) => Some(state), + _ => None, + })?; + unsampled("HLL", sample_p)?; + if state.precision == 0 || state.precision > 20 { + return Err(format!( + "HyperLogLogState precision {} out of range (expected 1..=20)", + state.precision + )); + } + let expected_len = 1usize << state.precision; + let registers = if state.registers.len() == expected_len { + state.registers.clone() + } else if !state.registers.is_empty() { + return Err(format!( + "HyperLogLogState registers has {} bytes, expected 2^precision = {}", + state.registers.len(), + expected_len + )); + } else if let Some(sparse) = state.registers_sparse.as_ref() { + expand_sparse_hll_registers(&sparse.packed, expected_len)? + } else { + vec![0u8; expected_len] + }; + let variant = match ProtoVariant::try_from(state.variant) + .map_err(|_| format!("HyperLogLogState has unknown variant tag {}", state.variant))? + { + ProtoVariant::Unspecified => HllVariant::Unspecified, + ProtoVariant::Regular => HllVariant::Regular, + ProtoVariant::ErtlMle => HllVariant::Datafusion, + ProtoVariant::Hip => HllVariant::Hip, + }; + Ok(HllSketch::from_raw( + variant, + state.precision, + registers, + state.hip_kxq0, + state.hip_kxq1, + state.hip_est, + )) } -/// Decode a `CountMinSketchWithHeap` from msgpack bytes — the OTLP -/// `CountMinSketch` wire bytes when the gateway/precompute layer -/// marked the sid as CmsWithHeap (heap embedded in the -/// `CountMinSketchWithHeapSerialized` outer wrapper). Delegates to -/// `asap_sketchlib::CountMinSketchWithHeap::deserialize_msgpack`. -pub fn decode_cms_with_heap_from_msgpack(buffer: &[u8]) -> Result { - CountMinSketchWithHeap::from_msgpack(buffer) - .map_err(|e| format!("deserialize CountMinSketchWithHeap msgpack: {e}")) +pub fn hll_from_msgpack(buffer: &[u8]) -> Result { + HllSketch::from_msgpack(buffer).map_err(|e| format!("deserialize HllSketch msgpack: {e}")) } -/// Decode a `CountSketchWithHeap` (median-estimator, Count Sketch family) -/// from msgpack bytes. Distinct wire type from `CountMinSketchWithHeap` -/// (min-estimator, Count-Min family) even though both are heap-bearing -/// frequency sketches — see `asap_sketchlib::CountSketchWithHeap`. -/// Delegates to `asap_sketchlib::CountSketchWithHeap::from_msgpack`. -pub fn decode_cs_with_heap_from_msgpack(buffer: &[u8]) -> Result { - CountSketchWithHeap::from_msgpack(buffer) - .map_err(|e| format!("deserialize CountSketchWithHeap msgpack: {e}")) +/// Apply an `HLLDelta` register frame (register-wise max) onto `sketch`. +pub fn apply_hll_proto_delta(sketch: &mut HllSketch, buffer: &[u8]) -> Result<(), String> { + sketch + .apply_delta_bytes(buffer) + .map_err(|e| format!("apply HLLDelta: {e}")) } // --------------------------------------------------------------------------- -// Delta decoders. Under the per-window-reset (PWR) contract -// (`asap-precompute-go/window.go`: a delta is that window's own state -// applied onto a freshly-reset per-series sketch), each stored *Delta -// frame reconstructs into the FULL window state when applied onto an -// EMPTY base of the frame's declared dimensions. The reducer's -// `FrequencyEstimate` / `FrequencyTopk` paths are per-window evaluations, -// so "empty + apply(this window's delta)" yields exactly the window's -// matrix/heap — no cross-window stitching needed (mirrors how the ingest -// accumulators reset_to_empty per window before applying). -// -// The proto path reuses the PUBLIC `asap_sketchlib::{CountSketch, -// CountMinSketch}::apply_delta`; the proto `*Delta` message is decoded via -// `asap_sketchlib::proto::sketchlib::{CountSketchDelta, CountMinDelta}`, -// exactly as `precompute_operators::{count_sketch, -// count_min_sketch}::apply_proto_delta_bytes` does. +// Count-Min and Count Sketch matrices // --------------------------------------------------------------------------- -/// Decode a `CountMinSketch` PROTO_DELTA frame into a FULL sketch by -/// applying the sparse cell delta onto an empty base of the frame's -/// declared dimensions. Mirrors -/// `precompute_operators::count_min_sketch::apply_proto_delta_bytes`. -pub fn decode_cms_from_proto_delta(buffer: &[u8]) -> Result { - use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; - use prost::Message; +/// Upper bound on wire-declared matrix cells, so a malformed frame cannot +/// force a huge allocation before its counts are checked. +const MAX_SKETCH_CELLS: usize = 8 * 1024 * 1024; - let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; - if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { +/// Reject degenerate, oversized, or packed-hash-incompatible dimensions: the +/// wire hasher reads `ceil(log2(cols))` bits per row from one 64-bit word. +fn validate_sketch_dims(what: &str, rows: usize, cols: usize) -> Result<(), String> { + if rows < 1 || cols < 1 { return Err(format!( - "CountMinDelta packed-array length mismatch: cell_rows={}, cell_cols={}, d_counts={}", - pb.cell_rows.len(), - pb.cell_cols.len(), - pb.d_counts.len() + "{what} has degenerate dims (rows={rows}, cols={cols}); rejecting" )); } - let rows = pb.rows as usize; - let cols = pb.cols as usize; - if rows == 0 || cols == 0 { + let mask_bits = cols.ilog2() as usize + usize::from(!cols.is_power_of_two()); + if rows.saturating_mul(mask_bits) > 64 { return Err(format!( - "CountMinDelta has zero dims (rows={rows}, cols={cols})" + "{what} dims (rows={rows}, cols={cols}) exceed the 64-bit packed-hash \ + column budget (rows * ceil(log2(cols)) = {} > 64)", + rows.saturating_mul(mask_bits) )); } - let cells = pb - .cell_rows - .iter() - .zip(pb.cell_cols.iter()) - .zip(pb.d_counts.iter()) - .map(|((r, c), dc)| (*r, *c, *dc)) - .collect(); - // hh_keys is parsed off the wire by the precompute accumulator but - // intentionally dropped (the vendored Go proto bindings don't yet - // populate it); match that to keep behavior identical. + if rows.saturating_mul(cols) > MAX_SKETCH_CELLS { + return Err(format!( + "{what} dims (rows={rows}, cols={cols}) exceed the {MAX_SKETCH_CELLS}-cell cap" + )); + } + Ok(()) +} + +/// Reshape a flat row-major counter array into the legacy matrix layout. +fn matrix( + what: &str, + rows: u32, + cols: u32, + counter_type: i32, + counts_int: &[i64], + counts_float: &[f64], +) -> Result<(Vec>, usize, usize), String> { + let (rows, cols) = (rows as usize, cols as usize); + validate_sketch_dims(what, rows, cols)?; + let expected = rows * cols; + let flat: Vec = match CounterType::try_from(counter_type) + .map_err(|_| format!("{what} has unknown counter_type tag {counter_type}"))? + { + CounterType::Int32 | CounterType::Int64 if counts_int.len() == expected => { + counts_int.iter().map(|&v| v as f64).collect() + } + CounterType::Float64 if counts_float.len() == expected => counts_float.to_vec(), + CounterType::Int32 | CounterType::Int64 | CounterType::Float64 => { + return Err(format!("{what} counts do not match rows*cols = {expected}")) + } + other => return Err(format!("{what} counter_type {other:?} is not supported")), + }; + Ok((flat.chunks(cols).map(<[f64]>::to_vec).collect(), rows, cols)) +} + +/// A `SketchEnvelope{CountMinState}` (or bare state) frame. +pub fn cms_from_proto(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::CountMinState; + let (state, sample_p) = enveloped::(buffer, "CountMinState", |s| match s { + sketch_envelope::SketchState::CountMin(state) => Some(state), + _ => None, + })?; + unsampled("CountMin", sample_p)?; + let (m, rows, cols) = matrix( + "CountMinState", + state.rows, + state.cols, + state.counter_type, + &state.counts_int, + &state.counts_float, + )?; + Ok(CountMinSketch::from_legacy_matrix(m, rows, cols)) +} + +pub fn cms_from_msgpack(buffer: &[u8]) -> Result { + CountMinSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountMinSketch msgpack: {e}")) +} + +/// Add a `CountMinDelta` cell frame onto `sketch`. +pub fn apply_cms_proto_delta(sketch: &mut CountMinSketch, buffer: &[u8]) -> Result<(), String> { + use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; + let cells = cells("CountMinDelta", &pb.cell_rows, &pb.cell_cols, &pb.d_counts)?; + // The vendored proto bindings do not decode `hh_keys`, and CountMin has + // no heap to rebuild from them. let delta = CountMinSketchDelta { rows: pb.rows, cols: pb.cols, @@ -240,43 +372,57 @@ pub fn decode_cms_from_proto_delta(buffer: &[u8]) -> Result Result { - use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; - use prost::Message; +/// A `CountMinDelta` frame applied onto an empty base of its declared shape +/// (the per-window-reset contract: a window's delta is its whole state). +pub fn cms_from_proto_delta(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; + let (rows, cols) = (pb.rows as usize, pb.cols as usize); + validate_sketch_dims("CountMinDelta", rows, cols)?; + let mut sketch = CountMinSketch::from_legacy_matrix(vec![vec![0.0; cols]; rows], rows, cols); + apply_cms_proto_delta(&mut sketch, buffer)?; + Ok(sketch) +} + +/// A `SketchEnvelope{CountSketchState}` (or bare state) frame. +pub fn cs_from_proto(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::CountSketchState; + let (state, sample_p) = + enveloped::(buffer, "CountSketchState", |s| match s { + sketch_envelope::SketchState::CountSketch(state) => Some(state), + _ => None, + })?; + unsampled("CountSketch", sample_p)?; + let (m, rows, cols) = matrix( + "CountSketchState", + state.rows, + state.cols, + state.counter_type, + &state.counts_int, + &state.counts_float, + )?; + Ok(CountSketch::from_legacy_matrix(m, rows, cols)) +} +pub fn cs_from_msgpack(buffer: &[u8]) -> Result { + CountSketch::from_msgpack(buffer).map_err(|e| format!("deserialize CountSketch msgpack: {e}")) +} + +/// Add a `CountSketchDelta` cell frame onto `sketch`. +pub fn apply_cs_proto_delta(sketch: &mut CountSketch, buffer: &[u8]) -> Result<(), String> { + use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountSketchDelta: {e}"))?; - if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { - return Err(format!( - "CountSketchDelta packed-array length mismatch: cell_rows={}, cell_cols={}, d_counts={}", - pb.cell_rows.len(), - pb.cell_cols.len(), - pb.d_counts.len() - )); - } - let rows = pb.rows as usize; - let cols = pb.cols as usize; - if rows == 0 || cols == 0 { - return Err(format!( - "CountSketchDelta has zero dims (rows={rows}, cols={cols})" - )); - } - let cells = pb - .cell_rows - .iter() - .zip(pb.cell_cols.iter()) - .zip(pb.d_counts.iter()) - .map(|((r, c), dc)| (*r, *c, *dc)) - .collect(); + let cells = cells( + "CountSketchDelta", + &pb.cell_rows, + &pb.cell_cols, + &pb.d_counts, + )?; let delta = CountSketchDelta { rows: pb.rows, cols: pb.cols, @@ -284,83 +430,388 @@ pub fn decode_cs_from_proto_delta(buffer: &[u8]) -> Result l2: pb.l2, hh_keys: Vec::new(), }; - let mut cs = CountSketch::from_legacy_matrix(vec![vec![0.0; cols]; rows], rows, cols); - cs.apply_delta(&delta) - .map_err(|e| format!("apply CountSketchDelta onto empty base: {e}"))?; - Ok(cs) -} - -/// Decode a heap-bearing CountSketch MSGPACK_DELTA frame into a FULL -/// `CountMinSketchWithHeap` by applying the sparse matrix delta + full -/// heap onto an empty base of the frame's declared dimensions. This -/// REUSES the ingest-side delta-heap apply logic -/// (`CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes` → -/// `apply_msgpack_heap_delta_bytes`), which decodes the frame generically -/// with `rmp_serde` — no `asap_sketchlib` delta API is added. -pub fn decode_cms_with_heap_from_msgpack_delta( - buffer: &[u8], -) -> Result { - let acc = CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes(buffer) - .map_err(|e| format!("reconstruct CountMinSketchWithHeap from delta: {e}"))?; - Ok(acc.inner) -} - -/// Decode a heap-bearing CountSketch (median-estimator) MSGPACK_DELTA frame -/// into a FULL `asap_sketchlib::CountSketchWithHeap` by applying the sparse -/// matrix delta + full heap onto an empty base of the frame's declared -/// dimensions. Same DELTA-HEAP wire shape as the CmsWithHeap delta frame -/// (see `HeapDeltaWire`/`MatrixDeltaWire` in -/// `count_min_sketch_with_heap.rs`), decoded here directly -/// with `rmp_serde` since there is no CountSketchWithHeap ingest -/// accumulator to delegate to. No `asap_sketchlib` delta API needed — the -/// public `from_legacy_matrix` rebuilds both the matrix and heap. -pub fn decode_cs_with_heap_from_msgpack_delta( - buffer: &[u8], -) -> Result { - #[derive(serde::Deserialize)] - struct HeapDeltaWire { - is_delta: bool, - matrix_delta: MatrixDeltaWire, - topk_heap: Vec<(String, f64)>, - heap_size: u64, - } - #[derive(serde::Deserialize)] - struct MatrixDeltaWire { - rows: u32, - cols: u32, - cells: Vec<(u32, u32, i64)>, - } + sketch + .apply_delta(&delta) + .map_err(|e| format!("apply CountSketchDelta: {e}")) +} - 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".to_string()); - } - let rows = wire.matrix_delta.rows as usize; - let cols = wire.matrix_delta.cols as usize; - if rows == 0 || cols == 0 { +/// A `CountSketchDelta` frame applied onto an empty base of its declared shape. +pub fn cs_from_proto_delta(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountSketchDelta: {e}"))?; + let (rows, cols) = (pb.rows as usize, pb.cols as usize); + validate_sketch_dims("CountSketchDelta", rows, cols)?; + let mut sketch = CountSketch::from_legacy_matrix(vec![vec![0.0; cols]; rows], rows, cols); + apply_cs_proto_delta(&mut sketch, buffer)?; + Ok(sketch) +} + +fn cells( + what: &str, + rows: &[u32], + cols: &[u32], + d_counts: &[i64], +) -> Result, String> { + if rows.len() != cols.len() || rows.len() != d_counts.len() { return Err(format!( - "CountSketchWithHeap delta frame has zero dims (rows={rows}, cols={cols})" + "{what} packed-array length mismatch: cell_rows={}, cell_cols={}, d_counts={}", + rows.len(), + cols.len(), + d_counts.len() )); } - let mut matrix = vec![vec![0.0; cols]; rows]; - for (r, c, dc) in &wire.matrix_delta.cells { - let (r, c) = (*r as usize, *c as usize); - if r >= rows || c >= cols { - continue; + Ok(rows + .iter() + .zip(cols) + .zip(d_counts) + .map(|((r, c), d)| (*r, *c, *d)) + .collect()) +} + +// --------------------------------------------------------------------------- +// Heap-bearing frequency sketches (legacy integer heaps) +// --------------------------------------------------------------------------- + +pub fn cms_with_heap_from_msgpack(buffer: &[u8]) -> Result { + CountMinSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountMinSketchWithHeap msgpack: {e}")) +} + +pub fn cs_with_heap_from_msgpack(buffer: &[u8]) -> Result { + CountSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketchWithHeap msgpack: {e}")) +} + +/// sketchlib-go's DELTA-HEAP frame (`MSGPACK_DELTA`): a positional array of +/// `[is_delta, [rows, cols, cells], topk_heap, heap_size]`. The matrix part is +/// a sparse signed delta; the heap is the window's full heap. +#[derive(serde::Deserialize)] +struct HeapDeltaWire { + is_delta: bool, + matrix_delta: MatrixDeltaWire, + topk_heap: Vec<(String, f64)>, + heap_size: u64, +} + +#[derive(serde::Deserialize)] +struct MatrixDeltaWire { + rows: u32, + cols: u32, + cells: Vec<(u32, u32, i64)>, +} + +fn heap_delta(buffer: &[u8]) -> Result { + let wire: HeapDeltaWire = + rmp_serde::from_slice(buffer).map_err(|e| format!("decode heap delta msgpack: {e}"))?; + if !wire.is_delta { + return Err("heap delta frame has is_delta=false".into()); + } + Ok(wire) +} + +/// Add a DELTA-HEAP frame's cells to `matrix`; cells outside it are ignored. +fn fold_heap_delta(matrix: &mut [Vec], wire: &HeapDeltaWire) { + for &(r, c, d) in &wire.matrix_delta.cells { + if let Some(cell) = matrix + .get_mut(r as usize) + .and_then(|row| row.get_mut(c as usize)) + { + *cell += d as f64; } - matrix[r][c] += *dc as f64; } - let heap: Vec = wire +} + +/// Apply a DELTA-HEAP frame onto a Count-Min heap state, replacing its heap. +pub fn apply_cms_heap_delta( + sketch: &mut CountMinSketchWithHeap, + buffer: &[u8], +) -> Result<(), String> { + let wire = heap_delta(buffer)?; + let (rows, cols, heap_size) = (sketch.rows(), sketch.cols(), sketch.heap_size); + let mut m = sketch.sketch_matrix(); + fold_heap_delta(&mut m, &wire); + let heap = wire + .topk_heap + .into_iter() + .map(|(key, value)| CmsHeapItem { key, value }) + .collect(); + *sketch = CountMinSketchWithHeap::from_legacy_matrix(m, heap, rows, cols, heap_size); + Ok(()) +} + +/// Apply a DELTA-HEAP frame onto a Count Sketch heap state, replacing its heap. +pub fn apply_cs_heap_delta(sketch: &mut CountSketchWithHeap, buffer: &[u8]) -> Result<(), String> { + let wire = heap_delta(buffer)?; + let (rows, cols, heap_size) = (sketch.rows(), sketch.cols(), sketch.heap_size); + let mut m = sketch.sketch_matrix(); + fold_heap_delta(&mut m, &wire); + let heap = wire .topk_heap .into_iter() .map(|(key, value)| CsHeapItem { key, value }) .collect(); - Ok(CountSketchWithHeap::from_legacy_matrix( - matrix, - heap, - rows, - cols, - wire.heap_size as usize, - )) + *sketch = CountSketchWithHeap::from_legacy_matrix(m, heap, rows, cols, heap_size); + Ok(()) +} + +fn heap_delta_shape(buffer: &[u8]) -> Result<(usize, usize, usize), String> { + let wire = heap_delta(buffer)?; + let (rows, cols) = ( + wire.matrix_delta.rows as usize, + wire.matrix_delta.cols as usize, + ); + if rows == 0 || cols == 0 { + return Err(format!( + "heap delta frame has zero dims (rows={rows}, cols={cols})" + )); + } + Ok((rows, cols, wire.heap_size as usize)) +} + +/// A DELTA-HEAP frame applied onto an empty Count-Min heap of its shape. +pub fn cms_with_heap_from_msgpack_delta(buffer: &[u8]) -> Result { + let (rows, cols, heap_size) = heap_delta_shape(buffer)?; + let mut sketch = CountMinSketchWithHeap::new(rows, cols, heap_size); + apply_cms_heap_delta(&mut sketch, buffer)?; + Ok(sketch) +} + +/// A DELTA-HEAP frame applied onto an empty Count Sketch heap of its shape. +pub fn cs_with_heap_from_msgpack_delta(buffer: &[u8]) -> Result { + let (rows, cols, heap_size) = heap_delta_shape(buffer)?; + let mut sketch = CountSketchWithHeap::new(rows, cols, heap_size); + apply_cs_heap_delta(&mut sketch, buffer)?; + Ok(sketch) +} + +// --------------------------------------------------------------------------- +// Exact Sum payload +// --------------------------------------------------------------------------- + +/// The edge Sum payload: little-endian `f64` sum followed by `u64` count. +pub fn sum_payload(buffer: &[u8]) -> Result { + let sum: [u8; 8] = buffer + .get(..8) + .filter(|_| buffer.len() >= 16) + .and_then(|b| b.try_into().ok()) + .ok_or_else(|| format!("Sum payload too short: {} bytes (want 16)", buffer.len()))?; + Ok(f64::from_le_bytes(sum)) +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_sketchlib::proto::sketchlib as pb; + + fn envelope(state: sketch_envelope::SketchState, sample_p: f64) -> Vec { + SketchEnvelope { + sketch_state: Some(state), + sample_p, + ..Default::default() + } + .encode_to_vec() + } + fn cms_state(rows: u32, cols: u32, counts: Vec) -> pb::CountMinState { + pb::CountMinState { + rows, + cols, + counter_type: CounterType::Int64 as i32, + counts_int: counts, + ..Default::default() + } + } + + // Envelope-wrapped and bare matrix frames decode to the same sketch. + #[test] + fn matrix_frames_decode_with_or_without_envelope() { + let state = cms_state(2, 4, vec![1, 0, 0, 0, 0, 2, 0, 0]); + let bare = cms_from_proto(&state.encode_to_vec()).unwrap(); + let wrapped = cms_from_proto(&envelope( + sketch_envelope::SketchState::CountMin(state), + 0.0, + )) + .unwrap(); + assert_eq!(bare.sketch(), wrapped.sketch()); + assert_eq!(bare.sketch()[1][1], 2.0); + let cs = pb::CountSketchState { + rows: 1, + cols: 2, + counter_type: CounterType::Float64 as i32, + counts_float: vec![-1.5, 2.0], + ..Default::default() + }; + assert_eq!( + cs_from_proto(&cs.encode_to_vec()).unwrap().sketch()[0], + vec![-1.5, 2.0] + ); + } + + // Malformed shapes, wrong families and sampled frames are rejected. + #[test] + fn invalid_or_sampled_frames_are_rejected() { + assert!(cms_from_proto(&cms_state(2, 4, vec![1]).encode_to_vec()).is_err()); + assert!(cms_from_proto(&cms_state(0, 4, vec![]).encode_to_vec()).is_err()); + assert!(cms_from_proto(&cms_state(64, 1 << 20, vec![]).encode_to_vec()).is_err()); + let kll = envelope( + sketch_envelope::SketchState::Kll(pb::KllState::default()), + 0.0, + ); + assert!(cms_from_proto(&kll).is_err()); + let sampled = envelope( + sketch_envelope::SketchState::CountMin(cms_state(1, 2, vec![1, 1])), + 0.5, + ); + assert!(cms_from_proto(&sampled) + .unwrap_err() + .contains("edge-sampled")); + let mut dd = DdSketch::new(0.01); + dd.update(1.0); + let mut dd_state = + asap_sketch_codec::ddsketch_state(&asap_sketch_codec::encode_ddsketch(&dd)) + .unwrap() + .0; + assert!(ddsketch_from_proto(&envelope( + sketch_envelope::SketchState::Ddsketch(dd_state.clone()), + 0.25 + )) + .unwrap_err() + .contains("edge-sampled")); + dd_state.alpha = 0.0; + assert!(ddsketch_from_proto(&envelope( + sketch_envelope::SketchState::Ddsketch(dd_state), + 0.0 + )) + .is_err()); + } + + // Cell deltas add onto the base; a standalone delta starts from empty. + #[test] + fn matrix_deltas_add_cells() { + let delta = pb::CountMinDelta { + rows: 2, + cols: 4, + cell_rows: vec![0, 1], + cell_cols: vec![1, 3], + d_counts: vec![5, 7], + ..Default::default() + } + .encode_to_vec(); + let mut base = + cms_from_proto(&cms_state(2, 4, vec![0, 1, 0, 0, 0, 0, 0, 0]).encode_to_vec()).unwrap(); + apply_cms_proto_delta(&mut base, &delta).unwrap(); + assert_eq!(base.sketch()[0][1], 6.0); + assert_eq!(base.sketch()[1][3], 7.0); + assert_eq!(cms_from_proto_delta(&delta).unwrap().sketch()[0][1], 5.0); + let mismatched = pb::CountSketchDelta { + rows: 1, + cols: 2, + cell_rows: vec![0], + ..Default::default() + } + .encode_to_vec(); + assert!(cs_from_proto_delta(&mismatched).is_err()); + } + + // Dense, sparse and absent HLL registers reconstruct the same register array. + #[test] + fn hll_register_forms_decode_and_delta_takes_register_max() { + let dense = pb::HyperLogLogState { + precision: 4, + variant: pb::HllVariant::Regular as i32, + registers: vec![0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], + ..Default::default() + }; + let sparse = pb::HyperLogLogState { + registers: Vec::new(), + registers_sparse: Some(pb::HllSparseRegisters { + packed: vec![1, 3, 14, 1], + num_registers: 16, + }), + ..dense.clone() + }; + let empty = pb::HyperLogLogState { + registers: Vec::new(), + ..dense.clone() + }; + let mut sketch = hll_from_proto(&dense.encode_to_vec()).unwrap(); + assert_eq!( + sketch.registers, + hll_from_proto(&sparse.encode_to_vec()).unwrap().registers + ); + assert_eq!( + hll_from_proto(&empty.encode_to_vec()).unwrap().registers, + vec![0; 16] + ); + let delta = pb::HllDelta { + packed_updates: vec![0, 2, 1, 1], + }; + apply_hll_proto_delta(&mut sketch, &delta.encode_to_vec()).unwrap(); + assert_eq!(&sketch.registers[..2], &[2, 3]); + } + + // KLL frames reconstruct from their level layout and reject bad layouts. + #[test] + fn kll_frames_follow_their_level_layout() { + let frame = |k, levels: Vec, items: Vec| { + envelope( + sketch_envelope::SketchState::Kll(pb::KllState { + k, + num_levels: levels.len().saturating_sub(1) as u32, + levels, + items, + ..Default::default() + }), + 0.0, + ) + }; + let sketch = kll_from_proto(&frame(200, vec![0, 3], vec![1.0, 2.0, 3.0])).unwrap(); + assert_eq!(sketch.count(), 3); + assert_eq!( + kll_from_proto(&frame(200, vec![], vec![4.0])) + .unwrap() + .count(), + 1 + ); + assert!(kll_from_proto(&frame(4, vec![0, 1], vec![1.0])).is_err()); + assert!(kll_from_proto(&frame(200, vec![0, 2], vec![1.0])).is_err()); + } + + // A heap delta adds its cells and replaces the heap with the frame's heap. + #[test] + fn heap_delta_adds_cells_and_replaces_heap() { + #[derive(serde::Serialize)] + struct Frame( + bool, + (u32, u32, Vec<(u32, u32, i64)>), + Vec<(String, f64)>, + u64, + ); + let frame = rmp_serde::to_vec(&Frame( + true, + (2, 8, vec![(0, 1, 4), (1, 2, 4)]), + vec![("b".into(), 4.0)], + 2, + )) + .unwrap(); + let mut base = CountMinSketchWithHeap::new(2, 8, 2); + base.update("a", 1.0); + let before = base.sketch_matrix(); + apply_cms_heap_delta(&mut base, &frame).unwrap(); + assert_eq!(base.sketch_matrix()[0][1], before[0][1] + 4.0); + let heap = base.topk_heap_items(); + assert_eq!((heap.len(), heap[0].key.as_str()), (1, "b")); + let standalone = cs_with_heap_from_msgpack_delta(&frame).unwrap(); + assert_eq!(standalone.sketch_matrix()[1][2], 4.0); + let not_delta = rmp_serde::to_vec(&Frame(false, (2, 8, vec![]), vec![], 2)).unwrap(); + assert!(cms_with_heap_from_msgpack_delta(¬_delta).is_err()); + } + + // The edge Sum payload is a little-endian f64 sum followed by a count. + #[test] + fn sum_payload_reads_the_sum_and_requires_the_count() { + let mut bytes = 2.5f64.to_le_bytes().to_vec(); + assert!(sum_payload(&bytes).is_err()); + bytes.extend(3u64.to_le_bytes()); + assert_eq!(sum_payload(&bytes).unwrap(), 2.5); + } } diff --git a/crates/asap_summary_state/src/stored_state/delta_apply.rs b/crates/asap_summary_state/src/stored_state/delta_apply.rs index a81e2b4f2..dc1df764d 100644 --- a/crates/asap_summary_state/src/stored_state/delta_apply.rs +++ b/crates/asap_summary_state/src/stored_state/delta_apply.rs @@ -1,30 +1,22 @@ -//! Shared sketch state reconstruction and decoding. -use asap_sketchlib::CountMinSketch; -use asap_sketchlib::CountMinSketchWithHeap; -use asap_sketchlib::CountSketch; -use asap_sketchlib::CountSketchWithHeap; -use asap_sketchlib::DdSketch; -use asap_sketchlib::HllSketch; -use asap_sketchlib::HllVariant; -use asap_sketchlib::KllSketch; -use asap_sketchlib::MessagePackCodec; - -use super::decoders::{ - decode_cms_from_msgpack, decode_cms_from_proto, decode_cms_from_proto_delta, - decode_cms_with_heap_from_msgpack, decode_cms_with_heap_from_msgpack_delta, - decode_cs_from_msgpack, decode_cs_from_proto, decode_cs_from_proto_delta, - decode_cs_with_heap_from_msgpack, decode_cs_with_heap_from_msgpack_delta, +//! Reconstruction of stored window states from full and delta frames. +use asap_physical_operators::summary_kernels as k; +use asap_physical_operators::summary_kernels::weighted_frequency::{ + FrequencyAlgorithm, WeightedFrequency, +}; +use asap_physical_operators::AggregateCore; +use asap_sketchlib::{ + CountMinSketch, CountMinSketchWithHeap, CountSketch, CountSketchWithHeap, DdSketch, HllSketch, + HllVariant, KllSketch, }; + +use super::decoders as d; use super::{SketchEncoding, SketchSampleState}; +use crate::univmon::UnivMonAccumulator; /// Which sketch family a candidate is, and the parameters needed to /// *bootstrap an empty state* — required by the per-window-reset (PWR) /// delta model where a window's FIRST frame is a delta-from-empty (no -/// carry-in Full). Most families' deltas embed their own params in the -/// wire fragment (decoded independently, then merged in — see -/// `SummaryState::apply_delta_bytes`); HLL register deltas and DD's -/// bucket-index deltas are applied onto a pre-sized structure instead, -/// so those two need the params known up front to allocate it. +/// carry-in Full). #[derive(Debug, Clone, Copy)] pub enum DeltaSketchKind { UnivMon { @@ -50,15 +42,9 @@ pub enum DeltaSketchKind { rows: usize, cols: usize, }, - /// `CmsWithHeap` wraps `asap_sketchlib::CountMinSketchWithHeap` - /// (min-over-rows estimator) and `CountSketchWithHeap` wraps the - /// distinct `asap_sketchlib::CountSketchWithHeap` (median-of-signed-rows - /// estimator) -- different algorithms that happen to share a storage - /// shape. Kept as two variants (not one shared `Heap`) so - /// `merge_same_family` rejects merging one into the other the same - /// way it already rejects e.g. merging a `Cms` into a `Kll`; now the - /// type system enforces it too, since the two variants hold different - /// Rust types. + /// Count-Min (min-over-rows) and Count Sketch (median-of-signed-rows) + /// heaps share a storage shape but are different algorithms, so they are + /// separate kinds and never merge into each other. CmsWithHeap { rows: usize, cols: usize, @@ -72,61 +58,62 @@ pub enum DeltaSketchKind { } impl DeltaSketchKind { - /// Construct an EMPTY state for this kind, used to seed a new window - /// when its first frame is a delta-from-empty (PWR). A delta applied - /// onto this empty base reconstructs exactly that window's state - /// (delta-from-empty ⊕ empty = window state). + /// An EMPTY state for this kind, seeding a window whose first frame is a + /// delta-from-empty (PWR): empty ⊕ delta = that window's state. fn bootstrap_empty(&self) -> SummaryState { - match self { + match *self { Self::UnivMon { heap_size, sketch_rows, sketch_cols, layers, } => SummaryState::UnivMon( - crate::summary_kernels::univmon::UnivMonAccumulator::new( - *heap_size as usize, - *sketch_rows as usize, - *sketch_cols as usize, - *layers as usize, + UnivMonAccumulator::new( + heap_size as usize, + sketch_rows as usize, + sketch_cols as usize, + layers as usize, ) .expect("validated UnivMon catalog dimensions"), ), - DeltaSketchKind::DDSketch { alpha } => SummaryState::Dd(DdSketch::new(*alpha)), - DeltaSketchKind::Kll { k } => SummaryState::Kll(KllSketch::new(*k as u16)), - DeltaSketchKind::Hll { precision } => { - SummaryState::Hll(HllSketch::new(HllVariant::Regular, *precision)) + Self::DDSketch { alpha } => SummaryState::Dd(k::DDSketchAccumulator::new(alpha)), + Self::Kll { k } => SummaryState::Kll(k::DatasketchesKLLAccumulator::new(k as u16)), + Self::Hll { precision } => { + SummaryState::Hll(k::HllSketchAccumulator::new(HllVariant::Regular, precision)) } - DeltaSketchKind::Cms { rows, cols } => { - SummaryState::Cms(CountMinSketch::new(*rows, *cols)) + Self::Cms { rows, cols } => { + SummaryState::Cms(k::CountMinSketchAccumulator::new(rows, cols)) } - DeltaSketchKind::CountSketch { rows, cols } => { - SummaryState::CountSketch(CountSketch::new(*rows, *cols)) + Self::CountSketch { rows, cols } => { + SummaryState::CountSketch(k::CountSketchAccumulator::new(rows, cols)) } - DeltaSketchKind::CmsWithHeap { + Self::CmsWithHeap { rows, cols, heap_size, - } => SummaryState::CmsWithHeap(CountMinSketchWithHeap::new(*rows, *cols, *heap_size)), - DeltaSketchKind::CountSketchWithHeap { + } => SummaryState::CmsWithHeap(k::CountMinSketchWithHeapAccumulator::new( + rows, cols, heap_size, + )), + Self::CountSketchWithHeap { rows, cols, heap_size, - } => SummaryState::CountSketchWithHeap(CountSketchWithHeap::new( - *rows, *cols, *heap_size, + } => SummaryState::CountSketchWithHeap(k::CountSketchWithHeapAccumulator::new( + rows, cols, heap_size, )), } } } -/// Try to decode a "full" sketch from the bytes (used by both -/// per-window and cumulative modes when the encoding is `*Full`). +/// Decode a "full" frame (used by both per-window and cumulative modes when +/// the encoding is `*Full`). fn decode_full( kind: &DeltaSketchKind, bytes: &[u8], encoding: SketchEncoding, ) -> Result { - match (kind, encoding) { + use SketchEncoding::{MsgpackFull, ProtoFull, WeightedFrequencyV1}; + Ok(match (kind, encoding) { ( DeltaSketchKind::UnivMon { heap_size, @@ -134,10 +121,9 @@ fn decode_full( sketch_cols, layers, }, - SketchEncoding::MsgpackFull, + MsgpackFull, ) => { - let state = crate::summary_kernels::univmon::UnivMonAccumulator::from_bytes(bytes) - .map_err(|e| e.to_string())?; + let state = UnivMonAccumulator::from_bytes(bytes).map_err(|e| e.to_string())?; if state.dimensions() != ( *heap_size as usize, @@ -148,91 +134,67 @@ fn decode_full( { return Err("UnivMon payload dimensions differ from installed catalog".into()); } - Ok(SummaryState::UnivMon(state)) - } - (DeltaSketchKind::DDSketch { .. }, SketchEncoding::ProtoFull) => { - let sk = dd_from_proto(bytes)?; - Ok(SummaryState::Dd(sk)) - } - (DeltaSketchKind::DDSketch { .. }, SketchEncoding::MsgpackFull) => { - let sk = DdSketch::from_msgpack(bytes) - .map_err(|e| format!("deserialize DDSketch msgpack: {e}"))?; - Ok(SummaryState::Dd(sk)) - } - (DeltaSketchKind::Hll { .. }, SketchEncoding::ProtoFull) => { - let sk = hll_from_proto(bytes)?; - Ok(SummaryState::Hll(sk)) - } - (DeltaSketchKind::Hll { .. }, SketchEncoding::MsgpackFull) => { - let sk = HllSketch::from_msgpack(bytes) - .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?; - Ok(SummaryState::Hll(sk)) - } - (DeltaSketchKind::Kll { .. }, SketchEncoding::ProtoFull) => { - let sk = kll_from_proto(bytes)?; - Ok(SummaryState::Kll(sk)) - } - (DeltaSketchKind::Kll { .. }, SketchEncoding::MsgpackFull) => { - let sk = KllSketch::from_msgpack(bytes) - .map_err(|e| format!("deserialize KllSketch msgpack: {e}"))?; - Ok(SummaryState::Kll(sk)) - } - (DeltaSketchKind::Cms { .. }, SketchEncoding::ProtoFull) => { - Ok(SummaryState::Cms(decode_cms_from_proto(bytes)?)) + SummaryState::UnivMon(state) } - (DeltaSketchKind::Cms { .. }, SketchEncoding::MsgpackFull) => { - Ok(SummaryState::Cms(decode_cms_from_msgpack(bytes)?)) + (DeltaSketchKind::DDSketch { .. }, ProtoFull) => dd(d::ddsketch_from_proto(bytes)?), + (DeltaSketchKind::DDSketch { .. }, MsgpackFull) => dd(d::ddsketch_from_msgpack(bytes)?), + (DeltaSketchKind::Hll { .. }, ProtoFull) => hll(d::hll_from_proto(bytes)?), + (DeltaSketchKind::Hll { .. }, MsgpackFull) => hll(d::hll_from_msgpack(bytes)?), + (DeltaSketchKind::Kll { .. }, ProtoFull) => kll(d::kll_from_proto(bytes)?), + (DeltaSketchKind::Kll { .. }, MsgpackFull) => kll(d::kll_from_msgpack(bytes)?), + (DeltaSketchKind::Cms { .. }, ProtoFull) => cms(d::cms_from_proto(bytes)?), + (DeltaSketchKind::Cms { .. }, MsgpackFull) => cms(d::cms_from_msgpack(bytes)?), + (DeltaSketchKind::CountSketch { .. }, ProtoFull) => cs(d::cs_from_proto(bytes)?), + (DeltaSketchKind::CountSketch { .. }, MsgpackFull) => cs(d::cs_from_msgpack(bytes)?), + // The legacy heap wire format is msgpack-only. + (DeltaSketchKind::CmsWithHeap { .. }, ProtoFull | MsgpackFull) => { + cms_heap_state(d::cms_with_heap_from_msgpack(bytes)?) } - (DeltaSketchKind::CountSketch { .. }, SketchEncoding::ProtoFull) => { - Ok(SummaryState::CountSketch(decode_cs_from_proto(bytes)?)) + (DeltaSketchKind::CountSketchWithHeap { .. }, ProtoFull | MsgpackFull) => { + cs_heap_state(d::cs_with_heap_from_msgpack(bytes)?) } - (DeltaSketchKind::CountSketch { .. }, SketchEncoding::MsgpackFull) => { - Ok(SummaryState::CountSketch(decode_cs_from_msgpack(bytes)?)) - } - // The heap-bearing wire format is msgpack-only in this - // deployment; `decode_cms_with_heap_from_msgpack` is the same - // "Full" decoder the reducer's existing per-frame dispatch falls - // through to for any non-MsgpackDelta encoding. - ( - DeltaSketchKind::CmsWithHeap { .. }, - SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull, - ) => Ok(SummaryState::CmsWithHeap( - decode_cms_with_heap_from_msgpack(bytes)?, - )), - ( - DeltaSketchKind::CountSketchWithHeap { .. }, - SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull, - ) => Ok(SummaryState::CountSketchWithHeap( - decode_cs_with_heap_from_msgpack(bytes)?, - )), ( - DeltaSketchKind::CmsWithHeap { .. } | DeltaSketchKind::CountSketchWithHeap { .. }, - SketchEncoding::WeightedFrequencyV1, + DeltaSketchKind::CmsWithHeap { rows, cols, .. } + | DeltaSketchKind::CountSketchWithHeap { rows, cols, .. }, + WeightedFrequencyV1, ) => { - use crate::summary_kernels::weighted_frequency::FrequencyAlgorithm; - let kernel = asap_sketchlib::WeightedFrequency::from_bytes(bytes) - .map_err(|e| format!("deserialize weighted frequency: {e:?}"))?; - let (expected, rows, cols) = match kind { - DeltaSketchKind::CmsWithHeap { rows, cols, .. } => { - (FrequencyAlgorithm::Cms, *rows, *cols) - } - DeltaSketchKind::CountSketchWithHeap { rows, cols, .. } => { - (FrequencyAlgorithm::CountSketch, *rows, *cols) - } - _ => unreachable!("matched heap kinds"), + let expected = match kind { + DeltaSketchKind::CmsWithHeap { .. } => FrequencyAlgorithm::Cms, + _ => FrequencyAlgorithm::CountSketch, }; + let state = super::codec::frequency_state(bytes).map_err(|e| e.to_string())?; + let kernel = super::codec::frequency_kernel(&state).map_err(|e| e.to_string())?; // The catalog's heap size is not carried here; matrix shape is. let (width, depth, _) = kernel.shape(); - if kernel.algorithm() != expected || (width, depth) != (cols, rows) { + if kernel.algorithm() != expected || (width, depth) != (*cols, *rows) { return Err("weighted frequency shape differs from installed catalog".into()); } - Ok(SummaryState::WeightedFrequency( - rmp_serde::from_slice(&rmp_serde::to_vec(&kernel).map_err(|e| e.to_string())?) - .map_err(|e| e.to_string())?, - )) + SummaryState::WeightedFrequency(state) } - (_, e) => Err(format!("decode_full called with non-Full encoding {e:?}")), - } + (_, e) => return Err(format!("decode_full called with non-Full encoding {e:?}")), + }) +} + +fn dd(inner: DdSketch) -> SummaryState { + SummaryState::Dd(k::DDSketchAccumulator { inner }) +} +fn hll(inner: HllSketch) -> SummaryState { + SummaryState::Hll(k::HllSketchAccumulator { inner }) +} +fn kll(inner: KllSketch) -> SummaryState { + SummaryState::Kll(k::DatasketchesKLLAccumulator { inner }) +} +fn cms(inner: CountMinSketch) -> SummaryState { + SummaryState::Cms(k::CountMinSketchAccumulator { inner }) +} +fn cs(inner: CountSketch) -> SummaryState { + SummaryState::CountSketch(k::CountSketchAccumulator { inner }) +} +fn cms_heap_state(inner: CountMinSketchWithHeap) -> SummaryState { + SummaryState::CmsWithHeap(k::CountMinSketchWithHeapAccumulator { inner }) +} +fn cs_heap_state(inner: CountSketchWithHeap) -> SummaryState { + SummaryState::CountSketchWithHeap(k::CountSketchWithHeapAccumulator { inner }) } /// Render a ranked heap item as the legacy heap key: item parts joined by @@ -294,238 +256,175 @@ fn series_key(labels: &std::collections::BTreeMap) -> String { } } -/// The reconstructed state one candidate sid contributes — either -/// folded across a window (or several) via delta application, or merged -/// in from another sid's own reconstruction. +/// The reconstructed Planner kernel state one candidate sid contributes — +/// either folded across a window (or several) via delta application, or +/// merged in from another sid's own reconstruction. pub enum SummaryState { - UnivMon(crate::summary_kernels::univmon::UnivMonAccumulator), - Dd(DdSketch), - Hll(HllSketch), - Kll(KllSketch), - Cms(CountMinSketch), - CountSketch(CountSketch), - /// See `DeltaSketchKind::CmsWithHeap`/`CountSketchWithHeap` for why - /// these are two variants holding two different sketchlib types. - CmsWithHeap(CountMinSketchWithHeap), - CountSketchWithHeap(CountSketchWithHeap), + UnivMon(UnivMonAccumulator), + Dd(k::DDSketchAccumulator), + Hll(k::HllSketchAccumulator), + Kll(k::DatasketchesKLLAccumulator), + Cms(k::CountMinSketchAccumulator), + CountSketch(k::CountSketchAccumulator), + /// See `DeltaSketchKind::CmsWithHeap` for why the two heaps differ. + CmsWithHeap(k::CountMinSketchWithHeapAccumulator), + CountSketchWithHeap(k::CountSketchWithHeapAccumulator), /// Planner's weighted-frequency heap, read through its ranked rows. - WeightedFrequency(crate::summary_kernels::weighted_frequency::PhysicalWeightedFrequency), + WeightedFrequency(WeightedFrequency), } impl SummaryState { - /// Apply a delta-encoded payload from a window sample. For DD / KLL, - /// the delta is interpreted as a "mergeable fragment" decoded - /// through the same full-state decoder and merged into the - /// rolling state. For HLL, the wire delta is a sparse register - /// update applied via the sketch's `apply_delta`. - /// - /// On encoding mismatch (e.g. trying to apply an HllDelta to a - /// DDSketch rolling state) returns Err. + /// The Planner kernel this state holds. + pub fn kernel(&self) -> &dyn AggregateCore { + match self { + Self::UnivMon(s) => s, + Self::Dd(s) => s, + Self::Hll(s) => s, + Self::Kll(s) => s, + Self::Cms(s) => s, + Self::CountSketch(s) => s, + Self::CmsWithHeap(s) => s, + Self::CountSketchWithHeap(s) => s, + Self::WeightedFrequency(s) => s, + } + } + + /// Rewrap a kernel of the same variant as `self`. + fn same_variant(&self, state: &dyn AggregateCore) -> Option { + let any = state.as_any(); + Some(match self { + Self::UnivMon(_) => Self::UnivMon(any.downcast_ref::()?.clone()), + Self::Dd(_) => Self::Dd(any.downcast_ref::()?.clone()), + Self::Hll(_) => Self::Hll(any.downcast_ref::()?.clone()), + Self::Kll(_) => Self::Kll(any.downcast_ref::()?.clone()), + Self::Cms(_) => Self::Cms(any.downcast_ref::()?.clone()), + Self::CountSketch(_) => { + Self::CountSketch(any.downcast_ref::()?.clone()) + } + Self::CmsWithHeap(_) => Self::CmsWithHeap( + any.downcast_ref::()? + .clone(), + ), + Self::CountSketchWithHeap(_) => Self::CountSketchWithHeap( + any.downcast_ref::()? + .clone(), + ), + Self::WeightedFrequency(_) => { + Self::WeightedFrequency(any.downcast_ref::()?.clone()) + } + }) + } + + /// Apply a delta-encoded frame. DDSketch bucket deltas and HLL register + /// deltas apply in place; every other frame decodes independently and + /// merges in. pub fn apply_delta_bytes( &mut self, bytes: &[u8], encoding: SketchEncoding, ) -> Result<(), String> { - if !matches!( - encoding, - SketchEncoding::ProtoDelta | SketchEncoding::MsgpackDelta - ) { + use SketchEncoding::{MsgpackDelta, MsgpackFull, ProtoDelta, ProtoFull}; + if !matches!(encoding, ProtoDelta | MsgpackDelta) { return Err(format!( "apply_delta_bytes called with non-Delta encoding {encoding:?}" )); } - match self { - SummaryState::UnivMon(_) => Err("UnivMon requires full pane snapshots".into()), - SummaryState::Dd(sk) => { - match encoding { - // PROTO_DELTA: dispatch on the payload SHAPE, mirroring the - // supported DDSketch frame decoder, which tries the - // full-envelope decode first, then falls back to - // the bucket-delta proto. Two wire shapes can arrive on the - // ProtoDelta channel: - // - // 1. `SketchEnvelope{DdSketchState}` — a full-state - // fragment, mergeable via `DdSketch::merge`. (The edge - // sends this when `compute_delta_against` hits the - // empty-current / undecodable-prior fallback and ships - // a full snapshot tagged as a delta.) - // 2. `DDSketchDelta { buckets: [{index, d_count}] }` — a - // bucket-index delta proto, applied additively. This is - // the COMMON delta_transmission frame the edge emits - // under per-window-reset (`compute_delta(&empty)`). - // - // Before this fix the reducer decoded ONLY shape (1) via - // `decode_full`. A real shape-(2) frame failed with a wire- - // type mismatch on field 1 (delta field 1 = repeated - // submessage; state field 1 = `double alpha`) → the whole - // `quantile_over_time` returned `No result` for every - // delta_transmission DDSketch stream. We wrap the rolling - // `DdSketch` in a transient accumulator so the bucket-delta - // apply lands on `sk` in place. - SketchEncoding::ProtoDelta => { - // Shape (1): full envelope fragment → merge. Try this - // first (cheap decode attempt; a bucket-delta proto - // fails it on the field-1 wire-type mismatch). - if let Ok(SummaryState::Dd(other)) = decode_full( - &DeltaSketchKind::DDSketch { alpha: 0.0 }, - bytes, - SketchEncoding::ProtoFull, - ) { - sk.merge(&other) - .map_err(|e| format!("merge DDSketch delta envelope: {e}"))?; - return Ok(()); - } - // Shape (2): bucket-delta proto → additive apply via the - // SAME decoder the ingest delta path uses. - use crate::summary_kernels::dd_sketch::DDSketchAccumulator; - let mut acc = DDSketchAccumulator { - inner: std::mem::replace(sk, DdSketch::new(sk.alpha)), - sample_p: 1.0, - }; - let res = acc.apply_proto_delta_bytes(bytes); - *sk = acc.inner; - res.map_err(|e| format!("apply DDSketch proto bucket-delta: {e}"))?; - Ok(()) - } - // MSGPACK_DELTA: a serialized full-sketch fragment, mergeable - // via the full-state decoder. Kept for completeness — the - // edge wires PROTO_DELTA for DDSketch today. - SketchEncoding::MsgpackDelta => { - let other = match decode_full( - &DeltaSketchKind::DDSketch { alpha: 0.0 }, - bytes, - SketchEncoding::MsgpackFull, - ) { - Ok(SummaryState::Dd(s)) => s, - Ok(_) => { - return Err( - "decode_full(DDSketch) returned non-DDSketch state".to_string() - ) - } - Err(e) => return Err(e), - }; - sk.merge(&other) - .map_err(|e| format!("merge DDSketch delta: {e}"))?; - Ok(()) - } - _ => unreachable!(), - } + let fragment = |kind: DeltaSketchKind, full| decode_full(&kind, bytes, full); + match (&mut *self, encoding) { + (Self::UnivMon(_), _) => Err("UnivMon requires full pane snapshots".into()), + (Self::WeightedFrequency(_), _) => { + Err("weighted frequency frames are complete states, not deltas".into()) } - SummaryState::Hll(sk) => { - // HLL has a true sparse register delta in the proto - // wire format. Use the same path the precompute - // accumulator uses (`apply_proto_delta_bytes`-style). - if encoding == SketchEncoding::ProtoDelta { - apply_hll_proto_delta(sk, bytes) + // Two shapes arrive on the DDSketch proto-delta channel: a full + // envelope fragment (merged) and the common bucket-index delta. + (Self::Dd(sketch), ProtoDelta) => { + if d::carries_sketch_state(bytes) { + let other = fragment(DeltaSketchKind::DDSketch { alpha: 0.0 }, ProtoFull)?; + self.merge_same_family(&other) } else { - // MsgpackDelta for HLL isn't a sparse encoding; - // it's a serialized HllSketch fragment, mergeable - // via `HllSketch::merge`. - let other = HllSketch::from_msgpack(bytes) - .map_err(|e| format!("deserialize HllSketch (delta-as-msgpack): {e}"))?; - sk.merge(&other) - .map_err(|e| format!("merge HLL delta: {e}"))?; - Ok(()) + d::apply_ddsketch_proto_delta(&mut sketch.inner, bytes) + .map_err(|e| format!("apply DDSketch proto bucket-delta: {e}")) } } - SummaryState::Kll(sk) => { - let full_enc = match encoding { - SketchEncoding::ProtoDelta => SketchEncoding::ProtoFull, - SketchEncoding::MsgpackDelta => SketchEncoding::MsgpackFull, - _ => unreachable!(), - }; - let other = match decode_full(&DeltaSketchKind::Kll { k: 0 }, bytes, full_enc) { - Ok(SummaryState::Kll(s)) => s, - Ok(_) => return Err("decode_full(Kll) returned non-Kll state".to_string()), - Err(e) => return Err(e), - }; - sk.merge(&other) - .map_err(|e| format!("merge KLL delta: {e}"))?; - Ok(()) - } - // CMS/CountSketch/Heap have no true sparse in-place delta - // (unlike DD's bucket-index proto or HLL's register proto, - // above) — every delta frame already decodes into a - // complete, standalone state on its own (the PWR wire - // contract resets to empty at the source), so applying one - // is always "decode independently, then merge". - SummaryState::Cms(sk) => { - if encoding != SketchEncoding::ProtoDelta { - return Err( - "CountMin (heap-less) MSGPACK_DELTA is not a valid producer encoding \ - (msgpack-delta is the heap-bearing form)" - .to_string(), - ); - } - let other = decode_cms_from_proto_delta(bytes)?; - sk.merge(&other) - .map_err(|e| format!("merge CountMinSketch delta: {e}")) + (Self::Dd(_), _) => { + let other = fragment(DeltaSketchKind::DDSketch { alpha: 0.0 }, MsgpackFull)?; + self.merge_same_family(&other) } - SummaryState::CountSketch(sk) => { - if encoding != SketchEncoding::ProtoDelta { - return Err( - "CountSketch (heap-less) MSGPACK_DELTA is not a valid producer encoding \ - (msgpack-delta is the heap-bearing form)" - .to_string(), - ); - } - let other = decode_cs_from_proto_delta(bytes)?; - sk.merge(&other) - .map_err(|e| format!("merge CountSketch delta: {e}")) + (Self::Hll(sketch), ProtoDelta) => d::apply_hll_proto_delta(&mut sketch.inner, bytes), + (Self::Hll(_), _) => { + let other = fragment(DeltaSketchKind::Hll { precision: 0 }, MsgpackFull)?; + self.merge_same_family(&other) } - SummaryState::CmsWithHeap(sk) => { - // Matches the existing per-frame reducer dispatch: only - // MsgpackDelta gets true delta treatment; ProtoDelta (not - // produced for this family in this deployment) falls - // through to the full-msgpack decoder, same as `decode_full`. - let other = if encoding == SketchEncoding::MsgpackDelta { - decode_cms_with_heap_from_msgpack_delta(bytes)? + (Self::Kll(_), _) => { + let full = if encoding == ProtoDelta { + ProtoFull } else { - decode_cms_with_heap_from_msgpack(bytes)? + MsgpackFull }; - sk.merge(&other) - .map_err(|e| format!("merge CmsWithHeap delta: {e}")) + let other = fragment(DeltaSketchKind::Kll { k: 0 }, full)?; + self.merge_same_family(&other) } - SummaryState::WeightedFrequency(_) => { - Err("weighted frequency frames are complete states, not deltas".into()) + (Self::Cms(_), ProtoDelta) => { + let other = cms(d::cms_from_proto_delta(bytes)?); + self.merge_same_family(&other) + } + (Self::CountSketch(_), ProtoDelta) => { + let other = cs(d::cs_from_proto_delta(bytes)?); + self.merge_same_family(&other) } - SummaryState::CountSketchWithHeap(sk) => { - let other = if encoding == SketchEncoding::MsgpackDelta { - decode_cs_with_heap_from_msgpack_delta(bytes)? + (Self::Cms(_) | Self::CountSketch(_), _) => Err( + "heap-less CountMin/CountSketch MSGPACK_DELTA is not a valid producer encoding \ + (msgpack-delta is the heap-bearing form)" + .into(), + ), + // A heap delta frame decodes to a standalone window state and + // merges in; heap proto "deltas" are full msgpack states. + (Self::CmsWithHeap(_), _) => { + let other = cms_heap_state(if encoding == MsgpackDelta { + d::cms_with_heap_from_msgpack_delta(bytes)? } else { - decode_cs_with_heap_from_msgpack(bytes)? - }; - sk.merge(&other) - .map_err(|e| format!("merge CountSketchWithHeap delta: {e}")) + d::cms_with_heap_from_msgpack(bytes)? + }); + self.merge_same_family(&other) + } + (Self::CountSketchWithHeap(_), _) => { + let other = cs_heap_state(if encoding == MsgpackDelta { + d::cs_with_heap_from_msgpack_delta(bytes)? + } else { + d::cs_with_heap_from_msgpack(bytes)? + }); + self.merge_same_family(&other) } } } pub fn quantile(&self, q: f64) -> f64 { + use planner_types::post_asap::SketchQuery; match self { - SummaryState::Dd(sk) => sk.quantile(q).unwrap_or(0.0), - SummaryState::Kll(sk) => sk.quantile(q), + Self::Dd(_) | Self::Kll(_) => self + .kernel() + .estimate(&SketchQuery::Quantile { q }) + .unwrap_or(0.0), _ => 0.0, } } pub fn cardinality(&self) -> f64 { + use planner_types::post_asap::SketchQuery; match self { - SummaryState::Hll(sk) => sk.estimate(), + Self::Hll(s) => s.estimate(&SketchQuery::Cardinality).unwrap_or(0.0), _ => 0.0, } } - /// The bucket TOTAL — sum of row 0 of the underlying matrix. What a - /// bare `count_over_time`/`sum by (item) (rate(...))`-shaped query - /// (no specific item key) reads out. `0.0` for non-Frequency-family - /// states. + /// The bucket TOTAL — sum of row 0 of the underlying matrix, what a bare + /// (no item key) frequency readout reads. `0.0` for other families. pub fn total(&self) -> f64 { let matrix = match self { - SummaryState::Cms(c) => c.sketch(), - SummaryState::CountSketch(c) => c.sketch().clone(), - SummaryState::CmsWithHeap(h) => h.sketch_matrix(), - SummaryState::CountSketchWithHeap(h) => h.sketch_matrix(), + Self::Cms(c) => c.inner.sketch(), + Self::CountSketch(c) => c.inner.sketch().clone(), + Self::CmsWithHeap(h) => h.inner.sketch_matrix(), + Self::CountSketchWithHeap(h) => h.inner.sketch_matrix(), _ => return 0.0, }; matrix @@ -534,44 +433,39 @@ impl SummaryState { .unwrap_or(0.0) } - /// Per-key point estimate — `count(metric{item="x"})`-shaped queries. - /// Unlike [`Self::topk_items`], no heap is needed: all four Frequency - /// variants (heap-bearing or not) already carry a keyed `estimate` - /// over their matrix. `None` for the quantile/cardinality states, - /// which have no item universe at all. + /// Per-item point estimate; `None` for families without an item universe. pub fn estimate(&self, key: &str) -> Option { match self { - SummaryState::Cms(c) => Some(c.estimate(key)), - SummaryState::CountSketch(c) => Some(c.estimate(key)), - SummaryState::CmsWithHeap(h) => Some(h.estimate(key)), - SummaryState::CountSketchWithHeap(h) => Some(h.estimate(key)), + Self::Cms(c) => Some(c.inner.estimate(key)), + Self::CountSketch(c) => Some(c.inner.estimate(key)), + Self::CmsWithHeap(h) => Some(h.inner.estimate(key)), + Self::CountSketchWithHeap(h) => Some(h.inner.estimate(key)), _ => None, } } - /// Top-k `(key, value)` pairs from the heap, descending by value. - /// `None` for anything other than a heap-bearing state — the - /// heap-less Frequency states (`Cms`/`CountSketch`) carry no item - /// universe to enumerate, and the quantile/cardinality states have - /// no heap at all. `group` is the stored group the state belongs to. + /// Heap `(key, value)` pairs, unordered; `None` for a heap-less state. + /// `group` is the stored group the state belongs to. pub fn topk_items( &self, group: &std::collections::BTreeMap, ) -> Option> { match self { - SummaryState::CmsWithHeap(h) => Some( - h.topk_heap_items() + Self::CmsWithHeap(h) => Some( + h.inner + .topk_heap_items() .into_iter() .map(|item| (item.key, item.value)) .collect(), ), - SummaryState::CountSketchWithHeap(h) => Some( - h.topk_heap_items() + Self::CountSketchWithHeap(h) => Some( + h.inner + .topk_heap_items() .into_iter() .map(|item| (item.key, item.value)) .collect(), ), - SummaryState::WeightedFrequency(h) => Some( + Self::WeightedFrequency(h) => Some( h.rows(usize::MAX) .into_iter() .filter_map(|mut row| { @@ -587,72 +481,31 @@ impl SummaryState { } } - /// Merge `other` into `self` in place — both must be the same sketch - /// family. Used to combine several sids' reconstructed states - /// (`cumulative_summary_state`/`per_window_summary_states`) into one - /// cross-sid answer. `CmsWithHeap`/`CountSketchWithHeap` fall through - /// to the catch-all mismatch arm below like any other mixed pair — - /// and since the two variants now hold distinct sketchlib types - /// (`CountMinSketchWithHeap` vs `CountSketchWithHeap`), there is no - /// arm that could accidentally match them together — see their doc - /// on `DeltaSketchKind`. + /// Merge `other` into `self` with the Planner kernel's merge; both must be + /// the same family and shape. pub fn merge_same_family(&mut self, other: &SummaryState) -> Result<(), String> { - match (self, other) { - (SummaryState::UnivMon(a), SummaryState::UnivMon(b)) => { - a.merge_in_place(b).map_err(|e| e.to_string()) - } - (SummaryState::Dd(a), SummaryState::Dd(b)) => { - a.merge(b).map_err(|e| format!("merge DDSketch: {e}")) - } - (SummaryState::Hll(a), SummaryState::Hll(b)) => { - a.merge(b).map_err(|e| format!("merge HLL: {e}")) - } - (SummaryState::Kll(a), SummaryState::Kll(b)) => { - a.merge(b).map_err(|e| format!("merge KLL: {e}")) - } - (SummaryState::Cms(a), SummaryState::Cms(b)) => { - a.merge(b).map_err(|e| format!("merge CountMinSketch: {e}")) - } - (SummaryState::CountSketch(a), SummaryState::CountSketch(b)) => { - a.merge(b).map_err(|e| format!("merge CountSketch: {e}")) - } - (SummaryState::CmsWithHeap(a), SummaryState::CmsWithHeap(b)) => { - a.merge(b).map_err(|e| format!("merge CmsWithHeap: {e}")) - } - (SummaryState::CountSketchWithHeap(a), SummaryState::CountSketchWithHeap(b)) => a - .merge(b) - .map_err(|e| format!("merge CountSketchWithHeap: {e}")), - (SummaryState::WeightedFrequency(a), SummaryState::WeightedFrequency(b)) => { - use asap_physical_operators::AggregateCore; - let merged = a - .merge_with(b) - .map_err(|e| format!("merge weighted frequency: {e}"))?; - *a = merged - .as_any() - .downcast_ref::() - .ok_or("weighted frequency merge changed state type")? - .clone(); - Ok(()) - } - (a, _) => Err(format!( - "SummaryState family mismatch in merge_same_family (self is {})", - a.family_name() - )), - } + let merged = self + .kernel() + .merge_with(other.kernel()) + .map_err(|e| format!("merge {} state: {e}", self.family_name()))?; + *self = self + .same_variant(merged.as_ref()) + .ok_or("Planner merge changed the state family")?; + Ok(()) } /// Diagnostic family name for error messages — not used for dispatch. fn family_name(&self) -> &'static str { match self { - SummaryState::UnivMon(_) => "UnivMon", - SummaryState::Dd(_) => "DDSketch", - SummaryState::Hll(_) => "Hll", - SummaryState::Kll(_) => "Kll", - SummaryState::Cms(_) => "Cms", - SummaryState::CountSketch(_) => "CountSketch", - SummaryState::CmsWithHeap(_) => "CmsWithHeap", - SummaryState::CountSketchWithHeap(_) => "CountSketchWithHeap", - SummaryState::WeightedFrequency(_) => "WeightedFrequency", + Self::UnivMon(_) => "UnivMon", + Self::Dd(_) => "DDSketch", + Self::Hll(_) => "Hll", + Self::Kll(_) => "Kll", + Self::Cms(_) => "Cms", + Self::CountSketch(_) => "CountSketch", + Self::CmsWithHeap(_) => "CmsWithHeap", + Self::CountSketchWithHeap(_) => "CountSketchWithHeap", + Self::WeightedFrequency(_) => "WeightedFrequency", } } } @@ -817,51 +670,6 @@ fn visit_window_summary_states( Ok(skipped) } -// --------------------------------------------------------------------------- -// Proto-envelope decoders — P2-4: ONE decoder per family. -// -// These delegate to the precompute-side accumulators' -// `from_sketchlib_proto_bytes`, which are the single source of truth for -// the modified-OTLP proto wire format (envelope unwrapping, alpha/k/ -// precision validation, and — critically for HLL — SPARSE -// `registers_sparse` expansion). Folding the warm read path onto the -// same decoder the ingest path uses means the sparse-register fix (and -// any future format change) can never drift between the two copies again -// — the bug class P2-3 / P2-4 closed. We extract the accumulator's -// public `inner` sketch for the rolling-state merge. -// --------------------------------------------------------------------------- - -fn dd_from_proto(buffer: &[u8]) -> Result { - use crate::summary_kernels::dd_sketch::DDSketchAccumulator; - DDSketchAccumulator::from_sketchlib_proto_bytes(buffer) - .map(|acc| acc.inner) - .map_err(|e| e.to_string()) -} - -fn kll_from_proto(buffer: &[u8]) -> Result { - use crate::summary_kernels::datasketches_kll::DatasketchesKLLAccumulator; - DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(buffer) - .map(|acc| acc.inner) - .map_err(|e| e.to_string()) -} - -fn hll_from_proto(buffer: &[u8]) -> Result { - use crate::summary_kernels::hll_sketch::HllSketchAccumulator; - HllSketchAccumulator::from_sketchlib_proto_bytes(buffer) - .map(|acc| acc.inner) - .map_err(|e| e.to_string()) -} - -/// Apply a proto-encoded `HllDelta` frame onto the HLL register vector — the -/// delta is a varint-packed (index_delta, value) blob; decode + apply -/// (register-wise max) via the shared sketch library so the unpacking stays a -/// single source of truth. -fn apply_hll_proto_delta(sk: &mut HllSketch, buffer: &[u8]) -> Result<(), String> { - sk.apply_delta_bytes(buffer) - .map_err(|e| format!("apply HLLDelta: {e}"))?; - Ok(()) -} - #[cfg(test)] mod tests { //! P2-3 / P2-4 regression tests for the consolidated single-decoder @@ -876,21 +684,18 @@ mod tests { // with another window, and ranks items under their series keys. #[test] fn weighted_frequency_frames_rank_items_by_series_key() { - use crate::summary_kernels::weighted_frequency::{ - FrequencyAlgorithm, PhysicalWeightedFrequency, WeightedFrequency, - }; - use crate::SerializableToSink; use asap_physical_operators::values::Value; let frame = |weight| { - let mut state = - PhysicalWeightedFrequency::new(FrequencyAlgorithm::Cms, 64, 3, 8).unwrap(); + let mut state = WeightedFrequency::new(FrequencyAlgorithm::Cms, 64, 3, 8).unwrap(); let identity = r#"{"__name__":"m","endpoint":"a\"b"}"#; state .update(&[Value::Utf8(identity.into())], weight) .unwrap(); state.update(&[Value::Utf8("plain".into())], 1.0).unwrap(); SketchSampleState { - bytes: WeightedFrequency(state).serialize_to_bytes(), + bytes: crate::stored_state::codec::frequency_kernel(&state) + .unwrap() + .to_bytes(), encoding: SketchEncoding::WeightedFrequencyV1, } }; @@ -930,11 +735,8 @@ mod tests { // restores non-empty ones from the stored group and keeps identity values. #[test] fn weighted_frequency_items_restore_group_labels() { - use crate::summary_kernels::weighted_frequency::{ - FrequencyAlgorithm, PhysicalWeightedFrequency, - }; use asap_physical_operators::values::Value; - let mut state = PhysicalWeightedFrequency::new(FrequencyAlgorithm::Cms, 64, 3, 8).unwrap(); + let mut state = WeightedFrequency::new(FrequencyAlgorithm::Cms, 64, 3, 8).unwrap(); let identity = r#"{"__name__":"m","endpoint":"a"}"#; state.update(&[Value::Utf8(identity.into())], 2.0).unwrap(); let state = SummaryState::WeightedFrequency(state); @@ -973,44 +775,6 @@ mod tests { .encode_to_vec() } - fn encode_kll(k: u16, items: &[f64]) -> Vec { - use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; - use prost::Message; - let state = KllState { - k: k as u32, - items: items.to_vec(), - levels: vec![], - num_levels: 0, - ..Default::default() - }; - SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Kll(state)), - ..Default::default() - } - .encode_to_vec() - } - - fn encode_hll_dense(sk: &HllSketch) -> Vec { - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, - }; - use prost::Message; - let state = HyperLogLogState { - variant: ProtoVariant::Regular as i32, - precision: sk.precision, - registers: sk.registers.clone(), - hip_kxq0: sk.hip_kxq0, - hip_kxq1: sk.hip_kxq1, - hip_est: sk.hip_est, - registers_sparse: None, - }; - SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Hll(state)), - ..Default::default() - } - .encode_to_vec() - } - /// Build a SPARSE HLL proto frame: dense `registers` left empty, /// `registers_sparse.packed` = varint (index_delta, value) pairs. /// This is exactly the wire form a low-cardinality producer emits @@ -1073,7 +837,7 @@ mod tests { let precision = 12u32; let nonzero = [(3u64, 5u8), (100, 2), (4000, 7)]; let bytes = encode_hll_sparse(precision, &nonzero); - let sk = hll_from_proto(&bytes).expect("sparse HLL frame must decode (P2-3 regression)"); + let sk = d::hll_from_proto(&bytes).expect("sparse HLL frame must decode (P2-3 regression)"); assert_eq!(sk.registers.len(), 1usize << precision); for (idx, val) in nonzero { assert_eq!( @@ -1083,56 +847,6 @@ mod tests { } } - #[test] - fn hll_from_proto_matches_accumulator_decoder() { - // P2-4: the warm read path and the ingest accumulator must decode - // the SAME bytes to the SAME sketch (one source of truth). - use crate::summary_kernels::hll_sketch::HllSketchAccumulator; - let mut sk = HllSketch::new(HllVariant::Regular, 12); - for i in 0..500u64 { - sk.update(format!("item-{i}").as_bytes()); - } - let bytes = encode_hll_dense(&sk); - let via_delta = hll_from_proto(&bytes).expect("delta_apply hll decode"); - let via_acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes) - .expect("accumulator hll decode") - .inner; - assert_eq!( - via_delta.registers, via_acc.registers, - "delta_apply and accumulator must produce identical HLL registers" - ); - assert!((via_delta.estimate() - via_acc.estimate()).abs() < 1e-9); - } - - #[test] - fn dd_from_proto_matches_accumulator_decoder() { - use crate::summary_kernels::dd_sketch::DDSketchAccumulator; - let mut sk = DdSketch::new(0.01); - for v in [1.0, 2.0, 5.0, 5.0, 9.0, 42.0] { - sk.update(v); - } - let bytes = encode_dd(&sk); - let via_delta = dd_from_proto(&bytes).expect("delta_apply dd decode"); - let via_acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes) - .expect("accumulator dd decode") - .inner; - // Same quantile answers from the same bytes through both paths. - assert_eq!(via_delta.quantile(0.5), via_acc.quantile(0.5)); - assert_eq!(via_delta.quantile(0.99), via_acc.quantile(0.99)); - } - - #[test] - fn kll_from_proto_matches_accumulator_decoder() { - use crate::summary_kernels::datasketches_kll::DatasketchesKLLAccumulator; - let items: Vec = (0..200).map(|i| i as f64).collect(); - let bytes = encode_kll(256, &items); - let via_delta = kll_from_proto(&bytes).expect("delta_apply kll decode"); - let via_acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes) - .expect("accumulator kll decode") - .inner; - assert_eq!(via_delta.quantile(0.5), via_acc.quantile(0.5)); - } - // ----------------------------------------------------------------- // Per-window-reset (PWR) delta-apply regression tests. // @@ -1166,6 +880,29 @@ mod tests { sk } + // A sampled DDSketch envelope on the delta channel is rejected rather + // than read as an empty bucket delta. + #[test] + fn sampled_ddsketch_envelope_on_delta_channel_is_rejected() { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; + use prost::Message; + let sketch = dd_over(0.01, &[1., 2.]); + let state = asap_sketch_codec::ddsketch_state(&encode_dd(&sketch)) + .unwrap() + .0; + let sampled = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Ddsketch(state)), + sample_p: 0.5, + ..Default::default() + } + .encode_to_vec(); + let mut rolling = DeltaSketchKind::DDSketch { alpha: 0.01 }.bootstrap_empty(); + assert!(rolling + .apply_delta_bytes(&sampled, SketchEncoding::ProtoDelta) + .unwrap_err() + .contains("edge-sampled")); + } + /// A full re-snapshot replaces its pane's earlier frames; cumulative /// readout must merge the finalized panes without counting updates twice. #[test] @@ -1180,7 +917,7 @@ mod tests { let SummaryState::Dd(state) = state else { panic!("expected DDSketch state"); }; - assert_eq!(state.store_counts.iter().sum::(), 4); + assert_eq!(state.inner.store_counts.iter().sum::(), 4); } /// PWR across 3 windows: window 1 is `[Full]`, windows 2 & 3 are @@ -1346,23 +1083,16 @@ mod tests { let mut cs_heap = CountSketchWithHeap::new(4, 256, 10); cs_heap.update("b", 1.0); - let mut a = SummaryState::CmsWithHeap( + let mut a = cms_heap_state( CountMinSketchWithHeap::from_msgpack(&cms_heap.to_msgpack().unwrap()).unwrap(), ); - let b = SummaryState::CountSketchWithHeap( + let b = cs_heap_state( CountSketchWithHeap::from_msgpack(&cs_heap.to_msgpack().unwrap()).unwrap(), ); - - match a.merge_same_family(&b) { - Err(msg) => assert!( - msg.contains("family mismatch"), - "expected a family-mismatch error, got: {msg}" - ), - Ok(()) => panic!( - "CmsWithHeap must not merge with CountSketchWithHeap -- \ - different algorithms sharing only a storage shape" - ), - } + assert!( + a.merge_same_family(&b).is_err(), + "CmsWithHeap must not merge with CountSketchWithHeap" + ); } fn encode_delta_heap( @@ -1419,8 +1149,8 @@ mod tests { .expect("decode_full CountSketchWithHeap"); match full_state { SummaryState::CountSketchWithHeap(inner) => { - assert_eq!(inner.sketch_matrix(), expected_matrix); - assert_eq!(inner.estimate("k"), expected_estimate); + assert_eq!(inner.inner.sketch_matrix(), expected_matrix); + assert_eq!(inner.inner.estimate("k"), expected_estimate); } other => panic!( "expected CountSketchWithHeap state, got {}", @@ -1465,11 +1195,11 @@ mod tests { match rolling { SummaryState::CountSketchWithHeap(inner) => { assert_eq!( - inner.sketch_matrix(), + inner.inner.sketch_matrix(), expected_matrix, "delta path must reconstruct the identical matrix" ); - assert_eq!(inner.estimate("k"), expected_estimate); + assert_eq!(inner.inner.estimate("k"), expected_estimate); } other => panic!( "expected CountSketchWithHeap state, got {}", diff --git a/crates/asap_summary_state/src/stored_state/mod.rs b/crates/asap_summary_state/src/stored_state/mod.rs index e45b851e0..0c7400fde 100644 --- a/crates/asap_summary_state/src/stored_state/mod.rs +++ b/crates/asap_summary_state/src/stored_state/mod.rs @@ -1,4 +1,5 @@ //! Portable stored-summary payloads and reconstruction, independent of storage engines. +pub mod codec; pub mod decoders; pub mod delta_apply; pub mod native; @@ -35,10 +36,8 @@ impl SketchEncoding { /// Encoding of a stored sketch state written as one complete window. pub fn full_frame_for(state: &dyn crate::AggregateCore) -> Self { - if state - .as_any() - .is::() - { + use asap_physical_operators::summary_kernels::weighted_frequency::WeightedFrequency; + if state.as_any().is::() { Self::WeightedFrequencyV1 } else { Self::MsgpackFull diff --git a/crates/asap_summary_state/src/stored_state/native.rs b/crates/asap_summary_state/src/stored_state/native.rs index d41133660..4f47c73ce 100644 --- a/crates/asap_summary_state/src/stored_state/native.rs +++ b/crates/asap_summary_state/src/stored_state/native.rs @@ -1,18 +1,20 @@ //! Versioned physical output batches. Deployment identities and coverage remain //! outside this payload and must be checked before decoding with the bound schema. -use crate::summary_kernels::{ - datasketches_kll::DatasketchesKLLAccumulator, dd_sketch::DDSketchAccumulator, - exact::ExactAccumulator, hll_sketch::HllSketchAccumulator, SumAccumulator, -}; +use super::codec::{self, StoredState}; +use crate::AggregationType; use asap_physical_operators::{ summary_kernels as physical, values::{Batch, Schema, Value}, - AggregateCore, Error, + AggregateCore, Error, KernelError, }; use planner_types::post_asap::{SummaryFamilyType, SummarySchema}; use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use std::sync::Arc; +/// Stored type tag of a published native output snapshot. +pub const NATIVE_OUTPUT_TYPE: &str = "NativePhysicalOutputV1"; + #[derive(Serialize, Deserialize)] #[serde(deny_unknown_fields)] struct StoredBatch { @@ -37,7 +39,7 @@ enum Cell { enum StateCodec { WeightedFrequencyV1, ExactAccumulatorV1, - /// Read-only: decodes to Planner's exact Sum state. + /// Retired: kept only so old batches fail with a named rejection. SumAccumulatorV1, KllMsgpackV1, DdMsgpackV1, @@ -46,25 +48,14 @@ enum StateCodec { fn invalid(message: impl ToString) -> Error { Error::Invalid(message.to_string()) } -type Kernel = asap_sketchlib::WeightedFrequency; impl StateCodec { fn encode(state: &dyn AggregateCore) -> Result<(Self, Vec), Error> { let any = state.as_any(); - if let Some(state) = any.downcast_ref::() { - return Ok(( - Self::WeightedFrequencyV1, - crate::physical::frequency_kernel(state) - .map_err(invalid)? - .to_bytes(), - )); - } - if let Some(state) = any.downcast_ref::() { - return Ok(( - Self::ExactAccumulatorV1, - rmp_serde::to_vec_named(state).map_err(invalid)?, - )); - } - let codec = if any.is::() { + let codec = if any.is::() { + Self::WeightedFrequencyV1 + } else if any.is::() { + Self::ExactAccumulatorV1 + } else if any.is::() { Self::KllMsgpackV1 } else if any.is::() { Self::DdMsgpackV1 @@ -73,35 +64,132 @@ impl StateCodec { } else { return Err(invalid("physical summary has no persisted native codec")); }; - let stored = crate::physical::from_physical(state).map_err(invalid)?; - Ok((codec, stored.serialize_to_bytes())) + Ok((codec, state.encode().map_err(invalid)?)) } fn decode(&self, bytes: &[u8]) -> Result, Error> { - let stored: Box = match self { - Self::WeightedFrequencyV1 => { - let kernel = Kernel::from_bytes(bytes).map_err(|e| invalid(format!("{e:?}")))?; - let state: physical::weighted_frequency::WeightedFrequency = - rmp_serde::from_slice(&rmp_serde::to_vec(&kernel).map_err(invalid)?) - .map_err(invalid)?; - return Ok(Arc::new(state)); - } - Self::ExactAccumulatorV1 => { - Box::new(ExactAccumulator::deserialize_from_bytes(bytes).map_err(invalid)?) - } + let tag = match self { + Self::WeightedFrequencyV1 => "WeightedFrequency", + Self::ExactAccumulatorV1 => codec::EXACT_V1, Self::SumAccumulatorV1 => { - Box::new(SumAccumulator::deserialize_from_bytes(bytes).map_err(invalid)?) + return Err(invalid( + "native codec SumAccumulatorV1 is retired and no longer decoded", + )) } - Self::KllMsgpackV1 => { - Box::new(DatasketchesKLLAccumulator::from_msgpack_bytes(bytes).map_err(invalid)?) + Self::KllMsgpackV1 => "DatasketchesKLLAccumulator", + Self::DdMsgpackV1 => "DDSketchAccumulator", + Self::HllMsgpackV1 => "HllSketchAccumulator", + }; + codec::decode(tag, bytes).map(Arc::from).map_err(invalid) + } +} + +/// A published native result held in the store as one opaque snapshot. It is +/// storage, not a kernel: it never merges and is read only by decoding its +/// batch against the installed physical DAG. +#[derive(Clone)] +pub struct NativeSummaryOutput { + batch: Batch, + bytes: Vec, + kind: AggregationType, +} + +impl NativeSummaryOutput { + /// Validate that the batch carries one summary family and fits the budget. + pub fn new(batch: Batch, max_bytes: usize) -> Result { + let families = batch + .schema() + .fields + .iter() + .filter_map(|field| { + (!matches!(field.dtype, SummaryFamilyType::Plain(_))).then_some(&field.dtype) + }) + .collect::>(); + let [family] = families.as_slice() else { + return Err("native stored batch requires one summary column".into()); + }; + use planner_types::post_asap::SketchAlgorithm; + let mut kind = match family { + SummaryFamilyType::Sketch(sketch, _) => match sketch.algorithm() { + SketchAlgorithm::CmsWithHeap => Some(AggregationType::CountMinSketchWithHeap), + SketchAlgorithm::CountSketchWithHeap => Some(AggregationType::CountSketchWithHeap), + _ => None, + }, + SummaryFamilyType::ExactAggregate(planner_types::post_asap::ExactKind::Sum, _) => { + Some(AggregationType::Sum) } - Self::DdMsgpackV1 => { - Box::new(DDSketchAccumulator::from_msgpack_bytes(bytes).map_err(invalid)?) + _ => None, + }; + for row in batch.rows() { + let states = row + .iter() + .filter_map(|value| match value { + Value::Summary { state, .. } => Some(state), + _ => None, + }) + .collect::>(); + let [state] = states.as_slice() else { + return Err("native stored row requires one summary state".into()); + }; + codec::check_storable(state.as_ref()).map_err(|error| error.to_string())?; + let row_kind = state.get_accumulator_type(); + if kind.is_some_and(|kind| kind != row_kind) { + return Err("native stored rows have different summary families".into()); } - Self::HllMsgpackV1 => { - Box::new(HllSketchAccumulator::from_msgpack_bytes(bytes).map_err(invalid)?) + kind = Some(row_kind); + } + let kind = kind.ok_or("empty native batch has no supported summary family")?; + let bytes = encode_batch(&batch).map_err(|error| error.to_string())?; + if bytes.len() > max_bytes || batch.bytes() > max_bytes { + return Err("native summary exceeds publication/read budget".into()); + } + Ok(Self { batch, bytes, kind }) + } + + pub fn batch(&self) -> &Batch { + &self.batch + } + pub fn bytes(&self) -> &[u8] { + &self.bytes + } + pub fn kind(&self) -> AggregationType { + self.kind + } + + /// Every stored group label must match the batch's own group column. + pub fn validate_group(&self, group: &BTreeMap) -> Result<(), String> { + for (key, value) in group { + let column = self + .batch + .schema() + .fields + .iter() + .position(|field| &field.name == key) + .ok_or("native output is missing its stored group key")?; + if self + .batch + .rows() + .iter() + .any(|row| !matches!(&row[column], Value::Utf8(actual) if actual.as_ref() == value)) + { + return Err("native output group differs from stored address".into()); } - }; - crate::physical::to_physical(stored.as_ref()).map_err(invalid) + } + Ok(()) + } +} + +impl AggregateCore for NativeSummaryOutput { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn merge_with(&self, _: &dyn AggregateCore) -> Result, KernelError> { + Err("native output snapshots require an explicit physical merge operator".into()) + } + fn approx_memory_bytes(&self) -> usize { + self.bytes.len() + self.batch.bytes() } } @@ -386,24 +474,4 @@ mod tests { ) .is_err()); } - - // A stored Sum payload is read back as Planner's exact Sum with the same value. - #[test] - fn stored_sum_payload_decodes_as_planner_exact_sum() { - use crate::SerializableToSink; - use planner_types::post_asap::{ExactKind, ExactParams}; - let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); - let state = StateCodec::SumAccumulatorV1 - .decode(&SumAccumulator::with_sum(5.5).serialize_to_bytes()) - .unwrap(); - let exact = state - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(exact.family(), &family); - assert_eq!( - exact.readout(crate::Statistic::Sum, None, None).unwrap(), - Some(5.5) - ); - } } diff --git a/crates/asap_summary_state/src/stored_state/readout.rs b/crates/asap_summary_state/src/stored_state/readout.rs index 6d29a2d0f..24796ed55 100644 --- a/crates/asap_summary_state/src/stored_state/readout.rs +++ b/crates/asap_summary_state/src/stored_state/readout.rs @@ -1,5 +1,6 @@ //! Planner-declared readouts over reconstructed summary states. use super::delta_apply::SummaryState; +use asap_physical_operators::AggregateCore; use planner_types::{post_asap::SketchQuery, pre_asap::ColumnRef}; #[derive(Debug, thiserror::Error)] pub enum Error { @@ -8,20 +9,9 @@ pub enum Error { } pub fn sketch_query_value(rs: &SummaryState, query: &SketchQuery) -> Result { if let SummaryState::UnivMon(state) = rs { - use crate::AggregateCore; - let statistic = match query { - SketchQuery::Cardinality => crate::Statistic::Cardinality, - SketchQuery::FrequencyL2 => crate::Statistic::FrequencyL2, - SketchQuery::FrequencyEntropy => crate::Statistic::FrequencyEntropy, - SketchQuery::PointCount { - key: ColumnRef::SampleValue, - value: None, - } => crate::Statistic::Count, - _ => return Err(Error::Unsupported("unsupported UnivMon readout")), - }; return state - .query_statistic(statistic, &None, &Default::default()) - .map_err(|_| Error::Unsupported("UnivMon readout failed")); + .estimate(query) + .map_err(|_| Error::Unsupported("unsupported UnivMon readout")); } if matches!(rs, SummaryState::WeightedFrequency(_)) { return Err(Error::Unsupported( @@ -35,9 +25,17 @@ pub fn sketch_query_value(rs: &SummaryState, query: &SketchQuery) -> Result match rs { // Typed PromQL/continuous-percentile readout uses interpolation; // portable DDS `quantile` deliberately retains lower-rank parity. - SummaryState::Dd(sketch) => sketch.quantile_interpolated(*q).ok_or(Error::Unsupported( - "DDS interpolated quantile is unavailable", - )), + SummaryState::Dd(sketch) => { + sketch + .inner + .quantile_interpolated(*q) + .ok_or(Error::Unsupported( + "DDS interpolated quantile is unavailable", + )) + } + SummaryState::Kll(sketch) => sketch + .estimate(query) + .map_err(|_| Error::Unsupported("KLL quantile must be in [0, 1]")), _ => Ok(rs.quantile(*q)), }, SketchQuery::Cardinality => Ok(rs.cardinality()), @@ -100,110 +98,68 @@ pub fn topk_ranked( Ok(items) } -/// Merge already selected exact panes and finalize using the shared accumulator contract. +type Parameters = std::collections::HashMap; + +/// Merge already selected exact panes and read one population; an absent +/// population (empty MIN/MAX, counter with too few samples) is an error. pub fn exact_readout( - states: impl IntoIterator>, + states: impl IntoIterator>, statistic: crate::Statistic, key: &Option, - parameters: &std::collections::HashMap, + parameters: &Parameters, ) -> Result { - let merged = merge_exact_states(states)?; - merged - .query_statistic(statistic, key, parameters) - .map_err(|e| e.to_string()) -} - -fn merge_exact_states( - states: impl IntoIterator>, -) -> Result, String> { - let mut states = states.into_iter(); - let first = states - .next() - .ok_or_else(|| "empty exact state input".to_string())?; - if let Some(first) = first - .as_any() - .downcast_ref::() - { - let mut merged = first.clone(); - for state in states { - let other = state - .as_any() - .downcast_ref::() - .ok_or_else(|| "merge requires Planner exact state".to_string())?; - merged - .merge_from(other) - .map_err(|error| error.to_string())?; - } - return Ok(Box::new(merged)); - } - let mut merged = first.clone_boxed_core(); - for state in states { - merged = merged - .merge_with(state.as_ref()) - .map_err(|e| e.to_string())?; - } - Ok(merged) + exact_readout_optional(states, statistic, key, parameters)? + .ok_or_else(|| "empty exact population".to_string()) } /// PromQL counter readouts omit a series with fewer than two samples. Other /// state/type/range failures remain errors rather than empty results. pub fn insufficient_counter_samples( - state: &dyn crate::AggregateCore, + state: &dyn AggregateCore, statistic: crate::Statistic, ) -> bool { - matches!( - statistic, - crate::Statistic::Rate | crate::Statistic::Increase - ) && (state - .as_any() - .downcast_ref::() - .is_some_and(|state| { - state.sample_count < 2 || state.last_seen_timestamp == state.starting_timestamp - }) - || state - .as_any() - .downcast_ref::() - .is_some_and(|state| state.insufficient_counter_samples(statistic, &None))) + asap_physical_operators::readout::insufficient_counter_samples(state, statistic) } +/// Merge already selected exact panes with Planner's exact readout. `None` +/// means a counter population with too few samples is absent; an empty +/// MIN/MAX population is an error so the query can fall back. pub fn exact_readout_optional( - states: impl IntoIterator>, + states: impl IntoIterator>, statistic: crate::Statistic, key: &Option, - parameters: &std::collections::HashMap, + parameters: &Parameters, ) -> Result, String> { - let merged = merge_exact_states(states)?; - let counter = key.as_ref().and_then(|key| { - merged - .as_any() - .downcast_ref::() - .and_then(|state| state.increases.get(key)) - }); - let exact_insufficient = merged - .as_any() - .downcast_ref::() - .is_some_and(|state| state.insufficient_counter_samples(statistic, key)); - if exact_insufficient - || insufficient_counter_samples(merged.as_ref(), statistic) - || counter.is_some_and(|counter| insufficient_counter_samples(counter, statistic)) - { - return Ok(None); + let range = super::codec::range_ms(parameters).map_err(|error| error.to_string())?; + let value = + asap_physical_operators::readout::exact_readout(states, statistic, range, key.as_ref())?; + if value.is_none() && matches!(statistic, crate::Statistic::Min | crate::Statistic::Max) { + return Err("empty exact population".into()); } - merged - .query_statistic(statistic, key, parameters) - .map(Some) - .map_err(|error| error.to_string()) + Ok(value) } #[cfg(test)] mod counter_tests { use super::*; - use crate::{summary_kernels::IncreaseAccumulator, AggregateCore, Measurement, Statistic}; + use crate::{KeyByLabelValues, Statistic}; + use asap_physical_operators::summary_kernels::exact::ExactAccumulator; + use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; use std::sync::Arc; + fn counter(kind: ExactKind, params: ExactParams, keyed: bool) -> ExactAccumulator { + ExactAccumulator::new(SummaryFamilyType::ExactAggregate(kind, params), keyed).unwrap() + } + fn range(start: &str, end: &str) -> Parameters { + Parameters::from([ + ("range_start_ms".into(), start.into()), + ("range_end_ms".into(), end.into()), + ]) + } + + // A counter population with a single sample is absent, keyed or not. #[test] fn planner_counter_population_omits_insufficient_samples() { - use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; for (kind, params, statistic) in [ (ExactKind::Rate, ExactParams::Rate, Statistic::Rate), ( @@ -213,13 +169,8 @@ mod counter_tests { ), ] { for keyed in [false, true] { - let mut state = crate::summary_kernels::exact::ExactAccumulator::new( - SummaryFamilyType::ExactAggregate(kind.clone(), params.clone()), - keyed, - ) - .unwrap(); - let key = keyed - .then(|| crate::KeyByLabelValues::new_with_labels(vec!["checkout".into()])); + let mut state = counter(kind.clone(), params.clone(), keyed); + let key = keyed.then(|| KeyByLabelValues::new_with_labels(vec!["checkout".into()])); state.update(key.as_ref(), 10., 10_000); assert_eq!( exact_readout_optional( @@ -235,78 +186,45 @@ mod counter_tests { } } + // Repeated or single samples are absent; unknown keys, inverted ranges + // and empty input remain errors. #[test] fn sparse_counter_is_absent_but_invalid_ranges_still_fail() { - let mut state = - IncreaseAccumulator::new(Measurement::new(10.), 10_000, Measurement::new(10.), 10_000); - let parameters = std::collections::HashMap::from([ - ("range_start_ms".into(), "0".into()), - ("range_end_ms".into(), "60000".into()), - ]); - assert_eq!( + let label = KeyByLabelValues::new_with_labels(vec!["checkout".into()]); + let mut keyed = counter(ExactKind::Rate, ExactParams::Rate, true); + keyed.update(Some(&label), 10., 10_000); + keyed.update(Some(&label), 10., 10_000); + let read = |state: &ExactAccumulator, key: Option, p: &Parameters| { exact_readout_optional( [Arc::new(state.clone()) as Arc], Statistic::Rate, - &None, - ¶meters - ) - .unwrap(), - None - ); - let mut repeated = state.clone(); - repeated.update(Measurement::new(10.), 10_000); - assert_eq!( - exact_readout_optional( - [Arc::new(repeated) as Arc], - Statistic::Rate, - &None, - ¶meters + &key, + p, ) - .unwrap(), - None - ); - let mut keyed = crate::summary_kernels::KeyedCounterState::new(); - let label = crate::KeyByLabelValues::new_with_labels(vec!["checkout".into()]); - keyed.update(label.clone(), state.clone()); - assert_eq!( - exact_readout_optional( - [Arc::new(keyed.clone()) as Arc], - Statistic::Rate, - &Some(label), - ¶meters - ) - .unwrap(), - None - ); - assert!(exact_readout_optional( - [Arc::new(keyed) as Arc], - Statistic::Rate, - &Some(crate::KeyByLabelValues::new_with_labels(vec![ - "missing".into() - ])), - ¶meters - ) - .is_err()); - state.update(Measurement::new(20.), 20_000); - assert!(exact_readout_optional( - [Arc::new(state.clone()) as Arc], - Statistic::Rate, - &None, - ¶meters - ) - .unwrap() - .is_some()); - let invalid = std::collections::HashMap::from([ - ("range_start_ms".into(), "60000".into()), - ("range_end_ms".into(), "0".into()), - ]); - assert!(exact_readout_optional( - [Arc::new(state) as Arc], - Statistic::Rate, - &None, - &invalid - ) - .is_err()); - assert!(exact_readout_optional([], Statistic::Rate, &None, ¶meters).is_err()); + }; + let window = range("0", "60000"); + assert_eq!(read(&keyed, Some(label.clone()), &window).unwrap(), None); + let mut repeated = counter(ExactKind::Rate, ExactParams::Rate, false); + repeated.update(None, 10., 10_000); + repeated.update(None, 10., 10_000); + assert_eq!(read(&repeated, None, &window).unwrap(), None); + let missing = KeyByLabelValues::new_with_labels(vec!["missing".into()]); + assert!(read(&keyed, Some(missing), &window).is_err()); + keyed.update(Some(&label), 20., 20_000); + assert!(read(&keyed, Some(label.clone()), &window) + .unwrap() + .is_some()); + assert!(read(&keyed, Some(label), &range("60000", "0")).is_err()); + assert!(exact_readout_optional([], Statistic::Rate, &None, &window).is_err()); + } + + // An empty MIN population is an error, not an absent series. + #[test] + fn empty_extremum_population_is_an_error() { + let empty = counter(ExactKind::Min, ExactParams::Min, false); + let states = || [Arc::new(empty.clone()) as Arc]; + let none = Parameters::new(); + assert!(exact_readout_optional(states(), Statistic::Min, &None, &none).is_err()); + assert!(exact_readout(states(), Statistic::Min, &None, &none).is_err()); } } diff --git a/crates/asap_summary_state/src/summary_kernels/count_min_sketch.rs b/crates/asap_summary_state/src/summary_kernels/count_min_sketch.rs deleted file mode 100644 index cffe90d28..000000000 --- a/crates/asap_summary_state/src/summary_kernels/count_min_sketch.rs +++ /dev/null @@ -1,1323 +0,0 @@ -use crate::summary_kernels::dd_sketch::normalize_sample_p; -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{CountMinSketch, CountMinSketchDelta, MessagePackCodec}; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -/// Count-Min Sketch accumulator — wraps asap_sketchlib::CountMinSketch. -/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. -/// This file retains QE-specific trait impls, legacy deserializers, and JSON output. -#[derive(Debug, Clone)] -pub struct CountMinSketchAccumulator { - pub inner: CountMinSketch, - /// Edge sampling probability `p ∈ (0,1]` carried on the producer's - /// `SketchEnvelope.sample_p`. The edge admits each insert with - /// probability `p`, so every stored cell count is ~`p`× the true count. - /// CMS is L1/additive and linear, so the unbiased rescale of BOTH a - /// point-frequency estimate (`query_key`) and the aggregate - /// total-event statistics (`Count`/`Sum`/`Increase`/`Rate`) is `×1/p`. - /// `1.0` (and the proto3 default `0.0`, dual-read as `1.0`) means no - /// sampling, so the rescale is a no-op and the behaviour is identical - /// to before. Mirrors `DDSketchAccumulator::sample_p`; set from the - /// envelope at the `from_sketchlib_proto_bytes` decode site and - /// preserved across `reset_to_empty` and `merge_with`. - pub sample_p: f64, -} - -impl CountMinSketchAccumulator { - pub fn new(row_num: usize, col_num: usize) -> Self { - Self { - inner: CountMinSketch::new(row_num, col_num), - sample_p: 1.0, - } - } - - // Marked as _update and kept private; only called internally. - fn _update(&mut self, key: &KeyByLabelValues, value: f64) { - self.inner.update(&key.to_semicolon_str(), value); - } - - pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { - // The edge sampled inserts with probability `sample_p`, so the - // stored point-frequency estimate is ~`p`× the true frequency. - // CMS is linear/additive, so `×1/p` is the unbiased rescale. - // `sample_p == 1.0` (unsampled / legacy) makes this a no-op. - self.inner.estimate(&key.to_semicolon_str()) / self.sample_p - } - - pub fn deserialize_from_json(data: &Value) -> Result> { - let row_num = data["row_num"] - .as_f64() - .ok_or("Missing or invalid 'row_num' field")? as usize; - let col_num = data["col_num"] - .as_f64() - .ok_or("Missing or invalid 'col_num' field")? as usize; - - let sketch_data = data["sketch"] - .as_array() - .ok_or("Missing or invalid 'sketch' field")?; - - let mut sketch = Vec::new(); - for row in sketch_data { - let row_array = row.as_array().ok_or("Invalid row in sketch data")?; - let mut sketch_row = Vec::new(); - for cell in row_array { - let value = cell.as_f64().ok_or("Invalid cell value in sketch data")?; - sketch_row.push(value); - } - sketch.push(sketch_row); - } - - Ok(Self { - inner: CountMinSketch::from_legacy_matrix(sketch, row_num, col_num), - sample_p: 1.0, - }) - } - - /// Decode from the modified OTLP wire format's - /// `CountMinSketchDataPoint.sketch` bytes when - /// `encoding = COUNT_MIN_SKETCH_ENCODING_MSGPACK`. The bytes are the - /// MessagePack serialization of the cross-language sketch-core - /// `CountMinSketch` wire struct (same format the legacy Arroyo path - /// uses — this method is the modified-OTLP entrypoint for PR I). - pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { - Ok(Self { - inner: CountMinSketch::from_msgpack(buffer) - .map_err(|e| -> Box { e.to_string().into() })?, - // The msgpack CountMinSketch struct carries no envelope/sample_p; - // the msgpack path is parity/test-only and is never edge-sampled. - sample_p: 1.0, - }) - } - - /// Decode from the modified OTLP wire format's - /// `CountMinSketchDataPoint.sketch` bytes — i.e. the protobuf-encoded - /// `asap_sketchlib::proto::sketchlib::CountMinState` message used by - /// DataCollector's `countminsketchprocessor` when emitting via - /// `Metric.data = CountMinSketch{…}` with - /// `encoding = COUNT_MIN_SKETCH_ENCODING_PROTO`. - /// - /// The resulting accumulator is constructed via - /// `CountMinSketch::from_legacy_matrix` after reshaping the flat - /// `counts_int` / `counts_float` field into a `Vec>`. - pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, CountMinState, CounterType, SketchEnvelope, - }; - use prost::Message; - - // DataCollector's countminsketchprocessor wraps the state in a - // `SketchEnvelope{count_min: CountMinState}` via - // `SerializePortableFO` + `proto.Marshal`. Try decoding as envelope - // first, fall back to bare `CountMinState` for callers (e.g. unit - // tests) that encode the state directly. Capture the envelope's - // `sample_p` alongside the state so the point-frequency - // (`query_key`) and aggregate statistics rescale by `1/p`. Bare - // `CountMinState` bytes (no envelope) carry no sampling info → - // `sample_p` 1.0 (no rescale). Mirrors `DDSketchAccumulator`. - let (state, sample_p) = match SketchEnvelope::decode(buffer) { - Ok(env) => { - let sp = env.sample_p; - match env.sketch_state { - Some(sketch_envelope::SketchState::CountMin(st)) => (st, sp), - Some(other) => { - return Err(format!( - "SketchEnvelope contains non-CountMin sketch: {:?}", - std::mem::discriminant(&other) - ) - .into()); - } - // Envelope decoded but was empty (e.g. the buffer is a - // bare CountMinState that happened to parse as a default - // envelope). Fall through to bare decode. - None => ( - CountMinState::decode(buffer) - .map_err(|e| format!("decode CountMinState: {e}"))?, - 1.0, - ), - } - } - Err(_) => ( - CountMinState::decode(buffer).map_err(|e| format!("decode CountMinState: {e}"))?, - 1.0, - ), - }; - let rows = state.rows as usize; - let cols = state.cols as usize; - // Defensive dim validation BEFORE reconstructing the matrix: - // reject degenerate / narrow-hash-budget-violating / absurdly - // oversized dims so a malformed payload fails gracefully (the - // ingest caller skips the data point) instead of building a - // degenerate or huge matrix. - validate_sketch_dims("CountMinState", rows, cols)?; - let expected_len = rows * cols; - let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { - format!( - "CountMinState has unknown counter_type tag {}", - state.counter_type - ) - })?; - let flat: Vec = match counter_type { - CounterType::Int32 | CounterType::Int64 => { - if state.counts_int.len() != expected_len { - return Err(format!( - "CountMinState counts_int has {} entries, expected rows*cols = {}", - state.counts_int.len(), - expected_len - ) - .into()); - } - state.counts_int.iter().map(|&v| v as f64).collect() - } - CounterType::Float64 => { - if state.counts_float.len() != expected_len { - return Err(format!( - "CountMinState counts_float has {} entries, expected rows*cols = {}", - state.counts_float.len(), - expected_len - ) - .into()); - } - state.counts_float.clone() - } - // INT128 stores (hi, lo) pairs and would have 2 * rows * cols - // entries in counts_int; defer to PR C if a producer ever uses it. - other => { - return Err(format!( - "CountMinState counter_type {other:?} not yet supported \ - (PR C will extend coverage)" - ) - .into()); - } - }; - let mut matrix = Vec::with_capacity(rows); - for r in 0..rows { - let start = r * cols; - matrix.push(flat[start..start + cols].to_vec()); - } - Ok(Self { - inner: CountMinSketch::from_legacy_matrix(matrix, rows, cols), - sample_p: normalize_sample_p(sample_p), - }) - } - - /// Apply a proto-encoded `CountMinDelta` frame to this - /// accumulator's inner sketch — the decode path for - /// `COUNT_MIN_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). - pub fn apply_proto_delta_bytes( - &mut self, - buffer: &[u8], - ) -> Result<(), Box> { - use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; - use prost::Message; - - let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; - - if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { - return Err(format!( - "CountMinDelta packed-array length mismatch: \ - cell_rows={}, cell_cols={}, d_counts={}", - pb.cell_rows.len(), - pb.cell_cols.len(), - pb.d_counts.len() - ) - .into()); - } - let cells = pb - .cell_rows - .iter() - .zip(pb.cell_cols.iter()) - .zip(pb.d_counts.iter()) - .map(|((r, c), dc)| (*r, *c, *dc)) - .collect(); - let delta = CountMinSketchDelta { - rows: pb.rows, - cols: pb.cols, - cells, - l1: pb.l1, - l2: pb.l2, - // The Go-side CountMinDelta proto now carries an hh_keys field - // (heavy-hitter candidates), mirrored on asap_sketchlib's - // CountMinSketchDelta. The vendored Rust proto bindings here don't - // decode it yet, and CountMin has no TopK to rebuild, so pass an - // empty set — same handling as CountSketch's hh_keys. - hh_keys: Vec::new(), - }; - self.inner - .apply_delta(&delta) - .map_err(|e| format!("apply CountMinDelta: {e}"))?; - Ok(()) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - if buffer.len() < 8 { - return Err("Buffer too short for row_num and col_num".into()); - } - - // TODO: this logic will need to be checked for i32 -> f64 - // Github Issue #11 - - let row_num = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize; - let col_num = u32::from_le_bytes([buffer[4], buffer[5], buffer[6], buffer[7]]) as usize; - - let expected_size = 8 + (row_num * col_num * 4); - if buffer.len() < expected_size { - return Err("Buffer too short for sketch data".into()); - } - - let mut sketch = Vec::new(); - let mut offset = 8; - - for _ in 0..row_num { - let mut row = Vec::new(); - for _ in 0..col_num { - let value = f64::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - buffer[offset + 4], - buffer[offset + 5], - buffer[offset + 6], - buffer[offset + 7], - ]); - row.push(value); - offset += 8; - } - sketch.push(row); - } - - Ok(Self { - inner: CountMinSketch::from_legacy_matrix(sketch, row_num, col_num), - sample_p: 1.0, - }) - } - - /// Merge multiple accumulators efficiently without cloning all of them. - pub fn merge_multiple( - accumulators: &[Box], - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut cms_accumulators = Vec::with_capacity(accumulators.len()); - for acc in accumulators { - if acc.get_accumulator_type() != AggregationType::CountMinSketch { - return Err(format!( - "Cannot merge CountMinSketchAccumulator with {:?}", - acc.get_accumulator_type() - ) - .into()); - } - let cms_acc = acc - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to CountMinSketchAccumulator")?; - cms_accumulators.push(cms_acc); - } - - // Check dimensions are consistent - let rows = cms_accumulators[0].inner.rows(); - let cols = cms_accumulators[0].inner.cols(); - for acc in &cms_accumulators { - if acc.inner.rows() != rows || acc.inner.cols() != cols { - return Err( - "Cannot merge CountMinSketch accumulators with different dimensions".into(), - ); - } - } - - let inner_refs: Vec<&CountMinSketch> = - cms_accumulators.iter().map(|acc| &acc.inner).collect(); - let merged_inner = CountMinSketch::merge_refs(&inner_refs)?; - // sample_p is a per-series config constant, so all operands carry the - // same value in practice. Mirror DDSketch's merge policy: prefer a - // sampled factor (< 1.0) over the no-sampling default so a merge with - // a freshly-reset (1.0) base keeps the series' sampling rate. - let sample_p = cms_accumulators - .iter() - .map(|acc| acc.sample_p) - .find(|&p| p < 1.0) - .unwrap_or(cms_accumulators[0].sample_p); - Ok(Self { - inner: merged_inner, - sample_p, - }) - } -} - -/// Defensive upper bound on the number of matrix cells (`rows * cols`) -/// we'll reconstruct from an inbound wire-declared CMS / CountSketch -/// dimension pair. A malformed / hostile payload could declare absurd -/// dims (e.g. `rows = cols = u32::MAX`) and trick the decoder into a -/// huge `Vec` allocation before the `counts_*.len() != rows*cols` -/// check ever runs. Realistic sketches are at most a few hundred rows -/// by tens-of-thousands of columns, so 8M cells (~64 MiB of f64) is a -/// generous ceiling that no legitimate producer reaches. -pub(crate) const MAX_SKETCH_CELLS: usize = 8 * 1024 * 1024; - -/// Validate an inbound, wire-declared `(rows, cols)` pair for a -/// matrix-backed frequency sketch (CMS / CountSketch) BEFORE any matrix -/// is reconstructed from it. Returns `Ok(())` for dimensions a -/// legitimate producer could have emitted, and an `Err` (never a panic) -/// for malformed / degenerate ones so the ingest path can skip the data -/// point and fall through to its existing decode-failure accounting. -/// -/// Rejections: -/// 1. `rows < 1` or `cols < 1` — a zero-dim matrix has no cells. -/// 2. Narrow-hash-budget violation. The cross-language wire hasher -/// (`sketchlib`'s `MatrixHashType::Packed64`) derives every row's -/// column index from disjoint bit-fields of a single 64-bit hash -/// word: row `r` reads `mask_bits = ceil(log2(cols))` bits at offset -/// `r * mask_bits`. Once `rows * mask_bits > 64` the per-row column -/// slices overflow / alias the 64-bit word and the matrix-cell -/// layout is no longer the one the producer hashed into — the sketch -/// is internally degenerate. This mirrors sketchlib's own -/// `MatrixFastHash::assert_compatible` budget (`rows * (mask_bits + 1) <= 64`); we check the column-index bits alone so realistic -/// configs (5x2048, 5x4096, 5x2000) — for which the sign bits share -/// the top of the word without affecting the cell layout — still -/// pass. -/// 3. Obviously-oversized dims: `rows * cols > MAX_SKETCH_CELLS`, -/// guarding against a huge allocation from a malformed payload. -/// -/// `what` names the wire struct for the error message (e.g. -/// `"CountMinState"`). -pub(crate) fn validate_sketch_dims(what: &str, rows: usize, cols: usize) -> Result<(), String> { - if rows < 1 || cols < 1 { - return Err(format!( - "{what} has degenerate dims (rows={rows}, cols={cols}); rejecting" - )); - } - // mask_bits = ceil(log2(cols)); cols >= 1 here. ilog2 is floor(log2). - let mask_bits = if cols.is_power_of_two() { - cols.ilog2() as usize - } else { - cols.ilog2() as usize + 1 - }; - if rows.saturating_mul(mask_bits) > 64 { - return Err(format!( - "{what} dims (rows={rows}, cols={cols}) exceed the 64-bit \ - packed-hash column budget (rows * ceil(log2(cols)) = {} > 64); \ - the sketch's matrix-cell layout is degenerate, rejecting", - rows.saturating_mul(mask_bits) - )); - } - if rows.saturating_mul(cols) > MAX_SKETCH_CELLS { - return Err(format!( - "{what} dims (rows={rows}, cols={cols}) declare {} cells, \ - exceeding the {MAX_SKETCH_CELLS}-cell ingest cap; rejecting to \ - avoid a huge allocation from a malformed payload", - rows.saturating_mul(cols) - )); - } - Ok(()) -} - -impl SerializableToSink for CountMinSketchAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ - "row_num": self.inner.rows(), - "col_num": self.inner.cols(), - "sketch": self.inner.sketch() - }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.inner.to_msgpack().unwrap_or_default() - } -} - -impl AggregateCore for CountMinSketchAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "CountMinSketchAccumulator" - } - - /// Per-window base rotation: rebuild an empty counter matrix with - /// the same (rows, cols) so the next window's additive cell deltas - /// align to the identical hash geometry. `sample_p` is a per-series - /// config constant (not per-window data), so it is intentionally - /// preserved across the rotation — mirrors `DDSketchAccumulator`. - fn reset_to_empty(&mut self) { - self.inner = CountMinSketch::new(self.inner.rows(), self.inner.cols()); - } - - 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 CountMinSketchAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - - let other_cms = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to CountMinSketchAccumulator")?; - - let merged_inner = CountMinSketch::merge_refs(&[&self.inner, &other_cms.inner])?; - // Mirror DDSketchAccumulator's merge policy exactly: sample_p is a - // per-series config constant, so both operands carry the same value - // in practice. Prefer a sampled factor over the no-sampling default - // so a merge with a freshly-reset (1.0) base keeps the series' - // sampling rate. - let sample_p = if self.sample_p < 1.0 { - self.sample_p - } else { - other_cms.sample_p - }; - Ok(Box::new(Self { - inner: merged_inner, - sample_p, - })) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::CountMinSketch - } - - fn approx_memory_bytes(&self) -> usize { - // Conservative constant for the CountMinSketch counter matrix. - // Real per-instance sizing would require exposing rows/cols on - // the inner sketch; 16 KiB is a reasonable v1 default. - 16 * 1024 - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::MultipleSubpopulationAggregate; - use crate::Statistic; - - // Key-provided path: route to MultipleSubpopulationAggregate::query - // (the canonical "what's the count of this key?" lookup). - 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 = crate::KeyByLabelValues::new_with_labels(vec![k.clone()]); - return self.query(statistic, &key_val, Some(query_kwargs)); - } - - // No-key path: return total event volume. The min-row-sum is the - // canonical CMS estimator for "how many inserts were observed" — - // each insert increments exactly one cell per row, so every row - // sums to the true insert count (modulo collisions, which CMS - // never *underestimates*; min is the tightest upper bound). - // - // When the edge sampled this series (sample_p < 1.0), each insert - // was admitted w.p. `p`, so the stored min-row-sum is ~`p`× the - // true event count. CMS is L1/additive and linear, so rescale by - // `1/sample_p` for an unbiased estimate. `sample_p == 1.0` - // (unsampled / legacy) makes this a no-op. This rescales BOTH the - // Count/Sum/Increase statistics and (via the same closure) the - // Rate per-second readout. - let total_events = || -> f64 { - let matrix = self.inner.sketch(); - if matrix.is_empty() || matrix[0].is_empty() { - return 0.0; - } - let row_totals = matrix.iter().map(|r| r.iter().sum::()); - let min_total = row_totals.fold(f64::INFINITY, f64::min); - if min_total.is_finite() { - min_total / self.sample_p - } else { - 0.0 - } - }; - match statistic { - Statistic::Count | Statistic::Sum => Ok(total_events()), - // PR #111 honest-gap closure (in-the-bag for ASAP tier). - // CMS records insert counts but not timestamps, so per-second - // `rate(metric[range])` requires the engine to push the - // range duration via `query_kwargs["range_ms"]`. When - // present, divide the min-row-sum by `range_ms / 1000`. When - // absent (the engine has not been wired to inject range_ms - // for this query, e.g. instant `rate` calls outside the - // PromQL range-vector pattern), fall back to the raw event - // count so the answer is at least non-empty — the caller's - // caveat is that the units are events/window rather than - // events/second. Increase carries the same caveat. - Statistic::Rate => { - let total = total_events(); - let range_ms_str = query_kwargs.get("range_ms").map(String::as_str); - let Some(s) = range_ms_str else { - return Ok(total); - }; - let range_ms: f64 = s - .parse() - .map_err(|e| format!("CountMinSketchAccumulator: bad range_ms='{s}': {e}"))?; - if range_ms <= 0.0 { - return Err("CountMinSketchAccumulator: range_ms must be positive".into()); - } - Ok(total * 1000.0 / range_ms) - } - Statistic::Increase => Ok(total_events()), - other => Err(format!( - "CountMinSketchAccumulator: statistic {:?} not supported \ - without a key (only Count / Sum / Rate / Increase aggregate \ - over the whole sketch)", - other, - ) - .into()), - } - } -} - -impl MultipleSubpopulationAggregate for CountMinSketchAccumulator { - 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 CountMinSketchAccumulator { - 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_min_sketch_creation() { - let cms = CountMinSketchAccumulator::new(4, 1000); - assert_eq!(cms.inner.rows(), 4); - assert_eq!(cms.inner.cols(), 1000); - let sketch = cms.inner.sketch(); - assert_eq!(sketch.len(), 4); - assert_eq!(sketch[0].len(), 1000); - - for row in &sketch { - for &value in row { - assert_eq!(value, 0.0); - } - } - } - - #[test] - fn test_count_min_sketch_update() { - let mut cms = CountMinSketchAccumulator::new(2, 10); - let key = KeyByLabelValues::new(); - cms._update(&key, 1.0); - let result = cms.query_key(&key); - assert!(result >= 1.0); - } - - #[test] - fn test_count_min_sketch_query() { - let cms = CountMinSketchAccumulator::new(2, 10); - let key = KeyByLabelValues::new(); - assert_eq!(cms.query_key(&key), 0.0); - - let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; - assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); - } - - #[test] - fn test_count_min_sketch_merge() { - // Build controlled state via from_legacy_matrix (works for both Legacy and Sketchlib backends). - let cms1 = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix( - vec![vec![5.0, 0.0, 0.0], vec![0.0, 0.0, 10.0]], - 2, - 3, - ), - sample_p: 1.0, - }; - let cms2 = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix( - vec![vec![3.0, 7.0, 0.0], vec![0.0, 0.0, 0.0]], - 2, - 3, - ), - sample_p: 1.0, - }; - - let merged = CountMinSketchAccumulator::merge_accumulators(vec![cms1, cms2]).unwrap(); - - let merged_sketch = merged.inner.sketch(); - assert_eq!(merged_sketch[0][0], 8.0); - assert_eq!(merged_sketch[0][1], 7.0); - assert_eq!(merged_sketch[1][2], 10.0); - } - - #[test] - fn test_count_min_sketch_merge_dimension_mismatch() { - let cms1 = CountMinSketchAccumulator::new(2, 3); - let cms2 = CountMinSketchAccumulator::new(3, 3); - let result = CountMinSketchAccumulator::merge_accumulators(vec![cms1, cms2]); - assert!(result.is_err()); - } - - #[test] - fn test_count_min_sketch_as_aggregate_core() { - let cms = CountMinSketchAccumulator::new(2, 3); - assert_eq!(cms.type_name(), "CountMinSketchAccumulator"); - } - - #[test] - fn test_trait_object() { - let cms = CountMinSketchAccumulator::new(2, 3); - let trait_obj: Box = Box::new(cms); - assert_eq!(trait_obj.type_name(), "CountMinSketchAccumulator"); - } - - #[test] - fn test_count_min_sketch_key_query() { - let mut cms = CountMinSketchAccumulator::new(4, 100); - let key = KeyByLabelValues::new(); - assert_eq!(cms.query_key(&key), 0.0); - cms._update(&key, 5.0); - let result = cms.query_key(&key); - assert!(result >= 5.0); - } - - #[test] - fn test_update_and_query_use_same_key_encoding() { - // Regression test: _update and query_key must hash the same key string. - // Previously _update went through serialize_to_json (which returns a JSON - // array, so as_object() is always None) and always stored under key "". - // query_key correctly used key.labels.join(";"), so they never matched. - let mut cms = CountMinSketchAccumulator::new(4, 1000); - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string(), "prod".to_string()]); - cms._update(&key, 5.0); - let result = cms.query_key(&key); - assert!( - result >= 5.0, - "_update and query_key used different key encodings: got {result}" - ); - - // Also verify a different key does not interfere. - let other_key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - // other_key was never updated; its estimate should be lower than key's. - let other_result = cms.query_key(&other_key); - // In a sketch this large there should be no collision, so other_result == 0. - assert_eq!( - other_result, 0.0, - "unrelated key returned non-zero: {other_result}" - ); - } - - #[test] - fn test_multiple_subpopulation_aggregate() { - let mut cms = CountMinSketchAccumulator::new(3, 50); - let key = KeyByLabelValues::new(); - cms._update(&key, 10.0); - - let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; - let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); - assert!(result >= 10.0); - - let keys = multi_trait.get_keys(); - assert!(keys.is_none()); - } - - #[test] - fn test_count_min_sketch_merge_multiple() { - // Build controlled state via from_legacy_matrix (works for both Legacy and Sketchlib backends). - let cms1 = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix( - vec![vec![5.0, 0.0, 0.0], vec![0.0, 0.0, 10.0]], - 2, - 3, - ), - sample_p: 1.0, - }; - let cms2 = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix( - vec![vec![3.0, 7.0, 0.0], vec![0.0, 0.0, 0.0]], - 2, - 3, - ), - sample_p: 1.0, - }; - let cms3 = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix( - vec![vec![2.0, 0.0, 0.0], vec![0.0, 0.0, 5.0]], - 2, - 3, - ), - sample_p: 1.0, - }; - - let boxed_accs: Vec> = - vec![Box::new(cms1), Box::new(cms2), Box::new(cms3)]; - - let merged = CountMinSketchAccumulator::merge_multiple(&boxed_accs).unwrap(); - - let merged_sketch = merged.inner.sketch(); - assert_eq!(merged_sketch[0][0], 10.0); - assert_eq!(merged_sketch[0][1], 7.0); - assert_eq!(merged_sketch[1][2], 15.0); - } - - #[test] - fn test_count_min_sketch_merge_multiple_error_cases() { - let empty: Vec> = vec![]; - assert!(CountMinSketchAccumulator::merge_multiple(&empty).is_err()); - - let cms1 = CountMinSketchAccumulator::new(2, 3); - let cms2 = CountMinSketchAccumulator::new(3, 3); - let boxed_accs: Vec> = vec![Box::new(cms1), Box::new(cms2)]; - assert!(CountMinSketchAccumulator::merge_multiple(&boxed_accs).is_err()); - - use crate::summary_kernels::sum::SumAccumulator; - let cms = CountMinSketchAccumulator::new(2, 3); - let sum = SumAccumulator::new(); - let mixed_accs: Vec> = vec![Box::new(cms), Box::new(sum)]; - assert!(CountMinSketchAccumulator::merge_multiple(&mixed_accs).is_err()); - } - - #[test] - fn test_from_sketchlib_proto_bytes_int64() { - // Hand-build a CountMinState proto with INT64 counters and verify - // round-tripping through from_sketchlib_proto_bytes yields the same - // matrix that the modified-OTLP wire format would carry. - use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; - use prost::Message; - - let rows = 2u32; - let cols = 3u32; - // Row-major: row 0 = [1,2,3], row 1 = [4,5,6] - let counts_int: Vec = vec![1, 2, 3, 4, 5, 6]; - let state = CountMinState { - rows, - cols, - counter_type: CounterType::Int64 as i32, - counts_int: counts_int.clone(), - counts_float: Vec::new(), - sum_counts: Vec::new(), - sum2_counts: Vec::new(), - l1: Vec::new(), - l2: Vec::new(), - }; - let bytes = state.encode_to_vec(); - - let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - let matrix = acc.inner.sketch(); - assert_eq!(matrix.len(), rows as usize); - assert_eq!(matrix[0], vec![1.0, 2.0, 3.0]); - assert_eq!(matrix[1], vec![4.0, 5.0, 6.0]); - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrapped() { - // Mirrors what DataCollector's countminsketchprocessor emits: - // the state is wrapped in a `SketchEnvelope{count_min: ...}` - // via sketchlib-go's `SerializePortableFO` + `proto.Marshal`. - // Before the fix, the Rust decoder decoded the envelope bytes as - // a bare CountMinState, which produced "invalid wire type" - // errors on field `cols` and silently fell through to §5.2. - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, CountMinState, CounterType, SketchEnvelope, - }; - use prost::Message; - - let state = CountMinState { - rows: 2, - cols: 3, - counter_type: CounterType::Int64 as i32, - counts_int: vec![7, 8, 9, 10, 11, 12], - counts_float: Vec::new(), - sum_counts: Vec::new(), - sum2_counts: Vec::new(), - l1: Vec::new(), - l2: Vec::new(), - }; - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::CountMin(state)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes) - .expect("envelope-wrapped decode should succeed"); - let matrix = acc.inner.sketch(); - assert_eq!(matrix[0], vec![7.0, 8.0, 9.0]); - assert_eq!(matrix[1], vec![10.0, 11.0, 12.0]); - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { - // An envelope carrying a non-CountMin sketch should be rejected - // with a clear error rather than silently producing garbage. - use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; - use prost::Message; - - let kll = KllState::default(); - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Kll(kll)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "wrong-sketch envelope should error"); - } - - #[test] - fn test_from_sketchlib_proto_bytes_float64() { - use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; - use prost::Message; - - let state = CountMinState { - rows: 2, - cols: 2, - counter_type: CounterType::Float64 as i32, - counts_int: Vec::new(), - counts_float: vec![1.5, 2.5, 3.5, 4.5], - sum_counts: Vec::new(), - sum2_counts: Vec::new(), - l1: Vec::new(), - l2: Vec::new(), - }; - let bytes = state.encode_to_vec(); - - let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - let matrix = acc.inner.sketch(); - assert_eq!(matrix[0], vec![1.5, 2.5]); - assert_eq!(matrix[1], vec![3.5, 4.5]); - } - - #[test] - fn test_from_sketchlib_proto_bytes_dimension_mismatch() { - // counts_int has 5 entries but rows*cols = 6 → expect error - use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; - use prost::Message; - - let state = CountMinState { - rows: 2, - cols: 3, - counter_type: CounterType::Int64 as i32, - counts_int: vec![1, 2, 3, 4, 5], - counts_float: Vec::new(), - sum_counts: Vec::new(), - sum2_counts: Vec::new(), - l1: Vec::new(), - l2: Vec::new(), - }; - let bytes = state.encode_to_vec(); - - let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - assert!( - result.unwrap_err().to_string().contains("counts_int"), - "error should mention counts_int dim mismatch" - ); - } - - #[test] - fn test_from_sketchlib_proto_bytes_zero_dims_rejected() { - use asap_sketchlib::proto::sketchlib::CountMinState; - use prost::Message; - - let state = CountMinState::default(); - let bytes = state.encode_to_vec(); - - let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("degenerate dims")); - } - - #[test] - fn test_apply_proto_delta_bytes_round_trip() { - use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; - use prost::Message; - - let mut acc = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix( - vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], - 2, - 3, - ), - sample_p: 1.0, - }; - let bytes = PbDelta { - rows: 2, - cols: 3, - cell_rows: vec![0, 1], - cell_cols: vec![0, 2], - d_counts: vec![10, 100], - l1: vec![], - l2: vec![], - ..Default::default() - } - .encode_to_vec(); - - acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); - assert_eq!( - acc.inner.sketch(), - vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 106.0]] - ); - } - - #[test] - fn test_apply_proto_delta_bytes_rejects_garbage() { - let mut acc = CountMinSketchAccumulator::new(2, 3); - assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); - } - - // ---------------------------------------------------------------- - // Statistic::Rate / Statistic::Increase — PR #111 honest-gap closure. - // CMS records insert counts but not timestamps. The Rate readout - // requires the engine to push `range_ms` via query_kwargs; without - // it the accumulator falls back to the raw event count (units of - // events/window) so the answer is at least non-empty. - // ---------------------------------------------------------------- - - #[test] - fn test_query_statistic_rate_with_range_ms() { - // Build a CMS whose min-row-sum is 100 events. With a 5-minute - // (300_000 ms) range, the per-second rate is 100 / 300 ≈ 0.333. - let cms = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix( - vec![vec![100.0, 0.0], vec![100.0, 0.0]], - 2, - 2, - ), - sample_p: 1.0, - }; - let mut kwargs = HashMap::new(); - kwargs.insert("range_ms".to_string(), "300000".to_string()); - let trait_obj: &dyn AggregateCore = &cms; - let v = trait_obj - .query_statistic(Statistic::Rate, &None, &kwargs) - .expect("Rate with range_ms is supported"); - assert!( - (v - (100.0 / 300.0)).abs() < 1e-9, - "expected 100/300 = {}, got {v}", - 100.0 / 300.0, - ); - } - - #[test] - fn test_query_statistic_rate_without_range_ms_falls_back_to_count() { - // Without `range_ms` in kwargs the accumulator returns the raw - // event volume (events/window units). Caller is responsible for - // surfacing that caveat to the user; this avoids `status=error` - // for instant rate-shape queries that bypass the matrix-selector - // code path. - let cms = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix(vec![vec![42.0, 0.0], vec![42.0, 0.0]], 2, 2), - sample_p: 1.0, - }; - let trait_obj: &dyn AggregateCore = &cms; - let v = trait_obj - .query_statistic(Statistic::Rate, &None, &HashMap::new()) - .expect("Rate without range_ms still answers (fallback)"); - assert_eq!(v, 42.0); - } - - #[test] - fn test_query_statistic_increase_returns_total_count() { - // Increase semantics on CMS: total events in the window — the - // same min-row-sum as Sum / Count. Differs from Rate only in - // that it never divides by range. - let cms = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix(vec![vec![5.0, 7.0], vec![3.0, 9.0]], 2, 2), - sample_p: 1.0, - }; - let trait_obj: &dyn AggregateCore = &cms; - let v = trait_obj - .query_statistic(Statistic::Increase, &None, &HashMap::new()) - .expect("Increase is supported"); - // min-row-sum: row0 = 12, row1 = 12, min = 12. - assert_eq!(v, 12.0); - } - - // ---------------------------------------------------------------- - // Defensive inbound-dimension validation (harden/sketch-dim-validation). - // Malformed / degenerate / narrow-hash-budget-violating CMS dims must - // be rejected gracefully (Err, never a panic); valid configs the - // backend actually uses (5x2048, 5x4096, 5x2000) must still decode. - // ---------------------------------------------------------------- - - /// Build a bare `CountMinState` proto carrying the given dims and a - /// row-major INT64 counts vector sized to `rows*cols` so that, IF the - /// dims pass validation, the reshape also succeeds. Used to prove a - /// malformed-dim payload is rejected at the dim gate, not later. - fn cms_state_bytes(rows: u32, cols: u32) -> Vec { - use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; - use prost::Message; - let n = (rows as usize).saturating_mul(cols as usize); - let state = CountMinState { - rows, - cols, - counter_type: CounterType::Int64 as i32, - counts_int: vec![0i64; n], - counts_float: Vec::new(), - sum_counts: Vec::new(), - sum2_counts: Vec::new(), - l1: Vec::new(), - l2: Vec::new(), - }; - state.encode_to_vec() - } - - #[test] - fn test_validate_sketch_dims_accepts_valid_configs() { - // The realistic configs the backend uses must pass unchanged. - for (r, c) in [(5usize, 2048usize), (5, 4096), (5, 2000), (4, 1000), (2, 3)] { - assert!( - validate_sketch_dims("CountMinState", r, c).is_ok(), - "valid config {r}x{c} was wrongly rejected" - ); - } - } - - #[test] - fn test_validate_sketch_dims_rejects_malformed() { - // Zero dims. - assert!(validate_sketch_dims("CountMinState", 0, 2048).is_err()); - assert!(validate_sketch_dims("CountMinState", 5, 0).is_err()); - // Narrow-hash-budget violation: 5 * ceil(log2(8192))=5*13=65 > 64. - let err = validate_sketch_dims("CountMinState", 5, 8192).unwrap_err(); - assert!(err.contains("budget"), "expected budget error, got: {err}"); - // Absurdly oversized: 1 x 16,777,216 = 16M cells > 8M cap. (1 row - // keeps the hash budget tiny — 1*24=24 — so the cap check, not the - // budget check, is what fires here.) - let err = validate_sketch_dims("CountMinState", 1, 16_777_216).unwrap_err(); - assert!(err.contains("cap"), "expected cell-cap error, got: {err}"); - // No panic on extreme dims (saturating_mul guards the products). - assert!(validate_sketch_dims("CountMinState", usize::MAX, usize::MAX).is_err()); - } - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { - // A data point declaring narrow-hash-budget-violating dims must be - // skipped (Err returned, NOT a panic). The ingest caller turns - // this Err into a dropped data point + WARN log. - let bytes = cms_state_bytes(5, 8192); - let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "budget-violating dims should be rejected"); - assert!(result.unwrap_err().to_string().contains("rejecting")); - - // A valid neighbour (5x4096) on the same path still decodes fine. - let ok_bytes = cms_state_bytes(5, 4096); - let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) - .expect("valid 5x4096 CMS should still decode"); - assert_eq!(acc.inner.rows(), 5); - assert_eq!(acc.inner.cols(), 4096); - } - - #[test] - fn test_query_statistic_rate_rejects_invalid_range_ms() { - let cms = CountMinSketchAccumulator::new(2, 2); - let mut kwargs = HashMap::new(); - kwargs.insert("range_ms".to_string(), "0".to_string()); - let trait_obj: &dyn AggregateCore = &cms; - let err = trait_obj - .query_statistic(Statistic::Rate, &None, &kwargs) - .expect_err("range_ms=0 should error"); - assert!(err.to_string().contains("positive")); - - let mut kwargs = HashMap::new(); - kwargs.insert("range_ms".to_string(), "not-a-number".to_string()); - let err = trait_obj - .query_statistic(Statistic::Rate, &None, &kwargs) - .expect_err("non-numeric range_ms should error"); - assert!(err.to_string().contains("bad range_ms")); - } - - // ---------------------------------------------------------------- - // sample_p rescale. The edge admits each insert with probability `p`, - // so every stored cell is ~p× the true count. CMS is L1/additive and - // linear, so BOTH the point-frequency (query_key) and the aggregate - // total-event statistics (Count/Sum/Increase/Rate) rescale by 1/p. - // ---------------------------------------------------------------- - - #[test] - fn test_query_key_rescaled_by_sample_p() { - // Same stored cell counts, two sample_p values: the p=0.25 sketch - // must report 4× the point-frequency of the unsampled one. - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - let mut unsampled = CountMinSketchAccumulator::new(4, 1000); - unsampled._update(&key, 10.0); - let mut sampled = CountMinSketchAccumulator::new(4, 1000); - sampled._update(&key, 10.0); - sampled.sample_p = 0.25; - - let raw = unsampled.query_key(&key); - let rescaled = sampled.query_key(&key); - assert!( - raw >= 10.0, - "raw estimate should be >= inserted 10, got {raw}" - ); - assert!( - (rescaled - raw * 4.0).abs() < 1e-9, - "expected point-frequency rescaled ≈ 4×raw ({}), got {rescaled}", - raw * 4.0 - ); - } - - #[test] - fn test_aggregate_statistics_rescaled_by_sample_p() { - use crate::Statistic; - // Build a CMS with a known min-row-sum of 12 events, sampled at - // p=0.25 → every aggregate statistic should report 12 / 0.25 = 48. - let cms = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix(vec![vec![5.0, 7.0], vec![3.0, 9.0]], 2, 2), - sample_p: 0.25, - }; - let trait_obj: &dyn AggregateCore = &cms; - for stat in [Statistic::Count, Statistic::Sum, Statistic::Increase] { - let v = trait_obj - .query_statistic(stat, &None, &HashMap::new()) - .unwrap_or_else(|e| panic!("{stat:?} should be supported: {e}")); - // min-row-sum = 12, rescaled by 1/0.25 = 48. - assert!( - (v - 48.0).abs() < 1e-9, - "{stat:?}: expected rescaled 48, got {v}" - ); - } - // Rate also divides through the rescaled total: 48 events over a - // 6-second (6000 ms) range = 8 events/s. - let mut kwargs = HashMap::new(); - kwargs.insert("range_ms".to_string(), "6000".to_string()); - let r = trait_obj - .query_statistic(Statistic::Rate, &None, &kwargs) - .expect("rate ok"); - assert!((r - 8.0).abs() < 1e-9, "expected rate 8.0, got {r}"); - } - - #[test] - fn test_sample_p_unset_behaves_as_one() { - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, CountMinState, CounterType, SketchEnvelope, - }; - use prost::Message; - // An envelope with no sample_p (proto3 default 0.0) must normalize - // to 1.0 (no rescale) — byte-compatible with legacy frames. - let state = CountMinState { - rows: 2, - cols: 2, - counter_type: CounterType::Int64 as i32, - counts_int: vec![1, 2, 3, 4], - counts_float: Vec::new(), - sum_counts: Vec::new(), - sum2_counts: Vec::new(), - l1: Vec::new(), - l2: Vec::new(), - }; - let env = SketchEnvelope { - // sample_p left at proto3 default 0.0. - sketch_state: Some(sketch_envelope::SketchState::CountMin(state)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.sample_p, 1.0, "unset sample_p must normalize to 1.0"); - } - - #[test] - fn test_from_sketchlib_proto_bytes_reads_envelope_sample_p() { - use crate::Statistic; - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, CountMinState, CounterType, SketchEnvelope, - }; - use prost::Message; - // min-row-sum = 12 raw; sample_p 0.25 → Count = 48. - let state = CountMinState { - rows: 2, - cols: 2, - counter_type: CounterType::Float64 as i32, - counts_int: Vec::new(), - counts_float: vec![5.0, 7.0, 3.0, 9.0], - sum_counts: Vec::new(), - sum2_counts: Vec::new(), - l1: Vec::new(), - l2: Vec::new(), - }; - let env = SketchEnvelope { - sample_p: 0.25, - sketch_state: Some(sketch_envelope::SketchState::CountMin(state)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.sample_p, 0.25); - let trait_obj: &dyn AggregateCore = &acc; - let v = trait_obj - .query_statistic(Statistic::Count, &None, &HashMap::new()) - .expect("count ok"); - assert!((v - 48.0).abs() < 1e-9, "expected rescaled 48, got {v}"); - } - - #[test] - fn test_reset_to_empty_preserves_sample_p() { - let mut acc = CountMinSketchAccumulator::new(2, 3); - acc.sample_p = 0.25; - acc.reset_to_empty(); - assert_eq!(acc.sample_p, 0.25, "window rotation must keep sample_p"); - } - - #[test] - fn test_merge_prefers_sampled_factor() { - let mut a = CountMinSketchAccumulator::new(2, 3); - a.sample_p = 0.25; - let b = CountMinSketchAccumulator::new(2, 3); // sample_p 1.0 - let merged = a.merge_with(&b).expect("merge ok"); - let merged = merged - .as_any() - .downcast_ref::() - .expect("downcast ok"); - assert_eq!(merged.sample_p, 0.25); - - // merge_multiple mirrors the same policy. - let mut c = CountMinSketchAccumulator::new(2, 3); - c.sample_p = 0.25; - let d = CountMinSketchAccumulator::new(2, 3); - let boxed: Vec> = vec![Box::new(d), Box::new(c)]; - let merged = CountMinSketchAccumulator::merge_multiple(&boxed).expect("merge ok"); - assert_eq!(merged.sample_p, 0.25); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/count_min_sketch_with_heap.rs b/crates/asap_summary_state/src/summary_kernels/count_min_sketch_with_heap.rs deleted file mode 100644 index f5d1369d1..000000000 --- a/crates/asap_summary_state/src/summary_kernels/count_min_sketch_with_heap.rs +++ /dev/null @@ -1,832 +0,0 @@ -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{CmsHeapItem, CountMinSketchWithHeap, MessagePackCodec}; -use serde::Deserialize; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -/// Local serde view of the DELTA-HEAP wire frame produced by sketchlib-go's -/// `CountSketch.SerializeMsgpackWithHeapDelta` (encoding `MSGPACK_DELTA`). -/// Decoded with `rmp_serde` directly in the backend so NO delta API needs to -/// be added to the public `asap_sketchlib`. -/// -/// rmp_serde compact layout — a 4-element positional array: -/// -/// [ -/// is_delta: bool (always true), -/// matrix_delta: ( rows:u32, cols:u32, cells: Vec<(u32,u32,i64)> ), -/// topk_heap: Vec<(String, f64)>, // FULL heap, [key, value] pairs -/// heap_size: u64, -/// ] -/// -/// Tuple structs deserialize from msgpack fixed arrays positionally, so this -/// matches the Go encoder's byte layout exactly (no field names on the wire). -#[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-Min Sketch with Heap accumulator — wraps `asap_sketchlib::CountMinSketchWithHeap`. -/// Core struct, update/merge/serde logic live in `asap_sketchlib::message_pack_format::portable::countminsketch_topk`. -/// This file retains QE-specific trait impls, legacy deserializers, and JSON output. -#[derive(Debug, Clone)] -pub struct CountMinSketchWithHeapAccumulator { - pub inner: CountMinSketchWithHeap, -} - -// Re-export HeapItem so existing code using CountMinSketchWithHeapAccumulator::HeapItem still works. -pub use asap_sketchlib::CmsHeapItem as HeapItemReexport; - -impl CountMinSketchWithHeapAccumulator { - pub fn new(row_num: usize, col_num: usize, heap_size: usize) -> Self { - Self { - inner: CountMinSketchWithHeap::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 - /// (`{sketch:[matrix,rows,cols], topk_heap, heap_size}`) into a heap - /// accumulator. This is the window-1 / full-frame base for the - /// DELTA-HEAP delta path: the backend caches THIS accumulator as the - /// per-series base so a later `MSGPACK_DELTA` frame applies its sparse - /// matrix delta onto a heap accumulator (not a plain CountSketch). - /// - /// Delegates to the PUBLIC `asap_sketchlib::CountMinSketchWithHeap:: - /// from_msgpack` (both heap-bearing frequency variants share the wire - /// shape; the CountSketch-with-heap promotion is decided by the ingest - /// router, not the bytes). - pub fn from_msgpack_with_heap_bytes(buffer: &[u8]) -> Result> { - Ok(Self { - inner: CountMinSketchWithHeap::from_msgpack(buffer) - .map_err(|e| format!("deserialize CountMinSketchWithHeap msgpack: {e}"))?, - }) - } - - /// Apply a DELTA-HEAP msgpack frame (encoding `MSGPACK_DELTA`) onto this - /// accumulator IN PLACE, WITHOUT any change to the public - /// `asap_sketchlib`: the frame is decoded generically with `rmp_serde` - /// into local serde structs, the sparse signed cell deltas are added to - /// the stored matrix (read back via the public `sketch_matrix()`), and - /// the top-k heap is REPLACED with the frame's full heap. The rebuilt - /// inner is produced via the public `from_legacy_matrix`, which rounds - /// cells to the i64 storage and re-seeds the heap. - /// - /// Under the per-window-reset model (`docs/delta-baseline-contract.md` - /// §3) the ingest caller resets this accumulator to empty at a window - /// boundary before applying, so the delta — which is the window's own - /// matrix against an empty base — reconstructs the window's state. - 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; - - // Read the current (post-reset, possibly empty) matrix and apply the - // sparse signed deltas additively. Cells outside the stored - // dimensions are skipped defensively (mirrors the plain-CountSketch - // delta apply). - 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; - } - - // Replace the heap with the frame's full heap. `from_legacy_matrix` - // re-seeds both the matrix and the heap from these inputs. - let heap: Vec = frame - .heap - .into_iter() - .map(|(key, value)| CmsHeapItem { key, value }) - .collect(); - - self.inner = - CountMinSketchWithHeap::from_legacy_matrix(matrix, heap, rows, cols, heap_size); - Ok(()) - } - - /// Reconstruct a heap accumulator STANDALONE from a single DELTA-HEAP - /// msgpack frame (encoding `MSGPACK_DELTA`), with NO cached per-series - /// base. Used by the read-side reducer's `FrequencyTopk` path, where — - /// unlike the ingest accumulator — there is no rolling base to apply - /// onto: under the per-window-reset contract - /// (`docs/delta-baseline-contract.md` §3) each window's delta encodes - /// that window's own state against an EMPTY base, so reconstruction is - /// "empty(dims) + apply(delta)". - /// - /// Reuses the exact ingest-side apply logic: read the (rows, cols, - /// heap_size) the frame declares, build an empty accumulator of those - /// dims (equivalent to `reset_to_empty` on a same-shape base), then - /// fold the frame in via `apply_msgpack_heap_delta_bytes`. No - /// `asap_sketchlib` change — the frame is decoded generically with - /// `rmp_serde`. - 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) - } - - /// This function seems will never be used anymore. Keep it for possible future use. - pub fn deserialize_from_json(data: &Value) -> Result> { - let row_num = data["row_num"] - .as_f64() - .ok_or("Missing or invalid 'row_num' field")? as usize; - let col_num = data["col_num"] - .as_f64() - .ok_or("Missing or invalid 'col_num' field")? as usize; - let heap_size = data["heap_size"] - .as_f64() - .ok_or("Missing or invalid 'heap_size' field")? as usize; - - let sketch_data = data["sketch"] - .as_array() - .ok_or("Missing or invalid 'sketch' field")?; - - let mut sketch = Vec::new(); - for row in sketch_data { - let row_array = row.as_array().ok_or("Invalid row in sketch data")?; - let mut sketch_row = Vec::new(); - for cell in row_array { - let value = cell.as_f64().ok_or("Invalid cell value in sketch data")?; - sketch_row.push(value); - } - sketch.push(sketch_row); - } - - let topk_heap_data = data["topk_heap"] - .as_array() - .ok_or("Missing or invalid 'topk_heap' field")?; - - let mut topk_heap = Vec::new(); - for item in topk_heap_data { - let key = item["key"] - .as_str() - .ok_or("Missing or invalid 'key' in heap item")? - .to_string(); - let value = item["value"] - .as_f64() - .ok_or("Missing or invalid 'value' in heap item")?; - topk_heap.push(CmsHeapItem { key, value }); - } - - Ok(Self { - inner: CountMinSketchWithHeap::from_legacy_matrix( - sketch, topk_heap, row_num, col_num, heap_size, - ), - }) - } - - pub fn deserialize_from_bytes(_buffer: &[u8]) -> Result> { - Err("deserialize_from_bytes for CountMinSketchWithHeapAccumulator not implemented".into()) - } - - /// VALUE-WEIGHTED heavy-hitter update (FIX: CountSketch/CMS topk - /// recall-0). The default ingest path inserts `+1` per occurrence keyed - /// by the raw `item`, so the heap ranks groups by OCCURRENCE COUNT — the - /// wrong answer for `topk(k, sum by (label) (metric))`, which asks for - /// the top groups by SUM OF VALUE. This update adds the sample `value` - /// (not `+1`) into both the CMS matrix and the top-k heap, keyed by the - /// GROUP LABEL (e.g. the `host` / `zone` value), so the heap's ranking is - /// by summed value. Repeated calls for the same `group_label` accumulate, - /// so after folding a window the heap holds Σvalue per group. - /// - /// Delegates to the library's value-weighted `CountMinSketchWithHeap:: - /// update(key, value)` (`sketchlib_cms_heap_update` → `insert_many(key, - /// round(value))`), which is the "separate update path" the evaluation - /// plan (Fig 3c) called for. - 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), keyed by - /// the group label. Pairs with [`Self::insert_value`]: the heap built by - /// value-weighted updates ranks by Σvalue, so this returns the - /// value-weighted top-k (not the occurrence-count top-k the raw `item` - /// heap would give). Sorted descending by value; ties broken by key for - /// determinism; truncated to `k`. - 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 CountMinSketchWithHeapAccumulator { - 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 CountMinSketchWithHeapAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "CountMinSketchWithHeapAccumulator" - } - - /// Per-window base rotation (`docs/delta-baseline-contract.md` §3): - /// rebuild an empty heap accumulator with the same (rows, cols, - /// heap_size) so the next window's DELTA-HEAP frame applies onto a clean, - /// same-shape base. Without this override the trait default is a no-op, - /// which would let the additive matrix delta accumulate across windows - /// (over-counting). Mirrors `CountSketchAccumulator::reset_to_empty`. - fn reset_to_empty(&mut self) { - self.inner = - CountMinSketchWithHeap::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 CountMinSketchWithHeapAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - - let other_cms = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to CountMinSketchWithHeapAccumulator")?; - - let merged = Self::merge_accumulators(vec![self.clone(), other_cms.clone()])?; - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::CountMinSketchWithHeap - } - - fn get_keys(&self) -> Option> { - Some(self.get_topk_keys()) - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::MultipleSubpopulationAggregate; - let key_val = key - .as_ref() - .ok_or("Key required for CountMinSketchWithHeapAccumulator")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for CountMinSketchWithHeapAccumulator { - 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 CountMinSketchWithHeapAccumulator { - 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_min_sketch_with_heap_creation() { - let cms = CountMinSketchWithHeapAccumulator::new(4, 1000, 20); - assert_eq!(cms.inner.rows(), 4); - assert_eq!(cms.inner.cols(), 1000); - assert_eq!(cms.inner.heap_size, 20); - assert_eq!(cms.inner.topk_heap_items().len(), 0); - } - - #[test] - fn test_count_min_sketch_with_heap_query() { - let cms = CountMinSketchWithHeapAccumulator::new(2, 10, 5); - let key = KeyByLabelValues::new(); - assert_eq!(cms.query_key(&key), 0.0); - - let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; - assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); - } - - #[test] - fn test_count_min_sketch_with_heap_merge() { - // Build controlled state via from_legacy_matrix (works regardless of backend config). - 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![ - CmsHeapItem { - key: "key1".to_string(), - value: 100.0, - }, - CmsHeapItem { - 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![ - CmsHeapItem { - key: "key3".to_string(), - value: 75.0, - }, - CmsHeapItem { - key: "key1".to_string(), - value: 80.0, - }, - ]; - - let cms1 = CountMinSketchWithHeapAccumulator { - inner: CountMinSketchWithHeap::from_legacy_matrix(sketch1, heap1, 2, 10, 5), - }; - let cms2 = CountMinSketchWithHeapAccumulator { - inner: CountMinSketchWithHeap::from_legacy_matrix(sketch2, heap2, 2, 10, 3), - }; - - let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); - 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_min_sketch_with_heap_merge_single() { - let cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); - let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms.clone()]); - assert!(result.is_ok()); - let merged = result.unwrap(); - assert_eq!(merged.inner.rows(), cms.inner.rows()); - assert_eq!(merged.inner.cols(), cms.inner.cols()); - assert_eq!(merged.inner.heap_size, cms.inner.heap_size); - } - - #[test] - fn test_count_min_sketch_with_heap_merge_dimension_mismatch() { - let cms1 = CountMinSketchWithHeapAccumulator::new(2, 10, 5); - let cms2 = CountMinSketchWithHeapAccumulator::new(3, 10, 5); - let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("dimension")); - } - - #[test] - fn test_count_min_sketch_with_heap_as_aggregate_core() { - let cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); - assert_eq!(cms.type_name(), "CountMinSketchWithHeapAccumulator"); - } - - #[test] - fn test_get_topk_keys() { - let mut cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); - cms.inner.update("label1;label2", 100.0); - cms.inner.update("label3;label4", 50.0); - - let keys = cms.get_topk_keys(); - assert_eq!(keys.len(), 2); - // Top-k order can differ between Legacy and Sketchlib backends (heap ordering / estimates). - 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 cms = CountMinSketchWithHeapAccumulator::new(3, 50, 10); - let key = KeyByLabelValues::new(); - - let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; - 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); - } - - // ---------------------------------------------------------------- - // DELTA-HEAP wire form (encoding MSGPACK_DELTA): apply a sparse matrix - // delta + replace the heap, decoded generically (rmp_serde) WITHOUT any - // asap_sketchlib delta API. The first test feeds a frame produced by the - // Go encoder (sketchlib-go `MarshalCountSketchWithHeapDelta`) to prove - // cross-language byte parity — mirrors how the full-heap parity is - // proven. The second proves PWR full -> delta -> delta reconstruction. - // ---------------------------------------------------------------- - - /// Cross-language byte-parity: this hex is the exact output of - /// sketchlib-go's `asapmsgpack.MarshalCountSketchWithHeapDelta(5, 1024, - /// cells=[(0,1,50),(1,3,-4),(4,1023,1_000_000)], - /// heap=[("/checkout",50),("/cart",20)], heap_size=20)` (captured via a - /// throw-away Go print test, identical methodology to the full-heap - /// golden in `sketchlib-go/.../count_sketch_with_heap_test.go`). If the - /// Go encoder or the rmp_serde layout ever shifts, this decode fails - /// loudly. - const GO_DELTA_HEAP_GOLDEN_HEX: &str = "94c39305cd04009393000132930103fc9304cd03ffce000f42409292a92f636865636b6f7574cb404900000000000092a52f63617274cb403400000000000014"; - - #[test] - fn test_apply_go_produced_delta_heap_frame_matrix_and_heap() { - let bytes = hex::decode(GO_DELTA_HEAP_GOLDEN_HEX).expect("hex"); - - // Base = empty heap accumulator with the frame's dims (what the - // ingest caller holds after the per-window base rotation). - let mut acc = CountMinSketchWithHeapAccumulator::new(5, 1024, 20); - acc.apply_msgpack_heap_delta_bytes(&bytes) - .expect("apply Go delta-heap frame"); - - // Matrix: the three sparse cells landed onto the empty base. - let m = acc.inner.sketch_matrix(); - assert_eq!(m.len(), 5); - assert_eq!(m[0].len(), 1024); - assert_eq!(m[0][1], 50.0, "cell (0,1)"); - assert_eq!(m[1][3], -4.0, "cell (1,3)"); - assert_eq!(m[4][1023], 1_000_000.0, "cell (4,1023)"); - // Everything else stays zero. - assert_eq!(m[2][2], 0.0); - assert_eq!(m[0][0], 0.0); - - // Heap: the frame's full heap, with /checkout ranked above /cart. - let mut items = acc.inner.topk_heap_items(); - items.sort_by(|a, b| b.value.partial_cmp(&a.value).unwrap()); - assert_eq!(items.len(), 2); - assert_eq!(items[0].key, "/checkout"); - assert_eq!(items[0].value, 50.0); - assert_eq!(items[1].key, "/cart"); - assert_eq!(items[1].value, 20.0); - } - - #[test] - fn test_pwr_full_then_delta_then_delta_reconstructs_per_window() { - use asap_sketchlib::MessagePackCodec; - - // Window 1 (full frame): build a heap-bearing CountSketch with mass - // and serialize the FULL `{sketch,topk_heap,heap_size}` frame, then - // decode it into a heap accumulator (the cached per-series base). - let w1 = CountMinSketchWithHeap::from_legacy_matrix( - vec![vec![300.0; 4]; 5], - vec![CmsHeapItem { - key: "k".into(), - value: 300.0, - }], - 5, - 4, - 20, - ); - let w1_bytes = w1.to_msgpack().expect("w1 full msgpack"); - let mut base = CountMinSketchWithHeapAccumulator::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); - - // Window 2 delta: this window's own state is matrix cells of value 50 - // against an EMPTY base + heap {k:50}. The DELTA-HEAP frame is encoded - // the same way the Go producer does (4-array, is_delta, sparse cells). - let w2_frame = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); - // PWR: rotate base to empty at the window boundary, then apply. - 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); - // No cross-window leakage from window 1's 300s. - 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); - - // Window 3 delta: 80s against empty + heap {k:80}. - 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_rmp_serde_layout_is_byte_identical_to_go_encoder() { - // The rmp_serde positional encoding of the delta-heap frame must be - // BYTE-IDENTICAL to sketchlib-go's hand-rolled - // `MarshalCountSketchWithHeapDelta`. This hex is the Go encoder's - // output for (5, 4, cells=[(0,0,50),(1,1,50)], heap=[("k",50)], - // heap_size=20) — the same inputs `encode_delta_heap` uses below. - // Equality here proves both encode AND decode are cross-language - // byte-compatible (the decode path is exercised by the Go-golden - // test above). - const GO_PARITY_HEX: &str = "94c39305049293000032930101329192a16bcb404900000000000014"; - let rust_bytes = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); - assert_eq!(hex::encode(&rust_bytes), GO_PARITY_HEX); - } - - #[test] - fn test_apply_delta_rejects_full_frame_and_garbage() { - use asap_sketchlib::MessagePackCodec; - let mut acc = CountMinSketchWithHeapAccumulator::new(2, 4, 5); - // A FULL frame (3-array, no is_delta marker) must NOT decode as a - // delta — the routing relies on the two shapes being distinct. - let full = CountMinSketchWithHeap::from_legacy_matrix( - vec![vec![1.0; 4]; 2], - vec![CmsHeapItem { - 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()); - } - - /// Encode a DELTA-HEAP frame the same way sketchlib-go's - /// `MarshalCountSketchWithHeapDelta` does (rmp_serde positional layout), - /// so the test exercises the real decode path. Tuple structs serialize - /// as msgpack fixed arrays — byte-identical to the Go hand-rolled writer. - 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") - } - - // ---------------------------------------------------------------- - // FIX 1 — VALUE-WEIGHTED top-k (recall 0 → correct). - // - // `topk(k, sum by (host) (cpu_load))` asks for the top-k hosts by - // SUM OF VALUE. The heavy-hitter heap built by the default `+1`-per- - // occurrence update ranks by COUNT keyed by `item`, so its recall - // against the value-weighted ground truth is 0 when the busiest host - // (most samples) is NOT the heaviest host (largest Σvalue). - // `insert_value(group_label, value)` adds the sample VALUE keyed by the - // GROUP LABEL, so `topk_by_value` ranks by Σvalue — correct recall. - // ---------------------------------------------------------------- - - /// Crafted adversarial dataset: the host with the MOST samples - /// (`h_chatty`, 100 tiny samples) is NOT the host with the largest - /// value-sum (`h_heavy`, a handful of huge samples). A COUNT-ranked - /// heap would surface `h_chatty`; the value-weighted top-k must surface - /// the true heavy hitters by Σvalue, giving recall 1.0 against the - /// ground-truth top-k-by-value-sum. - #[test] - fn value_weighted_topk_has_full_recall_vs_count_topk() { - // (host, per-sample value, sample count) → true Σvalue: - // h_heavy : 1000 × 3 = 3000 (few samples, huge value) - // h_mid : 200 × 5 = 1000 - // h_small : 50 × 6 = 300 - // h_chatty: 1 × 100 = 100 (MOST samples, tiny value) - let data: &[(&str, f64, usize)] = &[ - ("h_heavy", 1000.0, 3), - ("h_mid", 200.0, 5), - ("h_small", 50.0, 6), - ("h_chatty", 1.0, 100), - ]; - - // Wide CMS + heap large enough to hold every group exactly (4 groups) - // so the estimate equals the true Σvalue with no hash collisions. - let mut acc = CountMinSketchWithHeapAccumulator::new(5, 4096, 16); - let mut truth: std::collections::HashMap<&str, f64> = std::collections::HashMap::new(); - for (host, value, count) in data { - for _ in 0..*count { - acc.insert_value(host, *value); - } - *truth.entry(*host).or_insert(0.0) += value * (*count as f64); - } - - // Ground-truth top-2 by value-sum: h_heavy (3000), h_mid (1000). - let mut truth_ranked: Vec<(&str, f64)> = truth.into_iter().collect(); - truth_ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); - let truth_top2: std::collections::HashSet<&str> = - truth_ranked.iter().take(2).map(|(k, _)| *k).collect(); - assert!( - truth_top2.contains("h_heavy") && truth_top2.contains("h_mid"), - "ground-truth top-2 by value-sum should be h_heavy + h_mid" - ); - - // Value-weighted top-2 from the heap. - let got = acc.topk_by_value(2); - assert_eq!(got.len(), 2, "k=2 → two groups: {got:?}"); - let got_keys: std::collections::HashSet<&str> = - got.iter().map(|(k, _)| k.as_str()).collect(); - - // RECALL = |got ∩ truth| / |truth| must be 1.0. - let hits = got_keys.intersection(&truth_top2).count(); - let recall = hits as f64 / truth_top2.len() as f64; - assert_eq!( - recall, 1.0, - "value-weighted top-k recall must be 1.0 (count-ranked heap would \ - surface h_chatty and miss h_heavy → recall < 1): got={got:?}" - ); - - // The busiest-by-count host (h_chatty) must NOT be in the top-2, - // proving we rank by value-sum, not occurrence count. - assert!( - !got_keys.contains("h_chatty"), - "h_chatty (most samples, smallest value-sum) must be excluded: {got:?}" - ); - - // Estimates are exact here (no collisions, heap holds all groups): - // top-1 must be h_heavy with Σvalue 3000. - assert_eq!(got[0].0, "h_heavy"); - assert!( - (got[0].1 - 3000.0).abs() < 1e-6, - "h_heavy value-sum estimate ≈ 3000, got {}", - got[0].1 - ); - assert_eq!(got[1].0, "h_mid"); - assert!( - (got[1].1 - 1000.0).abs() < 1e-6, - "h_mid value-sum estimate ≈ 1000, got {}", - got[1].1 - ); - } - - /// A single value-weighted insert must put the full value (not +1) into - /// the heap, and repeated inserts for the same group must accumulate. - #[test] - fn insert_value_accumulates_summed_value_in_heap() { - let mut acc = CountMinSketchWithHeapAccumulator::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 - ); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/count_sketch.rs b/crates/asap_summary_state/src/summary_kernels/count_sketch.rs deleted file mode 100644 index 78c1f7d72..000000000 --- a/crates/asap_summary_state/src/summary_kernels/count_sketch.rs +++ /dev/null @@ -1,678 +0,0 @@ -//! CountSketch accumulator backed by `asap_sketchlib::CountSketch`. -//! -//! Supports worker merge, persistence serialization, and modified-OTLP proto -//! decoding. Per-key queries delegate to sketchlib's median-of-signed-rows -//! estimator so query and ingest use the same hash specification. Top-k -//! requires the separate heap-bearing accumulator. - -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{CountSketch, CountSketchDelta, MessagePackCodec}; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -/// Count Sketch accumulator — inner matrix of signed counts. -#[derive(Debug, Clone)] -pub struct CountSketchAccumulator { - pub inner: CountSketch, -} - -impl CountSketchAccumulator { - pub fn new(row_num: usize, col_num: usize) -> Self { - Self { - inner: CountSketch::new(row_num, col_num), - } - } - - /// 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 - /// MessagePack serialization of the cross-language sketch-core - /// `CountSketch` struct — PR I parity entrypoint. - pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { - Ok(Self { - inner: CountSketch::from_msgpack(buffer) - .map_err(|e| format!("deserialize CountSketch msgpack: {e}"))?, - }) - } - - /// Decode from the modified OTLP wire format's - /// `CountSketchDataPoint.sketch` bytes — the protobuf-encoded - /// `asap_sketchlib::proto::sketchlib::CountSketchState` message - /// that DataCollector's `countsketchprocessor` emits when - /// `encoding = COUNT_SKETCH_ENCODING_PROTO`. - /// - /// Mirrors `CountMinSketchAccumulator::from_sketchlib_proto_bytes` - /// but on the signed-counter `CountSketchState`. The resulting - /// accumulator is constructed via - /// `CountSketch::from_legacy_matrix` after reshaping the flat - /// `counts_int` / `counts_float` field into a `Vec>`. - pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, CountSketchState, CounterType, SketchEnvelope, - }; - use prost::Message; - - // DataCollector's countsketchprocessor wraps the state in a - // `SketchEnvelope{count_sketch: CountSketchState}` via - // sketchlib-go's `SerializePortableFO` + `proto.Marshal`. Try - // decoding as envelope first, fall back to bare - // `CountSketchState` for callers (e.g. unit tests) that - // encode the state directly. Mirrors the PR #14 fix on - // `CountMinSketchAccumulator::from_sketchlib_proto_bytes`. - let state = match SketchEnvelope::decode(buffer) { - Ok(env) => match env.sketch_state { - Some(sketch_envelope::SketchState::CountSketch(st)) => st, - Some(other) => { - return Err(format!( - "SketchEnvelope contains non-CountSketch sketch: {:?}", - std::mem::discriminant(&other) - ) - .into()); - } - None => CountSketchState::decode(buffer) - .map_err(|e| format!("decode CountSketchState: {e}"))?, - }, - Err(_) => CountSketchState::decode(buffer) - .map_err(|e| format!("decode CountSketchState: {e}"))?, - }; - let rows = state.rows as usize; - let cols = state.cols as usize; - // Defensive dim validation BEFORE reconstructing the matrix: - // reject degenerate / narrow-hash-budget-violating / absurdly - // oversized dims so a malformed payload fails gracefully (the - // ingest caller skips the data point) instead of building a - // degenerate or huge matrix. Shares the CMS validator since the - // CountSketch matrix uses the same packed-hash column layout. - crate::summary_kernels::count_min_sketch::validate_sketch_dims( - "CountSketchState", - rows, - cols, - )?; - let expected_len = rows * cols; - let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { - format!( - "CountSketchState has unknown counter_type tag {}", - state.counter_type - ) - })?; - let flat: Vec = match counter_type { - CounterType::Int32 | CounterType::Int64 => { - if state.counts_int.len() != expected_len { - return Err(format!( - "CountSketchState counts_int has {} entries, expected rows*cols = {}", - state.counts_int.len(), - expected_len - ) - .into()); - } - state.counts_int.iter().map(|&v| v as f64).collect() - } - CounterType::Float64 => { - if state.counts_float.len() != expected_len { - return Err(format!( - "CountSketchState counts_float has {} entries, expected rows*cols = {}", - state.counts_float.len(), - expected_len - ) - .into()); - } - state.counts_float.clone() - } - other => { - return Err(format!( - "CountSketchState counter_type {other:?} not yet supported \ - (INT128 stores interleaved hi/lo pairs; will be added when needed)" - ) - .into()); - } - }; - let mut matrix = Vec::with_capacity(rows); - for r in 0..rows { - let start = r * cols; - matrix.push(flat[start..start + cols].to_vec()); - } - Ok(Self { - inner: CountSketch::from_legacy_matrix(matrix, rows, cols), - }) - } - - /// Apply a proto-encoded `CountSketchDelta` frame to this - /// accumulator's inner sketch — the decode path for - /// `COUNT_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). - /// - /// Cells apply additively: `matrix[cell_rows[i]][cell_cols[i]] - /// += d_counts[i]`. Per-row L2 is parsed off the wire but - /// ignored at application time — it's a downstream error- - /// accounting signal, not a merge input. - pub fn apply_proto_delta_bytes( - &mut self, - buffer: &[u8], - ) -> Result<(), Box> { - use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; - use prost::Message; - - let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountSketchDelta: {e}"))?; - - if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { - return Err(format!( - "CountSketchDelta packed-array length mismatch: \ - cell_rows={}, cell_cols={}, d_counts={}", - pb.cell_rows.len(), - pb.cell_cols.len(), - pb.d_counts.len() - ) - .into()); - } - let cells = pb - .cell_rows - .iter() - .zip(pb.cell_cols.iter()) - .zip(pb.d_counts.iter()) - .map(|((r, c), dc)| (*r, *c, *dc)) - .collect(); - // This is the heap-less matrix kernel; ranked membership is handled - // by the explicit heap-bearing operator, not inferred from delta keys. - let delta = CountSketchDelta { - rows: pb.rows, - cols: pb.cols, - cells, - l2: pb.l2, - hh_keys: Vec::new(), - }; - self.inner - .apply_delta(&delta) - .map_err(|e| format!("apply CountSketchDelta: {e}"))?; - Ok(()) - } -} - -impl SerializableToSink for CountSketchAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ - "row_num": self.inner.rows, - "col_num": self.inner.cols, - "sketch": self.inner.sketch(), - }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.inner.to_msgpack().unwrap_or_default() - } -} - -impl AggregateCore for CountSketchAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "CountSketchAccumulator" - } - - /// Per-window base rotation: rebuild an empty signed-counter matrix - /// with the same (rows, cols) so the next window's additive cell - /// deltas align to the identical hash geometry. - fn reset_to_empty(&mut self) { - self.inner = CountSketch::new(self.inner.rows, self.inner.cols); - } - - 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 CountSketchAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - let other_cs = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to CountSketchAccumulator")?; - - let merged_inner = CountSketch::merge_refs(&[&self.inner, &other_cs.inner])?; - Ok(Box::new(Self { - inner: merged_inner, - })) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::CountSketch - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - key: &Option, - query_kwargs: &HashMap, - ) -> Result> { - use crate::Statistic; - // 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(); - 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 }) - } - Statistic::Sum => { - let matrix = self.inner.sketch(); - let total: f64 = matrix.iter().flatten().sum(); - let rows = matrix.len() as f64; - Ok(if rows > 0.0 { total / rows } else { 0.0 }) - } - other => Err(format!( - "CountSketchAccumulator: statistic {:?} not supported (only Topk / Count / Sum, with optional `key` in query_kwargs)", - other, - ) - .into()), - } - } -} - -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() { - // `query_key` must match sketchlib's estimator and hash specification. - 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, - counter_type: i32, - counts_int: Vec, - counts_float: Vec, - ) -> Vec { - use asap_sketchlib::proto::sketchlib::CountSketchState; - use prost::Message; - let state = CountSketchState { - rows, - cols, - counter_type, - counts_int, - counts_float, - l2: Vec::new(), - topk: None, - }; - state.encode_to_vec() - } - - #[test] - fn test_from_sketchlib_proto_bytes_int64() { - use asap_sketchlib::proto::sketchlib::CounterType; - // Signed 2x3 matrix: row 0 = [1,-2,3], row 1 = [-4,5,-6] - let bytes = encode_state( - 2, - 3, - CounterType::Int64 as i32, - vec![1, -2, 3, -4, 5, -6], - Vec::new(), - ); - let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - let matrix = acc.inner.sketch(); - assert_eq!(matrix[0], vec![1.0, -2.0, 3.0]); - assert_eq!(matrix[1], vec![-4.0, 5.0, -6.0]); - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrapped() { - // Mirrors what DataCollector's countsketchprocessor emits: - // the state wrapped in a `SketchEnvelope{count_sketch: ...}` - // via sketchlib-go's `SerializePortableFO` + `proto.Marshal`. - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, CountSketchState, CounterType, SketchEnvelope, - }; - use prost::Message; - - let state = CountSketchState { - rows: 2, - cols: 3, - counter_type: CounterType::Int64 as i32, - counts_int: vec![1, -2, 3, -4, 5, -6], - counts_float: Vec::new(), - ..Default::default() - }; - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::CountSketch(state)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes) - .expect("envelope-wrapped decode should succeed"); - let matrix = acc.inner.sketch(); - assert_eq!(matrix[0], vec![1.0, -2.0, 3.0]); - assert_eq!(matrix[1], vec![-4.0, 5.0, -6.0]); - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { - // An envelope carrying a non-CountSketch sketch should be - // rejected with a clear error rather than silently producing - // garbage. - use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; - use prost::Message; - - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "wrong-sketch envelope should error"); - } - - #[test] - fn test_from_sketchlib_proto_bytes_float64() { - use asap_sketchlib::proto::sketchlib::CounterType; - let bytes = encode_state( - 2, - 2, - CounterType::Float64 as i32, - Vec::new(), - vec![1.5, -2.5, 3.5, -4.5], - ); - let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - let matrix = acc.inner.sketch(); - assert_eq!(matrix[0], vec![1.5, -2.5]); - assert_eq!(matrix[1], vec![3.5, -4.5]); - } - - #[test] - fn test_from_sketchlib_proto_bytes_dimension_mismatch() { - use asap_sketchlib::proto::sketchlib::CounterType; - // 2x3 declared but only 5 int entries - let bytes = encode_state( - 2, - 3, - CounterType::Int64 as i32, - vec![1, 2, 3, 4, 5], - Vec::new(), - ); - let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - assert!( - result.unwrap_err().to_string().contains("counts_int"), - "error should mention counts_int dim mismatch" - ); - } - - #[test] - fn test_from_sketchlib_proto_bytes_zero_dims_rejected() { - use asap_sketchlib::proto::sketchlib::CountSketchState; - use prost::Message; - let state = CountSketchState::default(); - let bytes = state.encode_to_vec(); - let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("degenerate dims")); - } - - #[test] - fn test_aggregate_core_merge_matches_matrix_add() { - let a = CountSketchAccumulator { - inner: CountSketch::from_legacy_matrix(vec![vec![1.0, -2.0], vec![3.0, -4.0]], 2, 2), - }; - let b = CountSketchAccumulator { - inner: CountSketch::from_legacy_matrix(vec![vec![-1.0, 2.0], vec![-3.0, 4.0]], 2, 2), - }; - let merged_box = a.merge_with(&b).expect("merge ok"); - let merged = merged_box - .as_any() - .downcast_ref::() - .expect("downcast ok"); - let m = merged.inner.sketch(); - assert_eq!(m[0], vec![0.0, 0.0]); - assert_eq!(m[1], vec![0.0, 0.0]); - } - - #[test] - fn test_aggregate_core_merge_wrong_type_rejects() { - use crate::summary_kernels::count_min_sketch::CountMinSketchAccumulator; - let cs = CountSketchAccumulator::new(2, 3); - let cms = CountMinSketchAccumulator::new(2, 3); - let result = cs.merge_with(&cms); - assert!(result.is_err()); - } - - #[test] - fn test_from_msgpack_bytes_round_trip() { - let original = CountSketch::from_legacy_matrix( - vec![vec![1.0, -2.0, 3.0], vec![-4.0, 5.0, -6.0]], - 2, - 3, - ); - let bytes = original.to_msgpack().unwrap(); - let acc = CountSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.inner.rows, 2); - assert_eq!(acc.inner.cols, 3); - assert_eq!(acc.inner.sketch(), original.sketch()); - } - - #[test] - fn test_from_msgpack_bytes_rejects_garbage() { - let result = CountSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); - assert!(result.is_err()); - } - - #[test] - fn test_apply_proto_delta_bytes_round_trip() { - use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; - use prost::Message; - - let mut acc = CountSketchAccumulator { - inner: CountSketch::from_legacy_matrix( - vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], - 2, - 3, - ), - }; - let bytes = PbDelta { - rows: 2, - cols: 3, - cell_rows: vec![0, 1], - cell_cols: vec![0, 2], - d_counts: vec![10, -6], - l2: vec![], - ..Default::default() - } - .encode_to_vec(); - - acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); - assert_eq!( - acc.inner.sketch(), - &vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 0.0]] - ); - } - - #[test] - fn test_apply_proto_delta_bytes_rejects_garbage() { - let mut acc = CountSketchAccumulator::new(2, 3); - assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); - } - - // ---------------------------------------------------------------- - // Defensive inbound-dimension validation (harden/sketch-dim-validation). - // Malformed / narrow-hash-budget-violating CountSketch dims must be - // rejected gracefully (Err, never a panic); valid configs the backend - // actually uses (5x2048, 5x4096, 5x2000) must still decode. - // ---------------------------------------------------------------- - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { - use asap_sketchlib::proto::sketchlib::CounterType; - // 5 * ceil(log2(8192))=5*13=65 > 64 — narrow-hash-budget violation. - // counts sized to rows*cols so rejection is on dims, not length. - let n = 5usize * 8192usize; - let bytes = encode_state( - 5, - 8192, - CounterType::Int64 as i32, - vec![0i64; n], - Vec::new(), - ); - let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "budget-violating dims should be rejected"); - assert!(result.unwrap_err().to_string().contains("rejecting")); - - // A valid neighbour (5x4096) on the same path still decodes fine. - let n_ok = 5usize * 4096usize; - let ok_bytes = encode_state( - 5, - 4096, - CounterType::Int64 as i32, - vec![0i64; n_ok], - Vec::new(), - ); - let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) - .expect("valid 5x4096 CountSketch should still decode"); - assert_eq!(acc.inner.rows, 5); - assert_eq!(acc.inner.cols, 4096); - } - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_oversized_dims() { - use asap_sketchlib::proto::sketchlib::CounterType; - // Declare 1 x 16,777,216 = 16M cells (> 8M cap) but send an empty - // counts vector: validation must reject on the dim cap BEFORE the - // decoder tries to allocate/reshape a 16M-entry matrix. (1 row keeps - // the hash budget tiny so the cap check, not the budget check, fires.) - let bytes = encode_state( - 1, - 16_777_216, - CounterType::Int64 as i32, - Vec::new(), - Vec::new(), - ); - let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "oversized dims should be rejected"); - let msg = result.unwrap_err().to_string(); - assert!(msg.contains("cap"), "expected cell-cap error, got: {msg}"); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/count_sketch_with_heap.rs b/crates/asap_summary_state/src/summary_kernels/count_sketch_with_heap.rs deleted file mode 100644 index 1f614bcb5..000000000 --- a/crates/asap_summary_state/src/summary_kernels/count_sketch_with_heap.rs +++ /dev/null @@ -1,575 +0,0 @@ -//! Count Sketch with Heap accumulator — wraps -//! `asap_sketchlib::CountSketchWithHeap`. -//! -//! Port of `count_min_sketch_with_heap.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 `SketchAlgorithm::CountSketchWithHeap` sids -- the same conflation bug -//! already fixed on the read side, now closed on the write side too. - -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{CountSketchWithHeap, CsHeapItem, MessagePackCodec}; -use serde::Deserialize; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -/// Local serde view of the DELTA-HEAP wire frame (encoding `MSGPACK_DELTA`). -/// Identical shape to `count_min_sketch_with_heap.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: crate::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::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::summary_kernels::count_min_sketch_with_heap::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/crates/asap_summary_state/src/summary_kernels/datasketches_kll.rs b/crates/asap_summary_state/src/summary_kernels/datasketches_kll.rs deleted file mode 100644 index 1ab4df791..000000000 --- a/crates/asap_summary_state/src/summary_kernels/datasketches_kll.rs +++ /dev/null @@ -1,727 +0,0 @@ -use crate::{ - AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, -}; -use asap_sketchlib::{KllSketch, MessagePackCodec}; -use base64::{engine::general_purpose, Engine as _}; -use serde_json::Value; -use std::collections::HashMap; -#[cfg(feature = "extra_debugging")] -use std::time::Instant; -use tracing::debug; - -use crate::Statistic; - -/// KLL sketch accumulator — wraps asap_sketchlib::KllSketch. -/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. -/// This file retains QE-specific trait impls and JSON output. -pub struct DatasketchesKLLAccumulator { - pub inner: KllSketch, -} - -impl DatasketchesKLLAccumulator { - pub fn new(k: u16) -> Self { - Self { - inner: KllSketch::new(k), - } - } - - pub fn update(&mut self, value: f64) { - self.inner.update(value); - } - - pub fn get_quantile(&self, quantile: f64) -> f64 { - self.inner.quantile(quantile) - } - - /// Decode from the modified OTLP wire format's - /// `KLLSketchDataPoint.sketch` bytes when - /// `encoding = KLL_SKETCH_ENCODING_MSGPACK`. The bytes are the - /// MessagePack serialization of the cross-language sketch-core - /// `KllSketch` struct — PR I parity entrypoint. Unlike the - /// `_ENCODING_PROTO` path (which does lossy statistical - /// reconstruction via `update()` replay), the msgpack path is a - /// bit-identical round-trip because sketch-core's `KllSketch` - /// serializes its full internal state to msgpack. - pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { - Ok(Self { - inner: KllSketch::from_msgpack(buffer) - .map_err(|e| -> Box { e.to_string().into() })?, - }) - } - - /// Decode from the modified OTLP wire format's - /// `KLLSketchDataPoint.sketch` bytes — the protobuf-encoded - /// `asap_sketchlib::proto::sketchlib::KllState` message that - /// DataCollector's `kllprocessor` emits when - /// `encoding = KLL_SKETCH_ENCODING_PROTO`. - /// - /// The neutral codec decodes the sketchlib envelope. - /// The level-aware constructor below preserves the supplied retained - /// sample layout without replaying updates. - pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { - let state = asap_sketch_codec::kll_state(buffer)?; - if state.k < 8 { - return Err(format!("KllState.k must be >= 8 (got {})", state.k).into()); - } - if state.k > u16::MAX as u32 { - return Err(format!( - "KllState.k does not fit in u16 (got {}, max {})", - state.k, - u16::MAX - ) - .into()); - } - // Validate the levels[] boundary array if it is populated. The - // proto contract says `levels[0] == 0` and - // `levels[num_levels] == items.len()`. If the producer left - // levels empty (common when num_levels is zero), skip. - if !state.levels.is_empty() { - if state.levels.len() as u32 != state.num_levels + 1 { - return Err(format!( - "KllState levels length = {}, expected num_levels+1 = {}", - state.levels.len(), - state.num_levels + 1 - ) - .into()); - } - if state.levels[0] != 0 { - return Err(format!("KllState.levels[0] = {}, expected 0", state.levels[0]).into()); - } - if *state.levels.last().unwrap() as usize != state.items.len() { - return Err(format!( - "KllState.levels[{}] = {}, expected items.len() = {}", - state.num_levels, - state.levels.last().unwrap(), - state.items.len() - ) - .into()); - } - } - let k = state.k as u16; - // Direct, bit-exact reconstruction from the portable state (no per-item - // `update()` replay) whenever the producer supplied the `levels[]` - // boundary array — which it does for any non-empty sketch. Falls back to - // the statistical replay only when `levels` is absent (empty sketch). - if !state.levels.is_empty() { - // KllState is highest-level first; the in-memory constructor - // expects L0 first. Replaying or copying the wire order changes - // retained-item weights after the first compaction. - let mut items = Vec::with_capacity(state.items.len()); - let mut levels = vec![0]; - if state - .levels - .windows(2) - .any(|bounds| bounds[0] > bounds[1] || bounds[1] as usize > state.items.len()) - { - return Err("KllState levels must be monotonic and within items".into()); - } - for bounds in state.levels.windows(2).rev() { - items.extend_from_slice(&state.items[bounds[0] as usize..bounds[1] as usize]); - levels.push(items.len()); - } - return Ok(Self { - inner: KllSketch::from_portable_state( - k, - &items, - &levels, - state.num_levels as usize, - ) - .map_err(|e| -> Box { e.into() })?, - }); - } - let mut acc = Self::new(k); - for item in &state.items { - acc.update(*item); - } - Ok(acc) - } - - /// Merge multiple accumulators efficiently without cloning all of them. - pub fn merge_multiple( - accumulators: &[Box], - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut kll_accumulators = Vec::with_capacity(accumulators.len()); - for acc in accumulators { - if acc.get_accumulator_type() != AggregationType::DatasketchesKLL { - return Err(format!( - "Cannot merge DatasketchesKLLAccumulator with {:?}", - acc.get_accumulator_type() - ) - .into()); - } - let kll_acc = acc - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to DatasketchesKLLAccumulator")?; - kll_accumulators.push(kll_acc); - } - - let inner_refs: Vec<&KllSketch> = kll_accumulators.iter().map(|acc| &acc.inner).collect(); - let merged_inner = KllSketch::merge_refs(&inner_refs)?; - Ok(Self { - inner: merged_inner, - }) - } -} - -// Manual trait implementations since the C++ library doesn't provide them -impl Clone for DatasketchesKLLAccumulator { - fn clone(&self) -> Self { - Self { - inner: self.inner.clone(), - } - } -} - -impl std::fmt::Debug for DatasketchesKLLAccumulator { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DatasketchesKLLAccumulator") - .field("k", &self.inner.k) - .field("sketch_n", &self.inner.count()) - .finish() - } -} - -// TODO: verify this -// Thread safety: The C++ library is not thread-safe by default, but since we're using it -// in a single-threaded context per accumulator instance and only sharing read-only operations, -// this should be safe. -unsafe impl Send for DatasketchesKLLAccumulator {} -unsafe impl Sync for DatasketchesKLLAccumulator {} - -impl SerializableToSink for DatasketchesKLLAccumulator { - fn serialize_to_json(&self) -> Value { - // Mirror Python implementation: {"sketch": base64_encoded_string} - let sketch_bytes = self.inner.sketch_bytes(); - let sketch_b64 = general_purpose::STANDARD.encode(&sketch_bytes); - serde_json::json!({ "sketch": sketch_b64 }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.inner.to_msgpack().unwrap_or_default() - } -} - -impl AggregateCore for DatasketchesKLLAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "DatasketchesKLLAccumulator" - } - - 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> { - #[cfg(feature = "extra_debugging")] - let merge_with_start = Instant::now(); - #[cfg(feature = "extra_debugging")] - debug!( - "[PERF] DatasketchesKLLAccumulator::merge_with() started - self.k={}, self.n={}", - self.inner.k, - self.inner.count() - ); - - if other.get_accumulator_type() != self.get_accumulator_type() { - return Err(format!( - "Cannot merge DatasketchesKLLAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - - let other_kll = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to DatasketchesKLLAccumulator")?; - - let merged_inner = KllSketch::merge_refs(&[&self.inner, &other_kll.inner])?; - let merged = Self { - inner: merged_inner, - }; - - #[cfg(feature = "extra_debugging")] - debug!( - "[PERF] DatasketchesKLLAccumulator::merge_with() TOTAL TIME: {:?}", - merge_with_start.elapsed() - ); - - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::DatasketchesKLL - } - - fn approx_memory_bytes(&self) -> usize { - // KLL with default k=200 holds ~2*k items (~3 KiB). Round up - // for overhead. - 4 * 1024 - } - - fn aux_stats(&self) -> AuxStats { - // KLL natively tracks `count` (n, samples observed). min/max - // are available from the underlying sketch but only via a - // O(k) quantile extraction at quantile=0/1, which is not - // a cheap trait-method call. sum is not retained by KLL. - // - // Surface only count here; follow-up PR may add min/max via a - // dedicated accessor on sketch-core. `sum_over_time` queries - // on KLL fall back to query_statistic as they do today. - AuxStats { - count: Some(self.inner.count()), - ..AuxStats::empty() - } - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - _key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::SingleSubpopulationAggregate; - self.query(statistic, Some(query_kwargs)) - } -} - -impl SingleSubpopulationAggregate for DatasketchesKLLAccumulator { - fn query( - &self, - statistic: Statistic, - query_kwargs: Option<&HashMap>, - ) -> Result> { - match statistic { - Statistic::Quantile => { - debug!( - "Querying DatasketchesKLLAccumulator for quantile with kwargs: {:?}", - query_kwargs - ); - let quantile = query_kwargs - .and_then(|kwargs| kwargs.get("quantile")) - .ok_or("Missing quantile parameter for quantile query")? - .parse::() - .map_err(|_| "Invalid quantile parameter format")?; - - if !(0.0..=1.0).contains(&quantile) { - return Err("Quantile must be between 0.0 and 1.0".into()); - } - - Ok(self.get_quantile(quantile)) - } - _ => Err( - format!("Unsupported statistic in DatasketchesKLLAccumulator: {statistic:?}") - .into(), - ), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for DatasketchesKLLAccumulator { - 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::*; - use prost::Message; - - fn encode_state(state: asap_sketchlib::proto::sketchlib::KllState) -> Vec { - use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; - SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Kll(state)), - ..Default::default() - } - .encode_to_vec() - } - - #[test] - fn test_datasketches_kll_creation() { - let kll = DatasketchesKLLAccumulator::new(200); - assert!(kll.inner.count() == 0); - assert_eq!(kll.inner.k, 200); - } - - #[test] - fn test_datasketches_kll_update() { - let mut kll = DatasketchesKLLAccumulator::new(200); - kll.update(10.0); - kll.update(20.0); - kll.update(15.0); - assert_eq!(kll.inner.count(), 3); - } - - #[test] - fn test_datasketches_kll_quantile() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for i in 1..=10 { - kll.update(i as f64); - } - assert_eq!(kll.get_quantile(0.0), 1.0); - assert_eq!(kll.get_quantile(1.0), 10.0); - // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. - let q50 = kll.get_quantile(0.5); - assert!((q50 - 6.0).abs() <= 1.0, "expected median ~6, got {q50}"); - } - - #[test] - fn test_datasketches_kll_query() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for i in 1..=10 { - kll.update(i as f64); - } - - let mut query_kwargs = HashMap::new(); - query_kwargs.insert("quantile".to_string(), "0.5".to_string()); - let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); - // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. - assert!( - (result - 6.0).abs() <= 1.0, - "expected median ~6, got {result}" - ); - - assert!(kll.query(Statistic::Sum, Some(&query_kwargs)).is_err()); - } - - #[test] - fn test_datasketches_kll_merge() { - let mut kll1 = DatasketchesKLLAccumulator::new(200); - let mut kll2 = DatasketchesKLLAccumulator::new(200); - - for i in 1..=5 { - kll1.update(i as f64); - } - for i in 6..=10 { - kll2.update(i as f64); - } - - let merged = DatasketchesKLLAccumulator::merge_accumulators(vec![kll1, kll2]).unwrap(); - assert_eq!(merged.inner.count(), 10); - assert_eq!(merged.get_quantile(0.0), 1.0); - assert_eq!(merged.get_quantile(1.0), 10.0); - } - - #[test] - fn test_datasketches_kll_get_keys() { - let kll = DatasketchesKLLAccumulator::new(200); - assert_eq!(kll.type_name(), "DatasketchesKLLAccumulator"); - } - - #[test] - fn test_trait_object() { - let mut kll = DatasketchesKLLAccumulator::new(200); - kll.update(5.0); - let trait_obj: Box = Box::new(kll); - assert_eq!(trait_obj.type_name(), "DatasketchesKLLAccumulator"); - } - - #[test] - fn test_datasketches_kll_query_with_kwargs() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for i in 1..=10 { - kll.update(i as f64); - } - - let mut query_kwargs = HashMap::new(); - query_kwargs.insert("quantile".to_string(), "0.5".to_string()); - let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); - // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. - assert!( - (result - 6.0).abs() <= 1.0, - "expected median ~6, got {result}" - ); - - query_kwargs.insert("quantile".to_string(), "0.9".to_string()); - let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); - // Sketchlib KLL is approximate; 0.9 quantile of 1..10 may be 9 or 10. - assert!( - (9.0..=10.0).contains(&result), - "expected 0.9 quantile in [9,10], got {result}" - ); - - query_kwargs.insert("quantile".to_string(), "0.0".to_string()); - assert_eq!( - kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(), - 1.0 - ); - - query_kwargs.insert("quantile".to_string(), "1.0".to_string()); - assert_eq!( - kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(), - 10.0 - ); - - assert!(kll.query(Statistic::Quantile, None).is_err()); - - query_kwargs.insert("quantile".to_string(), "invalid".to_string()); - assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); - - query_kwargs.insert("quantile".to_string(), "1.5".to_string()); - assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); - - query_kwargs.insert("quantile".to_string(), "-0.1".to_string()); - assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); - - query_kwargs.insert("quantile".to_string(), "0.5".to_string()); - assert!(kll.query(Statistic::Sum, Some(&query_kwargs)).is_err()); - } - - #[test] - fn test_datasketches_kll_merge_multiple() { - let mut kll1 = DatasketchesKLLAccumulator::new(200); - let mut kll2 = DatasketchesKLLAccumulator::new(200); - let mut kll3 = DatasketchesKLLAccumulator::new(200); - - for i in 1..=5 { - kll1.update(i as f64); - } - for i in 6..=10 { - kll2.update(i as f64); - } - for i in 11..=15 { - kll3.update(i as f64); - } - - let boxed_accs: Vec> = - vec![Box::new(kll1), Box::new(kll2), Box::new(kll3)]; - - let merged = DatasketchesKLLAccumulator::merge_multiple(&boxed_accs).unwrap(); - assert_eq!(merged.inner.count(), 15); - assert_eq!(merged.get_quantile(0.0), 1.0); - assert_eq!(merged.get_quantile(1.0), 15.0); - assert_eq!(merged.get_quantile(0.5), 8.0); - } - - #[test] - fn test_datasketches_kll_merge_multiple_error_cases() { - let empty: Vec> = vec![]; - assert!(DatasketchesKLLAccumulator::merge_multiple(&empty).is_err()); - - let kll1 = DatasketchesKLLAccumulator::new(200); - let kll2 = DatasketchesKLLAccumulator::new(100); - let boxed_accs: Vec> = vec![Box::new(kll1), Box::new(kll2)]; - assert!(DatasketchesKLLAccumulator::merge_multiple(&boxed_accs).is_err()); - - use crate::summary_kernels::sum::SumAccumulator; - let kll = DatasketchesKLLAccumulator::new(200); - let sum = SumAccumulator::new(); - let mixed_accs: Vec> = vec![Box::new(kll), Box::new(sum)]; - assert!(DatasketchesKLLAccumulator::merge_multiple(&mixed_accs).is_err()); - } - - #[test] - fn test_from_sketchlib_proto_bytes_reconstructs_quantiles() { - // Build a KllState with 64 items in level order; the decoder - // replays every item through `update()` so the reconstructed - // sketch is statistically equivalent — quantile estimates - // match the ground truth (sorted items) within KLL's own - // rank-error bound for k=200. - use asap_sketchlib::proto::sketchlib::KllState; - - let items: Vec = (0..64).map(|i| i as f64).collect(); - let state = KllState { - k: 200, - m: 8, - num_levels: 1, - levels: vec![0, 64], - items: items.clone(), - coin: None, - offset: 0.0, - value_scale: 0, - residuals: Vec::new(), - }; - let bytes = encode_state(state); - - let acc = - DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.inner.count(), 64); - // For 64 values 0..63, the true median is 31.5 and quantile - // error is ~1% × range = 0.63. KLL's own point query can - // legally be off by up to ε × N ~= 0.01 × 64 = 0.64. Allow a - // generous tolerance since the important invariant is "the - // decoded sketch is queryable and returns a sensible value". - let median = acc.get_quantile(0.5); - assert!( - (median - 31.5).abs() <= 10.0, - "reconstructed median {median} is outside tolerance of true median 31.5" - ); - let q01 = acc.get_quantile(0.01); - let q99 = acc.get_quantile(0.99); - assert!( - q01 <= q99, - "quantile monotonicity violated: q01={q01}, q99={q99}" - ); - } - - // Compacted portable state is highest-level first, unlike the runtime buffer. - #[test] - fn compacted_wire_state_preserves_count_and_quantiles() { - use asap_sketchlib::{proto::sketchlib::KllState, sketches::KLL}; - let mut source = KLL::::init_kll_with_seed(32, 123); - for i in 0..1000 { - source.update(&(((i * 7919 + 17) % 1009) as f64 / 1009.0)); - } - assert!(source.wire_num_levels() > 1); - let state = KllState { - k: 32, - m: source.wire_m(), - num_levels: source.wire_num_levels(), - levels: source.wire_levels(), - items: source.wire_items(), - coin: None, - offset: 0.0, - value_scale: 0, - residuals: vec![], - }; - let decoded = - DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&encode_state(state)).unwrap(); - assert_eq!(decoded.inner.count(), source.count() as u64); - for q in [0.0, 0.1, 0.5, 0.9, 1.0] { - assert_eq!(decoded.inner.quantile(q), source.quantile(q), "q={q}"); - } - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrapped() { - // Mirrors what DataCollector's kllprocessor emits: the state - // wrapped in a `SketchEnvelope{kll: ...}` via sketchlib-go's - // `SerializePortableFO` + `proto.Marshal`. - use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; - - let items: Vec = (0..64).map(|i| i as f64).collect(); - let state = KllState { - k: 200, - m: 8, - num_levels: 1, - levels: vec![0, 64], - items, - coin: None, - offset: 0.0, - value_scale: 0, - residuals: Vec::new(), - }; - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Kll(state)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes) - .expect("envelope-wrapped decode should succeed"); - assert_eq!(acc.inner.count(), 64); - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { - use asap_sketchlib::proto::sketchlib::{sketch_envelope, CountMinState, SketchEnvelope}; - - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::CountMin( - CountMinState::default(), - )), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "wrong-sketch envelope should error"); - } - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_small_k() { - use asap_sketchlib::proto::sketchlib::KllState; - let state = KllState { - k: 4, // < minimum of 8 - m: 2, - num_levels: 0, - levels: Vec::new(), - items: Vec::new(), - coin: None, - offset: 0.0, - value_scale: 0, - residuals: Vec::new(), - }; - let bytes = encode_state(state); - let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("k must be >= 8")); - } - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_inconsistent_levels() { - use asap_sketchlib::proto::sketchlib::KllState; - // num_levels=1 but levels array has 3 entries instead of 2 - let state = KllState { - k: 200, - m: 8, - num_levels: 1, - levels: vec![0, 5, 10], - items: vec![1.0, 2.0, 3.0, 4.0, 5.0], - coin: None, - offset: 0.0, - value_scale: 0, - residuals: Vec::new(), - }; - let bytes = encode_state(state); - let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("levels length")); - } - - #[test] - fn aux_stats_exposes_count_via_kll_n() { - let mut acc = DatasketchesKLLAccumulator::new(200); - for i in 0..50 { - acc.update(i as f64); - } - let aux = acc.aux_stats(); - assert_eq!(aux.count, Some(50)); - // KLL doesn't natively expose min/max cheaply and doesn't - // track sum at all — those fields must be None so callers - // fall through to query_statistic. - assert_eq!(aux.sum, None); - assert_eq!(aux.min, None); - assert_eq!(aux.max, None); - } - - #[test] - fn aux_stats_empty_kll_has_zero_count() { - let acc = DatasketchesKLLAccumulator::new(200); - assert_eq!(acc.aux_stats().count, Some(0)); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/dd_sketch.rs b/crates/asap_summary_state/src/summary_kernels/dd_sketch.rs deleted file mode 100644 index 1c25d940c..000000000 --- a/crates/asap_summary_state/src/summary_kernels/dd_sketch.rs +++ /dev/null @@ -1,726 +0,0 @@ -//! DDSketch accumulator — wraps `asap_sketchlib::DdSketch`. -//! -//! Concrete accumulator reached from the modified-OTLP -//! `Metric.data = DDSketch{…}` hot path (PR C-CountSketch follow-up). -//! Merge via bucket-index alignment on the inner sketch, serialize as -//! MessagePack for the sink, and decode from the sketchlib -//! `DDSketchState` proto. -//! -//! Query semantics follow the STRICT policy after the DataPoint-level -//! METRIC scalars were dropped from the wire format -//! (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57): the sketch serves -//! Quantile (log-bucket estimation) and Count (sum of bucket counts). -//! Sum/Min/Max are no longer derivable from the wire bytes and are -//! served by controller-provisioned exact aggregations — `query_statistic` -//! returns the unavailable-statistic error for them. - -use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; -use asap_sketchlib::{DdSketch, DdSketchDelta, MessagePackCodec}; -use serde_json::Value; -use std::collections::HashMap; - -/// DDSketch accumulator — inner log-bucketed sketch. -#[derive(Debug, Clone)] -pub struct DDSketchAccumulator { - pub inner: DdSketch, - /// Edge sampling probability `p ∈ (0,1]` carried on the producer's - /// `SketchEnvelope.sample_p`. The edge admits each value with probability - /// `p` (NitroSketch geometric skip), so `inner.total_count()` is ~`p`× the - /// true count and a `Count` query must rescale by `1/p`. Quantiles are - /// rank-preserving and need NO rescale. `1.0` (and the proto3 default `0.0`, - /// dual-read as `1.0`) means no sampling, so the rescale is a no-op and the - /// behaviour is identical to before. The factor is a per-series config - /// constant: it is set from the first (always-full, otel.rs ingest - /// contract) frame and preserved across delta applies, window-boundary - /// `reset_to_empty`, and `merge_with`. - pub sample_p: f64, -} - -/// Normalize a wire `sample_p` to a usable rescale denominator. `0.0` (proto3 -/// default), `>= 1.0`, and non-finite all collapse to `1.0` (no sampling), so a -/// `Count` rescale by `1/p` is a no-op on unsampled / legacy frames. -pub(crate) fn normalize_sample_p(p: f64) -> f64 { - if p.is_finite() && p > 0.0 && p < 1.0 { - p - } else { - 1.0 - } -} - -impl DDSketchAccumulator { - pub fn new(alpha: f64) -> Self { - Self { - inner: DdSketch::new(alpha), - sample_p: 1.0, - } - } - - /// Read the normalized edge sampling probability from a full-frame - /// `SketchEnvelope`'s `sample_p`. Returns `1.0` (no sampling) for bare - /// `DdSketchState` bytes or any decode failure — the primary production - /// decode path (`reconstruct_via_runtime`) discards the envelope's - /// `sample_p`, so the ingest call site re-reads it from the same bytes. - pub fn sample_p_from_envelope_bytes(buffer: &[u8]) -> f64 { - use asap_sketchlib::proto::sketchlib::SketchEnvelope; - use prost::Message; - SketchEnvelope::decode(buffer) - .map(|env| normalize_sample_p(env.sample_p)) - .unwrap_or(1.0) - } - - /// Decode from the modified OTLP wire format's - /// `DDSketchDataPoint.sketch` bytes when - /// `encoding = DDSKETCH_ENCODING_MSGPACK`. The bytes are the - /// MessagePack serialization of the cross-language sketch-core - /// `DdSketch` struct — PR I parity entrypoint. - pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { - Ok(Self { - inner: DdSketch::from_msgpack(buffer) - .map_err(|e| format!("deserialize DdSketch msgpack: {e}"))?, - // The msgpack DdSketch struct carries no envelope/sample_p; the - // msgpack path is parity/test-only and is never edge-sampled. - sample_p: 1.0, - }) - } - - /// Decode from the modified OTLP wire format's - /// `DDSketchDataPoint.sketch` bytes — the protobuf-encoded - /// `asap_sketchlib::proto::sketchlib::DDSketchState` message that - /// DataCollector's `ddsketchprocessor` emits when - /// `encoding = DD_SKETCH_ENCODING_PROTO`. - pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { - let (state, sample_p) = asap_sketch_codec::ddsketch_state(buffer)?; - if !(state.alpha > 0.0 && state.alpha < 1.0) { - return Err(format!( - "DDSketchState alpha {} out of range (expected 0 < alpha < 1)", - state.alpha - ) - .into()); - } - // Preserve positive, negative and zero stores from the sketchlib wire state. - let inner = DdSketch::from_proto(state); - Ok(Self { - inner, - sample_p: normalize_sample_p(sample_p), - }) - } - - /// Apply a proto-encoded `DDSketchDelta` frame to this - /// accumulator's inner sketch — the decode path for - /// `DD_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). - /// - /// Called against an accumulator that already carries the base - /// sketch state; the caller is the per-series snapshot cache in - /// the ingest path. Bytes are the - /// `asap_sketchlib::proto::sketchlib::DdSketchDelta` message. - pub fn apply_proto_delta_bytes( - &mut self, - buffer: &[u8], - ) -> Result<(), Box> { - use asap_sketchlib::proto::sketchlib::DdSketchDelta as PbDelta; - use prost::Message; - - let pb = PbDelta::decode(buffer).map_err(|e| format!("decode DDSketchDelta: {e}"))?; - - // The delta no longer carries d_count/d_sum/min/max - // (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). Apply the - // bucket deltas only; `DdSketch` recomputes its total count from - // the merged bucket counts (`total_count()`). - let buckets = pb - .buckets - .into_iter() - .map(|b| (b.index, b.d_count)) - .collect(); - let delta = DdSketchDelta { - buckets, - negative_buckets: pb - .negative_buckets - .into_iter() - .map(|b| (b.index, b.d_count)) - .collect(), - zero_count: pb.zero_count, - ..Default::default() - }; - self.inner - .apply_delta(&delta) - .map_err(|error| format!("apply DDSketchDelta: {error}"))?; - Ok(()) - } -} - -impl SerializableToSink for DDSketchAccumulator { - fn serialize_to_json(&self) -> Value { - // The DataPoint-level scalars (sum/min/max) are no longer carried - // by `DdSketch` (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). - // `count` is the bucket-derived total via `total_count()`. - serde_json::json!({ - "alpha": self.inner.alpha, - "store_offset": self.inner.store_offset, - "bucket_count": self.inner.store_counts.len(), - // Raw bucket-derived count (admitted samples). `sample_p` is the - // scale factor a consumer applies (count / sample_p) to estimate - // the true count; `query_statistic(Count)` already does this. - "count": self.inner.total_count(), - "sample_p": self.sample_p, - }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.inner.to_msgpack().unwrap_or_default() - } -} - -impl AggregateCore for DDSketchAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "DDSketchAccumulator" - } - - /// Per-window base rotation: drop all bucket counts but keep the - /// relative-accuracy parameter so the next window's bucket deltas - /// index into the same log-bucket layout. `sample_p` is a per-series - /// config constant (not per-window data), so it is intentionally - /// preserved across the rotation — the next window's deltas are sampled - /// at the same rate and must rescale identically. - fn reset_to_empty(&mut self) { - self.inner = DdSketch::new(self.inner.alpha); - } - - 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 DDSketchAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - let other_dd = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to DDSketchAccumulator")?; - let merged_inner = DdSketch::merge_refs(&[&self.inner, &other_dd.inner])?; - // sample_p is a per-series config constant, so both operands carry the - // same value in practice. Prefer a sampled factor over the no-sampling - // default so a merge with a freshly-reset (1.0) base keeps the series' - // sampling rate. - let sample_p = if self.sample_p < 1.0 { - self.sample_p - } else { - other_dd.sample_p - }; - Ok(Box::new(Self { - inner: merged_inner, - sample_p, - })) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::DDSketch - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - _key: &Option, - query_kwargs: &HashMap, - ) -> Result> { - use crate::Statistic; - - match statistic { - Statistic::Quantile => { - // PromQL `histogram_quantile(q, …)` and - // `quantile_over_time(q, …)` both land here with - // `q` in `query_kwargs["quantile"]`. Default to - // 0.99 when the caller didn't provide one - // (defensive — pattern-matched queries in - // `inference_config.yaml` always populate it). - let q: f64 = query_kwargs - .get("quantile") - .and_then(|s| s.parse().ok()) - .unwrap_or(0.99); - if !(0.0..=1.0).contains(&q) { - return Err(format!("DDSketchAccumulator: quantile {q} out of [0,1]").into()); - } - self.inner.quantile(q).ok_or_else(|| { - "DDSketchAccumulator: quantile() returned None (sketch empty?)".into() - }) - } - // Count is derived by summing the bucket store counts — the only - // DataPoint-level scalar that survives the wire-format trim - // (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). When the edge - // sampled this series (sample_p < 1.0), the stored count is ~p× the - // true count, so rescale by 1/sample_p to recover an unbiased - // estimate. sample_p == 1.0 (unsampled / legacy) makes this a no-op. - Statistic::Count => Ok(self.inner.total_count() as f64 / self.sample_p), - // STRICT policy: the Sum/Min/Max scalars were removed from - // the DDSketch wire format. They are now served by the - // controller-provisioned exact aggregations (an exact `Sum` - // and an exact `MinMax`), NOT estimated from the buckets. - // Surface the unavailable-statistic error so the query path - // routes to those aggregations instead of returning a wrong - // (0 / panicked) value. - Statistic::Sum => Err( - "DDSketchAccumulator: Sum not available from DDSketch wire format \ - (ProjectASAP/sketchlib-go#243); use an exact Sum aggregation" - .into(), - ), - Statistic::Min | Statistic::Max => Err(format!( - "DDSketchAccumulator: {statistic:?} not available from DDSketch wire format \ - (ProjectASAP/sketchlib-go#243); use an exact MinMax aggregation", - ) - .into()), - other => Err(format!( - "DDSketchAccumulator: statistic {other:?} not supported (only Quantile / Count)", - ) - .into()), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - // The DataPoint-level METRIC scalars (count/sum/min/max) were dropped - // from `DdSketchState` (ProjectASAP/sketchlib-go#243 / - // asap_sketchlib#57); the proto now carries only - // `alpha`/`store_counts`/`store_offset`. - fn encode_state(alpha: f64, store_counts: Vec, store_offset: i32) -> Vec { - use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; - use prost::Message; - let state = DdSketchState { - alpha, - store_counts, - store_offset, - ..Default::default() - }; - SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Ddsketch(state)), - ..Default::default() - } - .encode_to_vec() - } - - #[test] - fn test_from_sketchlib_proto_bytes_round_trip() { - let bytes = encode_state(0.01, vec![1, 2, 3, 4], -2); - let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.inner.alpha, 0.01); - assert_eq!(acc.inner.store_counts, vec![1, 2, 3, 4]); - assert_eq!(acc.inner.store_offset, -2); - // `count` is recovered by summing the bucket store counts. - assert_eq!(acc.inner.total_count(), 10); - } - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_invalid_alpha() { - let bytes = encode_state(0.0, vec![1], 0); - let result = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("alpha")); - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrapped() { - // Mirrors what DataCollector's ddsketchprocessor emits: the - // state wrapped in a `SketchEnvelope{ddsketch: ...}` via - // sketchlib-go's `SerializePortableFO` + `proto.Marshal`. - use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; - use prost::Message; - - let state = DdSketchState { - alpha: 0.01, - store_counts: vec![1, 2, 3, 4], - store_offset: -2, - ..Default::default() - }; - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Ddsketch(state)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes) - .expect("envelope-wrapped decode should succeed"); - assert_eq!(acc.inner.alpha, 0.01); - assert_eq!(acc.inner.total_count(), 10); - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { - use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; - use prost::Message; - - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let result = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "wrong-sketch envelope should error"); - } - - #[test] - fn test_aggregate_core_merge_aligns_buckets() { - let a = DDSketchAccumulator { - inner: DdSketch::from_raw(0.01, vec![1, 1, 1], -1), - sample_p: 1.0, - }; - let b = DDSketchAccumulator { - inner: DdSketch::from_raw(0.01, vec![10, 10, 10], 0), - sample_p: 1.0, - }; - let merged_box = a.merge_with(&b).expect("merge ok"); - let merged = merged_box - .as_any() - .downcast_ref::() - .expect("downcast ok"); - assert_eq!(merged.inner.store_counts, vec![1, 11, 11, 10]); - assert_eq!(merged.inner.store_offset, -1); - assert_eq!(merged.inner.total_count(), 33); - } - - #[test] - fn test_aggregate_core_merge_wrong_type_rejects() { - use crate::summary_kernels::count_sketch::CountSketchAccumulator; - let dd = DDSketchAccumulator::new(0.01); - let cs = CountSketchAccumulator::new(2, 3); - assert!(dd.merge_with(&cs).is_err()); - } - - #[test] - fn test_from_msgpack_bytes_round_trip() { - let original = DdSketch::from_raw(0.01, vec![5, 10, 15, 20], -2); - let bytes = original.to_msgpack().unwrap(); - let acc = DDSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.inner.alpha, 0.01); - assert_eq!(acc.inner.store_counts, vec![5, 10, 15, 20]); - assert_eq!(acc.inner.store_offset, -2); - // `count` is recovered by summing the bucket store counts. - assert_eq!(acc.inner.total_count(), 50); - } - - #[test] - fn test_from_msgpack_bytes_rejects_garbage() { - let result = DDSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); - assert!(result.is_err()); - } - - #[test] - fn test_apply_proto_delta_bytes_round_trip() { - use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; - use prost::Message; - - let mut acc = DDSketchAccumulator::new(0.01); - acc.inner = DdSketch::from_raw(0.01, vec![1, 2, 3], 0); - - // The wire delta now carries only bucket deltas (tags 2-7 - // reserved); `DdSketchBucketDelta` has just `index` + `d_count`. - let bytes = PbDelta { - buckets: vec![ - DdSketchBucketDelta { - index: 0, - d_count: 10, - }, - DdSketchBucketDelta { - index: 2, - d_count: 20, - }, - ], - ..Default::default() - } - .encode_to_vec(); - - acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); - assert_eq!(acc.inner.store_counts, vec![11, 2, 23]); - // `count` recomputed from the merged buckets: 11 + 2 + 23 = 36. - assert_eq!(acc.inner.total_count(), 36); - } - - /// A valid protobuf with an inadmissible span must not acknowledge a dropped update. - #[test] - fn test_apply_proto_delta_rejects_span_without_mutating_state() { - use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; - use prost::Message; - let mut acc = DDSketchAccumulator::new(0.01); - acc.inner = DdSketch::from_raw(0.01, vec![1, 2, 3], 0); - let bytes = PbDelta { - buckets: vec![DdSketchBucketDelta { - index: i32::MAX, - d_count: 1, - }], - ..Default::default() - } - .encode_to_vec(); - assert!(acc.apply_proto_delta_bytes(&bytes).is_err()); - assert_eq!(acc.inner.store_counts, vec![1, 2, 3]); - assert_eq!(acc.inner.store_offset, 0); - } - - #[test] - fn test_apply_proto_delta_bytes_rejects_garbage() { - let mut acc = DDSketchAccumulator::new(0.01); - assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); - } - - // ----- query_statistic STRICT policy ----- - // - // After the DataPoint-level METRIC scalars were dropped from the - // DDSketch wire format (ProjectASAP/sketchlib-go#243 / - // asap_sketchlib#57), DDSketch serves only quantiles and Count. - // Sum/Min/Max move to controller-provisioned exact aggregations and - // MUST surface the unavailable-statistic error (never a panic / 0). - - fn sample_accumulator() -> DDSketchAccumulator { - // Build the in-memory sketch from bucket counts only — no scalars. - DDSketchAccumulator { - inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), - sample_p: 1.0, - } - } - - #[test] - fn test_query_statistic_quantile_is_sketch_derived() { - use crate::Statistic; - let acc = sample_accumulator(); - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), "0.5".to_string()); - let v = acc - .query_statistic(Statistic::Quantile, &None, &kwargs) - .expect("quantile should be served from the sketch buckets"); - assert!( - v.is_finite() && v > 0.0, - "quantile estimate should be positive finite, got {v}" - ); - } - - #[test] - fn test_query_statistic_count_is_bucket_derived() { - use crate::Statistic; - let acc = sample_accumulator(); - let v = acc - .query_statistic(Statistic::Count, &None, &HashMap::new()) - .expect("count should be derivable from the bucket store"); - // 1 + 2 + 3 + 4 = 10. - assert_eq!(v, 10.0); - } - - #[test] - fn test_query_statistic_sum_min_max_return_unavailable_error() { - use crate::Statistic; - let acc = sample_accumulator(); - for stat in [Statistic::Sum, Statistic::Min, Statistic::Max] { - let result = acc.query_statistic(stat, &None, &HashMap::new()); - assert!( - result.is_err(), - "{stat:?} must return the unavailable-statistic error (not a panic / 0)" - ); - let msg = result.unwrap_err().to_string(); - assert!( - msg.contains("not available"), - "{stat:?} error should explain the statistic is unavailable, got: {msg}" - ); - } - } - - // ----- sample_p count rescale ----- - // - // When the edge sampled a DDSketch (sample_p < 1.0), the stored count is - // ~p× the true count, so Count rescales by 1/p. Quantiles are - // rank-preserving and must NOT be rescaled. - - #[test] - fn test_count_is_rescaled_by_sample_p() { - use crate::Statistic; - let acc = DDSketchAccumulator { - inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), - sample_p: 0.1, - }; - let c = acc - .query_statistic(Statistic::Count, &None, &HashMap::new()) - .expect("count ok"); - // Raw bucket sum 10, rescaled by 1/0.1 = 100. - assert!((c - 100.0).abs() < 1e-9, "expected rescaled 100, got {c}"); - } - - #[test] - fn test_quantile_ignores_sample_p() { - use crate::Statistic; - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), "0.5".to_string()); - let unsampled = DDSketchAccumulator { - inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), - sample_p: 1.0, - }; - let sampled = DDSketchAccumulator { - inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), - sample_p: 0.1, - }; - let qu = unsampled - .query_statistic(Statistic::Quantile, &None, &kwargs) - .expect("q ok"); - let qs = sampled - .query_statistic(Statistic::Quantile, &None, &kwargs) - .expect("q ok"); - assert_eq!(qu, qs, "quantile must be sample_p-invariant"); - } - - #[test] - fn test_from_sketchlib_proto_bytes_reads_envelope_sample_p() { - use crate::Statistic; - use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; - use prost::Message; - - let env = SketchEnvelope { - sample_p: 0.25, - sketch_state: Some(sketch_envelope::SketchState::Ddsketch(DdSketchState { - alpha: 0.01, - store_counts: vec![2, 4, 6, 8], - store_offset: -2, - ..Default::default() - })), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.sample_p, 0.25); - // Raw 20, rescaled 20 / 0.25 = 80. - let c = acc - .query_statistic(Statistic::Count, &None, &HashMap::new()) - .expect("count ok"); - assert!((c - 80.0).abs() < 1e-9, "expected rescaled 80, got {c}"); - } - - #[test] - fn test_sample_p_normalization() { - // proto3 default (0.0), >=1.0, and non-finite all mean no sampling. - assert_eq!(normalize_sample_p(0.0), 1.0); - assert_eq!(normalize_sample_p(1.0), 1.0); - assert_eq!(normalize_sample_p(1.5), 1.0); - assert_eq!(normalize_sample_p(f64::NAN), 1.0); - assert_eq!(normalize_sample_p(-0.1), 1.0); - assert_eq!(normalize_sample_p(0.5), 0.5); - } - - #[test] - fn test_sample_p_from_envelope_bytes_defaults_to_one() { - use asap_sketchlib::proto::sketchlib::DdSketchState; - use prost::Message; - // Bare DdSketchState bytes (no envelope) → no sampling info → 1.0. - let bare = DdSketchState { - alpha: 0.01, - store_counts: vec![1, 2, 3], - store_offset: 0, - ..Default::default() - } - .encode_to_vec(); - assert_eq!( - DDSketchAccumulator::sample_p_from_envelope_bytes(&bare), - 1.0 - ); - } - - #[test] - fn test_reset_to_empty_preserves_sample_p() { - let mut acc = DDSketchAccumulator { - inner: DdSketch::from_raw(0.01, vec![1, 2, 3], 0), - sample_p: 0.2, - }; - acc.reset_to_empty(); - assert_eq!(acc.sample_p, 0.2, "window rotation must keep sample_p"); - assert_eq!(acc.inner.total_count(), 0, "buckets cleared"); - } - - #[test] - fn test_merge_prefers_sampled_factor() { - // A sampled base merged with a freshly-reset (1.0) operand keeps the - // series' sampling rate. - let a = DDSketchAccumulator { - inner: DdSketch::from_raw(0.01, vec![1, 1, 1], 0), - sample_p: 0.1, - }; - let b = DDSketchAccumulator { - inner: DdSketch::from_raw(0.01, vec![1, 1, 1], 0), - sample_p: 1.0, - }; - let merged = a.merge_with(&b).expect("merge ok"); - let merged = merged - .as_any() - .downcast_ref::() - .expect("downcast ok"); - assert_eq!(merged.sample_p, 0.1); - } -} - -#[cfg(test)] -mod dependency_upgrade_tests { - use super::*; - // The upgraded sketchlib state must retain negative and zero stores through both adapters. - #[test] - fn signed_state_survives_codec_and_accumulator_roundtrip() { - let mut inner = DdSketch::new(0.01); - for value in [-4.0, 0.0, 8.0] { - inner.update(value); - } - let bytes = asap_sketch_codec::encode_ddsketch(&inner); - let (wire, _) = asap_sketch_codec::ddsketch_state(&bytes).unwrap(); - assert_eq!(wire.zero_count, 1); - assert_eq!(wire.negative_store_counts.iter().sum::(), 1); - let restored = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes).unwrap(); - assert_eq!(restored.inner.total_count(), 3); - assert_eq!(restored.inner.alpha, inner.wire_alpha()); - assert_eq!(restored.inner.store_counts, inner.store_counts); - assert_eq!(restored.inner.store_offset, inner.store_offset); - assert_eq!( - restored.inner.negative_store_counts, - inner.negative_store_counts - ); - assert_eq!( - restored.inner.negative_store_offset, - inner.negative_store_offset - ); - assert_eq!(restored.inner.zero_count, inner.zero_count); - } - // Negative and zero delta fields added by sketchlib must not be discarded by the adapter. - #[test] - fn signed_delta_survives_adapter() { - use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; - use prost::Message; - let mut accumulator = DDSketchAccumulator::new(0.01); - let bytes = PbDelta { - negative_buckets: vec![DdSketchBucketDelta { - index: 0, - d_count: 2, - }], - zero_count: 3, - ..Default::default() - } - .encode_to_vec(); - accumulator.apply_proto_delta_bytes(&bytes).unwrap(); - assert_eq!(accumulator.inner.total_count(), 5); - assert_eq!(accumulator.inner.zero_count, 3); - assert_eq!( - accumulator.inner.negative_store_counts.iter().sum::(), - 2 - ); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/exact.rs b/crates/asap_summary_state/src/summary_kernels/exact.rs deleted file mode 100644 index 8652bec47..000000000 --- a/crates/asap_summary_state/src/summary_kernels/exact.rs +++ /dev/null @@ -1,350 +0,0 @@ -//! Exact summary state identified by Planner family, independent of keyed layout. -use super::increase::IncreaseAccumulator; -use crate::Statistic; -use crate::{ - AggregateCore, AggregationType, AuxStats, KeyByLabelValues, Measurement, SerializableToSink, -}; -use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; - -type Error = Box; - -#[derive(Debug, Clone, Serialize, Deserialize)] -enum ScalarState { - Sum(f64), - Count(u64), - Min(Option), - Max(Option), - Counter(Option), -} - -/// Both the family and population layout survive persistence. Sharing counter -/// arithmetic never authorizes a Rate state to answer an Increase readout. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ExactAccumulator { - family: SummaryFamilyType, - scalar: ScalarState, - keyed: Option>, -} - -impl ExactAccumulator { - /// Accumulate into run-local scratch state. Persistent input states remain - /// immutable; a failed merge discards this scratch state. - pub(crate) fn merge_from(&mut self, other: &Self) -> Result<(), Error> { - if self.family != other.family || self.is_keyed() != other.is_keyed() { - return Err("cannot merge different Planner families or layouts".into()); - } - if let (Some(target), Some(source)) = (&mut self.keyed, &other.keyed) { - for (key, state) in source { - let combined = match target.get(key) { - Some(old) => merge_scalar(old, state)?, - None => state.clone(), - }; - target.insert(key.clone(), combined); - } - } else { - self.scalar = merge_scalar(&self.scalar, &other.scalar)?; - } - Ok(()) - } - - pub fn new(family: SummaryFamilyType, keyed: bool) -> Result { - use ExactKind as K; - use ExactParams as P; - let scalar = match &family { - SummaryFamilyType::ExactAggregate(K::Sum, P::Sum) => ScalarState::Sum(0.0), - SummaryFamilyType::ExactAggregate(K::Count, P::Count) => ScalarState::Count(0), - SummaryFamilyType::ExactAggregate(K::Min, P::Min) => ScalarState::Min(None), - SummaryFamilyType::ExactAggregate(K::Max, P::Max) => ScalarState::Max(None), - SummaryFamilyType::ExactAggregate(K::Rate, P::Rate) - | SummaryFamilyType::ExactAggregate(K::Increase, P::Increase) => { - ScalarState::Counter(None) - } - _ => return Err(format!("unsupported exact Planner family: {family:?}")), - }; - Ok(Self { - family, - scalar, - keyed: keyed.then(HashMap::new), - }) - } - - pub fn family(&self) -> &SummaryFamilyType { - &self.family - } - pub(crate) fn insufficient_counter_samples( - &self, - statistic: Statistic, - key: &Option, - ) -> bool { - if statistic != self.statistic() { - return false; - } - let state = match (&self.keyed, key) { - (Some(states), Some(key)) => states.get(key), - (None, None) => Some(&self.scalar), - _ => None, - }; - match state { - Some(ScalarState::Counter(None)) => true, - Some(ScalarState::Counter(Some(counter))) => { - counter.sample_count < 2 - || counter.last_seen_timestamp == counter.starting_timestamp - } - _ => false, - } - } - pub fn is_keyed(&self) -> bool { - self.keyed.is_some() - } - - pub fn update(&mut self, key: Option<&KeyByLabelValues>, value: f64, timestamp: i64) { - let state = match (&mut self.keyed, key) { - (Some(states), Some(key)) => states - .entry(key.clone()) - .or_insert_with(|| self.scalar.clone()), - (None, None) => &mut self.scalar, - _ => panic!("exact update population layout differs from installed DAG"), - }; - match state { - ScalarState::Sum(sum) => *sum += value, - ScalarState::Count(count) => { - *count = count.checked_add(1).expect("exact count overflow") - } - ScalarState::Min(current) => { - *current = Some(current.map_or(value, |old| old.min(value))) - } - ScalarState::Max(current) => { - *current = Some(current.map_or(value, |old| old.max(value))) - } - ScalarState::Counter(current) => match current { - Some(counter) => counter.update(Measurement::new(value), timestamp), - None => { - *current = Some(IncreaseAccumulator::new( - Measurement::new(value), - timestamp, - Measurement::new(value), - timestamp, - )) - } - }, - } - } - - pub fn deserialize_from_bytes(bytes: &[u8]) -> Result { - let state: Self = rmp_serde::from_slice(bytes)?; - let expected = Self::new(state.family.clone(), state.is_keyed())?; - let same_variant = |value: &ScalarState| { - std::mem::discriminant(value) == std::mem::discriminant(&expected.scalar) - }; - if !same_variant(&state.scalar) - || state - .keyed - .as_ref() - .is_some_and(|states| states.values().any(|s| !same_variant(s))) - { - return Err("exact payload differs from declared Planner family".into()); - } - Ok(state) - } - - fn statistic(&self) -> Statistic { - match self.family { - SummaryFamilyType::ExactAggregate(ExactKind::Sum, _) => Statistic::Sum, - SummaryFamilyType::ExactAggregate(ExactKind::Count, _) => Statistic::Count, - SummaryFamilyType::ExactAggregate(ExactKind::Min, _) => Statistic::Min, - SummaryFamilyType::ExactAggregate(ExactKind::Max, _) => Statistic::Max, - SummaryFamilyType::ExactAggregate(ExactKind::Rate, _) => Statistic::Rate, - SummaryFamilyType::ExactAggregate(ExactKind::Increase, _) => Statistic::Increase, - _ => unreachable!("validated exact family"), - } - } -} - -fn merge_scalar(left: &ScalarState, right: &ScalarState) -> Result { - Ok(match (left, right) { - (ScalarState::Sum(a), ScalarState::Sum(b)) => ScalarState::Sum(a + b), - (ScalarState::Count(a), ScalarState::Count(b)) => { - ScalarState::Count(a.checked_add(*b).ok_or("exact count overflow")?) - } - (ScalarState::Min(a), ScalarState::Min(b)) => { - ScalarState::Min(a.iter().chain(b).copied().reduce(f64::min)) - } - (ScalarState::Max(a), ScalarState::Max(b)) => { - ScalarState::Max(a.iter().chain(b).copied().reduce(f64::max)) - } - (ScalarState::Counter(a), ScalarState::Counter(b)) => ScalarState::Counter(match (a, b) { - (Some(a), Some(b)) => Some(IncreaseAccumulator::merge_pair(a, b)), - (a, b) => a.clone().or_else(|| b.clone()), - }), - _ => return Err("exact scalar state families differ".into()), - }) -} - -impl SerializableToSink for ExactAccumulator { - fn serialize_to_json(&self) -> serde_json::Value { - serde_json::json!({"family": self.family, "scalar": self.scalar, "keyed": self.keyed.as_ref().map(|m|m.iter().collect::>())}) - } - fn serialize_to_bytes(&self) -> Vec { - rmp_serde::to_vec_named(self).expect("exact state encoding") - } -} - -impl AggregateCore for ExactAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - fn type_name(&self) -> &'static str { - "PlannerExactAccumulatorV1" - } - 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, Error> { - let other = other - .as_any() - .downcast_ref::() - .ok_or("merge requires Planner exact state")?; - let mut merged = self.clone(); - merged.merge_from(other)?; - Ok(Box::new(merged)) - } - fn get_accumulator_type(&self) -> AggregationType { - match self.statistic() { - Statistic::Sum => AggregationType::Sum, - Statistic::Count => AggregationType::Count, - Statistic::Min => AggregationType::Min, - Statistic::Max => AggregationType::Max, - Statistic::Rate => AggregationType::Rate, - Statistic::Increase => AggregationType::Increase, - _ => unreachable!(), - } - } - fn approx_memory_bytes(&self) -> usize { - std::mem::size_of::() - + self.keyed.as_ref().map_or(0, |m| { - m.keys() - .map(|k| { - std::mem::size_of::() - + k.labels.iter().map(String::len).sum::() - }) - .sum::() - }) - } - fn aux_stats(&self) -> AuxStats { - if self.is_keyed() { - return AuxStats::empty(); - } - match self.scalar { - ScalarState::Sum(value) => AuxStats { - sum: Some(value), - ..AuxStats::empty() - }, - ScalarState::Count(value) => AuxStats { - count: Some(value), - ..AuxStats::empty() - }, - ScalarState::Min(value) => AuxStats { - min: value, - ..AuxStats::empty() - }, - ScalarState::Max(value) => AuxStats { - max: value, - ..AuxStats::empty() - }, - ScalarState::Counter(_) => AuxStats::empty(), - } - } - fn get_keys(&self) -> Option> { - self.keyed.as_ref().map(|m| m.keys().cloned().collect()) - } - fn query_statistic( - &self, - statistic: Statistic, - key: &Option, - kwargs: &HashMap, - ) -> Result { - if statistic != self.statistic() { - return Err("readout differs from Planner exact family".into()); - } - let state = match (&self.keyed, key) { - (Some(states), Some(key)) => states.get(key).ok_or("unknown exact population")?, - (None, None) => &self.scalar, - _ => return Err("readout population differs from installed layout".into()), - }; - match state { - ScalarState::Sum(sum) => Ok(*sum), - ScalarState::Count(count) => Ok(*count as f64), - ScalarState::Min(value) | ScalarState::Max(value) => { - value.ok_or_else(|| "empty exact population".into()) - } - ScalarState::Counter(Some(counter)) => { - counter.query_statistic(statistic, &None, kwargs) - } - ScalarState::Counter(None) => Err("empty counter population".into()), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - // Identity, population isolation, and readout survive the persisted format. - #[test] - fn exact_families_roundtrip_and_reject_cross_family_operations() { - let families = [ - (ExactKind::Sum, ExactParams::Sum, Statistic::Sum, 16.0), - (ExactKind::Count, ExactParams::Count, Statistic::Count, 3.0), - (ExactKind::Min, ExactParams::Min, Statistic::Min, 2.0), - (ExactKind::Max, ExactParams::Max, Statistic::Max, 8.0), - (ExactKind::Rate, ExactParams::Rate, Statistic::Rate, 3.0), - ( - ExactKind::Increase, - ExactParams::Increase, - Statistic::Increase, - 6.0, - ), - ]; - for keyed in [false, true] { - let key = keyed.then(|| KeyByLabelValues::new_with_labels(vec!["a".into()])); - let mut states = Vec::new(); - for (kind, params, stat, value) in &families { - let mut state = ExactAccumulator::new( - SummaryFamilyType::ExactAggregate(kind.clone(), params.clone()), - keyed, - ) - .unwrap(); - for (ts, v) in [(1000, 8.0), (2000, 2.0), (3000, 6.0)] { - state.update(key.as_ref(), v, ts); - } - let restored = - ExactAccumulator::deserialize_from_bytes(&state.serialize_to_bytes()).unwrap(); - assert_eq!(restored.family(), state.family()); - assert_eq!( - restored - .query_statistic(*stat, &key, &HashMap::new()) - .unwrap(), - *value - ); - for (_, _, wrong, _) in &families { - if wrong != stat { - assert!(restored - .query_statistic(*wrong, &key, &HashMap::new()) - .is_err()); - } - } - states.push(restored); - } - for (i, a) in states.iter().enumerate() { - for (j, b) in states.iter().enumerate() { - assert_eq!(a.merge_with(b).is_ok(), i == j); - } - } - } - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/factory.rs b/crates/asap_summary_state/src/summary_kernels/factory.rs deleted file mode 100644 index c4e7ef829..000000000 --- a/crates/asap_summary_state/src/summary_kernels/factory.rs +++ /dev/null @@ -1,1176 +0,0 @@ -use crate::summary_kernels::{ - CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, CountSketchAccumulator, - CountSketchWithHeapAccumulator, DDSketchAccumulator, DatasketchesKLLAccumulator, - HydraKllSketchAccumulator, IncreaseAccumulator, KeyedCounterState, KeyedMaxState, - KeyedMinState, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, -}; -use crate::{AggregateCore, KeyByLabelValues, Measurement}; -// Production dispatch consumes Planner SummaryAgg payloads directly. The -// config adapter below is compiled only for isolated historical kernel tests. -use crate::summary_kernels::hll_sketch::HllSketchAccumulator; -use crate::summary_kernels::univmon::UnivMonAccumulator; -use planner_types::post_asap::{ExactKind, SketchAlgorithm, SketchParams, SummaryFamilyType}; - -/// Generate the two boilerplate clone-based `AccumulatorUpdater` methods -/// for updaters whose inner `acc` field implements `Clone + AggregateCore`. -/// Not applicable to `IncreaseAccumulatorUpdater` (its `acc` is `Option<_>` -/// with non-trivial `None` handling). -macro_rules! impl_clone_accumulator_methods { - ($acc_field:ident) => { - fn take_accumulator(&mut self) -> Box { - let result = Box::new(self.$acc_field.clone()); - self.reset(); - result - } - - fn snapshot_accumulator(&self) -> Box { - Box::new(self.$acc_field.clone()) - } - - fn into_accumulator(self: Box) -> Box { - // Consume the updater and MOVE the accumulator out — no clone. - // Avoids the expensive `Clone` (a full msgpack serialize/deserialize - // round-trip for sketch accumulators) when a pane is evicted at - // window close. - let this = *self; - Box::new(this.$acc_field) - } - }; -} - -/// Shared update interface for query-time and precompute-time accumulation. -/// -/// This provides a uniform interface over all accumulator types so that the -/// worker loop doesn't need to know which concrete type it's dealing with. -pub trait AccumulatorUpdater: Send { - /// Validate an immutable precompute input before an updater can silently - /// discard a value outside its representable domain. - fn validate_single_input(&self, value: f64) -> Result<(), String> { - if value.is_finite() { - Ok(()) - } else { - Err("accumulator input must be finite".into()) - } - } - - /// Feed a single (value, timestamp_ms) pair — for SingleSubpopulation types. - fn update_single(&mut self, value: f64, timestamp_ms: i64); - - /// Feed a keyed (key, value, timestamp_ms) triple — for MultipleSubpopulation types. - fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, timestamp_ms: i64); - - /// Extract the final accumulator as a boxed `AggregateCore`. - fn take_accumulator(&mut self) -> Box; - - /// Non-destructive read of the current accumulator state (clone without reset). - /// Used by pane-based sliding windows to read shared panes. - fn snapshot_accumulator(&self) -> Box; - - /// Consume the updater and return its accumulator BY MOVE, avoiding the - /// `Clone` that `take_accumulator`/`snapshot_accumulator` pay (for sketch - /// accumulators that clone is a full msgpack serialize/deserialize - /// round-trip). Used by `merge_panes_for_window` when a pane is evicted at - /// window close. Default falls back to a clone for updaters that can't - /// cheaply move their inner accumulator out. - fn into_accumulator(self: Box) -> Box { - self.snapshot_accumulator() - } - - /// Reset internal state for reuse (avoids re-allocation). - fn reset(&mut self); - - /// Whether this updater is keyed (MultipleSubpopulation). - fn is_keyed(&self) -> bool; - - /// Estimated memory usage in bytes. - fn memory_usage_bytes(&self) -> usize; -} - -// --------------------------------------------------------------------------- -// SumAccumulatorUpdater -// --------------------------------------------------------------------------- - -pub struct SumAccumulatorUpdater { - acc: SumAccumulator, -} - -impl SumAccumulatorUpdater { - pub fn new() -> Self { - Self { - acc: SumAccumulator::new(), - } - } -} - -impl Default for SumAccumulatorUpdater { - fn default() -> Self { - Self::new() - } -} - -impl AccumulatorUpdater for SumAccumulatorUpdater { - fn update_single(&mut self, value: f64, _timestamp_ms: i64) { - self.acc.update(value); - } - - fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { - self.update_single(value, timestamp_ms); - } - - impl_clone_accumulator_methods!(acc); - - fn reset(&mut self) { - self.acc = SumAccumulator::new(); - } - - fn is_keyed(&self) -> bool { - false - } - - fn memory_usage_bytes(&self) -> usize { - std::mem::size_of::() - } -} - -// --------------------------------------------------------------------------- -// MinAccumulatorUpdater / MaxAccumulatorUpdater -// --------------------------------------------------------------------------- - -macro_rules! extremum_updater { - ($updater:ident, $acc:ty) => { - #[derive(Default)] - pub struct $updater { - acc: $acc, - } - - impl $updater { - pub fn new() -> Self { - Self::default() - } - } - - impl AccumulatorUpdater for $updater { - fn update_single(&mut self, value: f64, _timestamp_ms: i64) { - self.acc.update(value); - } - - fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { - self.update_single(value, timestamp_ms); - } - - impl_clone_accumulator_methods!(acc); - - fn reset(&mut self) { - self.acc = <$acc>::new(); - } - - fn is_keyed(&self) -> bool { - false - } - - fn memory_usage_bytes(&self) -> usize { - std::mem::size_of::<$acc>() - } - } - }; -} - -extremum_updater!(MinAccumulatorUpdater, MinAccumulator); -extremum_updater!(MaxAccumulatorUpdater, MaxAccumulator); - -// --------------------------------------------------------------------------- -// IncreaseAccumulatorUpdater -// --------------------------------------------------------------------------- - -pub struct IncreaseAccumulatorUpdater { - acc: Option, -} - -impl IncreaseAccumulatorUpdater { - pub fn new() -> Self { - Self { acc: None } - } -} - -impl Default for IncreaseAccumulatorUpdater { - fn default() -> Self { - Self::new() - } -} - -impl AccumulatorUpdater for IncreaseAccumulatorUpdater { - fn update_single(&mut self, value: f64, timestamp_ms: i64) { - let measurement = Measurement::new(value); - match &mut self.acc { - Some(acc) => acc.update(measurement, timestamp_ms), - None => { - self.acc = Some(IncreaseAccumulator::new( - measurement.clone(), - timestamp_ms, - measurement, - timestamp_ms, - )); - } - } - } - - fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { - self.update_single(value, timestamp_ms); - } - - // Hand-written: acc is Option<_> with non-trivial None handling. - fn take_accumulator(&mut self) -> Box { - let acc = self.acc.take().unwrap_or_else(|| { - IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(0.0), 0) - }); - let result = Box::new(acc); - self.reset(); - result - } - - fn snapshot_accumulator(&self) -> Box { - match &self.acc { - Some(acc) => Box::new(acc.clone()), - None => Box::new(IncreaseAccumulator::new( - Measurement::new(0.0), - 0, - Measurement::new(0.0), - 0, - )), - } - } - - fn reset(&mut self) { - self.acc = None; - } - - fn is_keyed(&self) -> bool { - false - } - - fn memory_usage_bytes(&self) -> usize { - std::mem::size_of::>() - } -} - -// --------------------------------------------------------------------------- -// KllAccumulatorUpdater -// --------------------------------------------------------------------------- - -pub struct KllAccumulatorUpdater { - acc: DatasketchesKLLAccumulator, - k: u16, -} - -impl KllAccumulatorUpdater { - pub fn new(k: u16) -> Self { - Self { - acc: DatasketchesKLLAccumulator::new(k), - k, - } - } -} - -impl AccumulatorUpdater for KllAccumulatorUpdater { - fn update_single(&mut self, value: f64, _timestamp_ms: i64) { - self.acc.update(value); - } - - fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { - self.update_single(value, timestamp_ms); - } - - impl_clone_accumulator_methods!(acc); - - fn reset(&mut self) { - self.acc = DatasketchesKLLAccumulator::new(self.k); - } - - fn is_keyed(&self) -> bool { - false - } - - fn memory_usage_bytes(&self) -> usize { - // KLL sketch size is hard to estimate precisely; use a rough estimate - std::mem::size_of::() + 4096 - } -} - -// --------------------------------------------------------------------------- -// DDSketchAccumulatorUpdater — pendant to KllAccumulatorUpdater -// --------------------------------------------------------------------------- -// -// Drives the agent-aggregated DDSketch path: the worker either -// (a) merges an inbound `DDSketchAccumulator` from the -// modified-OTLP `Data::Ddsketch` ingest (via the worker's -// `merge_with`), or (b) consumes raw values via `update_single` -// when an OTLP scalar datapoint matches an aggregation typed as -// DDSketch. (b) is the less common path but it lets the same -// aggregation slot serve both pre-aggregated agent sketches and -// raw OTLP gauges. -pub struct DDSketchAccumulatorUpdater { - acc: DDSketchAccumulator, - alpha: f64, -} - -impl DDSketchAccumulatorUpdater { - pub fn new(alpha: f64) -> Self { - Self { - acc: DDSketchAccumulator::new(alpha), - alpha, - } - } -} - -impl AccumulatorUpdater for DDSketchAccumulatorUpdater { - fn validate_single_input(&self, value: f64) -> Result<(), String> { - let (minimum, maximum) = - asap_sketchlib::sketches::ddsketch::ddsketch_indexable_bounds(self.alpha); - if value.is_finite() && value > 0.0 && value >= minimum && value <= maximum { - Ok(()) - } else { - Err("DDS maintenance input is outside its positive representable domain".into()) - } - } - - fn update_single(&mut self, value: f64, _timestamp_ms: i64) { - // sketch-core's DdSketch (the inner of DDSketchAccumulator) - // exposes `update(f64)` for single-value ingestion. The - // worker calls this when a raw OTLP datapoint matches an - // aggregation typed as DDSketch — the sketch-merge path - // uses `merge_with` directly. - self.acc.inner.update(value); - } - - fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { - self.update_single(value, timestamp_ms); - } - - impl_clone_accumulator_methods!(acc); - - fn reset(&mut self) { - self.acc = DDSketchAccumulator::new(self.alpha); - } - - fn is_keyed(&self) -> bool { - false - } - - fn memory_usage_bytes(&self) -> usize { - // Bucket store is variable; rough estimate matches KLL. - std::mem::size_of::() + 4096 - } -} - -// --------------------------------------------------------------------------- -// KeyedSumCountAccumulatorUpdater -// --------------------------------------------------------------------------- - -pub struct KeyedSumCountAccumulatorUpdater { - acc: KeyedSumCountAccumulator, -} - -impl KeyedSumCountAccumulatorUpdater { - pub fn new() -> Self { - Self::for_family(ExactKind::Sum) - } - - pub fn for_family(family: ExactKind) -> Self { - Self { - acc: KeyedSumCountAccumulator::for_family(family), - } - } -} - -impl Default for KeyedSumCountAccumulatorUpdater { - fn default() -> Self { - Self::new() - } -} - -impl AccumulatorUpdater for KeyedSumCountAccumulatorUpdater { - 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.update(key.clone(), value); - } - - impl_clone_accumulator_methods!(acc); - - fn reset(&mut self) { - self.acc = KeyedSumCountAccumulator::for_family(self.acc.family.clone()); - } - - fn is_keyed(&self) -> bool { - true - } - - fn memory_usage_bytes(&self) -> usize { - std::mem::size_of::() - + self.acc.sums.len() * (std::mem::size_of::() + 16) - } -} - -// --------------------------------------------------------------------------- -// KeyedMinStateUpdater / KeyedMaxStateUpdater -// --------------------------------------------------------------------------- - -macro_rules! multiple_extremum_updater { - ($updater:ident, $acc:ty) => { - #[derive(Default)] - pub struct $updater { - acc: $acc, - } - - impl $updater { - pub fn new() -> Self { - Self::default() - } - } - - impl AccumulatorUpdater for $updater { - 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.update(key.clone(), value); - } - - impl_clone_accumulator_methods!(acc); - - fn reset(&mut self) { - self.acc = <$acc>::new(); - } - - fn is_keyed(&self) -> bool { - true - } - - fn memory_usage_bytes(&self) -> usize { - std::mem::size_of::<$acc>() - + self.acc.values.len() * (std::mem::size_of::() + 8) - } - } - }; -} - -multiple_extremum_updater!(KeyedMinStateUpdater, KeyedMinState); -multiple_extremum_updater!(KeyedMaxStateUpdater, KeyedMaxState); - -// --------------------------------------------------------------------------- -// KeyedCounterStateUpdater -// --------------------------------------------------------------------------- - -pub struct KeyedCounterStateUpdater { - acc: KeyedCounterState, -} - -impl KeyedCounterStateUpdater { - pub fn new() -> Self { - Self { - acc: KeyedCounterState::new(), - } - } -} - -impl Default for KeyedCounterStateUpdater { - fn default() -> Self { - Self::new() - } -} - -impl AccumulatorUpdater for KeyedCounterStateUpdater { - 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 measurement = Measurement::new(value); - match self.acc.increases.entry(key.clone()) { - std::collections::hash_map::Entry::Occupied(mut e) => { - e.get_mut().update(measurement, timestamp_ms); - } - std::collections::hash_map::Entry::Vacant(e) => { - e.insert(IncreaseAccumulator::new( - measurement.clone(), - timestamp_ms, - measurement, - timestamp_ms, - )); - } - } - } - - impl_clone_accumulator_methods!(acc); - - fn reset(&mut self) { - self.acc = KeyedCounterState::new(); - } - - fn is_keyed(&self) -> bool { - true - } - - fn memory_usage_bytes(&self) -> usize { - std::mem::size_of::() - + self.acc.increases.len() - * (std::mem::size_of::() - + std::mem::size_of::()) - } -} - -// --------------------------------------------------------------------------- -// CmsAccumulatorUpdater (CountMinSketch) -// --------------------------------------------------------------------------- - -/// Keyed weighted-frequency updater. -/// -/// A raw Prometheus sample represents the observed metric value, so a bare CMS -/// adds `value` for its key. Counting each received sample as one is a distinct -/// event-count operation and requires an explicit typed plan contract; it must -/// not be inferred from the sketch algorithm alone. -pub struct CmsAccumulatorUpdater { - acc: CountMinSketchAccumulator, - row_num: usize, - col_num: usize, -} - -impl CmsAccumulatorUpdater { - pub fn new(row_num: usize, col_num: usize) -> Self { - Self { - acc: CountMinSketchAccumulator::new(row_num, col_num), - row_num, - col_num, - } - } -} - -impl AccumulatorUpdater for CmsAccumulatorUpdater { - 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 = CountMinSketchAccumulator::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::() - } -} - -// --------------------------------------------------------------------------- -// CmsHeapAccumulatorUpdater — value-weighted / count-weighted top-k -// --------------------------------------------------------------------------- - -/// What quantity the top-k heap ranks keys by. -/// -/// These are DIFFERENT query semantics and must be chosen explicitly: -/// -/// * [`TopkWeight::Value`] — accumulate **Σ of the datapoint value** per key. -/// This answers "top-k by total " (e.g. "top-k hosts by -/// total CPU"). The heap value is the summed metric value, so the read-side -/// reducer's "sort heap descending by value" yields the correct ranking. -/// -/// * [`TopkWeight::Count`] — accumulate **+1 per event** per key (occurrence -/// frequency), the textbook heavy-hitter / frequency-top-k semantics -/// ("which keys appear most often"). -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum TopkWeight { - /// Σ datapoint value per key (value-weighted top-k). - Value, - /// +1 per event per key (count-weighted / frequency top-k). - Count, -} - -/// Keyed top-k updater backed by a real `CountMinSketchWithHeap` (a CMS -/// matrix PLUS a size-`heap_size` top-k heap). Unlike the heap-LESS -/// `CmsAccumulatorUpdater`, this enumerates top-k keys at read time -/// (`get_topk_keys` / `topk_heap_items`), which is what `topk(...)` queries -/// need. -/// -/// The key is the configured group-by (`aggregated_labels`) value vector — -/// e.g. `host` — formed by `extract_aggregated_key_from_series` in the worker, -/// NOT the hardcoded metric label `item`. The accumulated quantity is selected -/// by [`TopkWeight`]: -/// * `Value` → `inner.update(key, value)` adds the datapoint value (Σ value). -/// * `Count` → `inner.update(key, 1.0)` adds one per event (Σ count). -/// -/// Both `CountMinSketchWithHeap` and `CountSketchWithHeap` raw-input policies -/// route here; the heap is the shared distinguishing payload. -pub struct CmsHeapAccumulatorUpdater { - acc: CountMinSketchWithHeapAccumulator, - row_num: usize, - col_num: usize, - heap_size: usize, - weight: TopkWeight, -} - -impl CmsHeapAccumulatorUpdater { - pub fn new(row_num: usize, col_num: usize, heap_size: usize, weight: TopkWeight) -> Self { - Self { - acc: CountMinSketchWithHeapAccumulator::new(row_num, col_num, heap_size), - row_num, - col_num, - heap_size, - weight, - } - } -} - -impl AccumulatorUpdater for CmsHeapAccumulatorUpdater { - 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) { - // Heap key = the group-by label-value vector (e.g. `host`), joined the - // same way the read-side `get_topk_keys` splits it back apart (`;`). - let weighted = match self.weight { - // Σ value: feed the datapoint value. sketchlib's CMS-heap - // `update(key, w)` adds `w.round()` occurrences of `key`, so the - // heap value accumulates the (rounded) summed metric value. - TopkWeight::Value => value, - // Σ count: one occurrence per event, regardless of 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 = - CountMinSketchWithHeapAccumulator::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) - } -} - -// --------------------------------------------------------------------------- -// 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 `SketchAlgorithm::CountSketch` silently -/// shared `CmsAccumulatorUpdater` with bare CMS. -/// -/// As with bare CMS, each raw Prometheus sample contributes its `value`. -/// Unit event counting must be selected explicitly by a future typed plan -/// contract rather than being implied by `SketchAlgorithm::CountSketch`. -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 -// --------------------------------------------------------------------------- - -pub struct HydraKllAccumulatorUpdater { - acc: HydraKllSketchAccumulator, - row_num: usize, - col_num: usize, - k: u16, -} - -impl HydraKllAccumulatorUpdater { - pub fn new(row_num: usize, col_num: usize, k: u16) -> Self { - Self { - acc: HydraKllSketchAccumulator::new(row_num, col_num, k), - row_num, - col_num, - k, - } - } -} - -impl AccumulatorUpdater for HydraKllAccumulatorUpdater { - 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.update(key, value); - } - - impl_clone_accumulator_methods!(acc); - - fn reset(&mut self) { - self.acc = HydraKllSketchAccumulator::new(self.row_num, self.col_num, self.k); - } - - fn is_keyed(&self) -> bool { - true - } - - fn memory_usage_bytes(&self) -> usize { - // Rough estimate: each cell is a KLL sketch - std::mem::size_of::() + self.row_num * self.col_num * 4096 - } -} - -// --------------------------------------------------------------------------- -// Config helpers -// --------------------------------------------------------------------------- - -fn cms_dims(params: &SketchParams) -> (usize, usize) { - match params { - SketchParams::Cms { width, depth } | SketchParams::CountSketch { width, depth } => { - (*depth as usize, *width as usize) - } - other => unreachable!( - "accumulator_spec() paired SketchAlgorithm::Cms/CountSketch with unexpected params: {other:?}" - ), - } -} - -/// Read `(rows = depth, columns = width, heap_size)` out of `SketchParams::CmsWithHeap` -/// or `::CountSketchWithHeap`. -fn cms_heap_dims(params: &SketchParams) -> (usize, usize, usize) { - match params { - SketchParams::CmsWithHeap { - width, - depth, - heap_size, - } - | SketchParams::CountSketchWithHeap { - width, - depth, - heap_size, - } => (*depth as usize, *width as usize, *heap_size as usize), - other => unreachable!( - "accumulator_spec() paired a WithHeap SketchAlgorithm with unexpected params: {other:?}" - ), - } -} - -/// Construct the kernel declared by a Planner SummaryAgg. No backend config -/// tags participate in this dispatch and unsupported payloads are errors. -pub fn create_planner_accumulator( - family: &SummaryFamilyType, - input: &planner_types::post_asap::SummaryUpdate, - grouping: &planner_types::post_asap::GroupingStrategy, -) -> Result, String> { - if input.item.is_some() - && matches!( - input.weight_domain, - planner_types::post_asap::WeightDomain::NonNegative { - proof: - planner_types::post_asap::NonNegativeWeightProof::ResetAwareCounterDerivative - } - ) - { - return Err("window-weighted summaries require typed DAG binding; integer heap updaters cannot consume rates".into()); - } - - asap_physical_operators::capability::validate_summary_kernel(family, input, grouping)?; - use planner_types::post_asap::GroupingStrategy; - if grouping != &GroupingStrategy::PerSubpopulationInstance { - return Err("shared summary grouping requires a supported Planner Hydra kernel".into()); - } - if matches!(family, SummaryFamilyType::ExactAggregate(..)) { - return Ok(Box::new(PlannerExactUpdater { - acc: crate::summary_kernels::exact::ExactAccumulator::new( - family.clone(), - input.item.is_some(), - )?, - })); - } - let SummaryFamilyType::Sketch(kind, family_grouping) = family else { - return Err(format!("unsupported Planner summary family {family:?}")); - }; - if family_grouping != grouping { - return Err("Planner family and operator grouping disagree".into()); - } - let updater: Box = match (kind.algorithm(), kind.params()) { - (SketchAlgorithm::Kll, SketchParams::Kll { k }) => Box::new(KllAccumulatorUpdater::new( - u16::try_from(*k).map_err(|_| "KLL k exceeds runtime bound")?, - )), - (SketchAlgorithm::DDSketch, SketchParams::DDSketch { alpha }) => { - Box::new(DDSketchAccumulatorUpdater::new(*alpha)) - } - (SketchAlgorithm::Cms, params @ SketchParams::Cms { .. }) => { - let (r, c) = cms_dims(params); - Box::new(CmsAccumulatorUpdater::new(r, c)) - } - (SketchAlgorithm::CountSketch, params @ SketchParams::CountSketch { .. }) => { - let (r, c) = cms_dims(params); - Box::new(CountSketchAccumulatorUpdater::new(r, c)) - } - (SketchAlgorithm::CmsWithHeap, params @ SketchParams::CmsWithHeap { .. }) => { - let (r, c, h) = cms_heap_dims(params); - Box::new(CmsHeapAccumulatorUpdater::new(r, c, h, TopkWeight::Value)) - } - ( - SketchAlgorithm::CountSketchWithHeap, - params @ SketchParams::CountSketchWithHeap { .. }, - ) => { - let (r, c, h) = cms_heap_dims(params); - Box::new(CountSketchWithHeapAccumulatorUpdater::new( - r, - c, - h, - TopkWeight::Value, - )) - } - (SketchAlgorithm::Hll, SketchParams::Hll { precision }) => Box::new(HllUpdater { - acc: HllSketchAccumulator::new( - asap_sketchlib::HllVariant::Regular, - u32::from(*precision), - ), - }), - ( - SketchAlgorithm::UnivMon, - SketchParams::UnivMon { - heap_size, - sketch_rows, - sketch_cols, - layers, - }, - ) => Box::new(UnivMonUpdater { - acc: UnivMonAccumulator::new( - *heap_size as usize, - *sketch_rows as usize, - *sketch_cols as usize, - *layers as usize, - ) - .map_err(|e| e.to_string())?, - }), - _ => { - return Err(format!( - "unsupported Planner algorithm/parameters: {kind:?}" - )) - } - }; - if updater.is_keyed() != input.item.is_some() && !is_unit_sample_frequency(input) { - return Err("Planner item expression does not match the selected kernel layout".into()); - } - Ok(updater) -} - -/// A unit-weight sample-value frequency may use a keyed kernel without an item column. -fn is_unit_sample_frequency(update: &planner_types::post_asap::SummaryUpdate) -> bool { - use planner_types::post_asap::{NonNegativeWeightProof, SummaryInputExpr, WeightDomain}; - matches!( - update.item, - Some(SummaryInputExpr::Column( - planner_types::pre_asap::ColumnRef::SampleValue - )) - ) && matches!(update.weight, SummaryInputExpr::Constant(1.0)) - && matches!( - update.weight_domain, - WeightDomain::NonNegative { - proof: NonNegativeWeightProof::UnitCount - } - ) -} - -struct PlannerExactUpdater { - acc: crate::summary_kernels::exact::ExactAccumulator, -} -impl AccumulatorUpdater for PlannerExactUpdater { - fn update_single(&mut self, value: f64, timestamp: i64) { - self.acc.update(None, value, timestamp); - } - fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, timestamp: i64) { - self.acc.update(Some(key), value, timestamp); - } - impl_clone_accumulator_methods!(acc); - fn reset(&mut self) { - self.acc = crate::summary_kernels::exact::ExactAccumulator::new( - self.acc.family().clone(), - self.acc.is_keyed(), - ) - .expect("installed exact family"); - } - fn is_keyed(&self) -> bool { - self.acc.is_keyed() - } - fn memory_usage_bytes(&self) -> usize { - self.acc.approx_memory_bytes() - } -} - -struct UnivMonUpdater { - acc: UnivMonAccumulator, -} - -struct HllUpdater { - acc: HllSketchAccumulator, -} - -impl AccumulatorUpdater for HllUpdater { - fn is_keyed(&self) -> bool { - false - } - fn memory_usage_bytes(&self) -> usize { - self.acc.approx_memory_bytes() - } - fn update_single(&mut self, value: f64, _: i64) { - if !value.is_nan() { - let bits = if value == 0.0 { 0 } else { value.to_bits() }; - self.acc.inner.update(&bits.to_le_bytes()); - } - } - fn update_keyed(&mut self, _: &KeyByLabelValues, value: f64, timestamp_ms: i64) { - self.update_single(value, timestamp_ms); - } - impl_clone_accumulator_methods!(acc); - fn reset(&mut self) { - self.acc.reset_to_empty(); - } -} - -impl AccumulatorUpdater for UnivMonUpdater { - fn is_keyed(&self) -> bool { - false - } - fn memory_usage_bytes(&self) -> usize { - self.acc.approx_memory_bytes() - } - fn update_single(&mut self, value: f64, _: i64) { - self.acc - .insert_sample(value) - .expect("UnivMon sample counter overflow"); - } - fn update_keyed(&mut self, _: &KeyByLabelValues, value: f64, timestamp_ms: i64) { - self.update_single(value, timestamp_ms); - } - impl_clone_accumulator_methods!(acc); - fn reset(&mut self) { - self.acc.reset_to_empty(); - } -} - -#[cfg(test)] -mod planner_parameter_regression { - use super::*; - use planner_types::post_asap::{SketchKind, SummaryInputExpr, SummaryUpdate}; - - // Planner width is the bucket count; depth is the independent hash-row count. - #[test] - fn planner_sketch_dimensions_are_not_transposed() { - for (algorithm, params) in [ - ( - SketchAlgorithm::Cms, - SketchParams::Cms { - width: 128, - depth: 3, - }, - ), - ( - SketchAlgorithm::CountSketch, - SketchParams::CountSketch { - width: 128, - depth: 3, - }, - ), - ( - SketchAlgorithm::CmsWithHeap, - SketchParams::CmsWithHeap { - width: 128, - depth: 3, - heap_size: 8, - }, - ), - ( - SketchAlgorithm::CountSketchWithHeap, - SketchParams::CountSketchWithHeap { - width: 128, - depth: 3, - heap_size: 8, - }, - ), - ] { - let family = SummaryFamilyType::Sketch( - SketchKind::new(algorithm.clone(), params), - Default::default(), - ); - let update = SummaryUpdate { - item: Some(SummaryInputExpr::Column( - planner_types::pre_asap::ColumnRef::Named("host".into()), - )), - weight: SummaryInputExpr::Constant(1.0), - weight_domain: Default::default(), - }; - let state = create_planner_accumulator(&family, &update, &Default::default()) - .unwrap() - .snapshot_accumulator(); - let dims = match algorithm { - SketchAlgorithm::Cms => { - let s = state - .as_any() - .downcast_ref::() - .unwrap(); - (s.inner.rows(), s.inner.cols()) - } - SketchAlgorithm::CountSketch => { - let s = state - .as_any() - .downcast_ref::() - .unwrap(); - (s.inner.rows, s.inner.cols) - } - SketchAlgorithm::CmsWithHeap => { - let s = state - .as_any() - .downcast_ref::() - .unwrap(); - (s.inner.rows(), s.inner.cols()) - } - SketchAlgorithm::CountSketchWithHeap => { - let s = state - .as_any() - .downcast_ref::() - .unwrap(); - (s.inner.rows(), s.inner.cols()) - } - _ => unreachable!(), - }; - assert_eq!(dims, (3, 128), "{algorithm:?}"); - } - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/hll_sketch.rs b/crates/asap_summary_state/src/summary_kernels/hll_sketch.rs deleted file mode 100644 index 4c16156d6..000000000 --- a/crates/asap_summary_state/src/summary_kernels/hll_sketch.rs +++ /dev/null @@ -1,788 +0,0 @@ -//! HLL accumulator — wraps `asap_sketchlib::HllSketch`. -//! -//! Concrete accumulator reached from the modified-OTLP -//! `Metric.data = HLLSketch{…}` hot path (PR C-CountSketch follow-up). -//! Mirrors the CountSketch accumulator's shape: merge via register-wise -//! max on the inner sketch, serialize as MessagePack for the sink, and -//! decode from the sketchlib `HyperLogLogState` proto. -//! -//! Query semantics (cardinality estimation via the three HLL variants' -//! estimators) are intentionally deferred — the wire format carries the -//! registers + variant + HIP accumulators losslessly, so the merge + -//! store round-trip works end-to-end without that richer query surface. - -use crate::summary_kernels::dd_sketch::normalize_sample_p; -use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; -use asap_sketchlib::{HllSketch, HllVariant, MessagePackCodec}; -use serde_json::Value; -use std::collections::HashMap; - -/// Decode one protobuf base-128 varint (LEB128) from the front of `buf`. -/// Returns `(value, bytes_consumed)`, or `None` if the buffer is truncated -/// or the varint overflows u64. -pub(crate) fn read_uvarint(buf: &[u8]) -> Option<(u64, usize)> { - let mut result: u64 = 0; - let mut shift: u32 = 0; - for (i, &b) in buf.iter().enumerate() { - if shift >= 64 { - return None; - } - result |= u64::from(b & 0x7f) << shift; - if b & 0x80 == 0 { - return Some((result, i + 1)); - } - shift += 7; - } - None -} - -/// Expand sketchlib-go's sparse HLL register encoding -/// (`HLLSparseRegisters.packed`) into the dense `num_registers`-byte array. -/// -/// Layout (sketchlib-go `proto/hll/hll.proto`): varint-packed -/// `(index_delta, value)` pairs in ascending index order; `prev_index` -/// starts at 0, so each register's absolute index is the running sum of the -/// deltas. Mirrors the Go encoder in `sketches/HLL/sparse.go` -/// (`encodeSparseRegisters`). The reconstructed array is byte-identical to -/// the dense `registers` field a high-cardinality producer would have sent. -pub(crate) fn expand_sparse_hll_registers( - packed: &[u8], - num_registers: usize, -) -> Result, Box> { - let mut regs = vec![0u8; num_registers]; - let mut prev: u64 = 0; - let mut pos = 0usize; - while pos < packed.len() { - let (delta, n1) = read_uvarint(&packed[pos..]) - .ok_or("HLLSparseRegisters.packed: truncated index_delta varint")?; - pos += n1; - let (value, n2) = read_uvarint(&packed[pos..]) - .ok_or("HLLSparseRegisters.packed: truncated value varint")?; - pos += n2; - let idx = prev + delta; - let i = usize::try_from(idx) - .map_err(|_| format!("HLLSparseRegisters: index {idx} overflows usize"))?; - if i >= num_registers { - return Err(format!( - "HLLSparseRegisters: register index {i} >= num_registers {num_registers}" - ) - .into()); - } - regs[i] = u8::try_from(value) - .map_err(|_| format!("HLLSparseRegisters: register value {value} > 255"))?; - prev = idx; - } - Ok(regs) -} - -/// HLL accumulator — inner register array + variant metadata. -#[derive(Debug, Clone)] -pub struct HllSketchAccumulator { - pub inner: HllSketch, - /// Edge sampling probability `p ∈ (0,1]` carried on the producer's - /// `SketchEnvelope.sample_p`. HLL uses HASH-THRESHOLD sampling — each - /// DISTINCT key is admitted into the sketch with probability `p`, so the - /// register-derived distinct-count estimate is ~`p`× the true - /// cardinality and a `Cardinality`/`Count` query must rescale by `1/p`. - /// `1.0` (and the proto3 default `0.0`, dual-read as `1.0`) means no - /// sampling, so the rescale is a no-op and the behaviour is identical to - /// before. Mirrors `DDSketchAccumulator::sample_p`; set from the envelope - /// at the `from_sketchlib_proto_bytes` decode site and preserved across - /// `reset_to_empty` and `merge_with`. - /// - /// NOTE: HLL edge sampling is currently force-disabled in the edge - /// (`warm_sketch.go` HLL case always emits `sample_p = 1.0`), so in - /// practice `p = 1.0` today and this is a latent-correctness fix that - /// activates if HLL sampling is ever enabled. - pub sample_p: f64, -} - -impl HllSketchAccumulator { - pub fn new(variant: HllVariant, precision: u32) -> Self { - Self { - inner: HllSketch::new(variant, precision), - sample_p: 1.0, - } - } - - /// Decode from the modified OTLP wire format's - /// `HLLSketchDataPoint.sketch` bytes when - /// `encoding = HLL_SKETCH_ENCODING_MSGPACK`. The bytes are the - /// MessagePack serialization of the cross-language sketch-core - /// `HllSketch` struct — PR I parity entrypoint. - pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { - Ok(Self { - inner: HllSketch::from_msgpack(buffer) - .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?, - // The msgpack HllSketch struct carries no envelope/sample_p; the - // msgpack path is parity/test-only and is never edge-sampled. - sample_p: 1.0, - }) - } - - /// Decode from the modified OTLP wire format's - /// `HLLSketchDataPoint.sketch` bytes — the protobuf-encoded - /// `asap_sketchlib::proto::sketchlib::HyperLogLogState` message - /// that DataCollector's `hllprocessor` emits when - /// `encoding = HLL_SKETCH_ENCODING_PROTO`. - pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, - }; - use prost::Message; - - // DataCollector's hllprocessor wraps the state in a - // `SketchEnvelope{hll: HyperLogLogState}` via sketchlib-go's - // `SerializePortableFO` + `proto.Marshal`. Try envelope first, - // fall back to bare `HyperLogLogState` for callers (e.g. unit - // tests) that encode the state directly. Mirrors the PR #14 - // fix on `CountMinSketchAccumulator::from_sketchlib_proto_bytes`. - // Capture the envelope's `sample_p` alongside the state so a - // Cardinality query can rescale the distinct-count estimate by - // `1/p`. Bare `HyperLogLogState` bytes (no envelope) carry no - // sampling info → `sample_p` 1.0 (no rescale). Mirrors - // `DDSketchAccumulator`. - let (state, sample_p) = match SketchEnvelope::decode(buffer) { - Ok(env) => { - let sp = env.sample_p; - match env.sketch_state { - Some(sketch_envelope::SketchState::Hll(st)) => (st, sp), - Some(other) => { - return Err(format!( - "SketchEnvelope contains non-HLL sketch: {:?}", - std::mem::discriminant(&other) - ) - .into()); - } - None => ( - HyperLogLogState::decode(buffer) - .map_err(|e| format!("decode HyperLogLogState: {e}"))?, - 1.0, - ), - } - } - Err(_) => ( - HyperLogLogState::decode(buffer) - .map_err(|e| format!("decode HyperLogLogState: {e}"))?, - 1.0, - ), - }; - if state.precision == 0 || state.precision > 20 { - return Err(format!( - "HyperLogLogState precision {} out of range (expected 1..=20)", - state.precision - ) - .into()); - } - let expected_len = 1usize << state.precision; - // Register resolution. sketchlib-go emits the SPARSE - // `registers_sparse` (proto tag 7) form below its dense/sparse - // crossover (~6000 non-zero registers — see - // sketchlib-go/sketches/HLL/sparse.go); low-cardinality producers - // (the common case) therefore leave the dense `registers` (tag 3) - // field empty. The proto contract (hll.proto) is: read whichever of - // `registers` / `registers_sparse` is present; if both are empty the - // sketch is all-zero. Reconstruct the dense 2^precision array in all - // three cases so the inner `HllSketch` always gets a full register - // vector. - let dense_registers: Vec = if state.registers.len() == expected_len { - state.registers.clone() - } else if !state.registers.is_empty() { - // A non-empty dense field of the wrong length is a malformed frame. - return Err(format!( - "HyperLogLogState registers has {} bytes, expected 2^precision = {}", - state.registers.len(), - expected_len - ) - .into()); - } else if let Some(sparse) = state.registers_sparse.as_ref() { - expand_sparse_hll_registers(&sparse.packed, expected_len)? - } else { - // Neither representation populated → all-zero register array. - vec![0u8; expected_len] - }; - let proto_variant = ProtoVariant::try_from(state.variant) - .map_err(|_| format!("HyperLogLogState has unknown variant tag {}", state.variant))?; - let variant = match proto_variant { - ProtoVariant::Unspecified => HllVariant::Unspecified, - ProtoVariant::Regular => HllVariant::Regular, - ProtoVariant::ErtlMle => HllVariant::Datafusion, - ProtoVariant::Hip => HllVariant::Hip, - }; - let inner = HllSketch::from_raw( - variant, - state.precision, - dense_registers, - state.hip_kxq0, - state.hip_kxq1, - state.hip_est, - ); - Ok(Self { - inner, - sample_p: normalize_sample_p(sample_p), - }) - } - - /// Apply a proto-encoded `HLLDelta` frame to this accumulator's - /// inner sketch — the decode path for - /// `HLL_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). - /// - /// Called against an accumulator that already carries the base - /// sketch state; the caller is the per-series snapshot cache in - /// the ingest path. Bytes are the - /// `asap_sketchlib::proto::sketchlib::HllDelta` message. - pub fn apply_proto_delta_bytes( - &mut self, - buffer: &[u8], - ) -> Result<(), Box> { - // The HLLDelta wire format is a varint-packed (index_delta, value) blob; - // decode + apply (register-wise max) via the shared sketch library so - // the unpacking stays a single source of truth. - self.inner - .apply_delta_bytes(buffer) - .map_err(|e| format!("apply HLLDelta: {e}"))?; - Ok(()) - } -} - -impl SerializableToSink for HllSketchAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ - "variant": format!("{:?}", self.inner.variant), - "precision": self.inner.precision, - "register_bytes": self.inner.registers.len(), - "hip_kxq0": self.inner.hip_kxq0, - "hip_kxq1": self.inner.hip_kxq1, - "hip_est": self.inner.hip_est, - }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.inner.to_msgpack().unwrap_or_default() - } -} - -impl AggregateCore for HllSketchAccumulator { - fn approx_memory_bytes(&self) -> usize { - std::mem::size_of::().saturating_add(self.inner.registers.capacity()) - } - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "HllSketchAccumulator" - } - - /// Per-window base rotation: zero the registers but keep the variant - /// and precision. Critical for HLL — its register-wise `max` merge - /// has no inverse, so a never-reset base accumulates the all-time-max - /// across windows (`docs/delta-baseline-contract.md` §1.5); rotating - /// to an empty register array makes per-window cardinality correct. - /// `sample_p` is a per-series config constant (not per-window data), so - /// it is intentionally preserved across the rotation — mirrors - /// `DDSketchAccumulator`. - fn reset_to_empty(&mut self) { - self.inner = HllSketch::new(self.inner.variant, self.inner.precision); - } - - 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 HllSketchAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - let other_hll = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to HllSketchAccumulator")?; - let merged_inner = HllSketch::merge_refs(&[&self.inner, &other_hll.inner])?; - // Mirror DDSketchAccumulator's merge policy exactly: sample_p is a - // per-series config constant, so both operands carry the same value - // in practice. Prefer a sampled factor over the no-sampling default - // so a merge with a freshly-reset (1.0) base keeps the series' - // sampling rate. - let sample_p = if self.sample_p < 1.0 { - self.sample_p - } else { - other_hll.sample_p - }; - Ok(Box::new(Self { - inner: merged_inner, - sample_p, - })) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::HLL - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - _key: &Option, - _query_kwargs: &HashMap, - ) -> Result> { - use crate::Statistic; - match statistic { - // HLL's natural answer is unique-cardinality. PromQL's - // `count_over_time(...)` and `count(...)` both surface - // as `Statistic::Count` after pattern matching but - // semantically they mean "how many distinct values - // were observed in this window" when the underlying - // aggregator is HLL — that's the cardinality estimate, - // not a sample-count. Accept both. - Statistic::Cardinality | Statistic::Count => { - // HLL uses hash-threshold sampling — each distinct key is - // admitted with probability `sample_p`, so the register- - // derived distinct-count estimate is ~`p`× the true - // cardinality. Rescale by `1/sample_p` for an unbiased - // estimate. `sample_p == 1.0` (unsampled / legacy / edge - // HLL sampling currently force-disabled) makes this a no-op. - Ok(hll_cardinality_estimate(&self.inner.registers) / self.sample_p) - } - other => Err(format!( - "HllSketchAccumulator: statistic {:?} not supported (only Cardinality / Count)", - other, - ) - .into()), - } - } -} - -/// Standard HyperLogLog cardinality estimate with the canonical -/// `α_m × m² / Σ 2^(-register[i])` formula plus the small-range -/// (linear-counting) and large-range (32-bit space) corrections -/// from the original Flajolet et al. paper. -/// -/// Inlined here rather than added as a method on `asap_sketchlib::HllSketch` -/// because the existing `asap_sketchlib::asap` types only expose merge / -/// serialize today; adding a query method there would force a -/// cross-crate change. -fn hll_cardinality_estimate(registers: &[u8]) -> f64 { - let m = registers.len() as f64; - if m == 0.0 { - return 0.0; - } - let alpha = match registers.len() { - 16 => 0.673, - 32 => 0.697, - 64 => 0.709, - _ => 0.7213 / (1.0 + 1.079 / m), - }; - - let mut sum = 0.0f64; - let mut zero_registers = 0usize; - for &r in registers { - sum += 2f64.powi(-(r as i32)); - if r == 0 { - zero_registers += 1; - } - } - let raw = alpha * m * m / sum; - - // Small-range (linear-counting) correction. - if raw <= 2.5 * m && zero_registers > 0 { - return m * (m / zero_registers as f64).ln(); - } - - // Large-range correction (only meaningful with 32-bit register - // spaces; sketch-core uses up to 64-bit hashes so this branch - // rarely fires in practice — kept for completeness). - let two_pow_32 = 4_294_967_296f64; - if raw > two_pow_32 / 30.0 { - return -two_pow_32 * (1.0 - raw / two_pow_32).ln(); - } - raw -} - -#[cfg(test)] -mod tests { - use super::*; - - fn encode_state( - variant: i32, - precision: u32, - registers: Vec, - hip_kxq0: f64, - hip_kxq1: f64, - hip_est: f64, - ) -> Vec { - use asap_sketchlib::proto::sketchlib::HyperLogLogState; - use prost::Message; - let state = HyperLogLogState { - variant, - precision, - registers, - hip_kxq0, - hip_kxq1, - hip_est, - registers_sparse: None, - }; - state.encode_to_vec() - } - - #[test] - fn test_from_sketchlib_proto_bytes_regular() { - use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; - let bytes = encode_state( - ProtoVariant::Regular as i32, - 2, - vec![1, 2, 3, 4], - 0.0, - 0.0, - 0.0, - ); - let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.inner.variant, HllVariant::Regular); - assert_eq!(acc.inner.precision, 2); - assert_eq!(acc.inner.registers, vec![1, 2, 3, 4]); - } - - #[test] - fn test_from_sketchlib_proto_bytes_hip_preserves_accumulators() { - use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; - let bytes = encode_state( - ProtoVariant::Hip as i32, - 2, - vec![0, 0, 0, 0], - 1.5, - 2.5, - 42.0, - ); - let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.inner.variant, HllVariant::Hip); - assert_eq!(acc.inner.hip_kxq0, 1.5); - assert_eq!(acc.inner.hip_kxq1, 2.5); - assert_eq!(acc.inner.hip_est, 42.0); - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrapped() { - // Mirrors what DataCollector's hllprocessor emits: the state - // wrapped in a `SketchEnvelope{hll: ...}` via sketchlib-go's - // `SerializePortableFO` + `proto.Marshal`. - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, - }; - use prost::Message; - - let state = HyperLogLogState { - variant: ProtoVariant::Regular as i32, - precision: 2, - registers: vec![1, 2, 3, 4], - hip_kxq0: 0.0, - hip_kxq1: 0.0, - hip_est: 0.0, - registers_sparse: None, - }; - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Hll(state)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes) - .expect("envelope-wrapped decode should succeed"); - assert_eq!(acc.inner.variant, HllVariant::Regular); - assert_eq!(acc.inner.registers, vec![1, 2, 3, 4]); - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { - use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; - use prost::Message; - - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "wrong-sketch envelope should error"); - } - - #[test] - fn test_from_sketchlib_proto_bytes_register_length_mismatch() { - use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; - // precision=2 → expected 4 registers; supply only 3 - let bytes = encode_state( - ProtoVariant::Regular as i32, - 2, - vec![1, 2, 3], - 0.0, - 0.0, - 0.0, - ); - let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("registers")); - } - - #[test] - fn test_from_sketchlib_proto_bytes_zero_precision_rejected() { - use asap_sketchlib::proto::sketchlib::HyperLogLogState; - use prost::Message; - let state = HyperLogLogState::default(); - let bytes = state.encode_to_vec(); - let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - } - - #[test] - fn test_aggregate_core_merge_matches_register_max() { - let a = HllSketchAccumulator { - inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 5, 3, 7], 0.0, 0.0, 0.0), - sample_p: 1.0, - }; - let b = HllSketchAccumulator { - inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![4, 2, 6, 0], 0.0, 0.0, 0.0), - sample_p: 1.0, - }; - let merged_box = a.merge_with(&b).expect("merge ok"); - let merged = merged_box - .as_any() - .downcast_ref::() - .expect("downcast ok"); - assert_eq!(merged.inner.registers, vec![4, 5, 6, 7]); - } - - #[test] - fn test_aggregate_core_merge_wrong_type_rejects() { - use crate::summary_kernels::count_sketch::CountSketchAccumulator; - let hll = HllSketchAccumulator::new(HllVariant::Regular, 2); - let cs = CountSketchAccumulator::new(2, 3); - assert!(hll.merge_with(&cs).is_err()); - } - - #[test] - fn test_from_msgpack_bytes_round_trip() { - let original = HllSketch::from_raw( - HllVariant::Hip, - 3, - vec![0, 1, 2, 3, 4, 5, 6, 7], - 1.5, - 2.5, - 42.0, - ); - let bytes = original.to_msgpack().unwrap(); - let acc = HllSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.inner.variant, HllVariant::Hip); - assert_eq!(acc.inner.precision, 3); - assert_eq!(acc.inner.registers, vec![0, 1, 2, 3, 4, 5, 6, 7]); - assert_eq!(acc.inner.hip_kxq0, 1.5); - } - - #[test] - fn test_from_msgpack_bytes_rejects_garbage() { - let result = HllSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); - assert!(result.is_err()); - } - - #[test] - fn test_apply_proto_delta_bytes_round_trip() { - use asap_sketchlib::proto::sketchlib::HllDelta as PbDelta; - use prost::Message; - - let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 2); - acc.inner.registers = vec![1, 5, 3, 7]; - - // Packed (index_delta, value) blob for updates {0:4, 2:6}: - // varint(0),varint(4),varint(2),varint(6). - let delta_bytes = PbDelta { - packed_updates: vec![0, 4, 2, 6], - } - .encode_to_vec(); - - acc.apply_proto_delta_bytes(&delta_bytes).expect("apply ok"); - // Max semantics: reg[0]=max(1,4)=4, reg[2]=max(3,6)=6; others unchanged. - assert_eq!(acc.inner.registers, vec![4, 5, 6, 7]); - } - - #[test] - fn test_apply_proto_delta_bytes_rejects_garbage() { - let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 2); - assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); - } - - // ----- sample_p cardinality rescale ----- - // - // HLL uses hash-threshold sampling: each distinct key is admitted into - // the sketch with probability `p`, so the register-derived cardinality - // estimate is ~p× the true distinct count and must be rescaled by 1/p. - - #[test] - fn test_cardinality_is_rescaled_by_sample_p() { - use crate::Statistic; - // Build two accumulators with identical registers but different - // sample_p. The sampled one (p=0.25) must report ~4× the unsampled - // estimate. Use precision 8 (256 registers) with a spread of - // register values so the estimate is a non-trivial positive number. - let mut registers = vec![0u8; 256]; - for (i, r) in registers.iter_mut().enumerate() { - *r = ((i % 7) + 1) as u8; - } - let unsampled = HllSketchAccumulator { - inner: HllSketch::from_raw(HllVariant::Regular, 8, registers.clone(), 0.0, 0.0, 0.0), - sample_p: 1.0, - }; - let sampled = HllSketchAccumulator { - inner: HllSketch::from_raw(HllVariant::Regular, 8, registers, 0.0, 0.0, 0.0), - sample_p: 0.25, - }; - let raw = unsampled - .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) - .expect("cardinality ok"); - let rescaled = sampled - .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) - .expect("cardinality ok"); - assert!(raw > 0.0, "raw estimate should be positive, got {raw}"); - // Exact algebraic relationship: rescaled == raw / 0.25 == raw * 4. - assert!( - (rescaled - raw * 4.0).abs() < 1e-9, - "expected rescaled ≈ 4×raw ({}), got {rescaled}", - raw * 4.0 - ); - } - - #[test] - fn test_count_statistic_also_rescaled_by_sample_p() { - use crate::Statistic; - // Count maps to the same cardinality estimate for HLL, so it must - // rescale identically. - let registers = vec![3u8; 16]; - let unsampled = HllSketchAccumulator { - inner: HllSketch::from_raw(HllVariant::Regular, 4, registers.clone(), 0.0, 0.0, 0.0), - sample_p: 1.0, - }; - let sampled = HllSketchAccumulator { - inner: HllSketch::from_raw(HllVariant::Regular, 4, registers, 0.0, 0.0, 0.0), - sample_p: 0.25, - }; - let raw = unsampled - .query_statistic(Statistic::Count, &None, &HashMap::new()) - .expect("count ok"); - let rescaled = sampled - .query_statistic(Statistic::Count, &None, &HashMap::new()) - .expect("count ok"); - assert!((rescaled - raw * 4.0).abs() < 1e-9); - } - - #[test] - fn test_sample_p_unset_behaves_as_one() { - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, - }; - use prost::Message; - // An envelope with no sample_p set (proto3 default 0.0) must - // normalize to 1.0 (no rescale) — byte-compatible with legacy frames. - let state = HyperLogLogState { - variant: ProtoVariant::Regular as i32, - precision: 4, - registers: vec![2u8; 16], - hip_kxq0: 0.0, - hip_kxq1: 0.0, - hip_est: 0.0, - registers_sparse: None, - }; - let env = SketchEnvelope { - // sample_p left at proto3 default 0.0. - sketch_state: Some(sketch_envelope::SketchState::Hll(state)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.sample_p, 1.0, "unset sample_p must normalize to 1.0"); - } - - #[test] - fn test_from_sketchlib_proto_bytes_reads_envelope_sample_p() { - use crate::Statistic; - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, - }; - use prost::Message; - - let registers = vec![3u8; 16]; - let state = HyperLogLogState { - variant: ProtoVariant::Regular as i32, - precision: 4, - registers: registers.clone(), - hip_kxq0: 0.0, - hip_kxq1: 0.0, - hip_est: 0.0, - registers_sparse: None, - }; - let env = SketchEnvelope { - sample_p: 0.25, - sketch_state: Some(sketch_envelope::SketchState::Hll(state)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.sample_p, 0.25); - - // Compare against the unsampled estimate over the same registers. - let unsampled = HllSketchAccumulator { - inner: HllSketch::from_raw(HllVariant::Regular, 4, registers, 0.0, 0.0, 0.0), - sample_p: 1.0, - }; - let raw = unsampled - .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) - .expect("cardinality ok"); - let rescaled = acc - .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) - .expect("cardinality ok"); - assert!( - (rescaled - raw * 4.0).abs() < 1e-9, - "expected 4×raw rescale" - ); - } - - #[test] - fn test_reset_to_empty_preserves_sample_p() { - let mut acc = HllSketchAccumulator { - inner: HllSketch::from_raw(HllVariant::Regular, 4, vec![3u8; 16], 0.0, 0.0, 0.0), - sample_p: 0.25, - }; - acc.reset_to_empty(); - assert_eq!(acc.sample_p, 0.25, "window rotation must keep sample_p"); - assert_eq!(acc.inner.registers, vec![0u8; 16], "registers cleared"); - } - - #[test] - fn test_merge_prefers_sampled_factor() { - let a = HllSketchAccumulator { - inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 1, 1, 1], 0.0, 0.0, 0.0), - sample_p: 0.25, - }; - let b = HllSketchAccumulator { - inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 1, 1, 1], 0.0, 0.0, 0.0), - sample_p: 1.0, - }; - let merged = a.merge_with(&b).expect("merge ok"); - let merged = merged - .as_any() - .downcast_ref::() - .expect("downcast ok"); - assert_eq!(merged.sample_p, 0.25); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/hydra_kll.rs b/crates/asap_summary_state/src/summary_kernels/hydra_kll.rs deleted file mode 100644 index 167bcae38..000000000 --- a/crates/asap_summary_state/src/summary_kernels/hydra_kll.rs +++ /dev/null @@ -1,165 +0,0 @@ -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{HydraKllSketch, MessagePackCodec}; -use base64::{engine::general_purpose, Engine as _}; -use std::collections::HashMap; - -use crate::Statistic; - -/// HydraKLL sketch accumulator — wraps asap_sketchlib::HydraKllSketch. -/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. -/// This file retains QE-specific trait impls and JSON output. -#[derive(Debug, Clone)] -pub struct HydraKllSketchAccumulator { - pub inner: HydraKllSketch, -} - -impl HydraKllSketchAccumulator { - pub fn new(row_num: usize, col_num: usize, k: u16) -> Self { - Self { - inner: HydraKllSketch::new(row_num, col_num, k), - } - } - - pub fn update(&mut self, key: &KeyByLabelValues, value: f64) { - self.inner.update(&key.to_semicolon_str(), value); - } - - pub fn deserialize_from_bytes(_buffer: &[u8]) -> Result> { - Err("deserialize_from_bytes for HydraKllSketchAccumulator not implemented".into()) - } - - pub fn query_key(&self, key: &KeyByLabelValues, quantile: f64) -> f64 { - self.inner.quantile(&key.to_semicolon_str(), quantile) - } -} - -impl SerializableToSink for HydraKllSketchAccumulator { - fn serialize_to_json(&self) -> serde_json::Value { - // Mirror Python implementation: {"sketch": base64_encoded_string} - let sketch_bytes = self.inner.to_msgpack().unwrap_or_default(); - let sketch_b64 = general_purpose::STANDARD.encode(&sketch_bytes); - serde_json::json!({ "sketch": sketch_b64 }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.inner.to_msgpack().unwrap_or_default() - } -} - -impl MergeableAccumulator for HydraKllSketchAccumulator { - 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) - } -} - -impl AggregateCore for HydraKllSketchAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "HydraKllSketchAccumulator" - } - - 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 HydraKllSketchAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - - let hk = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to HydraKllSketchAccumulator")?; - - let merged = Self::merge_accumulators(vec![self.clone(), hk.clone()])?; - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::HydraKLL - } - - fn approx_memory_bytes(&self) -> usize { - // HydraKLL is a row*col grid of KLL sketches; typical instances - // are on the order of tens of KiB. 32 KiB is a conservative - // per-instance default. - 32 * 1024 - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::MultipleSubpopulationAggregate; - let key_val = key - .as_ref() - .ok_or("Key required for HydraKllSketchAccumulator")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for HydraKllSketchAccumulator { - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - query_kwargs: Option<&HashMap>, - ) -> Result> { - match statistic { - Statistic::Quantile => { - let quantile = query_kwargs - .and_then(|kwargs| kwargs.get("quantile")) - .ok_or("Missing quantile parameter for quantile query")? - .parse::() - .map_err(|_| "Invalid quantile parameter format")?; - - if !(0.0..=1.0).contains(&quantile) { - return Err("Quantile must be between 0.0 and 1.0".into()); - } - - Ok(self.query_key(key, quantile)) - } - _ => Err( - format!("Unsupported statistic in HydraKllSketchAccumulator: {statistic:?}").into(), - ), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/increase.rs b/crates/asap_summary_state/src/summary_kernels/increase.rs deleted file mode 100644 index c6b988325..000000000 --- a/crates/asap_summary_state/src/summary_kernels/increase.rs +++ /dev/null @@ -1,737 +0,0 @@ -use crate::codec::MeasurementCodec; -use crate::{ - AggregateCore, AggregationType, Measurement, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -const RESET_AWARE_WIRE_MAGIC: &[u8; 8] = b"ASAPINC2"; -const RESET_AWARE_WIRE_EXTENSION_LEN: usize = 8 + 8 + 8; - -/// Accumulator for tracking increases in counter metrics -/// Stores the starting and last seen measurements with timestamps -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct IncreaseAccumulator { - pub starting_measurement: Measurement, - pub starting_timestamp: i64, - pub last_seen_measurement: Measurement, - pub last_seen_timestamp: i64, - /// Sum of monotonic deltas, adding the post-reset value whenever the - /// counter decreases. This is the reset correction Prometheus applies. - #[serde(default)] - pub total_increase: f64, - #[serde(default)] - pub sample_count: u64, -} - -impl IncreaseAccumulator { - /// Merge two counter intervals without a temporary collection. Ties retain - /// the left input, matching the stable ordering of multi-pane merges. - pub(crate) fn merge_pair(left: &Self, right: &Self) -> Self { - let (first, second) = if left.starting_timestamp <= right.starting_timestamp { - (left, right) - } else { - (right, left) - }; - let mut merged = first.clone(); - if second.starting_timestamp > merged.last_seen_timestamp { - merged.total_increase += - if second.starting_measurement.value >= merged.last_seen_measurement.value { - second.starting_measurement.value - merged.last_seen_measurement.value - } else { - second.starting_measurement.value - }; - } - merged.total_increase += second.total_increase; - merged.sample_count = merged.sample_count.saturating_add(second.sample_count); - if second.last_seen_timestamp > merged.last_seen_timestamp { - merged.last_seen_measurement = second.last_seen_measurement.clone(); - merged.last_seen_timestamp = second.last_seen_timestamp; - } - - merged - } - - /// Return the number of bytes occupied by one accumulator at the start of - /// `buffer`. Old persisted values end after `last_seen_timestamp`; reset- - /// aware values carry a magic-prefixed extension. The magic makes this - /// safe when the buffer also contains the next keyed entry. - pub(crate) fn serialized_len_from_prefix( - buffer: &[u8], - ) -> Result> { - if buffer.len() < 4 { - return Err("Buffer too short for starting measurement length".into()); - } - let starting_len = u32::from_le_bytes(buffer[0..4].try_into()?) as usize; - let last_len_offset = 4usize - .checked_add(starting_len) - .and_then(|offset| offset.checked_add(8)) - .ok_or("IncreaseAccumulator length overflow")?; - if buffer.len() < last_len_offset + 4 { - return Err("Buffer too short for last seen measurement length".into()); - } - let last_len = - u32::from_le_bytes(buffer[last_len_offset..last_len_offset + 4].try_into()?) as usize; - let legacy_len = last_len_offset - .checked_add(4) - .and_then(|offset| offset.checked_add(last_len)) - .and_then(|offset| offset.checked_add(8)) - .ok_or("IncreaseAccumulator length overflow")?; - if buffer.len() < legacy_len { - return Err("Buffer too short for last seen timestamp".into()); - } - let has_extension = buffer.len() >= legacy_len + RESET_AWARE_WIRE_EXTENSION_LEN - && &buffer[legacy_len..legacy_len + RESET_AWARE_WIRE_MAGIC.len()] - == RESET_AWARE_WIRE_MAGIC; - Ok(legacy_len - + if has_extension { - RESET_AWARE_WIRE_EXTENSION_LEN - } else { - 0 - }) - } - - pub fn new( - starting_measurement: Measurement, - starting_timestamp: i64, - last_seen_measurement: Measurement, - last_seen_timestamp: i64, - ) -> Self { - let total_increase = if last_seen_timestamp <= starting_timestamp { - 0.0 - } else if last_seen_measurement.value >= starting_measurement.value { - last_seen_measurement.value - starting_measurement.value - } else { - last_seen_measurement.value - }; - let sample_count = if last_seen_timestamp > starting_timestamp { - 2 - } else { - 1 - }; - Self { - starting_measurement, - starting_timestamp, - last_seen_measurement, - last_seen_timestamp, - total_increase, - sample_count, - } - } - - pub fn update(&mut self, measurement: Measurement, timestamp: i64) { - if timestamp < self.last_seen_timestamp { - return; - } - if timestamp == self.last_seen_timestamp { - return; - } - if measurement.value >= self.last_seen_measurement.value { - self.total_increase += measurement.value - self.last_seen_measurement.value; - } else { - self.total_increase += measurement.value; - } - self.last_seen_measurement = measurement; - self.last_seen_timestamp = timestamp; - self.sample_count = self.sample_count.saturating_add(1); - } - - pub fn deserialize_from_json(data: &Value) -> Result> { - let starting_measurement = - Measurement::deserialize_from_json(&data["starting_measurement"])?; - let starting_timestamp = data["starting_timestamp"] - .as_i64() - .ok_or("Missing or invalid 'starting_timestamp' field")?; - let last_seen_measurement = - Measurement::deserialize_from_json(&data["last_seen_measurement"])?; - let last_seen_timestamp = data["last_seen_timestamp"] - .as_i64() - .ok_or("Missing or invalid 'last_seen_timestamp' field")?; - - let mut accumulator = Self::new( - starting_measurement, - starting_timestamp, - last_seen_measurement, - last_seen_timestamp, - ); - accumulator.total_increase = data["total_increase"] - .as_f64() - .unwrap_or(accumulator.total_increase); - accumulator.sample_count = data["sample_count"] - .as_u64() - .unwrap_or(accumulator.sample_count); - Ok(accumulator) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - let mut offset = 0; - - // Read starting measurement length and data - if buffer.len() < offset + 4 { - return Err("Buffer too short for starting measurement length".into()); - } - let starting_measurement_length = u32::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - ]) as usize; - offset += 4; - - if buffer.len() < offset + starting_measurement_length { - return Err("Buffer too short for starting measurement".into()); - } - let starting_measurement = Measurement::deserialize_from_bytes( - &buffer[offset..offset + starting_measurement_length], - )?; - offset += starting_measurement_length; - - // Read starting timestamp - if buffer.len() < offset + 8 { - return Err("Buffer too short for starting timestamp".into()); - } - let starting_timestamp = i64::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - buffer[offset + 4], - buffer[offset + 5], - buffer[offset + 6], - buffer[offset + 7], - ]); - offset += 8; - - // Read last seen measurement length and data - if buffer.len() < offset + 4 { - return Err("Buffer too short for last seen measurement length".into()); - } - let last_seen_measurement_length = u32::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - ]) as usize; - offset += 4; - - if buffer.len() < offset + last_seen_measurement_length { - return Err("Buffer too short for last seen measurement".into()); - } - let last_seen_measurement = Measurement::deserialize_from_bytes( - &buffer[offset..offset + last_seen_measurement_length], - )?; - offset += last_seen_measurement_length; - - // Read last seen timestamp - if buffer.len() < offset + 8 { - return Err("Buffer too short for last seen timestamp".into()); - } - let last_seen_timestamp = i64::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - buffer[offset + 4], - buffer[offset + 5], - buffer[offset + 6], - buffer[offset + 7], - ]); - - let mut accumulator = Self::new( - starting_measurement, - starting_timestamp, - last_seen_measurement, - last_seen_timestamp, - ); - offset += 8; - if buffer.len() >= offset + RESET_AWARE_WIRE_EXTENSION_LEN - && &buffer[offset..offset + RESET_AWARE_WIRE_MAGIC.len()] == RESET_AWARE_WIRE_MAGIC - { - offset += RESET_AWARE_WIRE_MAGIC.len(); - accumulator.total_increase = f64::from_le_bytes( - buffer[offset..offset + 8] - .try_into() - .expect("checked total-increase bytes"), - ); - offset += 8; - accumulator.sample_count = u64::from_le_bytes( - buffer[offset..offset + 8] - .try_into() - .expect("checked sample-count bytes"), - ); - } - Ok(accumulator) - } -} - -impl SerializableToSink for IncreaseAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ - "starting_measurement": self.starting_measurement.serialize_to_json(), - "starting_timestamp": self.starting_timestamp, - "last_seen_measurement": self.last_seen_measurement.serialize_to_json(), - "last_seen_timestamp": self.last_seen_timestamp, - "total_increase": self.total_increase, - "sample_count": self.sample_count, - }) - } - - fn serialize_to_bytes(&self) -> Vec { - let starting_measurement_bytes = self.starting_measurement.serialize_to_bytes(); - let last_seen_measurement_bytes = self.last_seen_measurement.serialize_to_bytes(); - - let mut buffer = Vec::new(); - - // Starting measurement length and data - buffer.extend_from_slice(&(starting_measurement_bytes.len() as u32).to_le_bytes()); - buffer.extend_from_slice(&starting_measurement_bytes); - - // Starting timestamp - buffer.extend_from_slice(&self.starting_timestamp.to_le_bytes()); - - // Last seen measurement length and data - buffer.extend_from_slice(&(last_seen_measurement_bytes.len() as u32).to_le_bytes()); - buffer.extend_from_slice(&last_seen_measurement_bytes); - - // Last seen timestamp - buffer.extend_from_slice(&self.last_seen_timestamp.to_le_bytes()); - buffer.extend_from_slice(RESET_AWARE_WIRE_MAGIC); - buffer.extend_from_slice(&self.total_increase.to_le_bytes()); - buffer.extend_from_slice(&self.sample_count.to_le_bytes()); - - buffer - } -} - -impl MergeableAccumulator for IncreaseAccumulator { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut accumulators = accumulators; - accumulators.sort_by_key(|accumulator| accumulator.starting_timestamp); - let mut result = accumulators[0].clone(); - - for acc in &accumulators[1..] { - result = Self::merge_pair(&result, acc); - } - - Ok(result) - } -} - -impl AggregateCore for IncreaseAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "IncreaseAccumulator" - } - - 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> { - // Check if other is also an IncreaseAccumulator - if other.get_accumulator_type() != self.get_accumulator_type() { - return Err(format!( - "Cannot merge IncreaseAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - - // Downcast to IncreaseAccumulator - let other_increase = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to IncreaseAccumulator")?; - - let merged = Self::merge_pair(self, other_increase); - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::Increase - } - - fn approx_memory_bytes(&self) -> usize { - // Two Measurements + two i64s. Measurements are a few f64 fields. - std::mem::size_of::() - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - _key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::SingleSubpopulationAggregate; - self.query( - statistic, - (!query_kwargs.is_empty()).then_some(query_kwargs), - ) - } -} - -impl SingleSubpopulationAggregate for IncreaseAccumulator { - fn query( - &self, - statistic: Statistic, - query_kwargs: Option<&HashMap>, - ) -> Result> { - match statistic { - Statistic::Increase => Ok(self.extrapolated_value(query_kwargs, false)?), - Statistic::Rate => Ok(self.extrapolated_value(query_kwargs, true)?), - // For instant `sum [by (...)] (counter_metric)` Prometheus - // sums the latest cumulative value of each matching series. - // The IncreaseAccumulator already tracks that latest value - // in `last_seen_measurement`, so per-series Sum is just - // that scalar; the engine's outer aggregation groups by the - // `by` labels and adds the per-series totals across keys. - // - // See PR #108 audit conclusion (commit 4359e10) and issue - // ProjectASAP/ASAPCollector#46: pre-fix the ASAP tier ingested - // counters as IncreaseAccumulator and bare `sum by (...) ()` - // capability-missed because this trait did not answer Sum. - Statistic::Sum => Ok(self.last_seen_measurement.value), - _ => Err(format!("Unsupported statistic in IncreaseAccumulator: {statistic:?}").into()), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl IncreaseAccumulator { - fn extrapolated_value( - &self, - query_kwargs: Option<&HashMap>, - is_rate: bool, - ) -> Result> { - if self.sample_count < 2 || self.last_seen_timestamp <= self.starting_timestamp { - return Err("at least two ordered counter samples are required".into()); - } - let sampled_interval = (self.last_seen_timestamp - self.starting_timestamp) as f64 / 1000.0; - let Some(kwargs) = query_kwargs else { - return Ok(if is_rate { - self.total_increase / sampled_interval - } else { - self.total_increase - }); - }; - let range_start = kwargs - .get("range_start_ms") - .ok_or("missing range_start_ms")? - .parse::()?; - let range_end = kwargs - .get("range_end_ms") - .ok_or("missing range_end_ms")? - .parse::()?; - if range_end <= range_start { - return Err("invalid counter evaluation range".into()); - } - - let mut duration_to_start = - (self.starting_timestamp.saturating_sub(range_start)) as f64 / 1000.0; - let duration_to_end = (range_end.saturating_sub(self.last_seen_timestamp)) as f64 / 1000.0; - let average_sample_interval = sampled_interval / (self.sample_count - 1) as f64; - let extrapolation_threshold = average_sample_interval * 1.1; - - if self.total_increase > 0.0 && self.starting_measurement.value >= 0.0 { - let duration_to_zero = - sampled_interval * (self.starting_measurement.value / self.total_increase); - duration_to_start = duration_to_start.min(duration_to_zero); - } - let mut extrapolate_to = sampled_interval; - extrapolate_to += if duration_to_start < extrapolation_threshold { - duration_to_start.max(0.0) - } else { - average_sample_interval / 2.0 - }; - extrapolate_to += if duration_to_end < extrapolation_threshold { - duration_to_end.max(0.0) - } else { - average_sample_interval / 2.0 - }; - let mut factor = extrapolate_to / sampled_interval; - if is_rate { - factor /= (range_end - range_start) as f64 / 1000.0; - } - Ok(self.total_increase * factor) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_increase_accumulator_creation() { - let starting_measurement = Measurement::new(10.0); - let last_seen_measurement = Measurement::new(25.0); - let acc = IncreaseAccumulator::new( - starting_measurement.clone(), - 1000, - last_seen_measurement.clone(), - 2000, - ); - - assert_eq!(acc.starting_measurement.value, 10.0); - assert_eq!(acc.starting_timestamp, 1000); - assert_eq!(acc.last_seen_measurement.value, 25.0); - assert_eq!(acc.last_seen_timestamp, 2000); - } - - #[test] - fn test_increase_accumulator_update() { - let starting_measurement = Measurement::new(10.0); - let mut acc = IncreaseAccumulator::new( - starting_measurement.clone(), - 1000, - starting_measurement.clone(), - 1000, - ); - - let new_measurement = Measurement::new(25.0); - acc.update(new_measurement.clone(), 2000); - - assert_eq!(acc.last_seen_measurement.value, 25.0); - assert_eq!(acc.last_seen_timestamp, 2000); - assert_eq!(acc.starting_measurement.value, 10.0); // Should remain unchanged - } - - #[test] - fn test_increase_accumulator_query() { - let starting_measurement = Measurement::new(10.0); - let last_seen_measurement = Measurement::new(25.0); - let acc = IncreaseAccumulator::new( - starting_measurement, - 1000, - last_seen_measurement, - 3000, // 2 second difference - ); - - // Test increase calculation - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, None).unwrap(), - 15.0 - ); - - // Test rate calculation (per second) - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, None).unwrap(), - 7.5 - ); // 15.0 / 2.0 - - // Statistic::Sum returns the latest cumulative counter value, - // matching Prometheus semantics for instant `sum()`. - // (Issue ProjectASAP/ASAPCollector#46, PR #108 diagnosis.) - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), - 25.0 - ); - - // Unsupported statistics still error. - assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); - } - - #[test] - fn prometheus_counter_reset_and_boundary_extrapolation() { - let mut acc = IncreaseAccumulator::new( - Measurement::new(10.0), - 10_000, - Measurement::new(10.0), - 10_000, - ); - acc.update(Measurement::new(20.0), 20_000); - acc.update(Measurement::new(3.0), 30_000); - acc.update(Measurement::new(13.0), 50_000); - assert_eq!(acc.total_increase, 23.0); - assert_eq!(acc.sample_count, 4); - - let kwargs = HashMap::from([ - ("range_start_ms".into(), "0".into()), - ("range_end_ms".into(), "60000".into()), - ]); - let increase = - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, Some(&kwargs)) - .unwrap(); - let rate = crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, Some(&kwargs)) - .unwrap(); - assert!((increase - 34.5).abs() < 1e-12); - assert!((rate - 0.575).abs() < 1e-12); - } - - #[test] - fn pane_merge_preserves_resets_and_prometheus_extrapolation() { - let mut left = IncreaseAccumulator::new( - Measurement::new(10.0), - 10_000, - Measurement::new(10.0), - 10_000, - ); - left.update(Measurement::new(20.0), 20_000); - let mut right = - IncreaseAccumulator::new(Measurement::new(3.0), 30_000, Measurement::new(3.0), 30_000); - right.update(Measurement::new(13.0), 50_000); - let merged = IncreaseAccumulator::merge_accumulators(vec![right, left]).unwrap(); - assert_eq!(merged.total_increase, 23.0); - assert_eq!(merged.sample_count, 4); - let kwargs = HashMap::from([ - ("range_start_ms".into(), "0".into()), - ("range_end_ms".into(), "60000".into()), - ]); - assert_eq!( - crate::SingleSubpopulationAggregate::query(&merged, Statistic::Increase, Some(&kwargs)) - .unwrap(), - 34.5 - ); - } - - #[test] - fn counter_sds_state_is_constant_size_per_pane() { - let mut acc = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(0.0), 0); - let initial = acc.serialize_to_bytes().len(); - for second in 1..=86_400 { - acc.update(Measurement::new(second as f64), second * 1_000); - } - assert_eq!(acc.serialize_to_bytes().len(), initial); - assert_eq!(acc.sample_count, 86_401); - assert_eq!( - acc.approx_memory_bytes(), - std::mem::size_of::() - ); - } - - #[test] - fn test_increase_accumulator_sum_is_latest_cumulative_value() { - // Instant `sum ()` semantics: the per-series summand is - // the latest cumulative counter value. Two series with latest - // values 100 and 50 (started at 10 and 5 respectively) should - // each report Sum = 100 and Sum = 50 — the engine's `sum by` - // outer aggregation does the cross-series total. - let acc_a = - IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(100.0), 2000); - let acc_b = - IncreaseAccumulator::new(Measurement::new(5.0), 1000, Measurement::new(50.0), 2000); - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc_a, Statistic::Sum, None).unwrap(), - 100.0 - ); - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc_b, Statistic::Sum, None).unwrap(), - 50.0 - ); - } - - #[test] - fn test_increase_accumulator_merge() { - let acc1 = - IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(20.0), 2000); - let acc2 = IncreaseAccumulator::new( - Measurement::new(5.0), - 500, // Earlier start - Measurement::new(15.0), - 1500, - ); - let acc3 = IncreaseAccumulator::new( - Measurement::new(20.0), - 2000, - Measurement::new(30.0), - 3000, // Later end - ); - - let merged = - >::merge_accumulators( - vec![acc1, acc2, acc3], - ) - .unwrap(); - - // Should use earliest start and latest end - assert_eq!(merged.starting_measurement.value, 5.0); - assert_eq!(merged.starting_timestamp, 500); - assert_eq!(merged.last_seen_measurement.value, 30.0); - assert_eq!(merged.last_seen_timestamp, 3000); - } - - #[test] - fn test_increase_accumulator_serialization() { - let acc = - IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(25.0), 2000); - - // Test JSON serialization - let json = acc.serialize_to_json(); - let deserialized = IncreaseAccumulator::deserialize_from_json(&json).unwrap(); - assert_eq!( - acc.starting_measurement.value, - deserialized.starting_measurement.value - ); - assert_eq!(acc.starting_timestamp, deserialized.starting_timestamp); - assert_eq!( - acc.last_seen_measurement.value, - deserialized.last_seen_measurement.value - ); - assert_eq!(acc.last_seen_timestamp, deserialized.last_seen_timestamp); - - // Test byte serialization - let bytes = acc.serialize_to_bytes(); - let deserialized_bytes = IncreaseAccumulator::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!( - acc.starting_measurement.value, - deserialized_bytes.starting_measurement.value - ); - assert_eq!( - acc.starting_timestamp, - deserialized_bytes.starting_timestamp - ); - assert_eq!( - acc.last_seen_measurement.value, - deserialized_bytes.last_seen_measurement.value - ); - assert_eq!( - acc.last_seen_timestamp, - deserialized_bytes.last_seen_timestamp - ); - assert_eq!(acc.total_increase, deserialized_bytes.total_increase); - assert_eq!(acc.sample_count, deserialized_bytes.sample_count); - - let legacy = &bytes[..bytes.len() - RESET_AWARE_WIRE_EXTENSION_LEN]; - let legacy_value = IncreaseAccumulator::deserialize_from_bytes(legacy).unwrap(); - assert_eq!(legacy_value.total_increase, 15.0); - assert_eq!(legacy_value.sample_count, 2); - } - - #[test] - fn test_trait_object() { - let acc: Box = Box::new(IncreaseAccumulator::new( - Measurement::new(10.0), - 1000, - Measurement::new(25.0), - 2000, - )); - - assert_eq!(acc.type_name(), "IncreaseAccumulator"); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/keyed_counter_state.rs b/crates/asap_summary_state/src/summary_kernels/keyed_counter_state.rs deleted file mode 100644 index d65860ad3..000000000 --- a/crates/asap_summary_state/src/summary_kernels/keyed_counter_state.rs +++ /dev/null @@ -1,530 +0,0 @@ -use crate::codec::KeyCodec; -use crate::summary_kernels::IncreaseAccumulator; -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, SingleSubpopulationAggregate, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -/// Accumulator that maintains separate increase accumulators for multiple keys -/// Allows tracking rate/increase for different label combinations -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct KeyedCounterState { - pub increases: HashMap, -} - -impl KeyedCounterState { - pub fn new() -> Self { - Self { - increases: HashMap::new(), - } - } - - pub fn update(&mut self, key: KeyByLabelValues, accumulator: IncreaseAccumulator) { - self.increases.insert(key, accumulator); - } - - pub fn deserialize_from_json(data: &Value) -> Result> { - let mut accumulator = Self::new(); - - if let Some(entries) = data["entries"].as_array() { - for entry in entries { - let key = KeyByLabelValues::deserialize_from_json(&entry["key"])?; - let increase_data = - IncreaseAccumulator::deserialize_from_json(&entry["increase_data"])?; - accumulator.increases.insert(key, increase_data); - } - } - - Ok(accumulator) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - let mut accumulator = Self::new(); - let mut offset = 0; - - // Read number of entries - if buffer.len() < 4 { - return Err("Buffer too short for entry count".into()); - } - let num_entries = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize; - offset += 4; - - for _ in 0..num_entries { - // Read key length and key - if offset + 4 > buffer.len() { - return Err("Buffer too short for key length".into()); - } - let key_length = u32::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - ]) as usize; - offset += 4; - - if offset + key_length > buffer.len() { - return Err("Buffer too short for key data".into()); - } - let key = - KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; - offset += key_length; - - // Read IncreaseAccumulator data - if offset >= buffer.len() { - return Err("Buffer too short for increase accumulator data".into()); - } - let consumed_bytes = - IncreaseAccumulator::serialized_len_from_prefix(&buffer[offset..])?; - let increase_data = IncreaseAccumulator::deserialize_from_bytes( - &buffer[offset..offset + consumed_bytes], - )?; - offset += consumed_bytes; - - accumulator.increases.insert(key, increase_data); - } - - Ok(accumulator) - } -} - -impl Default for KeyedCounterState { - fn default() -> Self { - Self::new() - } -} - -impl SerializableToSink for KeyedCounterState { - fn serialize_to_json(&self) -> Value { - let entries: Vec = self - .increases - .iter() - .map(|(key, data)| { - serde_json::json!({ - "key": key.serialize_to_json(), - "increase_data": data.serialize_to_json() - }) - }) - .collect(); - - serde_json::json!({ - "entries": entries - }) - } - - fn serialize_to_bytes(&self) -> Vec { - let mut buffer = Vec::new(); - - // Write number of entries - buffer.extend_from_slice(&(self.increases.len() as u32).to_le_bytes()); - - // Write each key-value pair - for (key, data) in &self.increases { - let key_bytes = key.serialize_to_bytes(); - buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); - buffer.extend_from_slice(&key_bytes); - - let data_bytes = data.serialize_to_bytes(); - buffer.extend_from_slice(&data_bytes); - } - - buffer - } -} - -impl AggregateCore for KeyedCounterState { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "KeyedCounterState" - } - - 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> { - // Check if other is also a KeyedCounterState - if other.get_accumulator_type() != self.get_accumulator_type() { - return Err(format!( - "Cannot merge KeyedCounterState with {}", - other.get_accumulator_type() - ) - .into()); - } - - // Downcast to KeyedCounterState - let other_multiple_increase = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to KeyedCounterState")?; - - // Clone self once, then merge each matching counter with the same - // reset-aware, boundary-aware implementation used by the unkeyed path. - let mut merged = self.clone(); - for (key, data) in &other_multiple_increase.increases { - if let Some(existing_data) = merged.increases.get_mut(key) { - *existing_data = IncreaseAccumulator::merge_accumulators(vec![ - existing_data.clone(), - data.clone(), - ])?; - } else { - merged.increases.insert(key.clone(), data.clone()); - } - } - - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::Increase - } - - fn approx_memory_bytes(&self) -> usize { - // HashMap. IncreaseAccumulator is ~64 B, - // per-entry key/overhead is ~96 B. - const BYTES_PER_ENTRY: usize = 160; - std::mem::size_of::() + self.increases.len() * BYTES_PER_ENTRY - } - - fn get_keys(&self) -> Option> { - Some(self.increases.keys().cloned().collect()) - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::MultipleSubpopulationAggregate; - let key_val = key.as_ref().ok_or("Key required for KeyedCounterState")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for KeyedCounterState { - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - query_kwargs: Option<&HashMap>, - ) -> Result> { - let data = self - .increases - .get(key) - .ok_or_else(|| format!("Key {key} not found in KeyedCounterState"))?; - - data.query(statistic, query_kwargs) - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for KeyedCounterState { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut result = KeyedCounterState::new(); - - for accumulator in accumulators { - for (key, data) in accumulator.increases { - if let Some(existing_data) = result.increases.get_mut(&key) { - *existing_data = - IncreaseAccumulator::merge_accumulators(vec![existing_data.clone(), data])?; - } else { - result.increases.insert(key, data); - } - } - } - - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::Measurement; - - fn create_test_increase_accumulator(start_val: f64, end_val: f64) -> IncreaseAccumulator { - IncreaseAccumulator::new( - Measurement::new(start_val), - 1000, - Measurement::new(end_val), - 2000, - ) - } - - fn create_test_increase_accumulator_with_time( - start_val: f64, - start_time: i64, - end_val: f64, - end_time: i64, - ) -> IncreaseAccumulator { - IncreaseAccumulator::new( - Measurement::new(start_val), - start_time, - Measurement::new(end_val), - end_time, - ) - } - - #[test] - fn test_keyed_counter_state_creation() { - let acc = KeyedCounterState::new(); - assert!(acc.increases.is_empty()); - } - - #[test] - fn test_keyed_counter_state_update() { - let mut acc = KeyedCounterState::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - let increase1 = create_test_increase_accumulator(10.0, 25.0); - let increase2 = create_test_increase_accumulator(5.0, 15.0); - - acc.update(key1.clone(), increase1); - acc.update(key2.clone(), increase2); - - assert_eq!(acc.increases.len(), 2); - assert!(acc.increases.contains_key(&key1)); - assert!(acc.increases.contains_key(&key2)); - } - - #[test] - fn test_keyed_counter_state_query() { - let mut acc = KeyedCounterState::new(); - - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let increase_acc = create_test_increase_accumulator(10.0, 25.0); - acc.update(key.clone(), increase_acc); - - // Test increase query - assert_eq!(acc.query(Statistic::Increase, &key, None).unwrap(), 15.0); - - // Test rate query (15.0 increase over 1 second = 15.0 per second) - assert_eq!(acc.query(Statistic::Rate, &key, None).unwrap(), 15.0); - - // Sum returns the latest cumulative counter value for the - // queried key (per-series Prometheus `sum()` semantics; - // see issue ProjectASAP/ASAPCollector#46 and PR #108 diagnosis). - // The series here was created with last_seen=25.0. - assert_eq!(acc.query(Statistic::Sum, &key, None).unwrap(), 25.0); - - // Unsupported statistic still errors. - assert!(acc.query(Statistic::Min, &key, None).is_err()); - - let unknown_key = KeyByLabelValues::new(); - assert!(acc.query(Statistic::Increase, &unknown_key, None).is_err()); - } - - #[test] - fn test_keyed_counter_state_sum_per_key() { - // `sum by (zone) (counter)` reaches KeyedCounterState - // only when the ASAP-tier ingest groups multiple series under - // a single accumulator (the `Multiple*` variant). In that case - // each per-key Sum should be the series' latest cumulative - // value; the engine's outer `by` aggregation does the cross-key - // grouping. (Issue ProjectASAP/ASAPCollector#46.) - let mut acc = KeyedCounterState::new(); - let east = KeyByLabelValues::new_with_labels(vec!["us-east-1".to_string()]); - let west = KeyByLabelValues::new_with_labels(vec!["us-west-2".to_string()]); - - acc.update( - east.clone(), - IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(100.0), 2000), - ); - acc.update( - west.clone(), - IncreaseAccumulator::new(Measurement::new(5.0), 1000, Measurement::new(50.0), 2000), - ); - - assert_eq!(acc.query(Statistic::Sum, &east, None).unwrap(), 100.0); - assert_eq!(acc.query(Statistic::Sum, &west, None).unwrap(), 50.0); - } - - #[test] - fn test_keyed_counter_state_merge() { - let mut acc1 = KeyedCounterState::new(); - let mut acc2 = KeyedCounterState::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - // Add different keys to each accumulator - acc1.update(key1.clone(), create_test_increase_accumulator(10.0, 20.0)); - acc2.update(key2.clone(), create_test_increase_accumulator(5.0, 15.0)); - - // Also add overlapping key with different time ranges (later timestamps) - acc2.update( - key1.clone(), - create_test_increase_accumulator_with_time(15.0, 2000, 30.0, 3000), - ); // Later time range - - let merged = KeyedCounterState::merge_accumulators(vec![acc1, acc2]).unwrap(); - - assert_eq!(merged.increases.len(), 2); - assert!(merged.increases.contains_key(&key1)); - assert!(merged.increases.contains_key(&key2)); - - // The merged key1 should have the full range (earliest start to latest end) - let merged_key1 = merged.increases.get(&key1).unwrap(); - assert_eq!(merged_key1.starting_measurement.value, 10.0); // Earlier start - assert_eq!(merged_key1.last_seen_measurement.value, 30.0); // Later end - } - - #[test] - fn test_keyed_counter_state_serialization() { - let mut acc = KeyedCounterState::new(); - - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - let second_key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - let mut reset_aware = create_test_increase_accumulator(10.0, 25.0); - reset_aware.update(Measurement::new(3.0), 3000); - acc.update(key.clone(), reset_aware); - acc.update( - second_key.clone(), - create_test_increase_accumulator(4.0, 9.0), - ); - - // Test JSON serialization - let json_value = acc.serialize_to_json(); - let deserialized = KeyedCounterState::deserialize_from_json(&json_value).unwrap(); - - assert_eq!(deserialized.increases.len(), 2); - let deserialized_acc = deserialized.increases.get(&key).unwrap(); - assert_eq!(deserialized_acc.starting_measurement.value, 10.0); - assert_eq!(deserialized_acc.last_seen_measurement.value, 3.0); - assert_eq!(deserialized_acc.total_increase, 18.0); - - // Test binary serialization - let bytes = acc.serialize_to_bytes(); - let deserialized_bytes = KeyedCounterState::deserialize_from_bytes(&bytes).unwrap(); - - assert_eq!(deserialized_bytes.increases.len(), 2); - let deserialized_acc_bytes = deserialized_bytes.increases.get(&key).unwrap(); - assert_eq!(deserialized_acc_bytes.starting_measurement.value, 10.0); - assert_eq!(deserialized_acc_bytes.last_seen_measurement.value, 3.0); - assert_eq!(deserialized_acc_bytes.total_increase, 18.0); - assert_eq!( - deserialized_bytes - .increases - .get(&second_key) - .unwrap() - .last_seen_measurement - .value, - 9.0 - ); - } - - #[test] - fn test_keyed_counter_state_get_keys() { - let mut acc = KeyedCounterState::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - acc.update(key1.clone(), create_test_increase_accumulator(10.0, 20.0)); - acc.update(key2.clone(), create_test_increase_accumulator(5.0, 15.0)); - - let keys = acc.get_keys().unwrap(); - assert_eq!(keys.len(), 2); - assert!(keys.contains(&key1)); - assert!(keys.contains(&key2)); - } - - #[test] - fn test_trait_object() { - let mut acc = KeyedCounterState::new(); - let key = KeyByLabelValues::new(); - acc.update(key.clone(), create_test_increase_accumulator(10.0, 25.0)); - - let trait_obj: Box = Box::new(acc); - assert_eq!( - trait_obj.query(Statistic::Increase, &key, None).unwrap(), - 15.0 - ); - - let keys = trait_obj.get_keys().unwrap(); - assert_eq!(keys.len(), 1); - } - - // #[test] - // fn test_keyed_counter_state_arroyo_deserialization() { - // // Create test data in Arroyo MessagePack format - // // Format: {key: [starting_value, starting_timestamp, last_seen_value, last_seen_timestamp]} - // let mut test_data = std::collections::HashMap::new(); - // test_data.insert("web;service".to_string(), vec![10.0, 1000.0, 25.0, 2000.0]); - // test_data.insert("api;service".to_string(), vec![5.0, 1500.0, 15.0, 2500.0]); - - // // Serialize to MessagePack - // let arroyo_buffer = rmp_serde::to_vec(&test_data).unwrap(); - - // // Test Arroyo deserialization - // let deserialized_acc = - // KeyedCounterState::deserialize_from_bytes_arroyo(&arroyo_buffer).unwrap(); - - // // Verify the deserialized accumulator has the correct data - // assert_eq!(deserialized_acc.increases.len(), 2); - - // // Check first key (web;service) - // let keys: Vec<_> = deserialized_acc.increases.keys().collect(); - // let key1 = keys - // .iter() - // .find(|k| k.labels.get("label_0").is_some_and(|v| v == "web")) - // .unwrap(); - - // let increase1 = deserialized_acc.increases.get(key1).unwrap(); - // assert_eq!(increase1.starting_measurement.value, 10.0); - // assert_eq!(increase1.starting_timestamp, 1000); - // assert_eq!(increase1.last_seen_measurement.value, 25.0); - // assert_eq!(increase1.last_seen_timestamp, 2000); - - // // Check second key (api;service) - // let key2 = keys - // .iter() - // .find(|k| k.labels.get("label_0").is_some_and(|v| v == "api")) - // .unwrap(); - - // let increase2 = deserialized_acc.increases.get(key2).unwrap(); - // assert_eq!(increase2.starting_measurement.value, 5.0); - // assert_eq!(increase2.starting_timestamp, 1500); - // assert_eq!(increase2.last_seen_measurement.value, 15.0); - // assert_eq!(increase2.last_seen_timestamp, 2500); - - // // Test querying - // assert_eq!( - // deserialized_acc.query(Statistic::Increase, key1).unwrap(), - // 15.0 - // ); // 25.0 - 10.0 - // assert_eq!( - // deserialized_acc.query(Statistic::Increase, key2).unwrap(), - // 10.0 - // ); // 15.0 - 5.0 - // } -} diff --git a/crates/asap_summary_state/src/summary_kernels/keyed_max_state.rs b/crates/asap_summary_state/src/summary_kernels/keyed_max_state.rs deleted file mode 100644 index 38b9094b4..000000000 --- a/crates/asap_summary_state/src/summary_kernels/keyed_max_state.rs +++ /dev/null @@ -1,336 +0,0 @@ -use crate::codec::KeyCodec; -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -/// Exact per-key maximum over many populations, mergeable by comparison. -/// -/// The minimum direction is -/// [`KeyedMinState`](super::keyed_min_state::KeyedMinState), -/// a separate type: these used to be one `MultipleMinMaxAccumulator` whose -/// direction lived in a `sub_type` string that every layer above had to carry -/// alongside the family. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct KeyedMaxState { - pub values: HashMap, -} - -impl KeyedMaxState { - pub fn new() -> Self { - Self::default() - } - - pub fn new_with_values(values: HashMap) -> Self { - Self { values } - } - - pub fn update(&mut self, key: KeyByLabelValues, value: f64) { - let current = self.values.entry(key).or_insert(f64::NEG_INFINITY); - if value > *current { - *current = value; - } - } - - pub fn add_value(&mut self, key: KeyByLabelValues, value: f64) { - self.values.insert(key, value); - } - - pub fn deserialize_from_json(data: &Value) -> Result> { - let values_data = data["values"] - .as_object() - .ok_or("Missing or invalid 'values' field")?; - - let mut values = HashMap::new(); - for (key_str, value) in values_data { - let key_json: Value = serde_json::from_str(key_str)?; - let key = KeyByLabelValues::deserialize_from_json(&key_json)?; - let val = value.as_f64().ok_or("Invalid value")?; - values.insert(key, val); - } - - Ok(Self { values }) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - let mut offset = 0; - - // Read number of entries - if buffer.len() < 4 { - return Err("Buffer too short for entry count".into()); - } - let num_entries = u32::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - ]) as usize; - offset += 4; - - let mut values = HashMap::new(); - - for _ in 0..num_entries { - // Read key length and data - if buffer.len() < offset + 4 { - return Err("Buffer too short for key length".into()); - } - let key_length = u32::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - ]) as usize; - offset += 4; - - if buffer.len() < offset + key_length { - return Err("Buffer too short for key data".into()); - } - let key = - KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; - offset += key_length; - - // Read value - if buffer.len() < offset + 8 { - return Err("Buffer too short for value".into()); - } - let value = f64::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - buffer[offset + 4], - buffer[offset + 5], - buffer[offset + 6], - buffer[offset + 7], - ]); - offset += 8; - - values.insert(key, value); - } - - Ok(Self { values }) - } -} - -impl SerializableToSink for KeyedMaxState { - fn serialize_to_json(&self) -> Value { - let mut values_obj = serde_json::Map::new(); - for (key, value) in &self.values { - let key_json = key.serialize_to_json(); - let key_str = serde_json::to_string(&key_json).unwrap(); - values_obj.insert( - key_str, - Value::Number(serde_json::Number::from_f64(*value).unwrap()), - ); - } - - serde_json::json!({ "values": values_obj }) - } - - fn serialize_to_bytes(&self) -> Vec { - let mut buffer = Vec::new(); - - // Write number of entries - buffer.extend_from_slice(&(self.values.len() as u32).to_le_bytes()); - - // Write each key-value pair - for (key, value) in &self.values { - let key_bytes = key.serialize_to_bytes(); - - // Write key length and data - buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); - buffer.extend_from_slice(&key_bytes); - - // Write value - buffer.extend_from_slice(&value.to_le_bytes()); - } - - buffer - } -} - -impl AggregateCore for KeyedMaxState { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "KeyedMaxState" - } - - 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 KeyedMaxState with {}", - other.get_accumulator_type() - ) - .into()); - } - - let other_multiple = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to KeyedMaxState")?; - - let merged = Self::merge_accumulators(vec![self.clone(), other_multiple.clone()])?; - - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::Max - } - - fn approx_memory_bytes(&self) -> usize { - const BYTES_PER_ENTRY: usize = 96; - std::mem::size_of::() + self.values.len() * BYTES_PER_ENTRY - } - - fn get_keys(&self) -> Option> { - Some(self.values.keys().cloned().collect()) - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::MultipleSubpopulationAggregate; - let key_val = key.as_ref().ok_or("Key required for KeyedMaxState")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for KeyedMaxState { - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, - ) -> Result> { - match statistic { - Statistic::Max => self - .values - .get(key) - .copied() - .ok_or_else(|| format!("Key {key} not found in KeyedMaxState").into()), - other => Err(format!("Unsupported statistic in KeyedMaxState: {other:?}").into()), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for KeyedMaxState { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut result = KeyedMaxState::new(); - - for acc in accumulators { - for (key, value) in acc.values { - result.update(key, value); - } - } - - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn key(value: &str) -> KeyByLabelValues { - KeyByLabelValues::new_with_labels(vec![value.to_string()]) - } - - #[test] - fn keeps_the_largest_per_key() { - let mut acc = KeyedMaxState::new(); - acc.update(key("a"), 10.0); - acc.update(key("a"), 5.0); - acc.update(key("a"), 15.0); - acc.update(key("b"), 7.0); - - assert_eq!(acc.query(Statistic::Max, &key("a"), None).unwrap(), 15.0); - assert_eq!(acc.query(Statistic::Max, &key("b"), None).unwrap(), 7.0); - } - - #[test] - fn refuses_the_opposite_statistic_and_unknown_keys() { - let mut acc = KeyedMaxState::new(); - acc.update(key("a"), 1.0); - assert!(acc.query(Statistic::Min, &key("a"), None).is_err()); - assert!(acc.query(Statistic::Max, &key("missing"), None).is_err()); - } - - #[test] - fn merges_per_key() { - let mut left = KeyedMaxState::new(); - left.update(key("a"), 10.0); - let mut right = KeyedMaxState::new(); - right.update(key("a"), 5.0); - right.update(key("b"), 3.0); - - let merged = - >::merge_accumulators(vec![ - left, right, - ]) - .unwrap(); - - assert_eq!(merged.query(Statistic::Max, &key("a"), None).unwrap(), 10.0); - assert_eq!(merged.query(Statistic::Max, &key("b"), None).unwrap(), 3.0); - } - - #[test] - fn refuses_to_merge_with_the_opposite_direction() { - use super::super::keyed_min_state::KeyedMinState; - let mine = KeyedMaxState::new(); - let theirs = KeyedMinState::new(); - assert!(mine.merge_with(&theirs).is_err()); - } - - #[test] - fn round_trips_through_both_serializations() { - let mut acc = KeyedMaxState::new(); - acc.update(key("a"), 4.0); - - let json = acc.serialize_to_json(); - let from_json = KeyedMaxState::deserialize_from_json(&json).unwrap(); - assert_eq!( - from_json.query(Statistic::Max, &key("a"), None).unwrap(), - 4.0 - ); - - let bytes = acc.serialize_to_bytes(); - let from_bytes = KeyedMaxState::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!( - from_bytes.query(Statistic::Max, &key("a"), None).unwrap(), - 4.0 - ); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/keyed_min_state.rs b/crates/asap_summary_state/src/summary_kernels/keyed_min_state.rs deleted file mode 100644 index b8fc2ddfb..000000000 --- a/crates/asap_summary_state/src/summary_kernels/keyed_min_state.rs +++ /dev/null @@ -1,336 +0,0 @@ -use crate::codec::KeyCodec; -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -/// Exact per-key minimum over many populations, mergeable by comparison. -/// -/// The maximum direction is -/// [`KeyedMaxState`](super::keyed_max_state::KeyedMaxState), -/// a separate type: these used to be one `MultipleMinMaxAccumulator` whose -/// direction lived in a `sub_type` string that every layer above had to carry -/// alongside the family. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct KeyedMinState { - pub values: HashMap, -} - -impl KeyedMinState { - pub fn new() -> Self { - Self::default() - } - - pub fn new_with_values(values: HashMap) -> Self { - Self { values } - } - - pub fn update(&mut self, key: KeyByLabelValues, value: f64) { - let current = self.values.entry(key).or_insert(f64::INFINITY); - if value < *current { - *current = value; - } - } - - pub fn add_value(&mut self, key: KeyByLabelValues, value: f64) { - self.values.insert(key, value); - } - - pub fn deserialize_from_json(data: &Value) -> Result> { - let values_data = data["values"] - .as_object() - .ok_or("Missing or invalid 'values' field")?; - - let mut values = HashMap::new(); - for (key_str, value) in values_data { - let key_json: Value = serde_json::from_str(key_str)?; - let key = KeyByLabelValues::deserialize_from_json(&key_json)?; - let val = value.as_f64().ok_or("Invalid value")?; - values.insert(key, val); - } - - Ok(Self { values }) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - let mut offset = 0; - - // Read number of entries - if buffer.len() < 4 { - return Err("Buffer too short for entry count".into()); - } - let num_entries = u32::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - ]) as usize; - offset += 4; - - let mut values = HashMap::new(); - - for _ in 0..num_entries { - // Read key length and data - if buffer.len() < offset + 4 { - return Err("Buffer too short for key length".into()); - } - let key_length = u32::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - ]) as usize; - offset += 4; - - if buffer.len() < offset + key_length { - return Err("Buffer too short for key data".into()); - } - let key = - KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; - offset += key_length; - - // Read value - if buffer.len() < offset + 8 { - return Err("Buffer too short for value".into()); - } - let value = f64::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - buffer[offset + 4], - buffer[offset + 5], - buffer[offset + 6], - buffer[offset + 7], - ]); - offset += 8; - - values.insert(key, value); - } - - Ok(Self { values }) - } -} - -impl SerializableToSink for KeyedMinState { - fn serialize_to_json(&self) -> Value { - let mut values_obj = serde_json::Map::new(); - for (key, value) in &self.values { - let key_json = key.serialize_to_json(); - let key_str = serde_json::to_string(&key_json).unwrap(); - values_obj.insert( - key_str, - Value::Number(serde_json::Number::from_f64(*value).unwrap()), - ); - } - - serde_json::json!({ "values": values_obj }) - } - - fn serialize_to_bytes(&self) -> Vec { - let mut buffer = Vec::new(); - - // Write number of entries - buffer.extend_from_slice(&(self.values.len() as u32).to_le_bytes()); - - // Write each key-value pair - for (key, value) in &self.values { - let key_bytes = key.serialize_to_bytes(); - - // Write key length and data - buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); - buffer.extend_from_slice(&key_bytes); - - // Write value - buffer.extend_from_slice(&value.to_le_bytes()); - } - - buffer - } -} - -impl AggregateCore for KeyedMinState { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "KeyedMinState" - } - - 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 KeyedMinState with {}", - other.get_accumulator_type() - ) - .into()); - } - - let other_multiple = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to KeyedMinState")?; - - let merged = Self::merge_accumulators(vec![self.clone(), other_multiple.clone()])?; - - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::Min - } - - fn approx_memory_bytes(&self) -> usize { - const BYTES_PER_ENTRY: usize = 96; - std::mem::size_of::() + self.values.len() * BYTES_PER_ENTRY - } - - fn get_keys(&self) -> Option> { - Some(self.values.keys().cloned().collect()) - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::MultipleSubpopulationAggregate; - let key_val = key.as_ref().ok_or("Key required for KeyedMinState")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for KeyedMinState { - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, - ) -> Result> { - match statistic { - Statistic::Min => self - .values - .get(key) - .copied() - .ok_or_else(|| format!("Key {key} not found in KeyedMinState").into()), - other => Err(format!("Unsupported statistic in KeyedMinState: {other:?}").into()), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for KeyedMinState { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut result = KeyedMinState::new(); - - for acc in accumulators { - for (key, value) in acc.values { - result.update(key, value); - } - } - - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn key(value: &str) -> KeyByLabelValues { - KeyByLabelValues::new_with_labels(vec![value.to_string()]) - } - - #[test] - fn keeps_the_smallest_per_key() { - let mut acc = KeyedMinState::new(); - acc.update(key("a"), 10.0); - acc.update(key("a"), 5.0); - acc.update(key("a"), 15.0); - acc.update(key("b"), 7.0); - - assert_eq!(acc.query(Statistic::Min, &key("a"), None).unwrap(), 5.0); - assert_eq!(acc.query(Statistic::Min, &key("b"), None).unwrap(), 7.0); - } - - #[test] - fn refuses_the_opposite_statistic_and_unknown_keys() { - let mut acc = KeyedMinState::new(); - acc.update(key("a"), 1.0); - assert!(acc.query(Statistic::Max, &key("a"), None).is_err()); - assert!(acc.query(Statistic::Min, &key("missing"), None).is_err()); - } - - #[test] - fn merges_per_key() { - let mut left = KeyedMinState::new(); - left.update(key("a"), 10.0); - let mut right = KeyedMinState::new(); - right.update(key("a"), 5.0); - right.update(key("b"), 3.0); - - let merged = - >::merge_accumulators(vec![ - left, right, - ]) - .unwrap(); - - assert_eq!(merged.query(Statistic::Min, &key("a"), None).unwrap(), 5.0); - assert_eq!(merged.query(Statistic::Min, &key("b"), None).unwrap(), 3.0); - } - - #[test] - fn refuses_to_merge_with_the_opposite_direction() { - use super::super::keyed_max_state::KeyedMaxState; - let mine = KeyedMinState::new(); - let theirs = KeyedMaxState::new(); - assert!(mine.merge_with(&theirs).is_err()); - } - - #[test] - fn round_trips_through_both_serializations() { - let mut acc = KeyedMinState::new(); - acc.update(key("a"), 4.0); - - let json = acc.serialize_to_json(); - let from_json = KeyedMinState::deserialize_from_json(&json).unwrap(); - assert_eq!( - from_json.query(Statistic::Min, &key("a"), None).unwrap(), - 4.0 - ); - - let bytes = acc.serialize_to_bytes(); - let from_bytes = KeyedMinState::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!( - from_bytes.query(Statistic::Min, &key("a"), None).unwrap(), - 4.0 - ); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/keyed_sum_count.rs b/crates/asap_summary_state/src/summary_kernels/keyed_sum_count.rs deleted file mode 100644 index 7ff044d7a..000000000 --- a/crates/asap_summary_state/src/summary_kernels/keyed_sum_count.rs +++ /dev/null @@ -1,559 +0,0 @@ -use crate::codec::KeyCodec; -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; -use planner_types::post_asap::ExactKind; - -fn sum_family() -> ExactKind { - ExactKind::Sum -} - -/// Accumulator that maintains separate sum values for multiple keys -/// Allows querying sums for specific label combinations -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct KeyedSumCountAccumulator { - #[serde(default = "sum_family")] - pub family: ExactKind, - pub sums: HashMap, - #[serde(default)] - pub counts: HashMap, -} - -impl KeyedSumCountAccumulator { - pub fn new() -> Self { - Self::for_family(ExactKind::Sum) - } - - pub fn for_family(family: ExactKind) -> Self { - assert!(matches!(family, ExactKind::Sum | ExactKind::Count)); - Self { - family, - sums: HashMap::new(), - counts: HashMap::new(), - } - } - - pub fn update(&mut self, key: KeyByLabelValues, value: f64) { - let is_new = !self.sums.contains_key(&key); - *self.sums.entry(key.clone()).or_insert(0.0) += value; - if let Some(count) = self.counts.get(&key).copied() { - if let Some(next) = count.checked_add(1).filter(|next| *next != u64::MAX) { - self.counts.insert(key, next); - } else { - self.counts.remove(&key); - } - } else if is_new { - self.counts.insert(key, 1); - } - } - - pub fn add_sum(&mut self, key: KeyByLabelValues, sum: f64) { - self.counts.remove(&key); - self.sums.insert(key, sum); - } - - pub fn deserialize_from_json(data: &Value) -> Result> { - let sums_data = data["sums"] - .as_object() - .ok_or("Missing or invalid 'sums' field")?; - - let mut sums = HashMap::new(); - for (key_str, value) in sums_data { - let key_json: Value = serde_json::from_str(key_str)?; - let key = KeyByLabelValues::deserialize_from_json(&key_json)?; - let sum = value.as_f64().ok_or("Invalid sum value")?; - sums.insert(key, sum); - } - - let mut counts = HashMap::new(); - if let Some(counts_data) = data.get("counts").and_then(Value::as_object) { - for (key_str, value) in counts_data { - let key_json: Value = serde_json::from_str(key_str)?; - let key = KeyByLabelValues::deserialize_from_json(&key_json)?; - let count = value.as_u64().ok_or("Invalid count value")?; - if !sums.contains_key(&key) { - return Err("Count key missing from sums".into()); - } - counts.insert(key, count); - } - } - let family = match data.get("family").and_then(Value::as_str) { - None | Some("Sum") => ExactKind::Sum, - Some("Count") => ExactKind::Count, - _ => return Err("Invalid keyed additive family".into()), - }; - Ok(Self { - family, - sums, - counts, - }) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - let mut offset = 0; - - // Read number of entries - if buffer.len() < 4 { - return Err("Buffer too short for entry count".into()); - } - let num_entries = u32::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - ]) as usize; - offset += 4; - - let mut sums = HashMap::new(); - let mut keys = Vec::new(); - - for _ in 0..num_entries { - // Read key length and data - if buffer.len() < offset + 4 { - return Err("Buffer too short for key length".into()); - } - let key_length = u32::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - ]) as usize; - offset += 4; - - if buffer.len() < offset + key_length { - return Err("Buffer too short for key data".into()); - } - let key = - KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; - offset += key_length; - - // Read sum value - if buffer.len() < offset + 8 { - return Err("Buffer too short for sum value".into()); - } - let sum = f64::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - buffer[offset + 4], - buffer[offset + 5], - buffer[offset + 6], - buffer[offset + 7], - ]); - offset += 8; - - keys.push(key.clone()); - sums.insert(key, sum); - } - let remaining = buffer.len() - offset; - let count_bytes = num_entries - .checked_mul(8) - .ok_or("Count section too large")?; - if remaining != 0 && remaining != count_bytes && remaining != count_bytes + 1 { - return Err("Invalid count section length".into()); - } - let mut counts = HashMap::new(); - if count_bytes != 0 && remaining >= count_bytes { - for key in keys { - let count = u64::from_le_bytes(buffer[offset..offset + 8].try_into()?); - offset += 8; - if count != u64::MAX { - counts.insert(key, count); - } - } - } - let family = if remaining == count_bytes + 1 { - match buffer[offset] { - 0 => ExactKind::Sum, - 1 => ExactKind::Count, - _ => return Err("Invalid keyed additive family tag".into()), - } - } else { - ExactKind::Sum - }; - Ok(Self { - family, - sums, - counts, - }) - } -} - -impl Default for KeyedSumCountAccumulator { - fn default() -> Self { - Self::new() - } -} - -impl SerializableToSink for KeyedSumCountAccumulator { - fn serialize_to_json(&self) -> Value { - let mut sums_obj = serde_json::Map::new(); - for (key, sum) in &self.sums { - let key_json = key.serialize_to_json(); - let key_str = serde_json::to_string(&key_json).unwrap(); - sums_obj.insert( - key_str, - Value::Number(serde_json::Number::from_f64(*sum).unwrap()), - ); - } - - let mut counts_obj = serde_json::Map::new(); - for (key, count) in &self.counts { - let key_str = serde_json::to_string(&key.serialize_to_json()).unwrap(); - counts_obj.insert(key_str, Value::from(*count)); - } - - serde_json::json!({ - "family": if self.family == ExactKind::Count { "Count" } else { "Sum" }, - "sums": sums_obj, - "counts": counts_obj - }) - } - - fn serialize_to_bytes(&self) -> Vec { - let mut buffer = Vec::new(); - - // Write number of entries - buffer.extend_from_slice(&(self.sums.len() as u32).to_le_bytes()); - - // Write each key-value pair - let mut ordered_keys = Vec::with_capacity(self.sums.len()); - for (key, sum) in &self.sums { - ordered_keys.push(key); - let key_bytes = key.serialize_to_bytes(); - - // Write key length and data - buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); - buffer.extend_from_slice(&key_bytes); - - // Write sum value - buffer.extend_from_slice(&sum.to_le_bytes()); - } - - for key in ordered_keys { - buffer.extend_from_slice( - &self - .counts - .get(key) - .copied() - .unwrap_or(u64::MAX) - .to_le_bytes(), - ); - } - - buffer.push(if self.family == ExactKind::Count { - 1 - } else { - 0 - }); - - buffer - } -} - -impl AggregateCore for KeyedSumCountAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "KeyedSumCountAccumulator" - } - - 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> { - // Check if other is also a KeyedSumCountAccumulator - if other.get_accumulator_type() != self.get_accumulator_type() { - return Err(format!( - "Cannot merge KeyedSumCountAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - - // Downcast to KeyedSumCountAccumulator - let other_multiple_sum = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to KeyedSumCountAccumulator")?; - - // Use the existing merge_accumulators method - let merged = Self::merge_accumulators(vec![self.clone(), other_multiple_sum.clone()])?; - - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - if self.family == ExactKind::Count { - AggregationType::Count - } else { - AggregationType::Sum - } - } - - fn approx_memory_bytes(&self) -> usize { - // HashMap. Label strings dominate; use a - // conservative per-entry estimate plus HashMap overhead. - const BYTES_PER_ENTRY: usize = 112; - std::mem::size_of::() + self.sums.len() * BYTES_PER_ENTRY - } - - fn get_keys(&self) -> Option> { - Some(self.sums.keys().cloned().collect()) - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::MultipleSubpopulationAggregate; - let key_val = key - .as_ref() - .ok_or("Key required for KeyedSumCountAccumulator")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for KeyedSumCountAccumulator { - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, - ) -> Result> { - match (&self.family, statistic) { - (ExactKind::Sum, Statistic::Sum) => self.sums.get(key).copied().ok_or_else(|| { - "Key not found in KeyedSumCountAccumulator" - .to_string() - .into() - }), - (ExactKind::Count, Statistic::Count) => self - .counts - .get(key) - .map(|count| *count as f64) - .ok_or_else(|| { - "Sample count unavailable in KeyedSumCountAccumulator" - .to_string() - .into() - }), - _ => Err( - format!("Unsupported statistic in KeyedSumCountAccumulator: {statistic:?}").into(), - ), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for KeyedSumCountAccumulator { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let family = accumulators[0].family.clone(); - if accumulators.iter().any(|acc| acc.family != family) { - return Err("Cannot merge different keyed additive families".into()); - } - let mut result = KeyedSumCountAccumulator::for_family(family); - - for acc in accumulators { - for key in acc.sums.keys() { - match ( - result.counts.get(key).copied(), - acc.counts.get(key).copied(), - ) { - (None, Some(count)) if !result.sums.contains_key(key) => { - result.counts.insert(key.clone(), count); - } - (Some(existing), Some(count)) => { - if let Some(total) = existing.checked_add(count) { - result.counts.insert(key.clone(), total); - } else { - result.counts.remove(key); - } - } - _ => { - result.counts.remove(key); - } - } - } - for (key, sum) in acc.sums { - *result.sums.entry(key).or_insert(0.0) += sum; - } - } - - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use std::vec; - - use super::*; - - #[test] - fn test_keyed_sum_count_accumulator_creation() { - let acc = KeyedSumCountAccumulator::new(); - assert!(acc.sums.is_empty()); - } - - #[test] - fn test_keyed_sum_count_accumulator_update() { - let mut acc = KeyedSumCountAccumulator::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - acc.update(key1.clone(), 10.0); - acc.update(key2.clone(), 20.0); - acc.update(key1.clone(), 5.0); // Should add to existing - - assert_eq!(acc.sums.get(&key1), Some(&15.0)); - assert_eq!(acc.sums.get(&key2), Some(&20.0)); - } - - #[test] - fn grouped_count_reads_sample_count_and_survives_merge_and_round_trip() { - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - let mut first = KeyedSumCountAccumulator::for_family(ExactKind::Count); - first.update(key.clone(), 10.0); - first.update(key.clone(), 20.0); - let mut second = KeyedSumCountAccumulator::for_family(ExactKind::Count); - second.update(key.clone(), 7.0); - let merged = KeyedSumCountAccumulator::merge_accumulators(vec![first, second]).unwrap(); - for acc in [ - merged.clone(), - KeyedSumCountAccumulator::deserialize_from_json(&merged.serialize_to_json()).unwrap(), - KeyedSumCountAccumulator::deserialize_from_bytes(&merged.serialize_to_bytes()).unwrap(), - ] { - assert_eq!(acc.family, ExactKind::Count); - assert!(acc.query(Statistic::Sum, &key, None).is_err()); - assert_eq!(acc.query(Statistic::Count, &key, None).unwrap(), 3.0); - } - } - - #[test] - fn keyed_additive_merge_rejects_different_planner_families() { - assert!(KeyedSumCountAccumulator::merge_accumulators(vec![ - KeyedSumCountAccumulator::for_family(ExactKind::Sum), - KeyedSumCountAccumulator::for_family(ExactKind::Count), - ]) - .is_err()); - } - - #[test] - fn test_keyed_sum_count_accumulator_query() { - let mut acc = KeyedSumCountAccumulator::new(); - - let key = KeyByLabelValues::new_with_labels(vec!["service".to_string()]); - - acc.add_sum(key.clone(), 42.0); - - // Test total queries (querying with the specific key) - assert_eq!( - crate::MultipleSubpopulationAggregate::query(&acc, Statistic::Sum, &key, None).unwrap(), - 42.0 - ); - - // Test error cases - assert!( - crate::MultipleSubpopulationAggregate::query(&acc, Statistic::Min, &key, None).is_err() - ); - } - - #[test] - fn test_keyed_sum_count_accumulator_get_keys() { - let mut acc = KeyedSumCountAccumulator::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - acc.add_sum(key1.clone(), 10.0); - acc.add_sum(key2.clone(), 20.0); - - let keys = crate::AggregateCore::get_keys(&acc).unwrap(); - assert_eq!(keys.len(), 2); - assert!(keys.contains(&key1)); - assert!(keys.contains(&key2)); - } - - #[test] - fn test_keyed_sum_count_accumulator_merge() { - let mut acc1 = KeyedSumCountAccumulator::new(); - let mut acc2 = KeyedSumCountAccumulator::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - acc1.add_sum(key1.clone(), 10.0); - acc1.add_sum(key2.clone(), 20.0); - - acc2.add_sum(key1.clone(), 5.0); // Same key, different accumulator - - let merged = >::merge_accumulators(vec![acc1, acc2]).unwrap(); - - assert_eq!(merged.sums.get(&key1), Some(&15.0)); // Should be merged - assert_eq!(merged.sums.get(&key2), Some(&20.0)); // Should be preserved - } - - #[test] - fn test_keyed_sum_count_accumulator_serialization() { - let mut acc = KeyedSumCountAccumulator::new(); - - let key = KeyByLabelValues::new_with_labels(vec!["service".to_string()]); - - acc.add_sum(key.clone(), 42.5); - - // Test JSON serialization - let json = acc.serialize_to_json(); - let deserialized = KeyedSumCountAccumulator::deserialize_from_json(&json).unwrap(); - assert_eq!(deserialized.sums.get(&key), Some(&42.5)); - - // Test byte serialization - let bytes = acc.serialize_to_bytes(); - let deserialized_bytes = KeyedSumCountAccumulator::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!(deserialized_bytes.sums.get(&key), Some(&42.5)); - } - - #[test] - fn test_trait_object() { - let mut acc = KeyedSumCountAccumulator::new(); - - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - acc.add_sum(key.clone(), 42.0); - - let trait_obj: Box = Box::new(acc); - - // Test type name through trait object - assert_eq!(trait_obj.type_name(), "KeyedSumCountAccumulator"); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/max.rs b/crates/asap_summary_state/src/summary_kernels/max.rs deleted file mode 100644 index 0c9547081..000000000 --- a/crates/asap_summary_state/src/summary_kernels/max.rs +++ /dev/null @@ -1,248 +0,0 @@ -use crate::{ - AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -/// Exact maximum over one population, mergeable by comparison. -/// -/// See [`MinAccumulator`](super::min::MinAccumulator) for why the -/// two directions are separate types rather than one accumulator carrying a -/// `sub_type` string. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct MaxAccumulator { - pub value: f64, -} - -impl Default for MaxAccumulator { - fn default() -> Self { - Self::new() - } -} - -impl MaxAccumulator { - pub fn new() -> Self { - Self { - value: f64::NEG_INFINITY, - } - } - - pub fn with_value(value: f64) -> Self { - Self { value } - } - - pub fn update(&mut self, value: f64) { - if value > self.value { - self.value = value; - } - } - - pub fn deserialize_from_json(data: &Value) -> Result> { - let value = data["value"] - .as_f64() - .ok_or("Missing or invalid 'value' field")?; - Ok(Self::with_value(value)) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - if buffer.len() < 8 { - return Err("Buffer too short".into()); - } - let value = f64::from_le_bytes([ - buffer[0], buffer[1], buffer[2], buffer[3], buffer[4], buffer[5], buffer[6], buffer[7], - ]); - Ok(Self::with_value(value)) - } -} - -impl SerializableToSink for MaxAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ "value": self.value }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.value.to_le_bytes().to_vec() - } -} - -impl MergeableAccumulator for MaxAccumulator { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - let mut result = MaxAccumulator::new(); - for acc in accumulators { - result.update(acc.value); - } - Ok(result) - } -} - -impl AggregateCore for MaxAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "MaxAccumulator" - } - - 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 MaxAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - let other_max = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to MaxAccumulator")?; - let mut merged = self.clone(); - merged.update(other_max.value); - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::Max - } - - fn approx_memory_bytes(&self) -> usize { - std::mem::size_of::() - } - - fn aux_stats(&self) -> AuxStats { - // The sentinel `f64::NEG_INFINITY` from `new()` is surfaced as-is; the - // query engine already treats it as "no data yet", the same way it - // does for `query_statistic`. - AuxStats { - max: Some(self.value), - ..AuxStats::empty() - } - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - _key: &Option, - _query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::SingleSubpopulationAggregate; - self.query(statistic, None) - } -} - -impl SingleSubpopulationAggregate for MaxAccumulator { - fn query( - &self, - statistic: Statistic, - query_kwargs: Option<&HashMap>, - ) -> Result> { - if query_kwargs.is_some() { - return Err("MaxAccumulator does not support query parameters".into()); - } - match statistic { - Statistic::Max => Ok(self.value), - other => Err(format!("Unsupported statistic in MaxAccumulator: {other:?}").into()), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn keeps_the_largest_update() { - let mut acc = MaxAccumulator::new(); - acc.update(10.0); - acc.update(5.0); - acc.update(15.0); - - assert_eq!(acc.value, 15.0); - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Max, None).unwrap(), - 15.0 - ); - } - - #[test] - fn refuses_to_answer_a_minimum_query() { - let acc = MaxAccumulator::with_value(15.0); - assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); - } - - #[test] - fn merges_by_taking_the_largest() { - let merged = - >::merge_accumulators(vec![ - MaxAccumulator::with_value(10.0), - MaxAccumulator::with_value(5.0), - MaxAccumulator::with_value(15.0), - ]) - .unwrap(); - assert_eq!(merged.value, 15.0); - } - - #[test] - fn refuses_to_merge_with_a_minimum() { - use super::super::min::MinAccumulator; - let max = MaxAccumulator::with_value(15.0); - let min = MinAccumulator::with_value(5.0); - assert!(max.merge_with(&min).is_err()); - } - - #[test] - fn round_trips_through_both_serializations() { - let acc = MaxAccumulator::with_value(42.5); - - let json = acc.serialize_to_json(); - assert_eq!( - MaxAccumulator::deserialize_from_json(&json).unwrap().value, - 42.5 - ); - - let bytes = acc.serialize_to_bytes(); - assert_eq!( - MaxAccumulator::deserialize_from_bytes(&bytes) - .unwrap() - .value, - 42.5 - ); - } - - #[test] - fn aux_stats_expose_max_only() { - let aux = MaxAccumulator::with_value(99.0).aux_stats(); - assert_eq!(aux.max, Some(99.0)); - assert_eq!(aux.min, None); - assert_eq!(aux.try_answer(Statistic::Max), Some(99.0)); - assert_eq!(aux.try_answer(Statistic::Min), None); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/min.rs b/crates/asap_summary_state/src/summary_kernels/min.rs deleted file mode 100644 index ff2ad548a..000000000 --- a/crates/asap_summary_state/src/summary_kernels/min.rs +++ /dev/null @@ -1,253 +0,0 @@ -use crate::{ - AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -/// Exact minimum over one population, mergeable by comparison. -/// -/// The sibling [`MaxAccumulator`](super::max::MaxAccumulator) is a -/// separate type on purpose: these two used to be one `MinMaxAccumulator` -/// whose direction lived in a `sub_type: String`, which meant every layer -/// above -- the wire `aggregationSubType`, the accumulator factory, the -/// summary catalog -- had to carry the direction alongside the family and -/// could silently answer a `min_over_time` read from maximum state. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct MinAccumulator { - pub value: f64, -} - -impl Default for MinAccumulator { - fn default() -> Self { - Self::new() - } -} - -impl MinAccumulator { - pub fn new() -> Self { - Self { - value: f64::INFINITY, - } - } - - pub fn with_value(value: f64) -> Self { - Self { value } - } - - pub fn update(&mut self, value: f64) { - if value < self.value { - self.value = value; - } - } - - pub fn deserialize_from_json(data: &Value) -> Result> { - let value = data["value"] - .as_f64() - .ok_or("Missing or invalid 'value' field")?; - Ok(Self::with_value(value)) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - if buffer.len() < 8 { - return Err("Buffer too short".into()); - } - let value = f64::from_le_bytes([ - buffer[0], buffer[1], buffer[2], buffer[3], buffer[4], buffer[5], buffer[6], buffer[7], - ]); - Ok(Self::with_value(value)) - } -} - -impl SerializableToSink for MinAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ "value": self.value }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.value.to_le_bytes().to_vec() - } -} - -impl MergeableAccumulator for MinAccumulator { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - let mut result = MinAccumulator::new(); - for acc in accumulators { - result.update(acc.value); - } - Ok(result) - } -} - -impl AggregateCore for MinAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "MinAccumulator" - } - - 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 MinAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - let other_min = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to MinAccumulator")?; - let mut merged = self.clone(); - merged.update(other_min.value); - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::Min - } - - fn approx_memory_bytes(&self) -> usize { - std::mem::size_of::() - } - - fn aux_stats(&self) -> AuxStats { - // The sentinel `f64::INFINITY` from `new()` is surfaced as-is; the - // query engine already treats it as "no data yet", the same way it - // does for `query_statistic`. - AuxStats { - min: Some(self.value), - ..AuxStats::empty() - } - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - _key: &Option, - _query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::SingleSubpopulationAggregate; - self.query(statistic, None) - } -} - -impl SingleSubpopulationAggregate for MinAccumulator { - fn query( - &self, - statistic: Statistic, - query_kwargs: Option<&HashMap>, - ) -> Result> { - if query_kwargs.is_some() { - return Err("MinAccumulator does not support query parameters".into()); - } - match statistic { - Statistic::Min => Ok(self.value), - other => Err(format!("Unsupported statistic in MinAccumulator: {other:?}").into()), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn keeps_the_smallest_update() { - let mut acc = MinAccumulator::new(); - acc.update(10.0); - acc.update(5.0); - acc.update(15.0); - - assert_eq!(acc.value, 5.0); - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).unwrap(), - 5.0 - ); - } - - #[test] - fn refuses_to_answer_a_maximum_query() { - let acc = MinAccumulator::with_value(5.0); - assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Max, None).is_err()); - } - - #[test] - fn merges_by_taking_the_smallest() { - let merged = - >::merge_accumulators(vec![ - MinAccumulator::with_value(10.0), - MinAccumulator::with_value(5.0), - MinAccumulator::with_value(15.0), - ]) - .unwrap(); - assert_eq!(merged.value, 5.0); - } - - #[test] - fn refuses_to_merge_with_a_maximum() { - use super::super::max::MaxAccumulator; - let min = MinAccumulator::with_value(5.0); - let max = MaxAccumulator::with_value(15.0); - assert!(min.merge_with(&max).is_err()); - } - - #[test] - fn round_trips_through_both_serializations() { - let acc = MinAccumulator::with_value(42.5); - - let json = acc.serialize_to_json(); - assert_eq!( - MinAccumulator::deserialize_from_json(&json).unwrap().value, - 42.5 - ); - - let bytes = acc.serialize_to_bytes(); - assert_eq!( - MinAccumulator::deserialize_from_bytes(&bytes) - .unwrap() - .value, - 42.5 - ); - } - - #[test] - fn aux_stats_expose_min_only() { - let aux = MinAccumulator::with_value(3.5).aux_stats(); - assert_eq!(aux.min, Some(3.5)); - assert_eq!(aux.max, None); - assert_eq!(aux.count, None); - assert_eq!(aux.sum, None); - assert_eq!(aux.try_answer(Statistic::Min), Some(3.5)); - assert_eq!(aux.try_answer(Statistic::Max), None); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/mod.rs b/crates/asap_summary_state/src/summary_kernels/mod.rs deleted file mode 100644 index 8ebc5d55d..000000000 --- a/crates/asap_summary_state/src/summary_kernels/mod.rs +++ /dev/null @@ -1,42 +0,0 @@ -//! ASAP state adapters and exact accumulators; sketch algorithms live in `asap_sketchlib`. -pub mod count_min_sketch; -pub mod count_min_sketch_with_heap; -pub mod count_sketch; -pub mod count_sketch_with_heap; -pub mod datasketches_kll; -pub mod dd_sketch; -pub mod exact; -pub mod hll_sketch; -pub mod hydra_kll; -pub mod increase; -pub mod keyed_counter_state; -pub mod keyed_max_state; -pub mod keyed_min_state; -pub mod keyed_sum_count; -pub mod max; -pub mod min; -pub mod sketch_envelope; -pub mod sum; -pub mod univmon; - -pub use count_min_sketch::*; -pub use count_min_sketch_with_heap::*; -pub use count_sketch::*; -pub use count_sketch_with_heap::*; -pub use datasketches_kll::*; -pub use dd_sketch::*; -pub use hll_sketch::*; -pub use hydra_kll::*; -pub use increase::*; -pub use keyed_counter_state::*; -pub use keyed_max_state::*; -pub use keyed_min_state::*; -pub use keyed_sum_count::*; -pub use max::*; -pub use min::*; -pub use sketch_envelope::*; -pub use sum::*; - -pub mod factory; -pub mod traits; -pub mod weighted_frequency; diff --git a/crates/asap_summary_state/src/summary_kernels/sketch_envelope.rs b/crates/asap_summary_state/src/summary_kernels/sketch_envelope.rs deleted file mode 100644 index 475d7f51d..000000000 --- a/crates/asap_summary_state/src/summary_kernels/sketch_envelope.rs +++ /dev/null @@ -1,154 +0,0 @@ -//! SketchEnvelopeAccumulator — wraps a raw SketchEnvelope protobuf payload -//! received via OTLP ingest so it can be stored through the `Store` trait. -//! -//! The accumulator preserves the opaque proto bytes and decodes them lazily -//! (via `SketchEnvelope::decode`) only when merge or query operations need -//! the inner sketch type. - -use crate::{AggregateCore, KeyByLabelValues, SerializableToSink}; -use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; -use prost::Message; -use serde_json::Value; -use std::collections::HashMap; - -use crate::AggregationType; -use crate::Statistic; - -/// Accumulator that stores a serialized `SketchEnvelope` protobuf. -/// -/// This is the simplest viable path for OTLP sketch ingest: the OTel Collector -/// has already computed the sketch, so the backend just stores the bytes and -/// serves them back at query time. -#[derive(Debug, Clone)] -pub struct SketchEnvelopeAccumulator { - /// Raw protobuf-encoded `SketchEnvelope`. - pub payload: Vec, - /// Sketch type string cached from decoding (e.g. "CountMin", "KLL"). - pub sketch_type: String, -} - -impl SketchEnvelopeAccumulator { - /// Create from raw protobuf bytes. Decodes the envelope once to cache - /// the sketch type; the full payload is kept for later use. - pub fn from_proto_bytes( - payload: Vec, - ) -> Result> { - let sketch_type = match SketchEnvelope::decode(payload.as_slice()) { - Ok(env) => match env.sketch_state { - Some(sketch_envelope::SketchState::CountMin(_)) => "CountMin".to_string(), - Some(sketch_envelope::SketchState::CountSketch(_)) => "CountSketch".to_string(), - Some(sketch_envelope::SketchState::Kll(_)) => "KLL".to_string(), - Some(sketch_envelope::SketchState::Hll(_)) => "HLL".to_string(), - Some(sketch_envelope::SketchState::Ddsketch(_)) => "DDSketch".to_string(), - Some(sketch_envelope::SketchState::Univmon(_)) => "UnivMon".to_string(), - Some(sketch_envelope::SketchState::Hydra(_)) => "Hydra".to_string(), - Some(sketch_envelope::SketchState::Coco(_)) => "CocoSketch".to_string(), - Some(sketch_envelope::SketchState::Elastic(_)) => "Elastic".to_string(), - None => "Unknown".to_string(), - }, - Err(e) => { - return Err(format!("Failed to decode SketchEnvelope: {}", e).into()); - } - }; - - Ok(Self { - payload, - sketch_type, - }) - } -} - -// --------------------------------------------------------------------------- -// Trait implementations -// --------------------------------------------------------------------------- - -impl SerializableToSink for SketchEnvelopeAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ - "type": "SketchEnvelopeAccumulator", - "sketch_type": self.sketch_type, - "payload_bytes": self.payload.len(), - }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.payload.clone() - } -} - -impl AggregateCore for SketchEnvelopeAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "SketchEnvelopeAccumulator" - } - - 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 SketchEnvelopeAccumulator with {:?}", - other.get_accumulator_type() - ) - .into()); - } - - // For now, merging opaque envelopes is not supported — each window is - // a self-contained sketch produced by the OTel Collector. Return self - // as-is so the store can still call merge_with without panicking. - Ok(Box::new(self.clone())) - } - - fn get_accumulator_type(&self) -> AggregationType { - // Opaque wrapper — report as the generic multi-subpopulation bucket. - // Direct dispatch is not supported; native sketch query path must - // decode the envelope and delegate to the correct accumulator. - AggregationType::MultipleSubpopulation - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - _statistic: Statistic, - _key: &Option, - _query_kwargs: &HashMap, - ) -> Result> { - Err( - "SketchEnvelopeAccumulator: query_statistic not supported; decode envelope first" - .into(), - ) - } -} - -impl crate::MultipleSubpopulationAggregate for SketchEnvelopeAccumulator { - fn query( - &self, - _statistic: Statistic, - _key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, - ) -> Result> { - Err( - "SketchEnvelopeAccumulator: direct query not supported; use native sketch query path" - .into(), - ) - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/sum.rs b/crates/asap_summary_state/src/summary_kernels/sum.rs deleted file mode 100644 index 4e74a45ae..000000000 --- a/crates/asap_summary_state/src/summary_kernels/sum.rs +++ /dev/null @@ -1,413 +0,0 @@ -use crate::{ - AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use crate::Statistic; - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SumAccumulator { - pub sum: f64, - /// None for scalar-only payloads; a sum does not establish a sample count. - #[serde(default)] - pub observation_count: Option, -} - -impl SumAccumulator { - pub fn new() -> Self { - Self { - sum: 0.0, - observation_count: Some(0), - } - } - - pub fn with_sum(sum: f64) -> Self { - Self { - sum, - observation_count: None, - } - } - - pub fn update(&mut self, value: f64) { - self.sum += value; - self.observation_count = self - .observation_count - .and_then(|count| count.checked_add(1)); - } - - pub fn deserialize_from_json(data: &Value) -> Result> { - let sum = data["sum"] - .as_f64() - .ok_or("Missing or invalid 'sum' field")?; - Ok(Self { - sum, - observation_count: data.get("observation_count").and_then(Value::as_u64), - }) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - match buffer.len() { - // Legacy Python scalar sums carry no sample-count evidence. - 4 => Ok(Self::with_sum(f32::from_le_bytes(buffer.try_into()?) as f64)), - // Counted sums use the same fixed layout as the Collector Sum payload. - 16 => Self::from_sum_bytes(buffer), - len => { - Err(format!("Invalid persisted Sum payload length: {len} (want 4 or 16)").into()) - } - } - } - - /// Decode the fixed Sum payload produced by the first-class Sum - /// AggregationType path (asap-precompute-go's SumWrapper): float64 sum - /// (little-endian) followed by uint64 count (little-endian), 16 bytes. - /// - /// Sum is an aggregation, NOT a sketch, so this deliberately does NOT - /// depend on the sketchlib sketch-envelope proto — the payload is a small - /// self-contained fixed layout. It decodes into the SAME - /// `AggregationType::Sum` accumulator as a plain-OTLP Sum, so the SumAgg - /// envelope and a plain Sum land on one identity (`exact_agg:Sum`) with no - /// new SketchAlgorithm. The supplied observation count is retained for - /// exact sample-count readouts; scalar-only legacy payloads leave it unknown. - pub fn from_sum_bytes(buffer: &[u8]) -> Result> { - if buffer.len() < 16 { - return Err(format!("Sum payload too short: {} bytes (want 16)", buffer.len()).into()); - } - let sum = f64::from_le_bytes(buffer[0..8].try_into().unwrap()); - let count = u64::from_le_bytes(buffer[8..16].try_into().unwrap()); - Ok(Self { - sum, - observation_count: Some(count), - }) - } -} - -impl Default for SumAccumulator { - fn default() -> Self { - Self::new() - } -} - -impl SerializableToSink for SumAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ - "sum": self.sum, - "observation_count": self.observation_count - }) - } - - fn serialize_to_bytes(&self) -> Vec { - match self.observation_count { - Some(count) => { - let mut bytes = Vec::with_capacity(16); - bytes.extend_from_slice(&self.sum.to_le_bytes()); - bytes.extend_from_slice(&count.to_le_bytes()); - bytes - } - None => (self.sum as f32).to_le_bytes().to_vec(), - } - } -} - -impl AggregateCore for SumAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "SumAccumulator" - } - - 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> { - // Check if other is also a SumAccumulator - if other.get_accumulator_type() != self.get_accumulator_type() { - return Err(format!( - "Cannot merge SumAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - - // Downcast to SumAccumulator - let other_sum = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to SumAccumulator")?; - - // Use the existing merge_accumulators method - let merged = Self::merge_accumulators(vec![self.clone(), other_sum.clone()])?; - - Ok(Box::new(merged)) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::Sum - } - - fn approx_memory_bytes(&self) -> usize { - // Single f64 + struct overhead. - std::mem::size_of::() - } - - fn aux_stats(&self) -> AuxStats { - AuxStats { - sum: Some(self.sum), - count: self.observation_count, - ..AuxStats::empty() - } - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: crate::Statistic, - _key: &Option, - _query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::SingleSubpopulationAggregate; - self.query(statistic, None) - } -} - -impl SingleSubpopulationAggregate for SumAccumulator { - fn query( - &self, - statistic: Statistic, - query_kwargs: Option<&HashMap>, - ) -> Result> { - // SumAccumulator doesn't use query_kwargs, assert it's None - if query_kwargs.is_some() { - return Err("SumAccumulator does not support query parameters".into()); - } - - match statistic { - Statistic::Sum => Ok(self.sum), - Statistic::Count => self - .observation_count - .map(|count| count as f64) - .ok_or_else(|| "sample count is unavailable for this Sum payload".into()), - _ => Err(format!("Unsupported statistic in SumAccumulator: {statistic:?}").into()), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for SumAccumulator { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - let total_sum = accumulators.iter().map(|acc| acc.sum).sum(); - let observation_count = accumulators - .iter() - .try_fold(0u64, |total, acc| total.checked_add(acc.observation_count?)); - Ok(SumAccumulator { - sum: total_sum, - observation_count, - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - // Sample counts must survive updates and merges independently of the sum. - #[test] - fn observation_count_survives_merge() { - let mut first = SumAccumulator::new(); - first.update(10.0); - first.update(20.0); - let mut second = SumAccumulator::new(); - second.update(100.0); - let merged = SumAccumulator::merge_accumulators(vec![first, second]).unwrap(); - assert_eq!(merged.sum, 130.0); - assert_eq!(merged.aux_stats().count, Some(3)); - } - - // A legacy scalar sum has no evidence of how many observations produced it. - #[test] - fn legacy_sum_does_not_invent_observation_count() { - let mut raw = SumAccumulator::new(); - raw.update(10.0); - let merged = - SumAccumulator::merge_accumulators(vec![raw, SumAccumulator::with_sum(20.0)]).unwrap(); - assert_eq!(merged.aux_stats().count, None); - } - - // Persistence retains known counts, including zero and the full u64 range. - #[test] - fn counted_sum_binary_round_trip() { - for count in [0, 3, u64::MAX] { - let acc = SumAccumulator { - sum: 1.0000000000001, - observation_count: Some(count), - }; - let bytes = acc.serialize_to_bytes(); - assert_eq!(bytes.len(), 16); - let restored = SumAccumulator::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!(restored.sum, acc.sum); - assert_eq!(restored.observation_count, Some(count)); - } - } - - // Existing scalar-only files remain readable without inventing counts. - #[test] - fn legacy_binary_sum_has_unknown_count() { - let bytes = 42.5f32.to_le_bytes(); - let restored = SumAccumulator::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!(restored.sum, 42.5); - assert_eq!(restored.observation_count, None); - assert_eq!(restored.serialize_to_bytes(), bytes); - } - - // Truncated counted payloads must not silently decode as scalar sums. - #[test] - fn persisted_sum_rejects_invalid_lengths() { - for len in [0, 3, 5, 8, 15, 17] { - assert!(SumAccumulator::deserialize_from_bytes(&vec![0; len]).is_err()); - } - } - - #[test] - fn test_sum_accumulator_creation() { - let acc = SumAccumulator::new(); - assert_eq!(acc.sum, 0.0); - - let acc2 = SumAccumulator::with_sum(42.5); - assert_eq!(acc2.sum, 42.5); - } - - #[test] - fn test_sum_accumulator_update() { - let mut acc = SumAccumulator::new(); - acc.update(10.0); - acc.update(20.0); - assert_eq!(acc.sum, 30.0); - } - - #[test] - fn test_sum_accumulator_query() { - let acc = SumAccumulator::with_sum(42.0); - - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), - 42.0 - ); - assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Count, None).is_err()); - - assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); - // SumAccumulator is a single subpopulation accumulator, doesn't need key-based queries - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), - 42.0 - ); - } - - #[test] - fn count_readout_uses_observation_count_not_sum() { - let mut acc = SumAccumulator::new(); - acc.update(10.0); - acc.update(20.0); - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Count, None).unwrap(), - 2.0 - ); - } - - #[test] - fn test_sum_accumulator_merge() { - let acc1 = SumAccumulator::with_sum(10.0); - let acc2 = SumAccumulator::with_sum(20.0); - let acc3 = SumAccumulator::with_sum(30.0); - - let merged = - >::merge_accumulators(vec![ - acc1, acc2, acc3, - ]) - .unwrap(); - assert_eq!(merged.sum, 60.0); - } - - #[test] - fn test_sum_accumulator_serialization() { - let acc = SumAccumulator::with_sum(42.5); - - // Test JSON serialization - let json = acc.serialize_to_json(); - let deserialized = SumAccumulator::deserialize_from_json(&json).unwrap(); - assert_eq!(acc.sum, deserialized.sum); - - // Test byte serialization - let bytes = acc.serialize_to_bytes(); - let deserialized_bytes = SumAccumulator::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!(acc.sum, deserialized_bytes.sum); - } - - #[test] - fn test_trait_object() { - let acc: Box = Box::new(SumAccumulator::with_sum(42.0)); - - assert_eq!(acc.type_name(), "SumAccumulator"); - } - - #[test] - fn from_sum_bytes_decodes_go_sum_payload() { - // GOLDEN: the 16-byte payload asap-precompute-go's - // SumWrapper{10,20,30,40}.Snapshot() emits — float64 sum (LE) followed - // by uint64 count (LE), sum=100, count=4. Proves the Rust backend - // decodes the first-class Sum payload the Go agent produces - // (cross-language wire parity, no sketchlib proto dependency). - let go_bytes: &[u8] = &[ - 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x59, 0x40, // 100.0 f64 LE - 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // 4 u64 LE - ]; - let acc = SumAccumulator::from_sum_bytes(go_bytes).expect("decode Go Sum payload"); - assert_eq!(acc.sum, 100.0, "decoded Go SumWrapper payload sum"); - } - - #[test] - fn from_sum_bytes_rejects_short_payload() { - // A short buffer is rejected (the ingest path then skips the point). - assert!(SumAccumulator::from_sum_bytes(&[]).is_err()); - assert!(SumAccumulator::from_sum_bytes(&[0u8; 8]).is_err()); - } - - #[test] - fn aux_stats_exposes_sum_only() { - let acc = SumAccumulator::with_sum(123.5); - let aux = acc.aux_stats(); - assert_eq!(aux.sum, Some(123.5)); - assert_eq!(aux.count, None); - assert_eq!(aux.min, None); - assert_eq!(aux.max, None); - } - - #[test] - fn aux_stats_try_answer_on_sum_statistic() { - use crate::Statistic; - let acc = SumAccumulator::with_sum(42.0); - // Sum statistic is covered by aux without deserialising. - assert_eq!(acc.aux_stats().try_answer(Statistic::Sum), Some(42.0)); - // Count is not tracked by SumAccumulator. - assert_eq!(acc.aux_stats().try_answer(Statistic::Count), None); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/traits.rs b/crates/asap_summary_state/src/summary_kernels/traits.rs deleted file mode 100644 index ae05ca06b..000000000 --- a/crates/asap_summary_state/src/summary_kernels/traits.rs +++ /dev/null @@ -1,357 +0,0 @@ -use crate::KeyByLabelValues; -use std::collections::HashMap; - -use crate::AggregationType; -use crate::Statistic; - -use serde_json::Value; - -/// Trait for objects that can be serialized to different formats -pub trait SerializableToSink { - fn serialize_to_json(&self) -> Value; - fn serialize_to_bytes(&self) -> Vec; -} - -/// Core trait for all aggregates containing shared functionality -/// This trait provides common operations like serialization, cloning, and type identification -pub trait AggregateCore: SerializableToSink + Send + Sync { - /// Clone this accumulator into a boxed trait object - fn clone_boxed_core(&self) -> Box; - - /// Get the type name of this accumulator - fn type_name(&self) -> &'static str; - - /// Downcast to Any for type checking - fn as_any(&self) -> &dyn std::any::Any; - - /// Mutable downcast to Any. Used by ingest paths that need to - /// mutate a boxed accumulator in place — e.g. the PROTO_DELTA - /// delta-merge applier in `drivers::ingest::otel::apply_modified_otlp_delta_bytes`. - fn as_any_mut(&mut self) -> &mut dyn std::any::Any; - - /// Merge this accumulator with another accumulator of the same type - /// Returns a new merged accumulator, leaving the original unchanged - fn merge_with( - &self, - other: &dyn AggregateCore, - ) -> Result, Box>; - - /// Get the accumulator type identifier for merge compatibility checking - fn get_accumulator_type(&self) -> AggregationType; - - /// Get all keys stored in this accumulator - fn get_keys(&self) -> Option>; - - /// Dispatch a statistic query without downcasting. - /// - /// Replaces the 12-arm `match get_accumulator_type()` in the engine. - /// Single-subpopulation types ignore `key`; multiple-subpopulation types - /// require it and return `Err` when it is `None`. - /// Special cases (DeltaSetAggregator, SetAggregator) fall back to a - /// cardinality value when `key` is `None`. - fn query_statistic( - &self, - statistic: Statistic, - key: &Option, - query_kwargs: &HashMap, - ) -> Result>; - - /// Approximate in-memory byte footprint of this accumulator. - /// - /// Used by the `SketchStore` persistence layer to drive its - /// memory-pressure trigger. Not required to be exact — the flusher - /// only needs rough proportionality. The default is a conservative - /// 4 KiB constant; concrete types should override it with a - /// type-aware estimate (e.g. KLL: `k * 8` plus overhead). - /// - /// Implementors must not call `serialize_to_bytes` here — this is - /// on the insert hot path. - fn approx_memory_bytes(&self) -> usize { - 4096 - } - - /// Typed auxiliary statistics — `count`, `sum`, `min`, `max` — - /// exposed as first-class scalars alongside the sketch payload. - /// - /// The overwhelming majority of production queries - /// (`count_over_time`, `sum_over_time`, `min_over_time`, - /// `max_over_time`, and the additive aggregations built on - /// them) only need these scalars. Returning them directly here - /// lets callers avoid deserialising the full sketch bytes. - /// - /// Returning fields as `None` means the accumulator doesn't - /// track that statistic exactly (e.g. a pure HLL doesn't carry - /// sum/min/max). Callers then fall back to the sketch's - /// `query_statistic` method. - /// - /// This is the phase-1 piece of the sketch DB design - /// (docs/design_docs/summary-storage.md). - fn aux_stats(&self) -> AuxStats { - AuxStats::empty() - } - - /// Reset the sketch state to empty **in place**, preserving its - /// shape / configuration (dimensions, relative accuracy, register - /// width, …) so a subsequent delta-apply lands on a clean, - /// same-shape base. - /// - /// Used by the OTLP ingest path's per-window base rotation: when a - /// delta frame opens a new tumbling window for a series, the cached - /// base is reset here before the new window's delta is applied, so - /// the reconstructed state reflects that window only rather than an - /// all-time accumulation across windows (see - /// `docs/delta-baseline-contract.md` §3). - /// - /// The default is a no-op: only the delta-capable, additive families - /// (DDSketch, CMS, CountSketch, HLL) ever reach the rotation path and - /// override this. KLL never deltas, and the non-sketch accumulators - /// are never cached as a delta base. - fn reset_to_empty(&mut self) {} -} - -/// Four typed auxiliary scalars tracked alongside every sketch entry: -/// `count`, `sum`, `min`, `max`. Exposed so the query engine can -/// serve Count / Sum / Min / Max statistics without touching sketch -/// bytes. -/// -/// Each field is `Option<…>` because not every accumulator tracks -/// every stat (e.g. HLL has cardinality but no meaningful -/// sum / min / max; DeltaSetAggregator tracks set transitions, not -/// numeric aggregates). -#[derive(Debug, Default, Clone, Copy, PartialEq)] -pub struct AuxStats { - pub count: Option, - pub sum: Option, - pub min: Option, - pub max: Option, -} - -impl AuxStats { - pub const fn empty() -> Self { - Self { - count: None, - sum: None, - min: None, - max: None, - } - } - - /// Attempt to fulfil a `Statistic` purely from the typed aux - /// columns, without needing to deserialise the sketch. Returns - /// `None` if the requested statistic isn't covered by aux - /// (e.g. Quantile, Cardinality, TopK) or if the corresponding - /// aux field is `None`. - pub fn try_answer(&self, statistic: Statistic) -> Option { - match statistic { - Statistic::Count => self.count.map(|c| c as f64), - Statistic::Sum => self.sum, - Statistic::Min => self.min, - Statistic::Max => self.max, - // Increase / Rate need two samples; aux columns carry - // window totals, so one entry's aux is insufficient. - // Cardinality / Quantile / Topk are sketch-native and - // must go through query_statistic. - _ => None, - } - } - - /// Merge two aux stats the way the corresponding sketch merge - /// would. Count / sum add, min / max take the extremum. When - /// either side is `None` the result is the other side (so a - /// window that only has partial aux still contributes). - pub fn merge(self, other: Self) -> Self { - fn add_opt_u(a: Option, b: Option) -> Option { - match (a, b) { - (Some(x), Some(y)) => Some(x.saturating_add(y)), - (x, None) => x, - (None, y) => y, - } - } - fn add_opt_f(a: Option, b: Option) -> Option { - match (a, b) { - (Some(x), Some(y)) => Some(x + y), - (x, None) => x, - (None, y) => y, - } - } - fn min_opt(a: Option, b: Option) -> Option { - match (a, b) { - (Some(x), Some(y)) => Some(x.min(y)), - (x, None) => x, - (None, y) => y, - } - } - fn max_opt(a: Option, b: Option) -> Option { - match (a, b) { - (Some(x), Some(y)) => Some(x.max(y)), - (x, None) => x, - (None, y) => y, - } - } - Self { - count: add_opt_u(self.count, other.count), - sum: add_opt_f(self.sum, other.sum), - min: min_opt(self.min, other.min), - max: max_opt(self.max, other.max), - } - } -} - -/// Trait for accumulators that support a single subpopulation -/// These accumulators store a single aggregate value (e.g., Sum, Increase) -pub trait SingleSubpopulationAggregate: AggregateCore { - /// Query the accumulator for a specific statistic - fn query( - &self, - statistic: Statistic, - query_kwargs: Option<&HashMap>, - ) -> Result>; - - /// Clone this accumulator into a boxed trait object - fn clone_boxed(&self) -> Box; -} - -/// Trait for accumulators that support multiple subpopulations identified by keys -/// These accumulators store separate values for different label combinations -pub trait MultipleSubpopulationAggregate: AggregateCore { - /// Query the accumulator for a specific statistic and key - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - query_kwargs: Option<&HashMap>, - ) -> Result>; - - /// Clone this accumulator into a boxed trait object - fn clone_boxed(&self) -> Box; -} - -/// Trait for merging multiple accumulators of the same type -pub trait MergeableAccumulator { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> - where - T: Sized; -} - -// Implement Clone for the new trait objects -impl Clone for Box { - fn clone(&self) -> Self { - self.clone_boxed_core() - } -} - -impl Clone for Box { - fn clone(&self) -> Self { - self.clone_boxed() - } -} - -impl Clone for Box { - fn clone(&self) -> Self { - self.clone_boxed() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn aux_stats_empty_answers_nothing() { - let e = AuxStats::empty(); - assert_eq!(e.try_answer(Statistic::Count), None); - assert_eq!(e.try_answer(Statistic::Sum), None); - assert_eq!(e.try_answer(Statistic::Min), None); - assert_eq!(e.try_answer(Statistic::Max), None); - } - - #[test] - fn aux_stats_try_answer_covers_typed_stats() { - let a = AuxStats { - count: Some(7), - sum: Some(42.0), - min: Some(1.5), - max: Some(9.25), - }; - assert_eq!(a.try_answer(Statistic::Count), Some(7.0)); - assert_eq!(a.try_answer(Statistic::Sum), Some(42.0)); - assert_eq!(a.try_answer(Statistic::Min), Some(1.5)); - assert_eq!(a.try_answer(Statistic::Max), Some(9.25)); - } - - #[test] - fn aux_stats_try_answer_skips_sketch_native_stats() { - let a = AuxStats { - count: Some(100), - sum: Some(500.0), - min: Some(1.0), - max: Some(10.0), - }; - assert_eq!(a.try_answer(Statistic::Quantile), None); - assert_eq!(a.try_answer(Statistic::Cardinality), None); - assert_eq!(a.try_answer(Statistic::Topk), None); - assert_eq!(a.try_answer(Statistic::Increase), None); - assert_eq!(a.try_answer(Statistic::Rate), None); - } - - #[test] - fn aux_stats_merge_adds_count_and_sum_takes_extrema() { - let a = AuxStats { - count: Some(10), - sum: Some(50.0), - min: Some(1.0), - max: Some(9.0), - }; - let b = AuxStats { - count: Some(5), - sum: Some(20.0), - min: Some(0.5), - max: Some(12.0), - }; - let merged = a.merge(b); - assert_eq!(merged.count, Some(15)); - assert_eq!(merged.sum, Some(70.0)); - assert_eq!(merged.min, Some(0.5)); - assert_eq!(merged.max, Some(12.0)); - } - - #[test] - fn aux_stats_merge_handles_partial_sides() { - // HLL-like (count only) merged with Sum-only side. - let hll_like = AuxStats { - count: Some(100), - ..AuxStats::empty() - }; - let sum_like = AuxStats { - sum: Some(500.0), - ..AuxStats::empty() - }; - let merged = hll_like.merge(sum_like); - assert_eq!(merged.count, Some(100)); - assert_eq!(merged.sum, Some(500.0)); - assert_eq!(merged.min, None); - assert_eq!(merged.max, None); - } - - #[test] - fn aux_stats_merge_is_empty_plus_empty() { - let merged = AuxStats::empty().merge(AuxStats::empty()); - assert_eq!(merged, AuxStats::empty()); - } - - #[test] - fn aux_stats_count_saturates_on_overflow() { - let a = AuxStats { - count: Some(u64::MAX - 1), - ..AuxStats::empty() - }; - let b = AuxStats { - count: Some(100), - ..AuxStats::empty() - }; - let merged = a.merge(b); - assert_eq!(merged.count, Some(u64::MAX)); - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/weighted_frequency.rs b/crates/asap_summary_state/src/summary_kernels/weighted_frequency.rs deleted file mode 100644 index 34e142d99..000000000 --- a/crates/asap_summary_state/src/summary_kernels/weighted_frequency.rs +++ /dev/null @@ -1,75 +0,0 @@ -//! Stored form of Planner's weighted frequency state. It is read only through -//! typed heap rows, never through a scalar statistic. -use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink, Statistic}; -pub use asap_physical_operators::summary_kernels::weighted_frequency::{ - FrequencyAlgorithm, WeightedFrequency as PhysicalWeightedFrequency, -}; -use asap_physical_operators::AggregateCore as PhysicalState; -use std::collections::HashMap; - -#[derive(Debug, Clone)] -pub struct WeightedFrequency(pub PhysicalWeightedFrequency); - -impl SerializableToSink for WeightedFrequency { - fn serialize_to_json(&self) -> serde_json::Value { - serde_json::to_value(&self.0).expect("finite validated frequency state") - } - fn serialize_to_bytes(&self) -> Vec { - crate::physical::frequency_kernel(&self.0) - .expect("weighted frequency kernel encoding") - .to_bytes() - } -} - -impl AggregateCore for WeightedFrequency { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - fn type_name(&self) -> &'static str { - "WeightedFrequency" - } - 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> { - let other = other - .as_any() - .downcast_ref::() - .ok_or("weighted frequency state type mismatch")?; - let merged = self.0.merge_with(&other.0)?; - let merged = merged - .as_any() - .downcast_ref::() - .ok_or("weighted frequency merge changed state type")?; - Ok(Box::new(Self(merged.clone()))) - } - fn get_accumulator_type(&self) -> AggregationType { - match crate::physical::frequency_kernel(&self.0) - .expect("weighted frequency kernel encoding") - .algorithm() - { - FrequencyAlgorithm::Cms => AggregationType::CountMinSketchWithHeap, - FrequencyAlgorithm::CountSketch => AggregationType::CountSketchWithHeap, - } - } - fn get_keys(&self) -> Option> { - None - } - fn query_statistic( - &self, - _: Statistic, - _: &Option, - _: &HashMap, - ) -> Result> { - Err("weighted frequency uses typed row readout".into()) - } - fn approx_memory_bytes(&self) -> usize { - self.0.approx_memory_bytes() - } -} diff --git a/crates/asap_summary_state/src/summary_kernels/univmon.rs b/crates/asap_summary_state/src/univmon.rs similarity index 51% rename from crates/asap_summary_state/src/summary_kernels/univmon.rs rename to crates/asap_summary_state/src/univmon.rs index 4a52bf1da..58546e961 100644 --- a/crates/asap_summary_state/src/summary_kernels/univmon.rs +++ b/crates/asap_summary_state/src/univmon.rs @@ -1,12 +1,12 @@ -//! One frequency state shared by count, distinct, L2 and entropy readouts. - -use crate::{AggregateCore, AuxStats, KeyByLabelValues, SerializableToSink}; -use crate::{AggregationType, Statistic}; +//! SHIM: stored UnivMon state. +//! +//! Planner's `UnivMonAccumulator` keeps its sketchlib `UnivMon` private and has +//! no byte form, so stored UnivMon bytes cannot be decoded into it. Delete this +//! module once Planner exposes `UnivMonAccumulator::from_sketch(UnivMon)` and +//! `UnivMonAccumulator::sketch(&self) -> &UnivMon` (or a byte codec). +use asap_physical_operators::{AggregateCore, KernelError}; use asap_sketchlib::{DataInput, UnivMon}; -use serde_json::Value; -use std::collections::HashMap; - -type Error = Box; +use planner_types::{post_asap::SketchQuery, pre_asap::ColumnRef}; #[derive(Debug, Clone)] pub struct UnivMonAccumulator { @@ -14,7 +14,12 @@ pub struct UnivMonAccumulator { } impl UnivMonAccumulator { - pub fn new(heap_size: usize, rows: usize, cols: usize, layers: usize) -> Result { + pub fn new( + heap_size: usize, + rows: usize, + cols: usize, + layers: usize, + ) -> Result { if heap_size == 0 || cols == 0 || !(1..=20).contains(&rows) || !(1..=64).contains(&layers) { return Err("invalid UnivMon dimensions".into()); } @@ -27,7 +32,7 @@ impl UnivMonAccumulator { } /// Each non-NaN sample is one occurrence. Signed zero has one identity. - pub fn insert_sample(&mut self, value: f64) -> Result<(), Error> { + pub fn insert_sample(&mut self, value: f64) -> Result<(), KernelError> { if value.is_nan() { return Ok(()); } @@ -40,7 +45,7 @@ impl UnivMonAccumulator { Ok(()) } - pub fn from_bytes(bytes: &[u8]) -> Result { + pub fn from_bytes(bytes: &[u8]) -> Result { let inner = UnivMon::deserialize_from_bytes(bytes) .map_err(|e| format!("invalid UnivMon state: {e}"))?; if !inner.accepts_standard_updates() { @@ -51,18 +56,8 @@ impl UnivMonAccumulator { Ok(Self { inner }) } - fn compatible(&self, other: &Self) -> bool { - ( - self.inner.heap_size, - self.inner.sketch_row, - self.inner.sketch_col, - self.inner.layer_size, - ) == ( - other.inner.heap_size, - other.inner.sketch_row, - other.inner.sketch_col, - other.inner.layer_size, - ) + pub fn to_bytes(&self) -> Result, KernelError> { + Ok(self.inner.serialize_to_bytes()?) } pub fn dimensions(&self) -> (usize, usize, usize, usize) { @@ -74,8 +69,13 @@ impl UnivMonAccumulator { ) } - pub fn merge_in_place(&mut self, other: &Self) -> Result<(), Error> { - if !self.compatible(other) { + /// Empty the sketch in place, keeping its shape. + pub fn clear(&mut self) { + self.inner.free(); + } + + pub fn merge_in_place(&mut self, other: &Self) -> Result<(), KernelError> { + if self.dimensions() != other.dimensions() { return Err("incompatible UnivMon dimensions".into()); } self.inner @@ -87,53 +87,14 @@ impl UnivMonAccumulator { } } -impl SerializableToSink for UnivMonAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({"count": self.inner.bucket_size}) - } - - fn serialize_to_bytes(&self) -> Vec { - self.inner - .serialize_to_bytes() - .expect("validated unit-frequency UnivMon state") - } -} - impl AggregateCore for UnivMonAccumulator { - fn approx_memory_bytes(&self) -> usize { - std::mem::size_of::().saturating_add( - self.inner.layer_size.saturating_mul( - self.inner - .sketch_row - .saturating_mul(self.inner.sketch_col) - .saturating_mul(16) - .saturating_add(self.inner.heap_size.saturating_mul(256)), - ), - ) - } fn clone_boxed_core(&self) -> Box { Box::new(self.clone()) } - fn type_name(&self) -> &'static str { - "UnivMonAccumulator" - } fn as_any(&self) -> &dyn std::any::Any { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::UnivMon - } - fn get_keys(&self) -> Option> { - None - } - fn reset_to_empty(&mut self) { - self.inner.free(); - } - - fn merge_with(&self, other: &dyn AggregateCore) -> Result, Error> { + fn merge_with(&self, other: &dyn AggregateCore) -> Result, KernelError> { let other = other .as_any() .downcast_ref::() @@ -142,30 +103,28 @@ impl AggregateCore for UnivMonAccumulator { merged.merge_in_place(other)?; Ok(Box::new(merged)) } - - fn query_statistic( - &self, - statistic: Statistic, - key: &Option, - _: &HashMap, - ) -> Result { - if key.is_some() { - return Err("UnivMon population is selected by the catalog binding".into()); - } - match statistic { - Statistic::Count => Ok(self.inner.calc_l1()), - Statistic::Cardinality => Ok(self.inner.calc_card()), - Statistic::FrequencyL2 => Ok(self.inner.calc_l2()), - Statistic::FrequencyEntropy => Ok(self.inner.calc_entropy()), - _ => Err("unsupported UnivMon readout".into()), - } + fn estimate(&self, query: &SketchQuery) -> Result { + Ok(match query { + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + value: None, + } => self.inner.calc_l1(), + SketchQuery::Cardinality => self.inner.calc_card(), + SketchQuery::FrequencyL2 => self.inner.calc_l2(), + SketchQuery::FrequencyEntropy => self.inner.calc_entropy(), + _ => return Err("unsupported UnivMon readout".into()), + }) } - - fn aux_stats(&self) -> AuxStats { - AuxStats { - count: Some(self.inner.bucket_size as u64), - ..AuxStats::empty() - } + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::().saturating_add( + self.inner.layer_size.saturating_mul( + self.inner + .sketch_row + .saturating_mul(self.inner.sketch_col) + .saturating_mul(16) + .saturating_add(self.inner.heap_size.saturating_mul(256)), + ), + ) } } @@ -173,33 +132,39 @@ impl AggregateCore for UnivMonAccumulator { mod tests { use super::*; - fn read(state: &dyn AggregateCore, stat: Statistic) -> f64 { - state.query_statistic(stat, &None, &HashMap::new()).unwrap() + fn read(state: &dyn AggregateCore, query: SketchQuery) -> f64 { + state.estimate(&query).unwrap() + } + fn count() -> SketchQuery { + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + value: None, + } } - /// Duplicate samples affect frequency but not cardinality, including signed zero. + // Duplicate samples affect frequency but not cardinality, including signed zero. #[test] fn shared_readouts_survive_serialization() { let mut state = UnivMonAccumulator::new(32, 5, 1024, 4).unwrap(); for value in [0.0, -0.0, 2.0, 2.0, f64::NAN] { state.insert_sample(value).unwrap(); } - let restored = UnivMonAccumulator::from_bytes(&state.serialize_to_bytes()).unwrap(); - for stat in [ - Statistic::Count, - Statistic::Cardinality, - Statistic::FrequencyL2, - Statistic::FrequencyEntropy, + let restored = UnivMonAccumulator::from_bytes(&state.to_bytes().unwrap()).unwrap(); + for query in [ + count(), + SketchQuery::Cardinality, + SketchQuery::FrequencyL2, + SketchQuery::FrequencyEntropy, ] { - assert_eq!(read(&state, stat), read(&restored, stat)); + assert_eq!(read(&state, query.clone()), read(&restored, query)); } - assert_eq!(read(&restored, Statistic::Count), 4.0); - assert!((read(&restored, Statistic::Cardinality) - 2.0).abs() < 0.01); - assert!((read(&restored, Statistic::FrequencyL2) - 8.0f64.sqrt()).abs() < 0.01); - assert!((read(&restored, Statistic::FrequencyEntropy) - 1.0).abs() < 0.01); + assert_eq!(read(&restored, count()), 4.0); + assert!((read(&restored, SketchQuery::Cardinality) - 2.0).abs() < 0.01); + assert!((read(&restored, SketchQuery::FrequencyL2) - 8.0f64.sqrt()).abs() < 0.01); + assert!((read(&restored, SketchQuery::FrequencyEntropy) - 1.0).abs() < 0.01); } - /// Terminal-mode serialization is valid sketchlib state but not this accumulator's update domain. + // Terminal-mode serialization is valid sketchlib state but not this accumulator's update domain. #[test] fn terminal_state_is_rejected_before_ingestion_or_merge() { let mut state = UnivMon::init_univmon(4, 3, 16, 2); @@ -210,9 +175,9 @@ mod tests { assert!(UnivMonAccumulator::from_bytes(&state.serialize_to_bytes().unwrap()).is_ok()); } - /// Pane merge preserves overlapping keys and reset removes the previous window. + // Pane merge preserves overlapping keys and clearing removes the previous window. #[test] - fn merge_and_reset_preserve_frequency_semantics() { + fn merge_and_clear_preserve_frequency_semantics() { let mut left = UnivMonAccumulator::new(32, 5, 1024, 4).unwrap(); let mut right = left.clone(); for value in [1.0, 2.0] { @@ -222,11 +187,11 @@ mod tests { right.insert_sample(value).unwrap(); } let merged = left.merge_with(&right).unwrap(); - assert_eq!(read(merged.as_ref(), Statistic::Count), 4.0); - assert!((read(merged.as_ref(), Statistic::Cardinality) - 3.0).abs() < 0.01); - left.reset_to_empty(); - assert_eq!(read(&left, Statistic::Count), 0.0); - assert_eq!(read(&left, Statistic::FrequencyEntropy), 0.0); + assert_eq!(read(merged.as_ref(), count()), 4.0); + assert!((read(merged.as_ref(), SketchQuery::Cardinality) - 3.0).abs() < 0.01); + left.clear(); + assert_eq!(read(&left, count()), 0.0); + assert_eq!(read(&left, SketchQuery::FrequencyEntropy), 0.0); assert!(left .merge_with(&UnivMonAccumulator::new(16, 5, 1024, 4).unwrap()) .is_err()); diff --git a/crates/asap_types/src/lib.rs b/crates/asap_types/src/lib.rs index abb83ec16..2275c0a4a 100644 --- a/crates/asap_types/src/lib.rs +++ b/crates/asap_types/src/lib.rs @@ -22,7 +22,6 @@ pub mod storage_backend; pub mod summary_catalog; pub mod summary_semantics; pub mod table_population; -pub mod traits; pub mod utils; pub use accumulator_spec::{AccumulatorSpec, AccumulatorSpecError, SampleUpdateRule}; diff --git a/crates/asap_types/src/traits.rs b/crates/asap_types/src/traits.rs deleted file mode 100644 index 4a91a97d5..000000000 --- a/crates/asap_types/src/traits.rs +++ /dev/null @@ -1 +0,0 @@ -pub use asap_summary_state::SerializableToSink; diff --git a/data_plane/benches/sketch_db.rs b/data_plane/benches/sketch_db.rs index 0797aaf44..48f02c7ed 100644 --- a/data_plane/benches/sketch_db.rs +++ b/data_plane/benches/sketch_db.rs @@ -33,7 +33,6 @@ use asap_sketchlib::DdSketch; use asap_sketchlib::{HllSketch, HllVariant}; use prost::Message; -use asap_summary_state::summary_kernels::SumAccumulator; use data_plane::storage_engines::sketch_db::data::{ AccuracyBound, AggKind, AggregationType, Capability, SketchAlgorithm, SketchConfig, SketchEncoding, @@ -280,7 +279,7 @@ fn bench_append_precompute(c: &mut Criterion) { sid, BTreeMap::new(), win, - Box::new(SumAccumulator::with_sum(c as f64)), + Box::new(exact_sum(c as f64)), ); black_box(&*store); }, @@ -349,7 +348,7 @@ fn build_precompute_store(num_sids: usize, windows_per_sid: u64, metric: &str) - sid, BTreeMap::new(), win, - Box::new(SumAccumulator::with_sum((sid + i) as f64)), + Box::new(exact_sum((sid + i) as f64)), ); } } @@ -513,3 +512,11 @@ criterion_group!( bench_group_key_projection, ); criterion_main!(benches); + +/// Planner's unkeyed exact Sum holding `sum`. +fn exact_sum(sum: f64) -> asap_physical_operators::summary_kernels::exact::ExactAccumulator { + asap_summary_state::stored_state::codec::exact_value( + planner_types::post_asap::ExactKind::Sum, + sum, + ) +} diff --git a/data_plane/examples/univmon_erp_artifact.rs b/data_plane/examples/univmon_erp_artifact.rs index 2ac5e07e1..a64610076 100644 --- a/data_plane/examples/univmon_erp_artifact.rs +++ b/data_plane/examples/univmon_erp_artifact.rs @@ -1,8 +1,8 @@ //! Measure readout-specific ERP evidence from finite JSONL evaluation data. //! This offline tool retains samples; the production backend does not. -use asap_summary_state::summary_kernels::hll_sketch::HllSketchAccumulator; -use asap_summary_state::summary_kernels::univmon::UnivMonAccumulator; -use data_plane::storage_engines::types::{AggregateCore, SerializableToSink}; +use asap_physical_operators::summary_kernels::hll_sketch::HllSketchAccumulator; +use asap_summary_state::univmon::UnivMonAccumulator; +use data_plane::storage_engines::types::{AggregateCore, StoredState}; use serde_json::{json, Value}; use std::collections::{BTreeMap, HashMap}; use std::io::{BufRead, BufReader}; @@ -109,18 +109,16 @@ fn main() -> Result<(), Box> { .map_err(|e| e.to_string())?; } left.merge_in_place(&right).map_err(|e| e.to_string())?; - bytes = bytes.max(left.serialize_to_bytes().len()); + bytes = bytes.max(left.to_bytes().map_err(|e| e.to_string())?.len()); for (i, stat) in [ - asap_types::Statistic::Cardinality, - asap_types::Statistic::FrequencyL2, - asap_types::Statistic::FrequencyEntropy, + planner_types::post_asap::SketchQuery::Cardinality, + planner_types::post_asap::SketchQuery::FrequencyL2, + planner_types::post_asap::SketchQuery::FrequencyEntropy, ] .into_iter() .enumerate() { - let answer = left - .query_statistic(stat, &None, &Default::default()) - .map_err(|e| e.to_string())?; + let answer = left.estimate(&stat).map_err(|e| e.to_string())?; if !answer.is_finite() { return Err("nonfinite estimate".into()); } @@ -171,15 +169,11 @@ fn main() -> Result<(), Box> { .downcast_ref::() .ok_or("HLL merge type")?; let estimate = merged - .query_statistic( - asap_types::Statistic::Cardinality, - &None, - &Default::default(), - ) + .estimate(&planner_types::post_asap::SketchQuery::Cardinality) .map_err(|e| e.to_string())?; max_error = max_error.max((estimate - distinct.len() as f64).abs() / distinct.len() as f64); - bytes = bytes.max(merged.serialize_to_bytes().len()); + bytes = bytes.max((merged as &dyn AggregateCore).serialize_to_bytes().len()); } records.push(json!({"id":format!("hll-p{precision}"),"sketch":"hll","implementation":"asap-sketchlib-hll-regular-v1", "parameters":{"precision":precision},"trials":populations.len(), diff --git a/data_plane/src/drivers/ingest/otel.rs b/data_plane/src/drivers/ingest/otel.rs index 92e470c41..1008da83d 100644 --- a/data_plane/src/drivers/ingest/otel.rs +++ b/data_plane/src/drivers/ingest/otel.rs @@ -36,7 +36,6 @@ use asap_otel_proto::tonic::common::v1::any_value::Value as AnyValueVariant; use asap_otel_proto::tonic::metrics::v1::number_data_point::Value as NumberValue; use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; use asap_sketchlib::MessagePackCodec; -use asap_summary_state::summary_kernels::sketch_envelope::SketchEnvelopeAccumulator; use axum::{body::Bytes, extract::State, routing::post, Json, Router}; use flate2::read::GzDecoder; use planner_types::post_asap::SketchAlgorithm; @@ -769,14 +768,10 @@ async fn route_otlp_to_precompute( continue; } let group_key = IngestState::extract_group_key_for(&series_key, config); - // Wrap the raw SketchEnvelope bytes in a SketchEnvelopeAccumulator - // so the precompute engine receives the opaque sketch as-is. This - // preserves all sketch state end-to-end; per-variant decoding - // (e.g. CountMin → CountMinSketchAccumulator) can layer on top - // later without changing the routing contract. + // Decode the envelope into the Planner kernel of its family. let accumulator: Box = - match SketchEnvelopeAccumulator::from_proto_bytes(point.payload.clone()) { - Ok(acc) => Box::new(acc), + match asap_summary_state::stored_state::codec::decode_envelope(&point.payload) { + Ok(state) => state, Err(e) => { warn!( "OTLP sketch decode failed for metric='{}' attr='{}': {}", @@ -1597,7 +1592,7 @@ async fn route_modified_otlp_sketches_to_precompute( new_window_start={}); rotating per-series base", metric.name, series_key, base_window_start, dp.start_time_unix_nano ); - merged.reset_to_empty(); + reset_to_empty(&mut merged); } if let Err(e) = apply_modified_otlp_delta_bytes( dp.algorithm.clone(), @@ -2110,9 +2105,8 @@ fn dp_carries_heap(dp: &ModifiedOtlpSketchDp) -> bool { .unwrap_or(false) } ENCODING_MSGPACK_DELTA => { - use asap_summary_state::summary_kernels::CountMinSketchWithHeapAccumulator; - CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes(&dp.sketch) - .map(|acc| !acc.inner.topk_heap_items().is_empty()) + asap_summary_state::stored_state::decoders::cms_with_heap_from_msgpack_delta(&dp.sketch) + .map(|heap| !heap.topk_heap_items().is_empty()) .unwrap_or(false) } _ => false, @@ -2311,7 +2305,7 @@ fn preflight_summary_frames( format!("delta frame for {metric_name} has no reconstructable base") })?; if base_window_start != dp.start_time_unix_nano { - base.reset_to_empty(); + reset_to_empty(&mut base); } apply_modified_otlp_delta_bytes( dp.algorithm.clone(), @@ -2539,10 +2533,8 @@ fn decode_modified_otlp_sketch_bytes( encoding: i32, bytes: &[u8], ) -> Result, Box> { - use asap_summary_state::summary_kernels::{ - CountMinSketchAccumulator, CountSketchAccumulator, DDSketchAccumulator, - DatasketchesKLLAccumulator, HllSketchAccumulator, - }; + use asap_physical_operators::summary_kernels as k; + use asap_summary_state::stored_state::decoders as d; // The encoding value is the raw i32 from the per-sketch encoding // enum. All five sketch variants share the same wire tag layout: @@ -2554,101 +2546,82 @@ fn decode_modified_otlp_sketch_bytes( // not standalone-decodable) // 3 — ENCODING_MSGPACK (full sketch-core msgpack state) // 4 — ENCODING_MSGPACK_DELTA (MSGPACK diff; not yet wired) - - match encoding { - ENCODING_PROTO => match algorithm { - // The neutral codec accepts both full envelopes and supported bare - // states. Query accumulators retain their family-specific readouts. - SketchAlgorithm::DDSketch => { - let (inner, sample_p) = asap_sketch_codec::reconstruct_ddsketch(bytes)?; - let sample_p = if sample_p.is_finite() && sample_p > 0.0 && sample_p < 1.0 { - sample_p - } else { - 1.0 - }; - Ok(Box::new(DDSketchAccumulator { inner, sample_p })) - } - SketchAlgorithm::Kll => Ok(Box::new( - DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(bytes)?, - )), - SketchAlgorithm::Cms => Ok(Box::new( - CountMinSketchAccumulator::from_sketchlib_proto_bytes(bytes)?, - )), - SketchAlgorithm::CountSketch => Ok(Box::new( - CountSketchAccumulator::from_sketchlib_proto_bytes(bytes)?, - )), - SketchAlgorithm::Hll => Ok(Box::new(HllSketchAccumulator::from_sketchlib_proto_bytes( - bytes, - )?)), - other => { - Err(format!("modified-OTLP PROTO decoding is not implemented for {other:?}").into()) - } - }, - ENCODING_MSGPACK => match algorithm { - SketchAlgorithm::Cms => Ok(Box::new(CountMinSketchAccumulator::from_msgpack_bytes( - bytes, - )?)), - SketchAlgorithm::CountSketch => { - // Heap-bearing CountSketch full frame: the bytes are the - // `{sketch,topk_heap,heap_size}` envelope (a DIFFERENT inner - // field order than the plain CountSketch msgpack), so - // `CountSketch::from_msgpack` can't parse it. Try the heap - // decode FIRST when the heap is non-empty (the same promotion - // gate `sketch_algorithm_for` uses); cache THAT heap - // accumulator as the per-series base so a later MSGPACK_DELTA - // 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). - // - // 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 asap_summary_state::summary_kernels::CountSketchWithHeapAccumulator; - return Ok(Box::new( - CountSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(bytes)?, - )); - } + Ok(match (encoding, algorithm) { + (ENCODING_PROTO, SketchAlgorithm::DDSketch) => Box::new(k::DDSketchAccumulator { + inner: d::ddsketch_from_proto(bytes)?, + }), + (ENCODING_PROTO, SketchAlgorithm::Kll) => Box::new(k::DatasketchesKLLAccumulator { + inner: d::kll_from_proto(bytes)?, + }), + (ENCODING_PROTO, SketchAlgorithm::Cms) => Box::new(k::CountMinSketchAccumulator { + inner: d::cms_from_proto(bytes)?, + }), + (ENCODING_PROTO, SketchAlgorithm::CountSketch) => Box::new(k::CountSketchAccumulator { + inner: d::cs_from_proto(bytes)?, + }), + (ENCODING_PROTO, SketchAlgorithm::Hll) => Box::new(k::HllSketchAccumulator { + inner: d::hll_from_proto(bytes)?, + }), + (ENCODING_MSGPACK, SketchAlgorithm::Cms) => Box::new(k::CountMinSketchAccumulator { + inner: d::cms_from_msgpack(bytes)?, + }), + (ENCODING_MSGPACK, SketchAlgorithm::CountSketch) => { + // A heap-bearing CountSketch full frame is the + // `{sketch,topk_heap,heap_size}` envelope, which the plain + // CountSketch decoder cannot parse. A non-empty heap selects the + // heap state (median-of-signed-rows, never the Count-Min heap), + // so a later MSGPACK_DELTA frame applies onto a heap base. + match d::cs_with_heap_from_msgpack(bytes) { + Ok(heap) if !heap.topk_heap_items().is_empty() => { + Box::new(k::CountSketchWithHeapAccumulator { inner: heap }) } - Ok(Box::new(CountSketchAccumulator::from_msgpack_bytes(bytes)?)) - } - SketchAlgorithm::Kll => Ok(Box::new(DatasketchesKLLAccumulator::from_msgpack_bytes( - bytes, - )?)), - SketchAlgorithm::DDSketch => { - Ok(Box::new(DDSketchAccumulator::from_msgpack_bytes(bytes)?)) + _ => Box::new(k::CountSketchAccumulator { + inner: d::cs_from_msgpack(bytes)?, + }), } - SketchAlgorithm::Hll => Ok(Box::new(HllSketchAccumulator::from_msgpack_bytes(bytes)?)), - other => Err(format!( - "modified-OTLP MSGPACK decoding is not implemented for {other:?}" + } + (ENCODING_MSGPACK, SketchAlgorithm::Kll) => Box::new(k::DatasketchesKLLAccumulator { + inner: d::kll_from_msgpack(bytes)?, + }), + (ENCODING_MSGPACK, SketchAlgorithm::DDSketch) => Box::new(k::DDSketchAccumulator { + inner: d::ddsketch_from_msgpack(bytes)?, + }), + (ENCODING_MSGPACK, SketchAlgorithm::Hll) => Box::new(k::HllSketchAccumulator { + inner: d::hll_from_msgpack(bytes)?, + }), + (ENCODING_PROTO, other) => { + return Err( + format!("modified-OTLP PROTO decoding is not implemented for {other:?}").into(), ) - .into()), - }, - ENCODING_PROTO_DELTA => Err(format!( - "sketch encoding PROTO_DELTA (2) is not standalone-decodable — \ + } + (ENCODING_MSGPACK, other) => { + return Err( + format!("modified-OTLP MSGPACK decoding is not implemented for {other:?}").into(), + ) + } + (ENCODING_PROTO_DELTA, _) => { + return Err( + "sketch encoding PROTO_DELTA (2) is not standalone-decodable — \ it carries only a diff against the caller's base snapshot. \ Caller must route these through \ `apply_modified_otlp_delta_bytes` with a cached accumulator; \ this decoder is for full-state frames only." - ) - .into()), - ENCODING_MSGPACK_DELTA => Err(format!( - "sketch encoding MSGPACK_DELTA (4) deferred — PR G wires \ + .into(), + ) + } + (ENCODING_MSGPACK_DELTA, _) => { + return Err("sketch encoding MSGPACK_DELTA (4) deferred — PR G wires \ PROTO_DELTA only; msgpack delta is a follow-up." - ) - .into()), - _ => Err(format!( - "unknown modified-OTLP sketch encoding {encoding} \ - (expected 1 / 2 / 3 / 4)" - ) - .into()), - } + .into()) + } + _ => { + return Err(format!( + "unknown modified-OTLP sketch encoding {encoding} \ + (expected 1 / 2 / 3 / 4)" + ) + .into()) + } + }) } /// P1-1/P1-2 — construct an EMPTY accumulator matching a sketch @@ -2676,41 +2649,32 @@ fn empty_accumulator_for_delta_bootstrap( encoding: i32, ) -> Option> { use crate::storage_engines::sketch_db::index::SketchConfig; - use asap_summary_state::summary_kernels::{ - CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, - HllSketchAccumulator, - }; + use asap_physical_operators::summary_kernels as k; match (algorithm, config) { (SketchAlgorithm::Hll, SketchConfig::Hll { precision }) => { - use asap_sketchlib::HllVariant; // Regular is the default agent variant; HLL's additive delta // merge tolerates an empty same-precision base. - Some(Box::new(HllSketchAccumulator::new( - HllVariant::Regular, + Some(Box::new(k::HllSketchAccumulator::new( + asap_sketchlib::HllVariant::Regular, *precision, ))) } (SketchAlgorithm::Cms, SketchConfig::CountMin { rows, cols }) => Some(Box::new( - CountMinSketchAccumulator::new(*rows as usize, *cols as usize), + k::CountMinSketchAccumulator::new(*rows as usize, *cols as usize), )), (SketchAlgorithm::CountSketch, SketchConfig::CountSketch { rows, cols }) => { - // A heap-bearing DELTA-HEAP frame must reconstruct onto a heap - // accumulator (the apply path downcasts to - // `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. + // A DELTA-HEAP frame reconstructs onto a heap base; a plain matrix + // delta onto a heap-less CountSketch. heap_size 0 is fine: the + // DELTA-HEAP apply replaces the heap wholesale. 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(CountSketchWithHeapAccumulator::new( + Some(Box::new(k::CountSketchWithHeapAccumulator::new( *rows as usize, *cols as usize, 0, ))) } else { - Some(Box::new(CountSketchAccumulator::new( + Some(Box::new(k::CountSketchAccumulator::new( *rows as usize, *cols as usize, ))) @@ -2749,79 +2713,44 @@ pub(crate) fn apply_modified_otlp_delta_bytes( existing: &mut Box, bytes: &[u8], ) -> Result<(), Box> { - use asap_summary_state::summary_kernels::{ - CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, - DDSketchAccumulator, HllSketchAccumulator, - }; + use asap_physical_operators::summary_kernels as k; + use asap_summary_state::stored_state::decoders as d; match (encoding, algorithm) { - (ENCODING_PROTO_DELTA, SketchAlgorithm::DDSketch) => { - let dd = existing - .as_any_mut() - .downcast_mut::() - .ok_or( - "apply_modified_otlp_delta_bytes: existing accumulator is \ - not a DDSketchAccumulator", - )?; - dd.apply_proto_delta_bytes(bytes) - } - (ENCODING_PROTO_DELTA, SketchAlgorithm::Hll) => { - let hll = existing - .as_any_mut() - .downcast_mut::() - .ok_or( - "apply_modified_otlp_delta_bytes: existing accumulator is \ - not an HllSketchAccumulator", - )?; - hll.apply_proto_delta_bytes(bytes) - } - (ENCODING_PROTO_DELTA, SketchAlgorithm::CountSketch) => { - let cs = existing - .as_any_mut() - .downcast_mut::() - .ok_or( - "apply_modified_otlp_delta_bytes: existing accumulator is \ - not a CountSketchAccumulator", - )?; - cs.apply_proto_delta_bytes(bytes) - } - (ENCODING_PROTO_DELTA, SketchAlgorithm::Cms) => { - let cms = existing - .as_any_mut() - .downcast_mut::() - .ok_or( - "apply_modified_otlp_delta_bytes: existing accumulator is \ - not a CountMinSketchAccumulator", - )?; - cms.apply_proto_delta_bytes(bytes) - } + (ENCODING_PROTO_DELTA, SketchAlgorithm::DDSketch) => edit( + existing, + "DDSketchAccumulator", + |s: &mut k::DDSketchAccumulator| d::apply_ddsketch_proto_delta(&mut s.inner, bytes), + ), + (ENCODING_PROTO_DELTA, SketchAlgorithm::Hll) => edit( + existing, + "HllSketchAccumulator", + |s: &mut k::HllSketchAccumulator| d::apply_hll_proto_delta(&mut s.inner, bytes), + ), + (ENCODING_PROTO_DELTA, SketchAlgorithm::CountSketch) => edit( + existing, + "CountSketchAccumulator", + |s: &mut k::CountSketchAccumulator| d::apply_cs_proto_delta(&mut s.inner, bytes), + ), + (ENCODING_PROTO_DELTA, SketchAlgorithm::Cms) => edit( + existing, + "CountMinSketchAccumulator", + |s: &mut k::CountMinSketchAccumulator| d::apply_cms_proto_delta(&mut s.inner, bytes), + ), (ENCODING_PROTO_DELTA, other) => Err(format!( "PROTO_DELTA for sketch kind {other:?} is not yet supported; \ DDSketch / HLL / CountSketch / CountMin are wired" ) .into()), - (ENCODING_MSGPACK_DELTA, SketchAlgorithm::CountSketch) => { - // DELTA-HEAP frame for the heap-bearing CountSketch: a sparse - // signed matrix delta + the full top-k heap. The cached base is - // a heap accumulator (window-1 full frame decoded via - // `from_msgpack_with_heap_bytes`); under the per-window-reset - // model the ingest caller has already reset it to empty at a - // 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). Real `CountSketchWithHeapAccumulator` - // (median-of-signed-rows), not the CMS-family wrapper. - let heap = existing - .as_any_mut() - .downcast_mut::() - .ok_or( - "apply_modified_otlp_delta_bytes: existing accumulator is \ - not a CountSketchWithHeapAccumulator (heap-bearing \ - CountSketch delta requires a heap base — the window-1 \ - full frame must have promoted the sid)", - )?; - heap.apply_msgpack_heap_delta_bytes(bytes) - } + // DELTA-HEAP frame for the heap-bearing CountSketch: a sparse signed + // matrix delta plus the full top-k heap, applied onto a heap base + // (already reset at a window boundary by the caller). + (ENCODING_MSGPACK_DELTA, SketchAlgorithm::CountSketch) => edit( + existing, + "CountSketchWithHeapAccumulator (heap-bearing CountSketch delta requires a \ + heap base — the window-1 full frame must have promoted the sid)", + |s: &mut k::CountSketchWithHeapAccumulator| d::apply_cs_heap_delta(&mut s.inner, bytes), + ), (ENCODING_MSGPACK_DELTA, other) => Err(format!( "MSGPACK_DELTA for sketch kind {other:?} is not yet wired; only \ the heap-bearing CountSketch DELTA-HEAP frame is supported" @@ -2837,6 +2766,33 @@ pub(crate) fn apply_modified_otlp_delta_bytes( } } +/// Apply `update` to the cached base of kernel type `T`. Planner kernels +/// expose no mutable downcast, so the base is copied, updated and replaced. +fn edit( + existing: &mut Box, + expected: &str, + update: impl FnOnce(&mut T) -> Result<(), String>, +) -> Result<(), Box> { + let mut state = existing + .as_any() + .downcast_ref::() + .ok_or_else(|| { + format!("apply_modified_otlp_delta_bytes: existing accumulator is not a {expected}") + })? + .clone(); + update(&mut state)?; + *existing = Box::new(state); + Ok(()) +} + +/// Rotate a cached delta base to an empty state of the same shape; families +/// that never delta keep their base. +fn reset_to_empty(state: &mut Box) { + if let Ok(empty) = asap_summary_state::stored_state::codec::empty_like(state.as_ref()) { + *state = empty; + } +} + /// Identify the concrete sketch type inside a `SketchEnvelope` payload, /// returning a human-readable name for logging. Returns `"Unknown"` if the /// payload does not decode or the `sketch_state` variant is unset. @@ -2995,13 +2951,16 @@ fn otlp_to_metric_points_and_sketches(request: &ExportMetricsServiceRequest) -> // ExactAgg(Sum) path as a plain delta Sum — the backend sums // the per-window/per-shard partials for the same sid. for dp in &sa.data_points { - let value = match asap_summary_state::summary_kernels::sum::SumAccumulator::from_sum_bytes(&dp.sketch) { - Ok(acc) => acc.sum, - Err(e) => { - debug!("asap_edge: SumAgg data point decode failed (skipping): {e}"); - continue; - } - }; + let value = + match asap_summary_state::stored_state::decoders::sum_payload( + &dp.sketch, + ) { + Ok(sum) => sum, + Err(e) => { + debug!("asap_edge: SumAgg data point decode failed (skipping): {e}"); + continue; + } + }; let labels = merge_point_attributes(&base_labels, &dp.attributes); points.push(MetricPoint { name: metric.name.clone(), @@ -3436,9 +3395,9 @@ mod policy_fp_lookup_tests { mod dispatcher_tests { use super::*; use crate::storage_engines::types::AggregateCore; + use asap_physical_operators::summary_kernels::{DDSketchAccumulator, HllSketchAccumulator}; use asap_sketchlib::DdSketch; use asap_sketchlib::HllVariant; - use asap_summary_state::summary_kernels::{DDSketchAccumulator, HllSketchAccumulator}; #[test] fn apply_modified_otlp_delta_bytes_ddsketch_round_trip() { @@ -3448,7 +3407,6 @@ mod dispatcher_tests { // Base sketch represents the last full snapshot the agent sent. let mut acc: Box = Box::new(DDSketchAccumulator { inner: DdSketch::from_raw(0.01, vec![1, 2, 3], 0), - sample_p: 1.0, }); // The wire delta now carries only bucket deltas (tags 2-7 @@ -3486,13 +3444,9 @@ mod dispatcher_tests { use asap_otel_proto::sketchlib::v1::HllDelta as PbDelta; use prost::Message; - let mut acc: Box = - Box::new(HllSketchAccumulator::new(HllVariant::Regular, 2)); - acc.as_any_mut() - .downcast_mut::() - .unwrap() - .inner - .registers = vec![1, 5, 3, 7]; + let mut base = HllSketchAccumulator::new(HllVariant::Regular, 2); + base.inner.registers = vec![1, 5, 3, 7]; + let mut acc: Box = Box::new(base); // Packed (index_delta, value) blob for updates {0:4, 2:6}. let bytes = PbDelta { @@ -3789,8 +3743,8 @@ mod sid_resolution_tests { #[tokio::test] async fn delta_apply_rotates_per_series_base_at_window_boundary() { use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use asap_physical_operators::summary_kernels::DDSketchAccumulator; use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; - use asap_summary_state::summary_kernels::DDSketchAccumulator; use prost::Message; let (state, drain) = make_state().await; @@ -4122,7 +4076,7 @@ mod sid_resolution_tests { #[tokio::test] async fn leading_cms_delta_bootstraps_onto_empty_base() { use asap_otel_proto::sketchlib::v1::CountMinDelta as PbDelta; - use asap_summary_state::summary_kernels::CountMinSketchAccumulator; + use asap_physical_operators::summary_kernels::CountMinSketchAccumulator; use prost::Message; let (state, drain) = make_state().await; @@ -4208,7 +4162,7 @@ mod sid_resolution_tests { #[tokio::test] async fn leading_hll_delta_bootstraps_onto_empty_base() { use asap_otel_proto::sketchlib::v1::HllDelta as PbDelta; - use asap_summary_state::summary_kernels::HllSketchAccumulator; + use asap_physical_operators::summary_kernels::HllSketchAccumulator; use prost::Message; let (state, drain) = make_state().await; diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index c7d3106df..94ce75212 100644 --- a/data_plane/src/lib.rs +++ b/data_plane/src/lib.rs @@ -37,13 +37,8 @@ pub mod utils; // Re-export commonly used types to avoid glob import conflicts pub use storage_engines::types::{ - AggregateCore, KeyByLabelValues, Measurement, MergeableAccumulator, - MultipleSubpopulationAggregate, PrecomputeMaterialization, PrecomputedOutput, - SerializableToSink, SingleSubpopulationAggregate, -}; - -pub use asap_summary_state::summary_kernels::{ - IncreaseAccumulator, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, + AggregateCore, KeyByLabelValues, Measurement, PrecomputeMaterialization, PrecomputedOutput, + StoredState, }; pub use storage_engines::StoreResult; diff --git a/data_plane/src/precompute_engine/config.rs b/data_plane/src/precompute_engine/config.rs index df6612be9..cdaf010f1 100644 --- a/data_plane/src/precompute_engine/config.rs +++ b/data_plane/src/precompute_engine/config.rs @@ -24,7 +24,7 @@ pub struct PrecomputeEngineConfig { /// Capacity of the MPSC channel between router and each worker. pub channel_buffer_size: usize, /// When true, skip all aggregation and pass each raw sample directly to the - /// output sink as a `SumAccumulator::with_sum(value)`. + /// output sink as an exact Sum state of that value. pub pass_raw_samples: bool, /// Aggregation ID to stamp on each raw-mode output. pub raw_mode_aggregation_id: u64, diff --git a/data_plane/src/precompute_engine/ingest_handler.rs b/data_plane/src/precompute_engine/ingest_handler.rs index 3d20001f4..a60e9fbd2 100644 --- a/data_plane/src/precompute_engine/ingest_handler.rs +++ b/data_plane/src/precompute_engine/ingest_handler.rs @@ -367,8 +367,8 @@ mod tests { async fn delta_path_reconstitutes_cumulative_state() { use crate::drivers::ingest::otel::apply_modified_otlp_delta_bytes; use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use asap_physical_operators::summary_kernels::DDSketchAccumulator; use asap_sketchlib::DdSketch; - use asap_summary_state::summary_kernels::DDSketchAccumulator; use planner_types::post_asap::SketchAlgorithm; use prost::Message; @@ -382,7 +382,6 @@ mod tests { let series_key = "__name__=latency_ms,inst=a"; let base = DDSketchAccumulator { inner: DdSketch::from_raw(0.01, vec![1, 2, 3], 0), - sample_p: 1.0, }; state.sketch_snapshots.insert( series_key.to_string(), @@ -472,8 +471,6 @@ mod tests { /// survive; a stale entry from far in the past must be swept. #[tokio::test] async fn stale_snapshot_entry_is_evicted_by_sweep() { - use asap_summary_state::summary_kernels::SumAccumulator; - let (state, drain) = setup_state(7, "evict_metric").await; // Pin a deterministic lag of 100ns so the test doesn't depend on @@ -489,14 +486,14 @@ mod tests { state.sketch_snapshots.insert( "stale".to_string(), SnapshotCacheEntry { - core: Box::new(SumAccumulator::with_sum(1.0)), + core: Box::new(crate::tests::accumulator_fixture::sum_state(1.0)), window_start: 10, }, ); state.sketch_snapshots.insert( "fresh".to_string(), SnapshotCacheEntry { - core: Box::new(SumAccumulator::with_sum(2.0)), + core: Box::new(crate::tests::accumulator_fixture::sum_state(2.0)), window_start: 1000, }, ); @@ -524,7 +521,7 @@ mod tests { state.sketch_snapshots.insert( "ancient".to_string(), SnapshotCacheEntry { - core: Box::new(SumAccumulator::with_sum(3.0)), + core: Box::new(crate::tests::accumulator_fixture::sum_state(3.0)), window_start: 1, }, ); diff --git a/data_plane/src/precompute_engine/maintenance_runtime.rs b/data_plane/src/precompute_engine/maintenance_runtime.rs index 560c5f627..fe42f2d2f 100644 --- a/data_plane/src/precompute_engine/maintenance_runtime.rs +++ b/data_plane/src/precompute_engine/maintenance_runtime.rs @@ -1,5 +1,6 @@ //! Bind immutable inputs to retained physical graphs and publish their stored outputs. use asap_summary_state::codec::KeyCodec; +use asap_summary_state::StoredState; #[cfg(test)] use asap_types::physical_plan_codec::PhysicalPlanCodec; @@ -1563,7 +1564,6 @@ pub(crate) fn affected_materializations( mod tests { use super::*; - use asap_summary_state::summary_kernels::SumAccumulator; use planner_types::post_asap::{ EdgeRole, GroupingEdgeCompatibility, PostAsapDag, PostAsapDagEdge, SummarySchema, WindowEdgeCompatibility, @@ -1619,8 +1619,8 @@ mod tests { fn cohort_lineage_is_order_independent_and_binds_every_input() { use crate::storage_engines::sketch_db::index::FrozenExactWindows; let make = |sid, id, value| { - let mut state = asap_summary_state::summary_kernels::SumAccumulator::new(); - state.update(value); + let mut state = crate::tests::accumulator_fixture::sum_state(0.0); + state.update(None, value, 0); FrozenExactWindows { stored_output_reference: asap_types::sds::StoredOutputReference::for_output( definition(id), @@ -1721,8 +1721,8 @@ mod tests { } fn sum(value: f64) -> SummaryState { - let mut accumulator = SumAccumulator::new(); - accumulator.update(value); + let mut accumulator = crate::tests::accumulator_fixture::sum_state(0.0); + accumulator.update(None, value, 0); Arc::new(accumulator) } @@ -1806,7 +1806,6 @@ mod tests { reduction: Reduction::by(vec![]), grouping: GroupingStrategy::default(), }; - let kwargs = std::collections::HashMap::from([("quantile".into(), "0.5".into())]); let mut source_node = node(1); source_node.output_schema.fields = vec![SummaryField { name: "state".into(), @@ -2056,7 +2055,7 @@ mod tests { .unwrap()); assert_eq!( result - .query_statistic(asap_types::Statistic::Quantile, &None, &kwargs) + .estimate(&planner_types::post_asap::SketchQuery::Quantile { q: 0.5 }) .unwrap(), 3.0 ); @@ -2461,11 +2460,7 @@ mod tests { assert!(group.is_empty()); assert_eq!( state - .query_statistic( - asap_types::Statistic::Quantile, - &None, - &std::collections::HashMap::from([("quantile".into(), "1.0".into())]) - ) + .estimate(&planner_types::post_asap::SketchQuery::Quantile { q: 1.0 }) .unwrap(), 29.0 ); @@ -2521,11 +2516,7 @@ mod tests { .unwrap(); assert_eq!( result - .query_statistic( - asap_types::Statistic::Quantile, - &None, - &std::collections::HashMap::from([("quantile".into(), "0.5".into())]) - ) + .estimate(&planner_types::post_asap::SketchQuery::Quantile { q: 0.5 }) .unwrap(), 9.0 ); diff --git a/data_plane/src/precompute_engine/native_precompute.rs b/data_plane/src/precompute_engine/native_precompute.rs index 9db27c275..78aced875 100644 --- a/data_plane/src/precompute_engine/native_precompute.rs +++ b/data_plane/src/precompute_engine/native_precompute.rs @@ -73,7 +73,7 @@ pub(super) fn execute( ), Value::Summary { family: family.clone(), - state: asap_summary_state::physical::to_physical(state.as_ref())?, + state: Arc::clone(state), }, ]); } @@ -93,8 +93,7 @@ pub(super) fn execute( _, ) => Ok(Value::Summary { family: field.dtype.clone(), - state: asap_summary_state::physical::to_physical(state.as_ref()) - .map_err(|e| e.to_string())?, + state: Arc::clone(state), }), SummaryFamilyType::Plain(DataType::Timestamp) => Ok(Value::Timestamp( i64::try_from(window.1).map_err(|_| "native window overflow")?, @@ -216,8 +215,8 @@ pub(super) fn population_states( return Err("duplicate population label".into()); } } - let state = asap_summary_state::physical::from_physical(state.as_ref())?; - if result.insert(group, Arc::from(state)).is_some() { + asap_summary_state::stored_state::codec::check_storable(state.as_ref())?; + if result.insert(group, Arc::clone(state)).is_some() { return Err("repeated precompute output population".into()); } } @@ -297,8 +296,8 @@ mod tests { snapshot_sha256: "0".repeat(64), }); let state = |value| { - let mut sum = asap_summary_state::summary_kernels::SumAccumulator::new(); - sum.update(value); + let mut sum = crate::tests::accumulator_fixture::sum_state(0.0); + sum.update(None, value, 0); Arc::new(sum) as Arc }; for (revision, value) in [(1, 2.0), (2, 7.0)] { @@ -335,10 +334,7 @@ mod tests { "shared merge must execute once, not once per output" ); assert_eq!( - asap_summary_state::physical::from_physical(state_at(3).as_ref()) - .unwrap() - .query_statistic(asap_types::Statistic::Sum, &None, &Default::default(),) - .unwrap(), + crate::tests::accumulator_fixture::sum_of(state_at(3).as_ref()), value + 3.0 ); let error = execute(&installed, &program, &inputs, (0, 2000), 1, revision).unwrap_err(); diff --git a/data_plane/src/precompute_engine/output_sink.rs b/data_plane/src/precompute_engine/output_sink.rs index 4b05edd1a..2a294ed9f 100644 --- a/data_plane/src/precompute_engine/output_sink.rs +++ b/data_plane/src/precompute_engine/output_sink.rs @@ -325,7 +325,7 @@ mod tests { use super::*; use crate::storage_engines::sketch_db::index::{AggKind, SeriesLookup}; use crate::storage_engines::types::{InstalledPrecomputePlan, KeyByLabelValues}; - use asap_summary_state::summary_kernels::{DDSketchAccumulator, SumAccumulator}; + use asap_physical_operators::summary_kernels::DDSketchAccumulator; use asap_types::aggregation_config::PrecomputeMaterialization; use asap_types::enums::WindowKind; use asap_types::AggregationType; @@ -428,7 +428,8 @@ mod tests { let key = KeyByLabelValues::new_with_labels(vec!["z0".to_string()]); let output = PrecomputedOutput::new(1000, 2000, Some(key), asap_types::PolicyFingerprint(agg_id)); - let acc: Box = Box::new(SumAccumulator::with_sum(42.0)); + let acc: Box = + Box::new(crate::tests::accumulator_fixture::sum_state(42.0)); sink.emit_batch(vec![(output, acc)]).expect("emit ok"); @@ -495,12 +496,18 @@ mod tests { output.catalog_generation = Some(Arc::clone(&original_generation)); output }; - sink.emit_batch(vec![(output(), Box::new(SumAccumulator::with_sum(7.0)))]) - .unwrap(); + sink.emit_batch(vec![( + output(), + Box::new(crate::tests::accumulator_fixture::sum_state(7.0)), + )]) + .unwrap(); let old_sid = store.series_ids_for_policy(fingerprint)[0]; store.remove_instance(old_sid).unwrap(); assert!(sink - .emit_batch(vec![(output(), Box::new(SumAccumulator::with_sum(11.0)))]) + .emit_batch(vec![( + output(), + Box::new(crate::tests::accumulator_fixture::sum_state(11.0)) + )]) .is_err()); let mut stale_output = output(); stale_output.storage_handle = Some(old_sid); @@ -512,33 +519,36 @@ mod tests { assert!(sink .emit_batch(vec![( stale_output.clone(), - Box::new(SumAccumulator::with_sum(99.0)) + Box::new(crate::tests::accumulator_fixture::sum_state(99.0)) )]) .is_err()); let mut next_output = output(); next_output.catalog_generation = Some(next_generation); sink.emit_batch(vec![( next_output, - Box::new(SumAccumulator::with_sum(11.0)), + Box::new(crate::tests::accumulator_fixture::sum_state(11.0)), )]) .unwrap(); assert!(sink .emit_batch(vec![( stale_output, - Box::new(SumAccumulator::with_sum(99.0)) + Box::new(crate::tests::accumulator_fixture::sum_state(99.0)) )]) .is_err()); let new_sid = store.series_ids_for_policy(fingerprint)[0]; // A derived/unbound stale output must not reuse an already rotated cache hit. assert!(sink - .emit_batch(vec![(output(), Box::new(SumAccumulator::with_sum(101.0)))]) + .emit_batch(vec![( + output(), + Box::new(crate::tests::accumulator_fixture::sum_state(101.0)) + )]) .is_err()); let mut stale_routed_output = output(); stale_routed_output.storage_handle = Some(new_sid); assert!(sink .emit_batch(vec![( stale_routed_output, - Box::new(SumAccumulator::with_sum(103.0)) + Box::new(crate::tests::accumulator_fixture::sum_state(103.0)) )]) .is_err()); let mut missing_generation = PrecomputedOutput::new(1000, 2000, None, fingerprint); @@ -546,14 +556,16 @@ mod tests { assert!(sink .emit_batch(vec![( missing_generation, - Box::new(SumAccumulator::with_sum(107.0)) + Box::new(crate::tests::accumulator_fixture::sum_state(107.0)) )]) .is_err()); assert_ne!(old_sid, new_sid); assert!(store.query_exact_agg_range(old_sid, 1000, 2000).is_empty()); let values = store.query_exact_agg_range(new_sid, 1000, 2000); assert_eq!( - values[0].1.values().next().unwrap().aux_stats().sum, + Some(crate::tests::accumulator_fixture::sum_of( + values[0].1.values().next().unwrap().as_ref() + )), Some(11.0) ); } @@ -620,7 +632,8 @@ mod tests { ); let output = PrecomputedOutput::new(1000, 2000, None, asap_types::PolicyFingerprint(99)); - let acc: Box = Box::new(SumAccumulator::with_sum(1.0)); + let acc: Box = + Box::new(crate::tests::accumulator_fixture::sum_state(1.0)); sink.emit_batch(vec![(output, acc)]) .expect_err("unpersisted output must not be acknowledged"); assert_eq!(summary_store.instance_count(), 0); @@ -644,7 +657,8 @@ mod tests { // policy_fp=42 is absent from the empty registry → registry miss. let output = PrecomputedOutput::new(1000, 2000, None, asap_types::PolicyFingerprint(42)); - let acc: Box = Box::new(SumAccumulator::with_sum(1.0)); + let acc: Box = + Box::new(crate::tests::accumulator_fixture::sum_state(1.0)); sink.emit_batch(vec![(output, acc)]) .expect_err("unpersisted output must not be acknowledged"); @@ -663,7 +677,8 @@ mod tests { // The UNSET sentinel is an expected raw-mode skip, NOT a policy // miss — it must not bump the counter. let unset = PrecomputedOutput::new(1000, 2000, None, asap_types::PolicyFingerprint::UNSET); - let acc2: Box = Box::new(SumAccumulator::with_sum(1.0)); + let acc2: Box = + Box::new(crate::tests::accumulator_fixture::sum_state(1.0)); sink.emit_batch(vec![(unset, acc2)]) .expect_err("unpersisted output must not be acknowledged"); assert_eq!( diff --git a/data_plane/src/precompute_engine/raw_dag.rs b/data_plane/src/precompute_engine/raw_dag.rs index e920bcfdc..fb46b81b2 100644 --- a/data_plane/src/precompute_engine/raw_dag.rs +++ b/data_plane/src/precompute_engine/raw_dag.rs @@ -2,13 +2,13 @@ //! The backend supplies one typed sample batch per pane; the Planner-compiled //! precompute graph owns every update, grouping and item computation. use crate::storage_engines::types::AggregateCore; +use asap_physical_operators::factory::create_planner_accumulator; use asap_physical_operators::{ operators::Operator, physical_planner::{precompute, CompiledPhysicalDag, Source as PhysicalSource}, runtime::{Limits, RunContext, Scope}, values::{Batch, Value}, }; -use asap_summary_state::factory::create_planner_accumulator; use asap_types::physical_plan_codec::PhysicalPlanCodec; use asap_types::{executable_plan::BackendNodeBinding, PrecomputeMaterialization}; use planner_types::post_asap::{ @@ -234,9 +234,10 @@ impl RawDagProgram { match self.execute(samples, pane, max_bytes)?.as_slice() { [] => Ok(None), [row] => match row.as_slice() { - [_, _, Value::Summary { state, .. }] => Ok(Some( - asap_summary_state::physical::from_physical(state.as_ref())?, - )), + [_, _, Value::Summary { state, .. }] => { + asap_summary_state::stored_state::codec::check_storable(state.as_ref())?; + Ok(Some(state.clone_boxed_core())) + } _ => Err("raw precompute output is not a population state".into()), }, _ => Err("one routed group produced several populations".into()), @@ -264,8 +265,8 @@ impl RawDagProgram { /// The family's empty state, for a pane known to have no samples. Heaps /// are Planner weighted-frequency states, as `build` produces. pub fn empty_state(&self) -> Result, String> { - use asap_summary_state::summary_kernels::weighted_frequency::{ - FrequencyAlgorithm, PhysicalWeightedFrequency, WeightedFrequency, + use asap_physical_operators::summary_kernels::weighted_frequency::{ + FrequencyAlgorithm, WeightedFrequency, }; use planner_types::post_asap::SketchParams; if let SummaryFamilyType::Sketch(kind, _) = &self.family { @@ -283,14 +284,32 @@ impl RawDagProgram { _ => None, }; if let Some((algorithm, width, depth, heap_size)) = heap { - let state = PhysicalWeightedFrequency::new( + let state = WeightedFrequency::new( algorithm, *width as usize, *depth as usize, *heap_size as usize, ) .map_err(|e| e.to_string())?; - return Ok(Box::new(WeightedFrequency(state))); + return Ok(Box::new(state)); + } + // Planner's UnivMon has no stored codec; the store keeps the + // backend UnivMon shim. + if let SketchParams::UnivMon { + heap_size, + sketch_rows, + sketch_cols, + layers, + } = kind.params() + { + return asap_summary_state::univmon::UnivMonAccumulator::new( + *heap_size as usize, + *sketch_rows as usize, + *sketch_cols as usize, + *layers as usize, + ) + .map(|state| Box::new(state) as Box) + .map_err(|e| e.to_string()); } } Ok( @@ -432,7 +451,7 @@ mod tests { source: 0, }; let empty = program.empty_state().unwrap(); - assert!(asap_summary_state::physical::to_physical(empty.as_ref()).is_ok()); + assert!(asap_summary_state::stored_state::codec::check_storable(empty.as_ref()).is_ok()); assert_eq!( asap_summary_state::stored_state::SketchEncoding::full_frame_for(empty.as_ref()), asap_summary_state::stored_state::SketchEncoding::WeightedFrequencyV1 diff --git a/data_plane/src/precompute_engine/revisions.rs b/data_plane/src/precompute_engine/revisions.rs index 3b2521329..9aecb674a 100644 --- a/data_plane/src/precompute_engine/revisions.rs +++ b/data_plane/src/precompute_engine/revisions.rs @@ -452,10 +452,7 @@ pub(crate) fn encode_state( ) -> Result, RevisionError> { Ok(native::encode_batch(&Batch::try_new( state_schema(family.clone()), - vec![vec![Value::Summary { - family, - state: asap_summary_state::physical::to_physical(state.as_ref())?, - }]], + vec![vec![Value::Summary { family, state }]], )?)?) } @@ -467,9 +464,10 @@ pub(crate) fn decode_state( let batch = native::decode_batch(&record.payload, state_schema(family), record.payload.len())?; match batch.rows() { [row] => match row.as_slice() { - [Value::Summary { state, .. }] => Ok(Arc::from( - asap_summary_state::physical::from_physical(state.as_ref())?, - )), + [Value::Summary { state, .. }] => { + asap_summary_state::stored_state::codec::check_storable(state.as_ref())?; + Ok(Arc::clone(state)) + } _ => Err("revision record must contain exactly one typed summary".into()), }, _ => Err("revision record must contain exactly one row".into()), diff --git a/data_plane/src/precompute_engine/series_router.rs b/data_plane/src/precompute_engine/series_router.rs index 0334f5fb2..9d2442e58 100644 --- a/data_plane/src/precompute_engine/series_router.rs +++ b/data_plane/src/precompute_engine/series_router.rs @@ -1,5 +1,6 @@ use crate::precompute_engine::group_key::GroupKey; use crate::storage_engines::types::AggregateCore; +use asap_summary_state::StoredState; use asap_types::PolicyFingerprint; use futures::future::try_join_all; use std::collections::HashMap; diff --git a/data_plane/src/precompute_engine/worker.rs b/data_plane/src/precompute_engine/worker.rs index 29a3b8bb5..010fa9111 100644 --- a/data_plane/src/precompute_engine/worker.rs +++ b/data_plane/src/precompute_engine/worker.rs @@ -9,9 +9,9 @@ use crate::storage_engines::types::{ }; #[cfg(test)] use crate::tests::accumulator_fixture::create_fixture_accumulator; +use asap_physical_operators::factory::AccumulatorUpdater; use asap_summary_state::codec::KeyCodec; -use asap_summary_state::factory::AccumulatorUpdater; -use asap_summary_state::summary_kernels::sum::SumAccumulator; +use asap_summary_state::StoredState; use asap_types::aggregation_config::PrecomputeMaterialization; use asap_types::PolicyFingerprint; use asap_types::SampleUpdateRule; @@ -1016,7 +1016,7 @@ impl Worker { Ok(()) } - /// Raw fast-path: emit each sample as a standalone `SumAccumulator`. + /// Raw fast-path: emit each sample as a standalone exact Sum state. pub fn process_samples_raw( &self, series_key: &str, @@ -1031,7 +1031,10 @@ impl Worker { let output = PrecomputedOutput::new(ts as u64, ts as u64, None, PolicyFingerprint::UNSET); let _ = self.raw_mode_aggregation_id; - let accumulator = SumAccumulator::with_sum(val); + let accumulator = asap_summary_state::stored_state::codec::exact_value( + planner_types::post_asap::ExactKind::Sum, + val, + ); emit_batch.push((output, Box::new(accumulator))); } @@ -1759,8 +1762,14 @@ mod tests { let mut output = PrecomputedOutput::new(0, 1000, None, PolicyFingerprint(1)); output.input_revision = Some(revision); let merged = coalesce_admitted_outputs(vec![ - (output.clone(), Box::new(SumAccumulator::with_sum(2.0))), - (output, Box::new(SumAccumulator::with_sum(3.0))), + ( + output.clone(), + Box::new(crate::tests::accumulator_fixture::sum_state(2.0)), + ), + ( + output, + Box::new(crate::tests::accumulator_fixture::sum_state(3.0)), + ), ]) .unwrap(); assert_eq!(merged.len(), 1); @@ -1768,7 +1777,8 @@ mod tests { merged[0] .1 .as_any() - .downcast_ref::() + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) .unwrap() .sum, 5.0 @@ -1888,10 +1898,9 @@ mod tests { use crate::precompute_engine::config::LateDataPolicy; use crate::precompute_engine::output_sink::CapturingOutputSink; use crate::storage_engines::types::InstalledPrecomputePlan; + use asap_physical_operators::summary_kernels::datasketches_kll::DatasketchesKLLAccumulator; + use asap_physical_operators::summary_kernels::exact::ExactAccumulator; use asap_sketchlib::KllSketch; - use asap_summary_state::summary_kernels::datasketches_kll::DatasketchesKLLAccumulator; - use asap_summary_state::summary_kernels::keyed_sum_count::KeyedSumCountAccumulator; - use asap_summary_state::summary_kernels::sum::SumAccumulator; use asap_types::enums::WindowKind; use asap_types::sds::StoredOutputId; use asap_types::AggregationType; @@ -2063,7 +2072,8 @@ mod tests { captured[0] .1 .as_any() - .downcast_ref::() + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) .unwrap() .sum, 0.0 @@ -2093,8 +2103,9 @@ mod tests { assert!(output.policy_fp.is_unset()); let sum_acc = acc .as_any() - .downcast_ref::() - .expect("should be SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("should be an exact Sum"); assert!( (sum_acc.sum - val).abs() < 1e-10, "sum should equal sample value" @@ -2176,8 +2187,9 @@ mod tests { let sum_acc = acc .as_any() - .downcast_ref::() - .expect("should be SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("should be an exact Sum"); assert!( (sum_acc.sum - 6.0).abs() < 1e-10, "sum should be 1+2+3=6, got {}", @@ -2247,8 +2259,9 @@ mod tests { let sum_acc = acc .as_any() - .downcast_ref::() - .expect("should be SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("should be an exact Sum"); assert!( (sum_acc.sum - 30.0).abs() < 1e-10, "sum should be 10+20=30, got {} (both series merged)", @@ -2324,7 +2337,11 @@ mod tests { let mut sums_by_key: HashMap = HashMap::new(); for (output, acc) in &captured { - let sum_acc = acc.as_any().downcast_ref::().unwrap(); + let sum_acc = acc + .as_any() + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .unwrap(); let key = output.key.as_ref().unwrap().labels.join(";"); sums_by_key.insert(key, sum_acc.sum); } @@ -2469,8 +2486,9 @@ mod tests { for (_output, acc) in &captured { let sum_acc = acc .as_any() - .downcast_ref::() - .expect("should be SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("should be an exact Sum"); assert!( (sum_acc.sum - 42.0).abs() < 1e-10, "pane should have sum=42, got {}", @@ -2531,7 +2549,8 @@ mod tests { assert_eq!( accumulator .as_any() - .downcast_ref::() + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) .unwrap() .sum, 42.0 @@ -2600,24 +2619,21 @@ mod tests { let (_output, acc) = &captured[0]; let ms_acc = acc .as_any() - .downcast_ref::() - .expect("should be KeyedSumCountAccumulator"); - - // The KeyedSumCountAccumulator should have two internal keys: "A" and "B" - assert_eq!(ms_acc.sums.len(), 2, "two host keys inside one accumulator"); - - let mut found_a = false; - let mut found_b = false; - for (key, &sum) in &ms_acc.sums { - if key.labels == vec!["A".to_string()] { - assert!((sum - 10.0).abs() < 1e-10); - found_a = true; - } - if key.labels == vec!["B".to_string()] { - assert!((sum - 20.0).abs() < 1e-10); - found_b = true; - } - } + .downcast_ref::() + .expect("should be a keyed exact Sum"); + assert!(ms_acc.is_keyed(), "both host keys inside one accumulator"); + let sum = |host: &str| { + ms_acc + .readout( + asap_types::Statistic::Sum, + None, + Some(&KeyByLabelValues::new_with_labels(vec![host.into()])), + ) + .unwrap() + .unwrap() + }; + let found_a = (sum("A") - 10.0).abs() < 1e-10; + let found_b = (sum("B") - 20.0).abs() < 1e-10; assert!(found_a, "expected key A inside accumulator"); assert!(found_b, "expected key B inside accumulator"); } @@ -2762,8 +2778,9 @@ mod tests { let sum_acc = acc .as_any() - .downcast_ref::() - .expect("should be SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("should be an exact Sum"); assert!( (sum_acc.sum - 55.0).abs() < 1e-10, "late sample sum should be 55.0, got {}", @@ -2955,8 +2972,9 @@ mod tests { let sum = group_b .1 .as_any() - .downcast_ref::() - .expect("must emit SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("must emit an exact Sum"); assert_eq!(sum.sum, 6.0, "group B's second sample must not be late"); } @@ -3007,8 +3025,9 @@ mod tests { let sum = emitted[0] .1 .as_any() - .downcast_ref::() - .expect("must emit SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("must emit an exact Sum"); assert_eq!(sum.sum, 3.0, "flush must not manufacture event time"); } @@ -3185,8 +3204,8 @@ mod tests { // OTLP ingest dispatch builds via `decode_modified_otlp_sketch_bytes`. // ----------------------------------------------------------------------- + use asap_physical_operators::summary_kernels::DDSketchAccumulator; use asap_sketchlib::DdSketch; - use asap_summary_state::summary_kernels::DDSketchAccumulator; /// Build a fresh DDSketch holding `vals` so each test has a real, /// non-empty sketch to push through `process_accumulator_input`. @@ -3198,10 +3217,7 @@ mod tests { // positive-only is the realistic shape. s.update(*v); } - DDSketchAccumulator { - inner: s, - sample_p: 1.0, - } + DDSketchAccumulator { inner: s } } /// Pinning test: a single-group, single-window sketch ingest must @@ -3635,8 +3651,9 @@ mod tests { let sum = captured[0] .1 .as_any() - .downcast_ref::() - .expect("must emit SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("must emit an exact Sum"); assert!((sum.sum - expected_sum).abs() < 1e-9); } @@ -3685,8 +3702,9 @@ mod tests { let initial = captured.pop().expect("deadline output").1; let sum = initial .as_any() - .downcast_ref::() - .expect("must emit SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("must emit an exact Sum"); assert_eq!(sum.sum, 28.0); // Continuing input for the already-closed event-time window becomes a @@ -3708,8 +3726,9 @@ mod tests { .expect("deadline output and correction must merge"); let merged_sum = merged .as_any() - .downcast_ref::() - .expect("merged output must remain SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("merged output must remain an exact Sum"); assert_eq!(merged_sum.sum, 36.0); } @@ -3950,7 +3969,6 @@ mod tests { // A pooled Sum is correct only for an explicit cross-entity reduction. #[test] fn pooled_sum_does_not_preserve_per_entity_output_rows() { - use asap_summary_state::summary_kernels::SumAccumulator; let config = make_agg_config( 1, "gauge", @@ -3987,7 +4005,8 @@ mod tests { captured[0] .1 .as_any() - .downcast_ref::() + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) .unwrap() .sum, 240.0 @@ -4004,7 +4023,6 @@ mod tests { // The physical compiler rejects raw counter producers until series state is preserved. #[test] fn pooled_counter_samples_lose_independent_same_timestamp_reset() { - use asap_summary_state::summary_kernels::IncreaseAccumulator; let config = make_agg_config( 1, "requests_total", @@ -4037,15 +4055,18 @@ mod tests { .unwrap(); worker.force_close_all().unwrap(); let captured = sink.drain(); - let accumulator = captured[0] + let total_increase = captured[0] .1 .as_any() - .downcast_ref::() + .downcast_ref::() + .unwrap() + .readout(asap_types::Statistic::Increase, None, None) + .unwrap() .unwrap(); - assert_eq!(accumulator.total_increase, 10.0); + assert_eq!(total_increase, 10.0); let independent_increases = (110.0 - 100.0) + 5.0; assert_eq!(independent_increases, 15.0); - assert_ne!(accumulator.total_increase, independent_increases); + assert_ne!(total_increase, independent_increases); } // Acknowledgement proves FIFO input processing and trailing-window publication. @@ -4089,7 +4110,8 @@ mod tests { outputs[0] .1 .as_any() - .downcast_ref::() + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) .unwrap() .sum, 5.0 @@ -4278,8 +4300,9 @@ mod tests { assert_eq!(output.end_timestamp, 10_000); let sum_acc = acc .as_any() - .downcast_ref::() - .expect("should be SumAccumulator"); + .downcast_ref::() + .map(crate::tests::accumulator_fixture::sum_view) + .expect("should be an exact Sum"); assert!( (sum_acc.sum - 5.0).abs() < 1e-10, "5 samples of 1.0 → sum 5, got {}", @@ -4368,7 +4391,7 @@ mod dag_execution_tests { use super::*; use crate::precompute_engine::output_sink::CapturingOutputSink; use crate::storage_engines::types::InstalledPrecomputePlan; - use asap_summary_state::summary_kernels::exact::ExactAccumulator; + use asap_physical_operators::summary_kernels::exact::ExactAccumulator; use asap_types::query_plan::ExactReadout; fn plan(query: &str) -> control_plane::physical::compiler::CompiledPhysicalPlan { @@ -4578,9 +4601,10 @@ mod dag_execution_tests { state.get_accumulator_type().planner_exact_family(), Some(readout.planner_family()) ); - let restored = - ExactAccumulator::deserialize_from_bytes(&state.serialize_to_bytes()) - .unwrap(); + let restored = asap_summary_state::stored_state::codec::decode_exact( + &state.serialize_to_bytes(), + ) + .unwrap(); states.insert( output.end_timestamp as i64, Arc::new(restored) as Arc, @@ -4774,7 +4798,7 @@ mod dag_execution_tests { end as u64, )) .or_insert_with(|| { - asap_summary_state::factory::create_planner_accumulator( + asap_physical_operators::factory::create_planner_accumulator( &program.family, &program.input, &program.grouping, @@ -5024,12 +5048,13 @@ mod dag_execution_tests { assert!(!corrections.is_empty(), "late sample must be forwarded"); for (_, state) in corrections { let exact = - ExactAccumulator::deserialize_from_bytes(&state.serialize_to_bytes()).unwrap(); + asap_summary_state::stored_state::codec::decode_exact(&state.serialize_to_bytes()) + .unwrap(); assert_eq!( exact - .query_statistic(asap_types::Statistic::Sum, &None, &Default::default()) + .readout(asap_types::Statistic::Sum, None, None) .unwrap(), - 7.0 + Some(7.0) ); } } diff --git a/data_plane/src/query_engines/asap_clickhouse_query_engine/accelerator.rs b/data_plane/src/query_engines/asap_clickhouse_query_engine/accelerator.rs index 3e91bd5e5..ee644c84a 100644 --- a/data_plane/src/query_engines/asap_clickhouse_query_engine/accelerator.rs +++ b/data_plane/src/query_engines/asap_clickhouse_query_engine/accelerator.rs @@ -2,7 +2,6 @@ #[cfg(test)] use asap_types::physical_plan_codec::PhysicalPlanCodec; -use asap_summary_state::summary_kernels::SumAccumulator; use async_trait::async_trait; use axum::{ body::Bytes, @@ -533,7 +532,6 @@ mod tests { } use crate::storage_engines::sketch_db::index::{AggKind, Capability, SummarySeriesMetadata}; - use asap_summary_state::summary_kernels::SumAccumulator; use asap_types::query_plan::{ ClickHousePlanningContext, ExactReadout, ExternalExactOutput, ExternalExactRequest, FallbackPolicy, FixedEvaluationRange, InstantExecution, MaterializationBinding, @@ -1003,13 +1001,13 @@ mod tests { 7, Default::default(), (0, 1_000), - Box::new(SumAccumulator::with_sum(2.0)), + Box::new(crate::tests::accumulator_fixture::sum_state(2.0)), ); store.append_precompute( 7, Default::default(), (1_000, 2_000), - Box::new(SumAccumulator::with_sum(3.0)), + Box::new(crate::tests::accumulator_fixture::sum_state(3.0)), ); } let envelope = asap_types::precompute_plan::PlanEnvelope { @@ -1178,7 +1176,7 @@ mod tests { 7, Default::default(), (2000, 3000), - Box::new(SumAccumulator::with_sum(5.0)), + Box::new(crate::tests::accumulator_fixture::sum_state(5.0)), ); for (sql, expected) in [ (request.sql.clone(), "1970-01-01T00:00:01\t20.0\n"), diff --git a/data_plane/src/query_engines/asap_query_engine/engine.rs b/data_plane/src/query_engines/asap_query_engine/engine.rs index 292fa4ca7..bc0b3976a 100644 --- a/data_plane/src/query_engines/asap_query_engine/engine.rs +++ b/data_plane/src/query_engines/asap_query_engine/engine.rs @@ -919,37 +919,6 @@ impl ASAPQueryEngine { }) } - #[cfg(test)] - fn query_precompute_for_statistic( - &self, - precompute: &dyn AggregateCore, - statistic: &Statistic, - key: &Option, - query_kwargs: &HashMap, - ) -> Result> { - // Phase 1b of the sketch DB design - // (docs/design_docs/summary-storage.md): - // for single-subpopulation queries on additive statistics - // (Count / Sum / Min / Max), serve from the typed aux - // columns without deserialising the sketch payload. - // - // Keyed queries (`key.is_some()`) still need the full - // `query_statistic` path — aux is per-accumulator, not - // per-subpopulation key. - // - // `try_answer` returns `None` when the statistic isn't - // covered by aux (Quantile / Cardinality / TopK / Increase / - // Rate) or when the accumulator doesn't track the requested - // aux field; both cases fall through to the existing path - // so the query result is semantically identical. - if key.is_none() { - if let Some(value) = precompute.aux_stats().try_answer(*statistic) { - return Ok(value); - } - } - precompute.query_statistic(*statistic, key, query_kwargs) - } - /// Modern warm-tier path for `/api/v1/query_range` — the range- /// query equivalent of the `QueryEngine::execute(&str)` trait /// surface. Used by the HTTP server as a fallback when the legacy @@ -1599,188 +1568,6 @@ mod sketch_query_tests { // } } -// Typed auxiliary statistics must answer covered queries without invoking -// the accumulator query method; uncovered statistics use that method. -#[cfg(test)] -mod aux_pushdown_tests { - use super::*; - use crate::storage_engines::types::AggregationType; - use asap_summary_state::summary_kernels::{ - max::MaxAccumulator, min::MinAccumulator, sum::SumAccumulator, - }; - use asap_types::Statistic; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - /// Accumulator that records how many times `query_statistic` - /// was invoked. Used to verify the aux fast path skips it. - struct SpyAccumulator { - inner_sum: f64, - query_calls: Arc, - } - - impl crate::storage_engines::types::SerializableToSink for SpyAccumulator { - fn serialize_to_bytes(&self) -> Vec { - Vec::new() - } - fn serialize_to_json(&self) -> serde_json::Value { - serde_json::Value::Null - } - } - - impl AggregateCore for SpyAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(SpyAccumulator { - inner_sum: self.inner_sum, - query_calls: self.query_calls.clone(), - }) - } - fn type_name(&self) -> &'static str { - "SpyAccumulator" - } - 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> { - unimplemented!() - } - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::Sum - } - fn get_keys(&self) -> Option> { - None - } - fn query_statistic( - &self, - _statistic: Statistic, - _key: &Option, - _query_kwargs: &HashMap, - ) -> Result> { - self.query_calls.fetch_add(1, Ordering::Relaxed); - Ok(-1.0) // sentinel: fast path should not return this - } - fn aux_stats(&self) -> crate::storage_engines::types::AuxStats { - crate::storage_engines::types::AuxStats { - sum: Some(self.inner_sum), - ..crate::storage_engines::types::AuxStats::empty() - } - } - } - - fn make_engine() -> ASAPQueryEngine { - use crate::storage_engines::types::{ - CleanupPolicy, InstalledPrecomputePlan, InstalledPrecomputePlanHandle, - }; - - let sc = Arc::new(InstalledPrecomputePlan::new(HashMap::new())); - let hr = InstalledPrecomputePlanHandle::from_arc(sc.clone()); - let _ = sc; - ASAPQueryEngine::new(60) - } - - #[test] - fn aux_covered_stat_skips_query_statistic() { - let engine = make_engine(); - let calls = Arc::new(AtomicUsize::new(0)); - let spy = SpyAccumulator { - inner_sum: 42.0, - query_calls: calls.clone(), - }; - let result = engine - .query_precompute_for_statistic(&spy, &Statistic::Sum, &None, &HashMap::new()) - .expect("query ok"); - assert_eq!(result, 42.0, "aux fast path should return aux value"); - assert_eq!( - calls.load(Ordering::Relaxed), - 0, - "query_statistic should NOT be called when aux covers the stat" - ); - } - - #[test] - fn aux_uncovered_stat_falls_through_to_query_statistic() { - let engine = make_engine(); - let calls = Arc::new(AtomicUsize::new(0)); - let spy = SpyAccumulator { - inner_sum: 42.0, - query_calls: calls.clone(), - }; - // Quantile is not covered by aux → must fall through. - let result = engine - .query_precompute_for_statistic(&spy, &Statistic::Quantile, &None, &HashMap::new()) - .expect("query ok"); - assert_eq!( - result, -1.0, - "should have returned query_statistic's sentinel" - ); - assert_eq!( - calls.load(Ordering::Relaxed), - 1, - "query_statistic should be called exactly once when aux misses" - ); - } - - #[test] - fn keyed_queries_always_use_query_statistic() { - let engine = make_engine(); - let calls = Arc::new(AtomicUsize::new(0)); - let spy = SpyAccumulator { - inner_sum: 42.0, - query_calls: calls.clone(), - }; - let key = Some(KeyByLabelValues::new()); - // Even for Sum (which aux covers), a keyed query must bypass aux - // — aux is per-accumulator, not per-subpopulation key. - let result = engine - .query_precompute_for_statistic(&spy, &Statistic::Sum, &key, &HashMap::new()) - .expect("query ok"); - assert_eq!(result, -1.0); - assert_eq!( - calls.load(Ordering::Relaxed), - 1, - "keyed queries must skip aux fast path" - ); - } - - #[test] - fn real_sum_accumulator_uses_aux_fast_path() { - // End-to-end: a real SumAccumulator goes through the fast path - // and returns its sum without ever hitting query_statistic. - let engine = make_engine(); - let acc = SumAccumulator::with_sum(7.5); - let result = engine - .query_precompute_for_statistic(&acc, &Statistic::Sum, &None, &HashMap::new()) - .expect("query ok"); - assert_eq!(result, 7.5); - } - - #[test] - fn real_min_max_accumulator_uses_aux_fast_path() { - let engine = make_engine(); - let min_acc = MinAccumulator::with_value(3.0); - let max_acc = MaxAccumulator::with_value(99.0); - assert_eq!( - engine - .query_precompute_for_statistic(&min_acc, &Statistic::Min, &None, &HashMap::new()) - .unwrap(), - 3.0 - ); - assert_eq!( - engine - .query_precompute_for_statistic(&max_acc, &Statistic::Max, &None, &HashMap::new()) - .unwrap(), - 99.0 - ); - } -} - // ── build_query_execution_context_promql_for_agg_id (forced-agg) tests ── // =========================================================================== @@ -1824,7 +1611,6 @@ mod asap_tier_classify_tests { async fn execute_sum_by_zone_dispatches_to_exact_agg_reducer() { use crate::query_engines::query_result::QueryResult; use crate::storage_engines::sketch_db::data::AggregationType; - use asap_summary_state::summary_kernels::sum::SumAccumulator; let idx = Arc::new(SketchStore::new()); // Mirror the acceptance-test setup: four ExactAgg(Sum) sids, one @@ -1862,7 +1648,7 @@ mod asap_tier_classify_tests { sid, lm, (window_start, window_end), - Box::new(SumAccumulator::with_sum(value)), + Box::new(crate::tests::accumulator_fixture::sum_state(value)), ); } @@ -2478,7 +2264,6 @@ mod asap_tier_classify_tests { async fn execute_instant_sum_accumulates_all_windows_not_last() { use crate::query_engines::query_result::QueryResult; use crate::storage_engines::sketch_db::data::AggregationType; - use asap_summary_state::summary_kernels::sum::SumAccumulator; let idx = Arc::new(SketchStore::new()); let now_ms = 600_000_u64; @@ -2512,7 +2297,7 @@ mod asap_tier_classify_tests { sid, lm, (ws, we), - Box::new(SumAccumulator::with_sum(*per_window)), + Box::new(crate::tests::accumulator_fixture::sum_state(*per_window)), ); } } diff --git a/data_plane/src/query_engines/asap_query_engine/exact_subqueries.rs b/data_plane/src/query_engines/asap_query_engine/exact_subqueries.rs index 0cc1facba..7ed9b1ae0 100644 --- a/data_plane/src/query_engines/asap_query_engine/exact_subqueries.rs +++ b/data_plane/src/query_engines/asap_query_engine/exact_subqueries.rs @@ -944,8 +944,6 @@ mod tests { data::AggKind, index::{Capability, SummarySeriesMetadata}, }; - use crate::storage_engines::types::Measurement; - use asap_summary_state::summary_kernels::IncreaseAccumulator; use asap_types::query_plan::{ExactReadout, MaterializationBinding, PhysicalGrouping}; use std::sync::{ atomic::{AtomicUsize, Ordering}, @@ -971,9 +969,10 @@ mod tests { expires_at_ms: None, policy_fp: MATERIALIZATION, }); - let mut denominator = - IncreaseAccumulator::new(Measurement::new(100.0), 0, Measurement::new(100.0), 0); - denominator.update(Measurement::new(400.0), AT as i64); + let denominator = crate::tests::accumulator_fixture::counter_state( + planner_types::post_asap::ExactKind::Rate, + &[(0, 100.0), (AT as i64, 400.0)], + ); store.append_precompute( 41, BTreeMap::from([("job".into(), "user-service".into())]), diff --git a/data_plane/src/query_engines/asap_query_engine/live_serve.rs b/data_plane/src/query_engines/asap_query_engine/live_serve.rs index a1c4c5aef..a584eb3ce 100644 --- a/data_plane/src/query_engines/asap_query_engine/live_serve.rs +++ b/data_plane/src/query_engines/asap_query_engine/live_serve.rs @@ -170,7 +170,7 @@ mod tests { 9, BTreeMap::new(), (start, end), - Box::new(asap_summary_state::summary_kernels::SumAccumulator::with_sum(value)), + Box::new(crate::tests::accumulator_fixture::sum_state(value)), ); } let entry = asap_types::query_plan::QueryPlanEntry { diff --git a/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs b/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs index 2e8609baa..7c079838b 100644 --- a/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs +++ b/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs @@ -1052,15 +1052,10 @@ mod tests { let groups = [1, 2] .into_iter() .map(|samples| { - let mut state = asap_summary_state::summary_kernels::IncreaseAccumulator::new( - Measurement::new(10.0), - 10_000, - Measurement::new(10.0), - 10_000, + let state = crate::tests::accumulator_fixture::counter_state( + planner_types::post_asap::ExactKind::Rate, + &[(10_000, 10.0), (20_000, 20.0)][..samples], ); - if samples == 2 { - state.update(Measurement::new(20.0), 20_000); - } ( BTreeMap::from([("instance".into(), samples.to_string())]), GroupState::ExactAgg { @@ -1126,11 +1121,9 @@ mod tests { policy_fp: policy, }); for (start, time, value) in [(20_000, 25_000, 10.), (40_000, 45_000, 20.)] { - let state = asap_summary_state::summary_kernels::IncreaseAccumulator::new( - Measurement::new(value), - time, - Measurement::new(value), - time, + let state = crate::tests::accumulator_fixture::counter_state( + planner_types::post_asap::ExactKind::Rate, + &[(time, value)], ); idx.append_precompute(7, BTreeMap::new(), (start, start + 10_000), Box::new(state)); } @@ -1197,7 +1190,7 @@ mod tests { sid, BTreeMap::from([("instance".into(), instance.into())]), (0, 1_000), - Box::new(asap_summary_state::summary_kernels::SumAccumulator::with_sum(1.)), + Box::new(crate::tests::accumulator_fixture::sum_state(1.)), ); } let binding = asap_types::query_plan::MaterializationBinding { @@ -1769,7 +1762,7 @@ mod tests { 1, BTreeMap::new(), (1_000, 2_000), - Box::new(asap_summary_state::summary_kernels::SumAccumulator::with_sum(42.0)), + Box::new(crate::tests::accumulator_fixture::sum_state(42.0)), ); let config = test_plan::materialization("bytes_total", "Sum", serde_json::json!({}), &[], 1000); @@ -1869,9 +1862,7 @@ mod tests { 7, BTreeMap::new(), bounds, - Box::new( - asap_summary_state::summary_kernels::SumAccumulator::with_sum(sum), - ), + Box::new(crate::tests::accumulator_fixture::sum_state(sum)), ); } let entry = plan.query_plan.entries.values().next().unwrap(); @@ -1938,11 +1929,9 @@ mod tests { 7, BTreeMap::new(), (pane * 60_000, (pane + 1) * 60_000), - Box::new( - asap_summary_state::summary_kernels::SumAccumulator::with_sum( - (pane + 1) as f64, - ), - ), + Box::new(crate::tests::accumulator_fixture::sum_state( + (pane + 1) as f64, + )), ); } let entry = plan.query_plan.entries.values().next().unwrap(); @@ -2013,11 +2002,9 @@ mod tests { 7, BTreeMap::new(), (pane * 10_000, (pane + 1) * 10_000), - Box::new( - asap_summary_state::summary_kernels::SumAccumulator::with_sum( - (pane + 1) as f64, - ), - ), + Box::new(crate::tests::accumulator_fixture::sum_state( + (pane + 1) as f64, + )), ); } @@ -2082,7 +2069,7 @@ mod tests { 7, BTreeMap::new(), (pane * 10_000, (pane + 1) * 10_000), - Box::new(asap_summary_state::summary_kernels::SumAccumulator::with_sum(1.0)), + Box::new(crate::tests::accumulator_fixture::sum_state(1.0)), ); } assert!( @@ -2111,16 +2098,15 @@ mod tests { expires_at_ms: None, policy_fp: policy, }); - use crate::storage_engines::types::Measurement; - let mut accumulator = asap_summary_state::summary_kernels::IncreaseAccumulator::new( - Measurement::new(10.0), - 10_000, - Measurement::new(10.0), - 10_000, + let accumulator = crate::tests::accumulator_fixture::counter_state( + planner_types::post_asap::ExactKind::Rate, + &[ + (10_000, 10.0), + (20_000, 20.0), + (30_000, 3.0), + (50_000, 13.0), + ], ); - accumulator.update(Measurement::new(20.0), 20_000); - accumulator.update(Measurement::new(3.0), 30_000); - accumulator.update(Measurement::new(13.0), 50_000); idx.append_precompute(7, BTreeMap::new(), (0, 60_000), Box::new(accumulator)); let entry = asap_types::query_plan::QueryPlanEntry { diff --git a/data_plane/src/query_engines/asap_query_engine/request_tests.rs b/data_plane/src/query_engines/asap_query_engine/request_tests.rs index 3f876c26a..1da971ed9 100644 --- a/data_plane/src/query_engines/asap_query_engine/request_tests.rs +++ b/data_plane/src/query_engines/asap_query_engine/request_tests.rs @@ -201,7 +201,7 @@ async fn installed_local_range_accounts_for_accumulated_results() { 1, BTreeMap::new(), (pane * 1000, (pane + 1) * 1000), - Box::new(asap_summary_state::summary_kernels::SumAccumulator::with_sum(2.0)), + Box::new(crate::tests::accumulator_fixture::sum_state(2.0)), ); } engine diff --git a/data_plane/src/query_engines/asap_query_engine/summary_executor.rs b/data_plane/src/query_engines/asap_query_engine/summary_executor.rs index 3560f017b..8db617a03 100644 --- a/data_plane/src/query_engines/asap_query_engine/summary_executor.rs +++ b/data_plane/src/query_engines/asap_query_engine/summary_executor.rs @@ -1,5 +1,4 @@ //! Deployment adapters for resolving, decoding, and reading installed materializations. -use asap_summary_state::summary_kernels::{MaxAccumulator, MinAccumulator}; use std::collections::{BTreeMap, BTreeSet}; use std::rc::Rc; use std::sync::Arc; @@ -469,11 +468,15 @@ impl QueryExecutionContext<'_> { project_group_key(keys, &labels) } }; - let accumulator: Arc = if is_min { - Arc::new(MinAccumulator::with_value(value)) + let kind = if is_min { + planner_types::post_asap::ExactKind::Min } else { - Arc::new(MaxAccumulator::with_value(value)) + planner_types::post_asap::ExactKind::Max }; + let accumulator: Arc = + Arc::new(asap_summary_state::stored_state::codec::exact_value( + kind, value, + )); source_order.entry(key.clone()).or_insert(sid); by_group.entry(key).or_default().push(GroupState::ExactAgg { entries: vec![Rc::new(BTreeMap::from([( @@ -1016,13 +1019,20 @@ mod tests { #[test] fn keyed_count_state_follows_planner_family_and_query_readout() { - use asap_summary_state::summary_kernels::KeyedSumCountAccumulator; + use asap_physical_operators::summary_kernels::exact::ExactAccumulator; use asap_types::query_plan::ExactReadout; let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - let mut payload = KeyedSumCountAccumulator::for_family(ExactKind::Count); - payload.update(key.clone(), 10.0); - payload.update(key.clone(), 20.0); + let mut payload = ExactAccumulator::new( + planner_types::post_asap::SummaryFamilyType::ExactAggregate( + ExactKind::Count, + planner_types::post_asap::ExactParams::Count, + ), + true, + ) + .unwrap(); + payload.update(Some(&key), 10.0, 0); + payload.update(Some(&key), 20.0, 0); let state = GroupState::ExactAgg { entries: vec![Rc::new(BTreeMap::from([( 60_000, @@ -1213,8 +1223,7 @@ mod tests { #[test] fn bound_univmon_merges_panes_for_four_readouts() { use crate::storage_engines::sketch_db::index::SketchEncoding; - use crate::storage_engines::types::SerializableToSink; - use asap_summary_state::summary_kernels::univmon::UnivMonAccumulator; + use asap_summary_state::univmon::UnivMonAccumulator; use asap_types::query_plan::{MaterializationBinding, PhysicalGrouping}; let index = SketchStore::new(); let fp = asap_types::PolicyFingerprint(701); @@ -1243,7 +1252,7 @@ mod tests { BTreeMap::from([("job".into(), "a".into())]), (start, start + 1000), SketchSampleState { - bytes: state.serialize_to_bytes(), + bytes: state.to_bytes().unwrap(), encoding: SketchEncoding::MsgpackFull, }, ); @@ -1300,16 +1309,21 @@ mod tests { #[test] fn typed_dds_quantile_interpolates_without_changing_portable_rank_semantics() { + use asap_physical_operators::summary_kernels::DDSketchAccumulator; let mut sketch = asap_sketchlib::DdSketch::new(0.01); assert!(sketch_query_value( - &SummaryState::Dd(sketch.clone()), + &SummaryState::Dd(DDSketchAccumulator { + inner: sketch.clone(), + }), &SketchQuery::Quantile { q: 0.9 } ) .is_err()); sketch.update(20.0); for q in [0.0, 0.5, 0.9, 1.0] { let value = sketch_query_value( - &SummaryState::Dd(sketch.clone()), + &SummaryState::Dd(DDSketchAccumulator { + inner: sketch.clone(), + }), &SketchQuery::Quantile { q }, ) .unwrap(); @@ -1319,14 +1333,16 @@ mod tests { assert!(sketch.quantile(0.9).unwrap() < 21.0); for (q, expected) in [(0.0, 20.0), (0.5, 30.0), (0.9, 38.0), (1.0, 40.0)] { let value = sketch_query_value( - &SummaryState::Dd(sketch.clone()), + &SummaryState::Dd(DDSketchAccumulator { + inner: sketch.clone(), + }), &SketchQuery::Quantile { q }, ) .unwrap(); assert!((value - expected).abs() <= expected * 0.01); } assert!(sketch_query_value( - &SummaryState::Dd(sketch), + &SummaryState::Dd(DDSketchAccumulator { inner: sketch }), &SketchQuery::Quantile { q: f64::NAN } ) .is_err()); diff --git a/data_plane/src/storage_engines/sketch_db/backfill/processor.rs b/data_plane/src/storage_engines/sketch_db/backfill/processor.rs index d408a54db..ee6c355ee 100644 --- a/data_plane/src/storage_engines/sketch_db/backfill/processor.rs +++ b/data_plane/src/storage_engines/sketch_db/backfill/processor.rs @@ -414,6 +414,7 @@ mod tests { use crate::storage_engines::sketch_db::backfill::worker::BackfillWorker; use crate::storage_engines::sketch_db::backfill::BackfillSource; use crate::storage_engines::types::InstalledPrecomputePlan; + use asap_summary_state::StoredState; use asap_types::enums::WindowKind; use asap_types::AggregationType; use asap_types::KeyByLabelNames; @@ -985,7 +986,7 @@ mod tests { ), cfg.policy_fingerprint(), ); - let acc = asap_summary_state::summary_kernels::sum::SumAccumulator::with_sum(1.0); + let acc = crate::tests::accumulator_fixture::sum_state(1.0); let live_sid = store .ingest_precompute_for_agg_config( |_metric, attrs, _kind| { diff --git a/data_plane/src/storage_engines/sketch_db/backfill/window_builder.rs b/data_plane/src/storage_engines/sketch_db/backfill/window_builder.rs index 54db73627..d419cb1fc 100644 --- a/data_plane/src/storage_engines/sketch_db/backfill/window_builder.rs +++ b/data_plane/src/storage_engines/sketch_db/backfill/window_builder.rs @@ -11,7 +11,7 @@ use crate::storage_engines::types::AggregateCore; #[cfg(test)] use crate::tests::accumulator_fixture::create_fixture_accumulator; #[cfg(test)] -use asap_summary_state::factory::AccumulatorUpdater; +use asap_physical_operators::factory::AccumulatorUpdater; #[cfg(test)] use asap_types::aggregation_config::PrecomputeMaterialization; @@ -86,7 +86,7 @@ mod tests { // Replay must preserve each series and rank by the selected update mode. #[test] fn backfilled_topk_preserves_series_and_weight_mode() { - use asap_summary_state::summary_kernels::{ + use asap_physical_operators::summary_kernels::{ CountMinSketchWithHeapAccumulator, CountSketchWithHeapAccumulator, }; for kind in [ @@ -145,19 +145,15 @@ mod tests { raw("m{svc=\"a\"}", 30, 3.0), ]; let acc = build_backfilled_accumulator(&config, &samples); - // SumAccumulator's AuxStats exposes the sum. - let aux = acc.aux_stats(); - assert_eq!(aux.sum, Some(6.0)); + assert_eq!(crate::tests::accumulator_fixture::sum_of(acc.as_ref()), 6.0); } #[test] fn empty_samples_produce_empty_accumulator() { let config = sum_config(); let acc = build_backfilled_accumulator(&config, &[]); - let aux = acc.aux_stats(); - // A fresh SumAccumulator has sum = Some(0.0) per its AuxStats - // implementation (identity element). - assert!(aux.sum == Some(0.0) || aux.sum.is_none()); + // An empty exact Sum is the additive identity. + assert_eq!(crate::tests::accumulator_fixture::sum_of(acc.as_ref()), 0.0); } } diff --git a/data_plane/src/storage_engines/sketch_db/data/mod.rs b/data_plane/src/storage_engines/sketch_db/data/mod.rs index 5c8ed04e3..f515c9e14 100644 --- a/data_plane/src/storage_engines/sketch_db/data/mod.rs +++ b/data_plane/src/storage_engines/sketch_db/data/mod.rs @@ -524,7 +524,10 @@ impl std::fmt::Debug for AggPayload { AggPayload::Sketch(s) => f.debug_tuple("Sketch").field(s).finish(), AggPayload::ExactAgg(p) => f .debug_struct("ExactAgg") - .field("type_name", &p.type_name()) + .field( + "type_name", + &asap_summary_state::StoredState::type_name(p.as_ref()), + ) .finish(), } } diff --git a/data_plane/src/storage_engines/sketch_db/index/maintenance.rs b/data_plane/src/storage_engines/sketch_db/index/maintenance.rs index c99a20734..b57e679de 100644 --- a/data_plane/src/storage_engines/sketch_db/index/maintenance.rs +++ b/data_plane/src/storage_engines/sketch_db/index/maintenance.rs @@ -5,6 +5,7 @@ //! summary. Completion is an admission barrier, not merely an emitted flag. use super::*; use crate::storage_engines::types::AggregateCore; +use asap_summary_state::StoredState; pub(crate) struct FrozenExactWindows { pub(crate) stored_output_reference: asap_types::sds::StoredOutputReference, @@ -283,7 +284,7 @@ impl SketchStore { expected_windows, group, |name, _, bytes| { - reconstruct_exact_agg(name, bytes) + reconstruct_exact_agg(name, bytes)? .ok_or_else(|| "immutable input accumulator cannot be decoded".to_string()) }, ) @@ -797,8 +798,6 @@ impl SketchStore { mod tests { use super::*; use crate::storage_engines::types::PrecomputedOutput; - use asap_summary_state::summary_kernels::SumAccumulator; - use asap_types::traits::SerializableToSink; #[test] fn complete_population_keeps_every_sid_and_rejects_missing_live_binding() { @@ -855,8 +854,8 @@ mod tests { let mut output = PrecomputedOutput::new(0, 1000, None, config.policy_fingerprint()); output.catalog_generation = Some(Arc::clone(&generation)); output.population_labels = Some(population); - let mut sum = SumAccumulator::new(); - sum.update(5.0 + index as f64); + let mut sum = crate::tests::accumulator_fixture::sum_state(0.0); + sum.update(None, 5.0 + index as f64, 0); let sid = 900 + index as u64; store .publish_admitted_summary_update( @@ -970,8 +969,8 @@ mod tests { .unwrap(); let mut output = PrecomputedOutput::new(0, 1000, None, config.policy_fingerprint()); output.catalog_generation = Some(Arc::clone(&generation)); - let mut sum = SumAccumulator::new(); - sum.update(5.0 + index as f64); + let mut sum = crate::tests::accumulator_fixture::sum_state(0.0); + sum.update(None, 5.0 + index as f64, 0); let sid = 900 + index as u64; store .publish_admitted_summary_update( @@ -1006,8 +1005,8 @@ mod tests { let mut extra_window = PrecomputedOutput::new(1000, 2000, None, configs[0].policy_fingerprint()); extra_window.catalog_generation = Some(Arc::clone(&generation)); - let mut extra_sum = SumAccumulator::new(); - extra_sum.update(99.0); + let mut extra_sum = crate::tests::accumulator_fixture::sum_state(0.0); + extra_sum.update(None, 99.0, 0); store .publish_admitted_summary_update( &generation, @@ -1053,8 +1052,8 @@ mod tests { let target = &configs[2]; let mut output = PrecomputedOutput::new(0, 1000, None, target.policy_fingerprint()); output.catalog_generation = Some(Arc::clone(&generation)); - let mut sum = SumAccumulator::new(); - sum.update(11.0); + let mut sum = crate::tests::accumulator_fixture::sum_state(0.0); + sum.update(None, 11.0, 0); assert!(store .publish_complete_raw_maintenance_output( 902, target, &output, &sum, &complete, [42; 32] @@ -1196,8 +1195,8 @@ mod tests { let mut output = PrecomputedOutput::new(0, 60_000, None, source.policy_fingerprint()); output.population_labels = Some(population); output.catalog_generation = Some(Arc::clone(&generation)); - let mut sum = SumAccumulator::new(); - sum.update(value); + let mut sum = crate::tests::accumulator_fixture::sum_state(0.0); + sum.update(None, value, 0); store .publish_admitted_summary_update( &generation, @@ -1305,7 +1304,7 @@ mod tests { ); let assert_complete_output = |store: &SketchStore| { use crate::storage_engines::sketch_db::data::SketchEncoding; - use asap_summary_state::summary_kernels::DDSketchAccumulator; + use asap_physical_operators::summary_kernels::DDSketchAccumulator; let rows = store.query_range(target_sid, 0, 60_000); assert_eq!(rows.len(), 1); assert!(rows[0].series_label_values.is_empty()); @@ -1313,12 +1312,18 @@ mod tests { let frames = &rows[0].samples[&60_000]; assert_eq!(frames.len(), 1); let sketch = match frames[0].encoding { - SketchEncoding::MsgpackFull => { - DDSketchAccumulator::from_msgpack_bytes(&frames[0].bytes).unwrap() - } - SketchEncoding::ProtoFull => { - DDSketchAccumulator::from_sketchlib_proto_bytes(&frames[0].bytes).unwrap() - } + SketchEncoding::MsgpackFull => DDSketchAccumulator { + inner: asap_summary_state::stored_state::decoders::ddsketch_from_msgpack( + &frames[0].bytes, + ) + .unwrap(), + }, + SketchEncoding::ProtoFull => DDSketchAccumulator { + inner: asap_summary_state::stored_state::decoders::ddsketch_from_proto( + &frames[0].bytes, + ) + .unwrap(), + }, other => panic!("unexpected derived encoding: {other:?}"), }; assert_eq!(sketch.inner.total_count(), 2); @@ -1395,8 +1400,8 @@ mod tests { let mut output = PrecomputedOutput::new(0, 60_000, None, source.policy_fingerprint()); output.population_labels = Some(population.clone()); output.catalog_generation = Some(Arc::clone(&next_generation)); - let mut sum = SumAccumulator::new(); - sum.update(9.0); + let mut sum = crate::tests::accumulator_fixture::sum_state(0.0); + sum.update(None, 9.0, 0); restarted .publish_admitted_summary_update( &next_generation, @@ -1513,7 +1518,7 @@ mod tests { let record = store .metadata_record(&store.instances.read().unwrap()[&601]) .unwrap(); - let state = SumAccumulator::new(); + let state = crate::tests::accumulator_fixture::sum_state(0.0); let snapshot = persistence::source::EpochSnapshot { agg_id: 601, epoch_id: 0, @@ -1524,9 +1529,9 @@ mod tests { start_ts: 0, end_ts: 1000, label: None, - sketch_type_name: state.type_name().into(), + sketch_type_name: StoredState::type_name(&state as &dyn AggregateCore).into(), encoding_tag: 0, - sketch_bytes: state.serialize_to_bytes(), + sketch_bytes: StoredState::serialize_to_bytes(&state as &dyn AggregateCore), }], }; persistence @@ -1563,7 +1568,7 @@ mod tests { 601, BTreeMap::new(), (2000, 3000), - Box::new(SumAccumulator::new()) + Box::new(crate::tests::accumulator_fixture::sum_state(0.0)) )); assert!(!store.append_sample( 601, diff --git a/data_plane/src/storage_engines/sketch_db/index/mod.rs b/data_plane/src/storage_engines/sketch_db/index/mod.rs index c116a379d..b476f1430 100644 --- a/data_plane/src/storage_engines/sketch_db/index/mod.rs +++ b/data_plane/src/storage_engines/sketch_db/index/mod.rs @@ -19,6 +19,7 @@ //! See design doc §4.6 ("OTLP metadata model + backend store layout") at //! `docs/design_docs/series-identity.md`. +use asap_summary_state::StoredState; use asap_types::sds::StoredOutputReference; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::sync::{Arc, RwLock}; @@ -83,57 +84,44 @@ fn tag_to_encoding(tag: u8) -> SketchEncoding { } /// Reconstruct an exact-aggregation accumulator from its on-disk -/// `(type_name, bytes)` pair so the durable tier can serve the -/// exact-agg query path (`query_exact_agg_range` / `sum by (...)`) after -/// flush+evict. Covers the deterministic scalar accumulators the live -/// marquee `sum by (zone)` path uses; the sketch-backed accumulator forms -/// (DDSketch/KLL/HLL/CountSketch — registered as `AggKind::Sketch`) are -/// served as opaque bytes via [`SketchStore::query_range`] and are NOT -/// reconstructed here. Returns `None` for an unrecognized `type_name` -/// (the caller skips the disk entry rather than fabricating a wrong -/// payload) — see the remaining-follow-up note in the PR. +/// `(type_name, bytes)` pair so the durable tier can serve the exact-agg query +/// path after flush+evict. `Ok(None)` for entries that are not exact state +/// (sketch-backed forms are served as opaque bytes via +/// [`SketchStore::query_range`]); an error names an unreadable or retired +/// exact format instead of misreading it. fn reconstruct_exact_agg( type_name: &str, bytes: &[u8], -) -> Option> { - use crate::storage_engines::types::AggregateCore; - use asap_summary_state::summary_kernels::{ - IncreaseAccumulator, KeyedCounterState, KeyedSumCountAccumulator, MaxAccumulator, - MinAccumulator, SumAccumulator, +) -> Result>, String> { + use asap_summary_state::stored_state::codec; + if type_name != codec::EXACT_V1 && !codec::is_retired_exact(type_name) { + return Ok(None); + } + codec::decode(type_name, bytes) + .map(Some) + .map_err(|error| error.to_string()) +} + +/// The rollup value of an unkeyed exact MIN/MAX pane; an empty pane has none. +fn extremum_rollup( + payload: &dyn crate::storage_engines::types::AggregateCore, +) -> Option<(RollupReduction, f64)> { + use asap_physical_operators::summary_kernels::exact::ExactAccumulator; + use planner_types::post_asap::{ExactKind, SummaryFamilyType}; + let state = payload + .as_any() + .downcast_ref::() + .filter(|state| !state.is_keyed())?; + let (reduction, statistic) = match state.family() { + SummaryFamilyType::ExactAggregate(ExactKind::Min, _) => { + (RollupReduction::Min, asap_types::Statistic::Min) + } + SummaryFamilyType::ExactAggregate(ExactKind::Max, _) => { + (RollupReduction::Max, asap_types::Statistic::Max) + } + _ => return None, }; - match type_name { - "PlannerExactAccumulatorV1" => { - asap_summary_state::summary_kernels::exact::ExactAccumulator::deserialize_from_bytes( - bytes, - ) - .ok() - .map(|a| Box::new(a) as Box) - } - "SumAccumulator" => SumAccumulator::deserialize_from_bytes(bytes) - .ok() - .map(|a| Box::new(a) as Box), - "IncreaseAccumulator" => IncreaseAccumulator::deserialize_from_bytes(bytes) - .ok() - .map(|a| Box::new(a) as Box), - "MinAccumulator" => MinAccumulator::deserialize_from_bytes(bytes) - .ok() - .map(|a| Box::new(a) as Box), - "MaxAccumulator" => MaxAccumulator::deserialize_from_bytes(bytes) - .ok() - .map(|a| Box::new(a) as Box), - "KeyedSumCountAccumulator" => KeyedSumCountAccumulator::deserialize_from_bytes(bytes) - .ok() - .map(|a| Box::new(a) as Box), - "KeyedCounterState" => KeyedCounterState::deserialize_from_bytes(bytes) - .ok() - .map(|a| Box::new(a) as Box), - // The keyed `MultipleMin`/`MultipleMax` forms and the - // sketch-backed accumulators have no generic byte factory — left - // to the deferred exact-agg/sketch precompute read-back work (see - // PR follow-up note). They are still served from memory; only the - // evicted-to-disk portion is skipped for these types. - _ => None, - } + Some((reduction, state.readout(statistic, None, None).ok()??)) } /// Joint helper shared by [`SketchStore::ingest_precompute_for_agg_config`] @@ -1706,6 +1694,12 @@ impl SketchStore { window: TimestampRange, payload: Box, ) -> bool { + if let Err(error) = + asap_summary_state::stored_state::codec::check_storable(payload.as_ref()) + { + tracing::warn!(sid, %error, "rejecting summary state without a stored codec"); + return false; + } let completed = self.completed_windows.read().unwrap(); if completed.get(&sid).is_some_and(|end| window.1 <= *end) { return false; @@ -1716,16 +1710,7 @@ impl SketchStore { // Both directions are their own accumulator type, so the reduction // follows from the payload's type rather than from a `sub_type` // string that had to agree with it. - let rollup_value = payload - .as_any() - .downcast_ref::() - .map(|acc| (RollupReduction::Min, acc.value)) - .or_else(|| { - payload - .as_any() - .downcast_ref::() - .map(|acc| (RollupReduction::Max, acc.value)) - }); + let rollup_value = extremum_rollup(payload.as_ref()); let store = self .series .entry(sid) @@ -2328,9 +2313,14 @@ impl SketchStore { let Ok(entry) = reader.load_entry(&rec) else { continue; }; - let Some(acc) = reconstruct_exact_agg(&entry.sketch_type_name, &entry.sketch_bytes) - else { - continue; + let acc = match reconstruct_exact_agg(&entry.sketch_type_name, &entry.sketch_bytes) + { + Ok(Some(acc)) => acc, + Ok(None) => continue, + Err(error) => { + tracing::warn!(sid, %error, "exact-agg disk read: skipping unreadable entry"); + continue; + } }; let label_map = Self::rebuild_label_map(&keys, &entry.label); by_label_map @@ -2420,8 +2410,10 @@ impl SketchStore { let Ok(entry) = reader.load_entry(&rec) else { continue; }; - if reconstruct_exact_agg(&entry.sketch_type_name, &entry.sketch_bytes).is_none() - { + if !matches!( + reconstruct_exact_agg(&entry.sketch_type_name, &entry.sketch_bytes), + Ok(Some(_)) + ) { continue; } any = true; @@ -3329,6 +3321,10 @@ impl SketchStore { output: &crate::storage_engines::types::PrecomputedOutput, accumulator: &dyn crate::storage_engines::types::AggregateCore, ) -> Option { + if let Err(error) = asap_summary_state::stored_state::codec::check_storable(accumulator) { + tracing::warn!(sid, %error, "rejecting summary state without a stored codec"); + return None; + } let expected = agg_cfg.accumulator_spec().ok()?.family; if matches!( expected, @@ -4619,10 +4615,43 @@ mod tests { assert_eq!(canonical_parameters(&p_ab), canonical_parameters(&p_ba)); } + // Unkeyed Planner exact MIN panes feed the MIN rollup. #[test] - fn precompute_payload_round_trips_through_storage() { - use asap_summary_state::summary_kernels::SumAccumulator; + fn planner_exact_min_panes_feed_the_min_rollup() { + use planner_types::post_asap::ExactKind; + let idx = SketchStore::new(); + let mut min_meta = meta(43); + min_meta.capability = None; + min_meta.accuracy = None; + min_meta.agg_kind = AggKind::ExactAgg { + agg_type: AggregationType::Min, + parameters_canonical: String::new(), + spatial_filter_canonical: String::new(), + }; + idx.register(min_meta); + for (start, value) in [(0, 7.0), (1000, 3.0)] { + idx.append_precompute( + 43, + BTreeMap::new(), + (start, start + 1000), + Box::new(asap_summary_state::stored_state::codec::exact_value( + ExactKind::Min, + value, + )), + ); + } + assert_eq!( + idx.query_rollup_range(RollupReduction::Min, 43, 0, 2000), + Some(vec![(BTreeMap::new(), 3.0)]) + ); + assert_eq!( + idx.query_rollup_range(RollupReduction::Max, 43, 0, 2000), + None + ); + } + #[test] + fn precompute_payload_round_trips_through_storage() { let idx = SketchStore::new(); let cfg = SketchConfig::DDSketch { relative_accuracy: 0.01, @@ -4642,7 +4671,7 @@ mod tests { 42, BTreeMap::new(), (1000, 1010), - Box::new(SumAccumulator::with_sum(5.0)), + Box::new(crate::tests::accumulator_fixture::sum_state(5.0)), ); // Sketch-side query_range filters out precompute payloads, so @@ -4662,8 +4691,6 @@ mod tests { #[test] fn query_precomputes_by_agg_returns_data_grouped_by_label_values() { - use asap_summary_state::summary_kernels::SumAccumulator; - let idx = SketchStore::new(); let cfg = SketchConfig::DDSketch { relative_accuracy: 0.01, @@ -4689,13 +4716,13 @@ mod tests { 99, lv.clone(), (1000, 2000), - Box::new(SumAccumulator::with_sum(1.0)), + Box::new(crate::tests::accumulator_fixture::sum_state(1.0)), ); idx.append_precompute( 99, lv, (2000, 3000), - Box::new(SumAccumulator::with_sum(2.0)), + Box::new(crate::tests::accumulator_fixture::sum_state(2.0)), ); let result = idx.query_precomputes_by_agg("cpu_seconds", AggregationType::Sum, 0, 10_000); @@ -4815,9 +4842,8 @@ mod tests { }); assert!(sketch.as_sketch().is_some()); assert!(sketch.as_exact_agg().is_none()); - - use asap_summary_state::summary_kernels::SumAccumulator; - let exact_agg = AggPayload::ExactAgg(Arc::new(SumAccumulator::with_sum(1.0))); + let exact_agg = + AggPayload::ExactAgg(Arc::new(crate::tests::accumulator_fixture::sum_state(1.0))); assert!(exact_agg.as_sketch().is_none()); assert!(exact_agg.as_exact_agg().is_some()); } @@ -5628,7 +5654,7 @@ mod tests { 850, BTreeMap::new(), (0, 30_000), - Box::new(asap_summary_state::summary_kernels::SumAccumulator::new()) + Box::new(crate::tests::accumulator_fixture::sum_state(0.0)) )); // A flusher that captured metadata before completion cannot reopen it. writer.upsert_all(&[stale_record]).unwrap(); @@ -6096,11 +6122,7 @@ mod tests { 8100, lv_zone("z0"), (s, s + 30_000), - Box::new( - asap_summary_state::summary_kernels::SumAccumulator::with_sum( - (i + 1) as f64, - ), - ), + Box::new(crate::tests::accumulator_fixture::sum_state((i + 1) as f64)), ); } assert!( @@ -6315,9 +6337,7 @@ mod tests { 8001, lv_zone("z0"), (s, s + 30_000), - Box::new( - asap_summary_state::summary_kernels::SumAccumulator::with_sum((i + 1) as f64), - ), + Box::new(crate::tests::accumulator_fixture::sum_state((i + 1) as f64)), ); } assert!( @@ -6380,9 +6400,9 @@ mod tests { lv_zone("z0"), (s, s + 30_000), Box::new({ - let mut acc = asap_summary_state::summary_kernels::SumAccumulator::new(); - acc.update((i + 1) as f64); - acc.update(10.0); + let mut acc = crate::tests::accumulator_fixture::sum_state(0.0); + acc.update(None, (i + 1) as f64, 0); + acc.update(None, 10.0, 0); acc }), ); @@ -6405,10 +6425,10 @@ mod tests { samples.contains_key(&30_000), "evicted exact-agg window missing from disk" ); - let stats = samples[&30_000].aux_stats(); - assert_eq!(stats.count, Some(2)); - assert_eq!(stats.sum, Some(11.0)); - assert_eq!(stats.sum.unwrap() / stats.count.unwrap() as f64, 5.5); + assert_eq!( + crate::tests::accumulator_fixture::sum_of(samples[&30_000].as_ref()), + 11.0 + ); drop(p); } @@ -6574,7 +6594,7 @@ mod tests { #[test] fn planner_exact_families_survive_disk_eviction_and_restart() { use crate::storage_engines::types::{AggregateCore, AggregationType}; - use asap_summary_state::summary_kernels::exact::ExactAccumulator; + use asap_physical_operators::summary_kernels::exact::ExactAccumulator; let kinds = [ AggregationType::Sum, AggregationType::Count, @@ -6641,17 +6661,14 @@ mod tests { assert_eq!(series.len(), 1, "{kind:?}"); let state = &series[0].1[&30000]; assert_eq!(state.get_accumulator_type(), *kind); + let exact = state.as_any().downcast_ref::().unwrap(); assert_eq!( - state - .query_statistic(stats[i], &None, &HashMap::new()) - .unwrap(), - expected[i] + exact.readout(stats[i], None, None).unwrap(), + Some(expected[i]) ); for (j, stat) in stats.iter().enumerate() { if i != j { - assert!(state - .query_statistic(*stat, &None, &HashMap::new()) - .is_err()); + assert!(exact.readout(*stat, None, None).is_err()); } } } diff --git a/data_plane/src/storage_engines/sketch_db/index/native.rs b/data_plane/src/storage_engines/sketch_db/index/native.rs index 7488c7860..636be0961 100644 --- a/data_plane/src/storage_engines/sketch_db/index/native.rs +++ b/data_plane/src/storage_engines/sketch_db/index/native.rs @@ -4,13 +4,12 @@ use super::*; use crate::drivers::ingest::series_resolver::SeriesIdResolver; use asap_physical_operators::values::{Batch, Schema, Value}; use asap_summary_state::{ - stored_state::native::{decode_batch, encode_batch}, - AggregateCore, SerializableToSink, + stored_state::native::{decode_batch, encode_batch, NativeSummaryOutput, NATIVE_OUTPUT_TYPE}, + AggregateCore, StoredState, }; #[cfg(test)] use asap_types::physical_plan_codec::PhysicalPlanCodec; -const NATIVE_OUTPUT_TYPE: &str = "NativePhysicalOutputV1"; const NATIVE_OUTPUT_TAG: u8 = persistence::part::encoding_tag::NATIVE_BATCH_V1; /// Preserve execution failures across the storage boundary; missing state remains @@ -33,134 +32,6 @@ impl From<&'static str> for NativeReadError { } } -#[derive(Clone)] -struct NativeSummaryOutput { - batch: Batch, - bytes: Vec, - kind: AggregationType, -} -impl NativeSummaryOutput { - fn new(batch: Batch, max_bytes: usize) -> Result { - let families = batch - .schema() - .fields - .iter() - .filter_map(|field| { - (!matches!( - field.dtype, - planner_types::post_asap::SummaryFamilyType::Plain(_) - )) - .then_some(&field.dtype) - }) - .collect::>(); - let [family] = families.as_slice() else { - return Err("native stored batch requires one summary column".into()); - }; - use planner_types::post_asap::{SketchAlgorithm, SummaryFamilyType}; - let schema_kind = match family { - SummaryFamilyType::Sketch(sketch, _) => match sketch.algorithm() { - SketchAlgorithm::CmsWithHeap => Some(AggregationType::CountMinSketchWithHeap), - SketchAlgorithm::CountSketchWithHeap => Some(AggregationType::CountSketchWithHeap), - _ => None, - }, - SummaryFamilyType::ExactAggregate(planner_types::post_asap::ExactKind::Sum, _) => { - Some(AggregationType::Sum) - } - _ => None, - }; - let mut kind = schema_kind; - for row in batch.rows() { - let states = row - .iter() - .filter_map(|value| match value { - Value::Summary { state, .. } => Some(state), - _ => None, - }) - .collect::>(); - let [state] = states.as_slice() else { - return Err("native stored row requires one summary state".into()); - }; - let row_kind = asap_summary_state::physical::aggregation_type(state.as_ref()) - .map_err(|error| error.to_string())?; - if kind.is_some_and(|kind| kind != row_kind) { - return Err("native stored rows have different summary families".into()); - } - kind = Some(row_kind); - } - let kind = kind.ok_or("empty native batch has no supported summary family")?; - let bytes = encode_batch(&batch).map_err(|error| error.to_string())?; - if bytes.len() > max_bytes || batch.bytes() > max_bytes { - return Err("native summary exceeds publication/read budget".into()); - } - Ok(Self { batch, bytes, kind }) - } - fn validate_group(&self, group: &BTreeMap) -> Result<(), String> { - for (key, value) in group { - let column = self - .batch - .schema() - .fields - .iter() - .position(|field| &field.name == key) - .ok_or("native output is missing its stored group key")?; - if self - .batch - .rows() - .iter() - .any(|row| !matches!(&row[column], Value::Utf8(actual) if actual.as_ref() == value)) - { - return Err("native output group differs from stored address".into()); - } - } - Ok(()) - } -} -impl SerializableToSink for NativeSummaryOutput { - fn serialize_to_bytes(&self) -> Vec { - self.bytes.clone() - } - fn serialize_to_json(&self) -> serde_json::Value { - serde_json::json!({"format": NATIVE_OUTPUT_TYPE, "bytes": self.bytes}) - } -} -impl AggregateCore for NativeSummaryOutput { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - fn type_name(&self) -> &'static str { - NATIVE_OUTPUT_TYPE - } - fn as_any(&self) -> &dyn std::any::Any { - self - } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } - fn get_accumulator_type(&self) -> AggregationType { - self.kind - } - fn get_keys(&self) -> Option> { - None - } - fn approx_memory_bytes(&self) -> usize { - self.bytes.len() + self.batch.bytes() - } - fn merge_with( - &self, - _: &dyn AggregateCore, - ) -> Result, Box> { - Err("native output snapshots require an explicit physical merge operator".into()) - } - fn query_statistic( - &self, - _: asap_types::Statistic, - _: &Option, - _: &HashMap, - ) -> Result> { - Err("native output readout requires the installed physical DAG".into()) - } -} - impl SketchStore { /// Publish a finalized native result only after the complete raw input /// cohort is durable. Existing publication fences prevent duplicate commits. @@ -188,7 +59,7 @@ impl SketchStore { let state = NativeSummaryOutput::new(batch, max_bytes)?; let family = config.accumulator_spec().map_err(|e| e.to_string())?.family; if state - .batch + .batch() .schema() .fields .iter() @@ -202,7 +73,7 @@ impl SketchStore { { return Err("native output schema differs from installed definition".into()); } - for value in state.batch.rows().iter().flatten() { + for value in state.batch().rows().iter().flatten() { if let Value::Summary { family: actual, .. } = value { if actual != &family { return Err("native output family differs from installed definition".into()); @@ -385,7 +256,7 @@ impl SketchStore { .as_any() .downcast_ref::() .ok_or("native output decoder mismatch")?; - Ok(state.batch.clone()) + Ok(state.batch().clone()) } } @@ -541,8 +412,8 @@ mod tests { let mut output = PrecomputedOutput::new(0, 60_000, None, source.policy_fingerprint()); output.population_labels = Some(group); output.catalog_generation = Some(generation.clone()); - let mut state = asap_summary_state::summary_kernels::SumAccumulator::new(); - state.update(value); + let mut state = crate::tests::accumulator_fixture::sum_state(0.0); + state.update(None, value, 0); store .publish_admitted_summary_update( &generation, @@ -571,11 +442,9 @@ mod tests { .iter() .flat_map(|input| input.windows.values()) .map(|state| { - vec![Value::Float64( - state - .query_statistic(asap_types::Statistic::Sum, &None, &HashMap::new()) - .unwrap(), - )] + vec![Value::Float64(crate::tests::accumulator_fixture::sum_of( + state.as_ref(), + ))] }) .collect(); let raw_schema = Arc::new(SummarySchema { diff --git a/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs b/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs index fd8e8b63e..93b0ed5c8 100644 --- a/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs +++ b/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs @@ -194,7 +194,6 @@ pub fn warn_if_retention_inverted( mod tests { use super::*; use crate::storage_engines::types::{AggregationType, InstalledPrecomputePlan}; - use asap_summary_state::summary_kernels::SumAccumulator; use asap_types::aggregation_config::PrecomputeMaterialization; use asap_types::enums::WindowKind; use asap_types::KeyByLabelNames; @@ -258,7 +257,7 @@ mod tests { ) -> u64 { use crate::drivers::ingest::series_resolver::SeriesIdResolver; use std::sync::Arc; - let acc = SumAccumulator::with_sum(1.0); + let acc = crate::tests::accumulator_fixture::sum_state(1.0); let output = crate::storage_engines::types::PrecomputedOutput::new( ts, ts + 1000, diff --git a/data_plane/src/storage_engines/sketch_db/query/window_merger.rs b/data_plane/src/storage_engines/sketch_db/query/window_merger.rs index 5df9853d2..7da26ae24 100644 --- a/data_plane/src/storage_engines/sketch_db/query/window_merger.rs +++ b/data_plane/src/storage_engines/sketch_db/query/window_merger.rs @@ -107,8 +107,6 @@ pub fn create_window_merger(_accumulator_type: AggregationType) -> Box Value { - serde_json::json!({"value": self.value}) - } - - fn serialize_to_bytes(&self) -> Vec { - self.value.to_le_bytes().to_vec() - } - } - impl AggregateCore for MockSumAccumulator { fn clone_boxed_core(&self) -> Box { Box::new(self.clone()) } - fn type_name(&self) -> &'static str { - "MockSumAccumulator" - } - fn as_any(&self) -> &dyn Any { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } - fn merge_with( &self, other: &dyn AggregateCore, @@ -162,23 +142,6 @@ mod tests { Err("Cannot merge with different accumulator type".into()) } } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::Sum - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - _statistic: asap_types::Statistic, - _key: &Option, - _query_kwargs: &std::collections::HashMap, - ) -> Result> { - Err("MockSumAccumulator does not support query_statistic".into()) - } } // Basic structure tests diff --git a/data_plane/src/storage_engines/types/mod.rs b/data_plane/src/storage_engines/types/mod.rs index 235b006ee..53397f964 100644 --- a/data_plane/src/storage_engines/types/mod.rs +++ b/data_plane/src/storage_engines/types/mod.rs @@ -30,4 +30,4 @@ pub use crate::query_engines::routing::{ QueryOperatorShape, RoutingTarget, }; -pub use asap_summary_state::{traits::*, KeyByLabelValues, Measurement}; +pub use asap_summary_state::{AggregateCore, KeyByLabelValues, Measurement, StoredState}; diff --git a/data_plane/src/tests/accumulator_fixture.rs b/data_plane/src/tests/accumulator_fixture.rs index 9a2d92686..4df1d3018 100644 --- a/data_plane/src/tests/accumulator_fixture.rs +++ b/data_plane/src/tests/accumulator_fixture.rs @@ -1,6 +1,7 @@ //! Config fixtures for backend integration tests; production binds Planner payloads. -use asap_summary_state::factory::*; -use asap_summary_state::{AggregateCore, AggregationType}; +use asap_physical_operators::factory::*; +use asap_physical_operators::summary_kernels::exact::ExactAccumulator; +use asap_summary_state::{AggregateCore, KeyByLabelValues}; use asap_types::{accumulator_spec::cms_params, PrecomputeMaterialization}; use planner_types::post_asap::{ExactKind, SketchAlgorithm, SketchParams, SummaryFamilyType}; #[cfg(test)] @@ -109,40 +110,9 @@ pub fn create_fixture_accumulator( let keyed = spec.grouping.is_some(); match (&spec.family, keyed) { - (SummaryFamilyType::ExactAggregate(ExactKind::Sum | ExactKind::Count, _), false) => { - Box::new(SumAccumulatorUpdater::new()) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Sum, _), true) => { - Box::new(KeyedSumCountAccumulatorUpdater::for_family(ExactKind::Sum)) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Count, _), true) => Box::new( - KeyedSumCountAccumulatorUpdater::for_family(ExactKind::Count), - ), - - // Direction comes off the family itself now. It used to be read - // back out of `aggregation_sub_type` because Planner had one - // `MinMax` accumulator for both directions, which meant a config - // whose sub_type was lost or misspelled silently built the wrong - // extremum. - (SummaryFamilyType::ExactAggregate(ExactKind::Min, _), false) => { - Box::new(MinAccumulatorUpdater::new()) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Min, _), true) => { - Box::new(KeyedMinStateUpdater::new()) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Max, _), false) => { - Box::new(MaxAccumulatorUpdater::new()) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Max, _), true) => { - Box::new(KeyedMaxStateUpdater::new()) - } - - (SummaryFamilyType::ExactAggregate(ExactKind::Increase | ExactKind::Rate, _), false) => { - Box::new(IncreaseAccumulatorUpdater::new()) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Increase | ExactKind::Rate, _), true) => { - Box::new(KeyedCounterStateUpdater::new()) - } + (SummaryFamilyType::ExactAggregate(..), keyed) => Box::new(ExactUpdater { + acc: ExactAccumulator::new(spec.family.clone(), keyed).expect("exact fixture family"), + }), (SummaryFamilyType::Sketch(kind, _), false) if kind.algorithm() == &SketchAlgorithm::Kll => @@ -232,7 +202,7 @@ pub fn create_fixture_accumulator( else { unreachable!("validated UnivMon family parameters") }; - asap_summary_state::factory::create_planner_accumulator( + create_planner_accumulator( &spec.family, &planner_types::post_asap::SummaryUpdate::column( planner_types::pre_asap::ColumnRef::SampleValue, @@ -248,7 +218,7 @@ pub fn create_fixture_accumulator( let SketchParams::Hll { precision } = kind.params() else { unreachable!("validated HLL family parameters") }; - asap_summary_state::factory::create_planner_accumulator( + create_planner_accumulator( &spec.family, &planner_types::post_asap::SummaryUpdate::column( planner_types::pre_asap::ColumnRef::SampleValue, @@ -263,3 +233,90 @@ pub fn create_fixture_accumulator( } } } + +/// Feeds Planner's exact state directly; Planner keeps its own exact updater private. +#[cfg(test)] +struct ExactUpdater { + acc: ExactAccumulator, +} + +#[cfg(test)] +impl AccumulatorUpdater for ExactUpdater { + fn update_single(&mut self, value: f64, timestamp_ms: i64) { + self.acc.update(None, value, timestamp_ms); + } + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.acc.update(Some(key), value, timestamp_ms); + } + fn take_accumulator(&mut self) -> Box { + let taken = Box::new(self.acc.clone()); + self.reset(); + taken + } + fn snapshot_accumulator(&self) -> Box { + Box::new(self.acc.clone()) + } + fn reset(&mut self) { + self.acc = ExactAccumulator::new(self.acc.family().clone(), self.acc.is_keyed()) + .expect("exact fixture family"); + } + fn is_keyed(&self) -> bool { + self.acc.is_keyed() + } + fn memory_usage_bytes(&self) -> usize { + self.acc.approx_memory_bytes() + } +} + +/// Planner's unkeyed exact Sum holding `sum`. +#[cfg(test)] +pub fn sum_state(sum: f64) -> ExactAccumulator { + asap_summary_state::stored_state::codec::exact_value(ExactKind::Sum, sum) +} + +/// The value of an unkeyed exact Sum state. +#[cfg(test)] +pub fn sum_of(state: &dyn AggregateCore) -> f64 { + state + .as_any() + .downcast_ref::() + .expect("exact Sum state") + .readout(asap_types::Statistic::Sum, None, None) + .expect("exact Sum readout") + .expect("present Sum population") +} + +/// Test view of an unkeyed exact Sum state's value. +#[cfg(test)] +pub struct SumView { + pub sum: f64, +} + +/// Panics unless `state` is an unkeyed exact Sum. +#[cfg(test)] +pub fn sum_view(state: &ExactAccumulator) -> SumView { + SumView { + sum: state + .readout(asap_types::Statistic::Sum, None, None) + .expect("exact Sum readout") + .expect("present exact population"), + } +} + +/// Planner's unkeyed exact counter state (Rate or Increase) over `samples` +/// of `(timestamp_ms, value)`. +#[cfg(test)] +pub fn counter_state(kind: ExactKind, samples: &[(i64, f64)]) -> ExactAccumulator { + use planner_types::post_asap::ExactParams; + let params = match kind { + ExactKind::Rate => ExactParams::Rate, + ExactKind::Increase => ExactParams::Increase, + other => panic!("{other:?} is not a counter family"), + }; + let mut state = ExactAccumulator::new(SummaryFamilyType::ExactAggregate(kind, params), false) + .expect("counter family"); + for (timestamp, value) in samples { + state.update(None, *value, *timestamp); + } + state +} diff --git a/data_plane/src/tests/mod.rs b/data_plane/src/tests/mod.rs index e016a48f2..288bf2651 100644 --- a/data_plane/src/tests/mod.rs +++ b/data_plane/src/tests/mod.rs @@ -1,7 +1,6 @@ pub mod accuracy_empirical_validation_tests; pub mod accuracy_in_promql_response_tests; pub mod prometheus_forwarding_tests; -pub mod trait_design_tests; #[cfg(test)] pub mod test_utilities; diff --git a/data_plane/src/tests/trait_design_tests.rs b/data_plane/src/tests/trait_design_tests.rs deleted file mode 100644 index 8e2fafae6..000000000 --- a/data_plane/src/tests/trait_design_tests.rs +++ /dev/null @@ -1,81 +0,0 @@ -#[cfg(test)] -use crate::storage_engines::types::{ - KeyByLabelValues, MultipleSubpopulationAggregate, SingleSubpopulationAggregate, -}; -use asap_summary_state::summary_kernels::{KeyedSumCountAccumulator, SumAccumulator}; -use asap_types::Statistic; - -#[test] -fn test_single_subpopulation_interface() { - // Single accumulator - matches Python behavior exactly - let acc: Box = Box::new(SumAccumulator::with_sum(42.0)); - - // ✅ Query without key - this is the correct interface for Single accumulators - let result = acc.query(Statistic::Sum, None).unwrap(); - assert_eq!(result, 42.0); -} - -#[test] -fn test_multiple_subpopulation_interface() { - // Multiple accumulator - matches Python behavior exactly - let mut multi_acc = KeyedSumCountAccumulator::new(); - - let mut key = KeyByLabelValues::new(); - key.insert("web".to_string()); - multi_acc.add_sum(key.clone(), 100.0); - - let acc: Box = Box::new(multi_acc); - - // ✅ Query with key - this is the correct interface for Multiple accumulators - let result = acc.query(Statistic::Sum, &key, None).unwrap(); - assert_eq!(result, 100.0); - - // ✅ Get all keys - let keys = acc.get_keys().unwrap(); - assert_eq!(keys.len(), 1); - assert_eq!(keys[0], key); -} - -#[test] -fn test_interface_prevents_misuse() { - // This test documents what WON'T compile - which is exactly what we want! - - let single_acc: Box = - Box::new(SumAccumulator::with_sum(42.0)); - let multi_acc: Box = - Box::new(KeyedSumCountAccumulator::new()); - - // ✅ These work - correct usage - let _result1 = single_acc.query(Statistic::Sum, None); - let key = KeyByLabelValues::new(); - let _result2 = multi_acc.query(Statistic::Sum, &key, None); - - // ❌ These would be compile-time errors (commented out): - // let _result3 = single_acc.query(Statistic::Sum, &key); // Too many args for Single - // let _result4 = multi_acc.query(Statistic::Sum); // Too few args for Multiple - - // This is exactly the type safety we wanted to achieve! -} - -#[test] -fn test_python_alignment() { - // Demonstrate that the Rust interface now matches Python exactly - - // Python: sum_accumulator.query(Statistic.SUM) - // Rust: sum_accumulator.query(Statistic::Sum) - let sum_acc: Box = Box::new(SumAccumulator::with_sum(42.0)); - assert_eq!(sum_acc.query(Statistic::Sum, None).unwrap(), 42.0); - - // Python: multiple_accumulator.query(Statistic.SUM, key) - // Rust: multiple_accumulator.query(Statistic::Sum, &key) - let mut multi_acc = KeyedSumCountAccumulator::new(); - let key = KeyByLabelValues::new(); - multi_acc.add_sum(key.clone(), 100.0); - let multi_trait: Box = Box::new(multi_acc); - assert_eq!( - multi_trait.query(Statistic::Sum, &key, None).unwrap(), - 100.0 - ); - - // Perfect alignment with Python behavior! -} diff --git a/data_plane/tests/edge_sketch_codec.rs b/data_plane/tests/edge_sketch_codec.rs index f4d758b98..306e8bec5 100644 --- a/data_plane/tests/edge_sketch_codec.rs +++ b/data_plane/tests/edge_sketch_codec.rs @@ -35,18 +35,10 @@ fn ddsketch_bare_state_is_rejected_and_envelope_supports_query_readout() { let bare = prost::Message::encode_to_vec(&state); assert!(asap_sketch_codec::reconstruct_ddsketch(&bare).is_err()); let (decoded, _) = asap_sketch_codec::reconstruct_ddsketch(&envelope).unwrap(); - let accumulator = asap_summary_state::summary_kernels::DDSketchAccumulator { - inner: decoded, - sample_p: 1.0, - }; + let accumulator = + asap_physical_operators::summary_kernels::DDSketchAccumulator { inner: decoded }; let median = accumulator - .query_statistic( - asap_types::Statistic::Quantile, - &None, - &[("quantile".to_string(), "0.5".to_string())] - .into_iter() - .collect(), - ) + .estimate(&planner_types::post_asap::SketchQuery::Quantile { q: 0.5 }) .unwrap(); assert!((median - 50.0).abs() / 50.0 < 0.05); } @@ -62,7 +54,9 @@ fn kll_envelope_keeps_level_layout_for_backend_readout() { assert_eq!(state.k, 200); assert_eq!(state.items.len(), 50); let snapshot_bytes = bytes; - let accumulator = asap_summary_state::summary_kernels::DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&snapshot_bytes).unwrap(); + let accumulator = asap_physical_operators::summary_kernels::DatasketchesKLLAccumulator { + inner: asap_summary_state::stored_state::decoders::kll_from_proto(&snapshot_bytes).unwrap(), + }; assert!(accumulator.get_quantile(0.5).is_finite()); } diff --git a/data_plane/tests/support/univmon_erp_process.rs b/data_plane/tests/support/univmon_erp_process.rs index 73330a7b8..9ca9ea769 100644 --- a/data_plane/tests/support/univmon_erp_process.rs +++ b/data_plane/tests/support/univmon_erp_process.rs @@ -1,7 +1,7 @@ use super::*; -use asap_summary_state::summary_kernels::univmon::UnivMonAccumulator; +use asap_summary_state::univmon::UnivMonAccumulator; use control_plane::physical::erp::ErpShapeObserver; -use data_plane::storage_engines::types::{AggregateCore, SerializableToSink}; +use data_plane::storage_engines::types::AggregateCore; fn values(offset: usize) -> Vec { (1..=128) @@ -48,18 +48,16 @@ fn measured_artifact() -> Value { } let other = panes[1].clone(); panes[0].merge_in_place(&other).unwrap(); - bytes = bytes.max(panes[0].serialize_to_bytes().len()); + bytes = bytes.max(panes[0].to_bytes().unwrap().len()); for (i, stat) in [ - asap_types::Statistic::Cardinality, - asap_types::Statistic::FrequencyL2, - asap_types::Statistic::FrequencyEntropy, + planner_types::post_asap::SketchQuery::Cardinality, + planner_types::post_asap::SketchQuery::FrequencyL2, + planner_types::post_asap::SketchQuery::FrequencyEntropy, ] .into_iter() .enumerate() { - let estimate = panes[0] - .query_statistic(stat, &None, &Default::default()) - .unwrap(); + let estimate = panes[0].estimate(&stat).unwrap(); assert!(estimate.is_finite()); let error = (estimate - exact[i]).abs() / if i == 2 { 1.0 } else { exact[i] }; errors[i] = errors[i].max(error); From 787573cfa192f25bf13ef1afd83cc4e328fe0e25 Mon Sep 17 00:00:00 2001 From: zzylol Date: Wed, 30 Sep 2026 13:23:37 +0000 Subject: [PATCH 2/4] fix(store): stop serving a MIN/MAX rollup after a non-contiguous pane A late correction, replacement or gap was dropped from the rollup, and the next contiguous pane re-validated it, so the rollup could serve an extremum that ignored the correction. Such a pane now invalidates the series' rollup permanently and readers fall back to the exact panes. Co-Authored-By: Claude Opus 5.5 --- .../storage_engines/sketch_db/index/mod.rs | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/data_plane/src/storage_engines/sketch_db/index/mod.rs b/data_plane/src/storage_engines/sketch_db/index/mod.rs index b476f1430..bb47a012d 100644 --- a/data_plane/src/storage_engines/sketch_db/index/mod.rs +++ b/data_plane/src/storage_engines/sketch_db/index/mod.rs @@ -300,6 +300,10 @@ struct ReductionRollupSeries { base_width_ms: Option, next_start_ms: Option, valid: bool, + /// A pane that does not extend the series contiguously (a late correction, + /// a replacement or a gap) is not folded in, so the rollup can no longer + /// answer and stays invalid; readers fall back to the exact panes. + broken: bool, levels: Vec>, } @@ -311,6 +315,7 @@ impl ReductionRollupSeries { base_width_ms: None, next_start_ms: None, valid: false, + broken: false, levels: Vec::new(), } } @@ -318,9 +323,13 @@ impl ReductionRollupSeries { impl ReductionRollupSeries { fn append(&mut self, window: TimestampRange, value: f64, retention_horizon_ms: Option) { + if self.broken { + return; + } let width = window.1.saturating_sub(window.0); if width == 0 || self.next_start_ms.is_some_and(|next| next != window.0) { self.valid = false; + self.broken = true; return; } let anchor = *self.anchor_start_ms.get_or_insert(window.0); @@ -330,6 +339,7 @@ impl ReductionRollupSeries { || !(window.0 - anchor).is_multiple_of(base_width) { self.valid = false; + self.broken = true; return; } self.valid = true; @@ -3880,6 +3890,19 @@ pub use crate::storage_engines::sketch_db::persistence; mod tests { use super::*; + // A late correction to a closed pane makes the rollup stop answering + // instead of serving the uncorrected extremum. + #[test] + fn min_rollup_stops_answering_after_a_late_correction() { + let mut rollup = ReductionRollupSeries::new(RollupReduction::Min); + rollup.append((0, 1000), 7.0, None); + rollup.append((1000, 2000), 3.0, None); + assert_eq!(rollup.query(0, 2000), Some(3.0)); + rollup.append((0, 1000), 1.0, None); + rollup.append((2000, 3000), 5.0, None); + assert_eq!(rollup.query(0, 3000), None); + } + #[test] fn max_rollup_answers_aligned_and_partial_ranges_and_prunes_history() { let mut rollup = ReductionRollupSeries::new(RollupReduction::Max); From 2a2934e5befd1077891f9e617f7cedcf13149252 Mon Sep 17 00:00:00 2001 From: zzylol Date: Wed, 30 Sep 2026 13:23:37 +0000 Subject: [PATCH 3/4] docs: describe asap_summary_state as codecs over Planner kernels Co-Authored-By: Claude Opus 5.5 --- docs/design_docs/physical-operators.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/design_docs/physical-operators.md b/docs/design_docs/physical-operators.md index a9f514763..5d00727e2 100644 --- a/docs/design_docs/physical-operators.md +++ b/docs/design_docs/physical-operators.md @@ -13,10 +13,10 @@ changed and tested together in Planner. Planner does not own storage formats. The backend owns them: -- `crates/asap_summary_state`: the summary kernels that ingest and the sketch - store keep, their stored byte encodings, delta reconstruction, native batch - frames, and conversion to Planner physical states where stored state enters - or leaves a physical DAG. +- `crates/asap_summary_state`: the stored byte encodings of Planner kernel + states (the store keeps Planner's kernels directly), edge wire decoding, + delta reconstruction, native batch frames, and statistic readout binding. + It implements no summary algorithm. - `crates/asap_sketch_codec`: the sketchlib `SketchEnvelope` codec. - `asap_types::physical_plan_codec`: the versioned envelope of persisted physical plans and candidates. From 173734852d1d8f3e59d695d0c5d3384da22ccf0f Mon Sep 17 00:00:00 2001 From: zzylol Date: Wed, 30 Sep 2026 13:24:05 +0000 Subject: [PATCH 4/4] docs: stop naming deleted kernels in data-plane comments Co-Authored-By: Claude Opus 5.5 --- data_plane/src/drivers/ingest/otel.rs | 2 +- data_plane/src/precompute_engine/worker.rs | 2 +- data_plane/src/query_engines/asap_query_engine/engine.rs | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/data_plane/src/drivers/ingest/otel.rs b/data_plane/src/drivers/ingest/otel.rs index 1008da83d..5a3ef0eb4 100644 --- a/data_plane/src/drivers/ingest/otel.rs +++ b/data_plane/src/drivers/ingest/otel.rs @@ -1512,7 +1512,7 @@ async fn route_modified_otlp_sketches_to_precompute( // starts fresh, so the reconstructed `state(N)` is // window N only. Sketch-agnostic: the reset is the // additive families' (DDSketch / CMS / CountSketch / - // HLL) `AggregateCore::reset_to_empty`; KLL never + // HLL) `codec::empty_like`; KLL never // deltas. Full frames keep REPLACE semantics and set // the stored `window_start`. let accumulator: Box = if dp.encoding == ENCODING_PROTO_DELTA diff --git a/data_plane/src/precompute_engine/worker.rs b/data_plane/src/precompute_engine/worker.rs index 010fa9111..d1b1591bb 100644 --- a/data_plane/src/precompute_engine/worker.rs +++ b/data_plane/src/precompute_engine/worker.rs @@ -2026,7 +2026,7 @@ mod tests { } // ----------------------------------------------------------------------- - // Test: raw mode — each sample forwarded as SumAccumulator with sum==value + // Test: raw mode — each sample forwarded as an exact Sum with sum==value // ----------------------------------------------------------------------- #[test] diff --git a/data_plane/src/query_engines/asap_query_engine/engine.rs b/data_plane/src/query_engines/asap_query_engine/engine.rs index bc0b3976a..8b57629bc 100644 --- a/data_plane/src/query_engines/asap_query_engine/engine.rs +++ b/data_plane/src/query_engines/asap_query_engine/engine.rs @@ -1575,8 +1575,8 @@ mod sketch_query_tests { // // Pins that the warm engine answers `count(metric)` from an HLL-backed // aggregation: capability matching picks HLL (per -// `compatible_agg_types(Statistic::Count)`), and the HLL accumulator's -// `query_statistic` returns the cardinality estimate. This is the +// `compatible_agg_types(Statistic::Count)`), and the HLL state's readout +// returns the cardinality estimate. This is the // runtime contract the wire-side _hll alias resolver above relies on. // =========================================================================== // =========================================================================== @@ -1615,7 +1615,7 @@ mod asap_tier_classify_tests { let idx = Arc::new(SketchStore::new()); // Mirror the acceptance-test setup: four ExactAgg(Sum) sids, one // per zone (z0..z3), registered with `group_by_keys=["zone"]` - // and carrying a `SumAccumulator` per window. + // and carrying an exact Sum state per window. let zones = ["z0", "z1", "z2", "z3"]; // Anchor windows so the engine's instant-query default // lookback (5 min) reaches them. @@ -1684,7 +1684,7 @@ mod asap_tier_classify_tests { other => panic!("expected Vector, got {other:?}"), }; assert_eq!(vector.values.len(), 4, "one entry per zone"); - // Per-zone values match what each SumAccumulator carries. + // Per-zone values match what each exact Sum state carries. // KeyByLabelValues stores values only; the override carries // the corresponding keys. let mut by_zone: std::collections::HashMap = std::collections::HashMap::new();