diff --git a/Cargo.lock b/Cargo.lock index 28106fd5..af73a82c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -987,7 +987,9 @@ dependencies = [ "anyhow", "arc-swap", "arrow", + "asap-ir", "asap-precompute-rs", + "asap-sketch", "asap_otel_proto", "asap_sketchlib", "asap_types", diff --git a/data_plane/Cargo.toml b/data_plane/Cargo.toml index 985ea52f..1714c464 100644 --- a/data_plane/Cargo.toml +++ b/data_plane/Cargo.toml @@ -13,6 +13,24 @@ asap_types.workspace = true # compile in the workspace and importable from main.rs. control_plane = { path = "../control_plane" } +# Step C (data_plane/docs/l4node-plan-executor-design.md): data_plane +# implements `asap_sketch::exec::SummaryExecutor`, which means naming +# `asap_sketch::{L4Node, SketchQuery, exec::*}` types directly -- not +# reachable through `control_plane::asap_tier_implement`'s public +# `Vec>` return type alone (Rust requires importing a type's +# defining crate to implement traits on it or match its variants, even +# when a dependency's function already returns that type). Pin MUST +# match control_plane's `asap-sketch` pin exactly -- same rule as +# `crates/asap_types/Cargo.toml`'s existing pin comment: two revs of the +# same git dependency in one workspace resolve to two distinct Rust +# types that won't unify. +asap-sketch = { git = "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/ProjectASAP/ASAPController", rev = "d4c175633f6ce46b801ca2115c00fa757d5b6240" } +# `asap_sketch::L4Node`'s own fields (`SummaryExpr::Logical(Box)`, +# `SummaryAgg { col: ColumnRef, by: Vec, .. }`) are `asap-ir` +# types, not re-exported by `asap-sketch` -- `find_candidates` needs to +# walk/match them directly. Same pin-must-match rule as `asap-sketch` above. +asap-ir = { git = "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/ProjectASAP/ASAPController", rev = "d4c175633f6ce46b801ca2115c00fa757d5b6240" } + # Shared external (workspace) serde.workspace = true serde_json.workspace = true diff --git a/data_plane/src/query_engines/asap_query_engine/mod.rs b/data_plane/src/query_engines/asap_query_engine/mod.rs index 8ec546b9..f4b202d8 100644 --- a/data_plane/src/query_engines/asap_query_engine/mod.rs +++ b/data_plane/src/query_engines/asap_query_engine/mod.rs @@ -10,6 +10,7 @@ //! the JSONL leg has been deleted). pub mod engine; +pub mod summary_executor; // Phase-5 reorg: ASAP-tier reducer moved to `sketch_db::query`. The // engine still consumes it via that canonical path. 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 new file mode 100644 index 00000000..e96fd185 --- /dev/null +++ b/data_plane/src/query_engines/asap_query_engine/summary_executor.rs @@ -0,0 +1,1268 @@ +//! `data_plane`'s implementation of `asap_sketch::exec::SummaryExecutor` +//! — the serving-time interface that resolves an `L4Node` plan tree +//! against whatever is actually materialized right now. See +//! `data_plane/docs/l4node-plan-executor-design.md` for the surrounding +//! design. +//! +//! ## Scope +//! +//! Covers **quantile/cardinality queries** (DDSketch/Kll/Hll) **and the +//! Frequency family's bare total** (CMS/CountSketch/CMS-with-heap/ +//! CountSketch-with-heap, `count`/`sum` with no specific item key), both +//! cumulative (instant) and per-window (matrix/range). All modes do real +//! cross-sid merging via `delta_apply::SummaryState`: reconstruct each +//! candidate sid's own state over the range (or per window), then merge +//! same-window/same-range states *across* sids before reading out one +//! answer per group (or per group per window). +//! +//! Not covered, and reported as an explicit `Unsupported` error rather +//! than silently mishandled: +//! - `SketchQuery::TopK`. Its answer is fundamentally shaped differently +//! from everything else this executor reads out — K `(item, count)` +//! pairs per timestamp, not one scalar — so it doesn't fit +//! `Self::Value = Vec<(i64, f64)>`. Forcing it into that shape (e.g. +//! returning only the top item) would silently drop data rather than +//! error; that needs a `Value` type redesign, a separate decision from +//! this executor's rollout. +//! - `SketchQuery::PointCount` with a *named* item key (a point lookup +//! for one specific item, e.g. `count(cms_metric{item="x"})`). The +//! *value* to look up isn't carried by `SketchQuery` or available in +//! `readout`'s signature — `PointCount{key: ColumnRef}` names which +//! *column* is being queried, not the value to filter for, which would +//! come from a `Filter` predicate elsewhere in the tree. Resolving +//! that is a separate problem from this trait's scope. +//! `PointCount{key: ColumnRef::SampleValue}` (no specific item — the +//! bare bucket total) is covered. +//! - `ExactAgg` intents (`Sum`/`Rate`/`Increase`/`MinMax`/exact `Count`). +//! These don't reach `readout` at all — `asap_plan::bind` never wraps +//! an `ExactAccumulator` implementation in a `SummaryEstimate` +//! (`estimate = false` in `bind_summary_agg`), so `execute()` on such a +//! tree returns `ExecOutcome::State` at the root; the caller must read +//! the final value out of that `State` itself, not through this trait. + +use std::collections::{BTreeMap, BTreeSet}; +use std::rc::Rc; + +use asap_ir::intent_algebra::{ColumnId, ColumnRef, QueryExpr, Source}; +use asap_sketch::exec::SummaryExecutor; +use asap_sketch::{L4Node, SketchQuery, SummaryExpr, SummaryKind, SummaryParams}; + +use control_plane::sketch_algebra::capability::SketchKindHandle; + +use crate::storage_engines::sketch_db::data::{SketchConfig, SketchTimeSeries}; +use crate::storage_engines::sketch_db::index::{SketchSampleState, SketchStore}; +use crate::storage_engines::sketch_db::query::delta_apply::{ + cumulative_summary_state, per_window_summary_states, DeltaSketchKind, SummaryState, +}; + +/// Per-query, per-call execution context — constructed fresh for each +/// incoming query (never shared across concurrent queries, never +/// mutated after construction). This is what carries the time range and +/// cumulative-vs-per-window mode: `SummaryExecutor`'s trait methods take +/// no such parameters, and `ASAPQueryEngine` itself is called +/// concurrently (`Arc`), so threading the range through +/// shared mutable state on the engine would be a race — a fresh, +/// stack-local context per call is the safe alternative. +pub struct QueryExecutionContext<'a> { + pub index: &'a SketchStore, + pub t0_ms: u64, + pub t1_ms: u64, + /// `true` for `quantile_over_time`/`count_distinct_over_time`-shaped + /// instant queries (fold the whole range into one answer, via + /// `readout_cumulative`); `false` for a per-window matrix (one merged + /// answer per window, via `readout_per_window`). + pub is_cumulative: bool, +} + +/// One candidate sid, already carrying its `[t0, t1]` data and decode +/// parameters. `find_candidates` fetches this once per candidate (it +/// needs the sid's label values to build the group key anyway); folding +/// it into the handle means `fetch_state`/`readout` reuse it instead of +/// re-querying the same `(sid, t0, t1)` range a second time. `Rc` keeps +/// clones of the handle cheap (a refcount bump, not a re-fetch or a +/// re-clone of the sample bytes). +#[derive(Debug, Clone)] +pub struct SidHandle { + series: Rc, + kind: DeltaSketchKind, +} + +/// One group's accumulated candidates, all sharing one `DeltaSketchKind` +/// (guaranteed by `find_candidates`'s exact-match contract). +#[derive(Debug, Clone)] +pub struct GroupState { + entries: Vec, + kind: DeltaSketchKind, +} + +#[derive(Debug)] +pub enum SummaryExecutorError { + /// No sid in the catalog matches the requested `(metric, by, + /// SummaryKind, SummaryParams)` — mirrors today's `CapabilityMiss` + /// contract; the caller fails over to archive. + NoCandidates, + /// Couldn't recover a metric name by walking the `SummaryAgg`'s + /// child subtree — an unsupported/CSE-`Ref`-shaped `QueryExpr` this + /// executor doesn't walk through. + NoMetricFound, + /// A requested `by` `ColumnId` doesn't resolve to a name against the + /// child's schema. + UnresolvedColumn(ColumnId), + /// A candidate sid claims a `SummaryKind` this executor doesn't + /// implement cross-sid merge for, or the sid's on-disk + /// `SketchConfig` didn't decode into a `DeltaSketchKind`. + UnsupportedFamily, + /// Decode/merge failure surfaced from `delta_apply`/`asap_sketchlib`. + Decode(String), + /// A `SummaryExpr::Logical` node — nothing committed at L4. Same + /// meaning as today's "no candidate bound"; the caller fails over. + Logical, + /// A query shape not covered by this executor — see the module doc. + Unsupported(&'static str), +} + +impl<'a> SummaryExecutor for QueryExecutionContext<'a> { + type Handle = SidHandle; + type State = GroupState; + type Value = Vec<(i64, f64)>; + type Error = SummaryExecutorError; + type GroupKey = BTreeMap; + + fn find_candidates( + &self, + sketch: &SummaryKind, + params: &SummaryParams, + _col: &ColumnRef, + by: &[ColumnId], + child: &L4Node, + ) -> Result, Self::Error> { + let metric = find_metric(child).ok_or(SummaryExecutorError::NoMetricFound)?; + + let mut by_names: Vec = Vec::with_capacity(by.len()); + for &col_id in by { + let name = child + .schema + .fields + .get(col_id) + .map(|f| f.name.clone()) + .ok_or(SummaryExecutorError::UnresolvedColumn(col_id))?; + by_names.push(name); + } + let required_keys: BTreeSet = by_names.iter().cloned().collect(); + + let candidate_sids = self.index.instances_matching(&metric, &required_keys); + let mut out = Vec::new(); + for sid in candidate_sids { + let candidate_kind = self.index.with_instance(sid, |m| { + let kind = m.sketch_kind()?; + let config = m.sketch_config()?; + summary_params_match(sketch, params, kind, config) + .then(|| to_delta_kind(kind, config)) + .flatten() + }); + let Some(candidate_kind) = candidate_kind.flatten() else { + continue; + }; + + // Fetching the series here (rather than just checking + // membership) is what lets `fetch_state`/`readout` skip a + // second identical `query_range` call later -- see + // `SidHandle`'s doc. The label values it carries are also + // the only place a group's actual values live (metadata only + // has the group-by KEY names, not values). + let series = self.index.query_range(sid, self.t0_ms, self.t1_ms); + let Some(series) = series.into_iter().next() else { + continue; + }; + let group_key: BTreeMap = by_names + .iter() + .map(|k| { + let v = series + .series_label_values + .get(k) + .cloned() + .unwrap_or_default(); + (k.clone(), v) + }) + .collect(); + out.push(( + group_key, + SidHandle { + series: Rc::new(series), + kind: candidate_kind, + }, + )); + } + // Empty is NOT an error here -- `asap_sketch::exec::execute()` + // itself checks `find_candidates`'s result for emptiness and + // raises the canonical `ExecError::NoCandidates`; erroring here + // too would just wrap that in `ExecError::Executor(..)` instead, + // losing the distinction callers match on. + Ok(out) + } + + fn fetch_state(&self, handle: &Self::Handle) -> Result { + Ok(GroupState { + kind: handle.kind, + entries: vec![handle.clone()], + }) + } + + fn merge_states(&self, states: Vec) -> Result { + // The actual decode/merge math (`cumulative_summary_state`/ + // `merge_same_family`) happens in `readout`, not here: it needs + // to distinguish cumulative vs. per-window mode + // (`self.is_cumulative`), which only `readout` is positioned to + // do generically for both callers. `merge_states` and + // `fetch_state` just assemble the group's full candidate list. + let mut states = states.into_iter(); + let mut acc = states.next().ok_or(SummaryExecutorError::NoCandidates)?; + for s in states { + acc.entries.extend(s.entries); + } + Ok(acc) + } + + fn readout( + &self, + state: &Self::State, + query: &SketchQuery, + ) -> Result { + if self.is_cumulative { + readout_cumulative(state, query, self.t1_ms as i64) + } else { + readout_per_window(state, query, self.t0_ms as i64) + } + } + + fn logical(&self, _expr: &QueryExpr) -> Result { + Err(SummaryExecutorError::Logical) + } +} + +/// Fold a group's whole `[t0, t1]` into one merged state and read out one +/// scalar -- `quantile_over_time`/`count_distinct_over_time`-shaped +/// instant queries. +fn readout_cumulative( + state: &GroupState, + query: &SketchQuery, + t1_ms: i64, +) -> Result, SummaryExecutorError> { + let mut merged: Option = None; + let mut latest_window_end: Option = None; + for entry in &state.entries { + let samples_vec: Vec<(i64, &SketchSampleState)> = entry + .series + .samples + .iter() + .flat_map(|(t, frames)| frames.iter().map(move |s| (*t, s))) + .collect(); + if let Some((w, _)) = samples_vec.last() { + latest_window_end = Some(latest_window_end.map_or(*w, |prev| prev.max(*w))); + } + let rs = cumulative_summary_state(&samples_vec, state.kind) + .map_err(SummaryExecutorError::Decode)?; + if let Some(rs) = rs { + merged = Some(match merged.take() { + None => rs, + Some(mut acc) => { + acc.merge_same_family(&rs) + .map_err(SummaryExecutorError::Decode)?; + acc + } + }); + } + } + let Some(merged) = merged else { + return Err(SummaryExecutorError::NoCandidates); + }; + let value = sketch_query_value(&merged, query)?; + Ok(vec![(latest_window_end.unwrap_or(t1_ms), value)]) +} + +/// Per-window matrix/range-query readout: reconstruct each of the +/// group's sids' own per-window states, then merge same-window states +/// *across* sids before evaluating each window -- one merged answer per +/// window, not one merged answer for the whole range. Windows are +/// unioned across sids: a sid that's missing a particular window just +/// doesn't contribute to it, rather than the whole window being dropped. +fn readout_per_window( + state: &GroupState, + query: &SketchQuery, + t0_ms: i64, +) -> Result, SummaryExecutorError> { + let mut by_window: BTreeMap = BTreeMap::new(); + for entry in &state.entries { + let samples_vec: Vec<(i64, &SketchSampleState)> = entry + .series + .samples + .iter() + .flat_map(|(t, frames)| frames.iter().map(move |s| (*t, s))) + .collect(); + let (per_window, _skipped) = per_window_summary_states(&samples_vec, state.kind) + .map_err(SummaryExecutorError::Decode)?; + for (w_end, rs) in per_window { + // `SketchStore::query_range` may splice in a carry-in Full + // snapshot ending before the requested range so the + // delta-apply walk can establish a rolling base for a + // delta-only leading window; that base must not surface as + // an output point. + if w_end < t0_ms { + continue; + } + match by_window.get_mut(&w_end) { + Some(acc) => acc + .merge_same_family(&rs) + .map_err(SummaryExecutorError::Decode)?, + None => { + by_window.insert(w_end, rs); + } + } + } + } + if by_window.is_empty() { + return Err(SummaryExecutorError::NoCandidates); + } + by_window + .into_iter() + .map(|(w_end, rs)| sketch_query_value(&rs, query).map(|v| (w_end, v))) + .collect() +} + +/// Read one scalar out of a merged `SummaryState` for the requested +/// `SketchQuery` -- shared by both the cumulative and per-window readout +/// paths. +fn sketch_query_value(rs: &SummaryState, query: &SketchQuery) -> Result { + match query { + SketchQuery::Quantile { q } => Ok(rs.quantile(*q)), + SketchQuery::Cardinality => Ok(rs.cardinality()), + // `key: ColumnRef::SampleValue` means "no specific item" -- the + // bare bucket total. Any other column names an item to look up + // by VALUE, which isn't carried by `SketchQuery` -- see the + // module doc. + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + } => Ok(rs.total()), + SketchQuery::PointCount { .. } => Err(SummaryExecutorError::Unsupported( + "PointCount for a named item key needs a filter value this trait doesn't carry", + )), + SketchQuery::TopK { .. } => Err(SummaryExecutorError::Unsupported( + "TopK's answer shape (K items per timestamp) doesn't fit this executor's \ + Value = Vec<(i64, f64)> -- needs a Value type redesign", + )), + } +} + +/// Exact `(SummaryKind, SummaryParams)` match against a sid's own +/// `(SketchKindHandle, SketchConfig)` -- the check `find_candidates`'s +/// trait contract requires (not the looser family-only +/// `Capability::is_satisfied_by` check the legacy analyzer path uses), +/// so a `SummaryMerge`'s precondition (every child agrees on kind AND +/// params) is guaranteed by construction for anything routed through +/// this executor. +fn summary_params_match( + sketch: &SummaryKind, + params: &SummaryParams, + kind: SketchKindHandle, + config: &SketchConfig, +) -> bool { + // `SummaryParams::{Cms,CmsWithHeap,CountSketch,CountSketchWithHeap}` + // use width=cols/depth=rows (matches the control-plane wire + // convention -- see `sketch_config_to_json`'s comment). `SketchConfig` + // has no `heap_size` field at all (heap-bearing kinds reuse their + // heap-less base's config shape for identity -- see + // `base_sketch_kind_handle`'s doc in `drivers/ingest/otel.rs`), so + // heap_size can't be part of this match; width/depth are. + match (sketch, params, kind, config) { + ( + SummaryKind::DDSketch, + SummaryParams::DDSketch { alpha }, + SketchKindHandle::DDSketch, + SketchConfig::DDSketch { relative_accuracy }, + ) => alpha == relative_accuracy, + ( + SummaryKind::Kll, + SummaryParams::Kll { k }, + SketchKindHandle::Kll, + SketchConfig::Kll { k: sid_k }, + ) => k == sid_k, + ( + SummaryKind::Hll, + SummaryParams::Hll { precision }, + SketchKindHandle::Hll, + SketchConfig::Hll { precision: sid_p }, + ) => u32::from(*precision) == *sid_p, + ( + SummaryKind::Cms, + SummaryParams::Cms { width, depth }, + SketchKindHandle::CountMin, + SketchConfig::CountMin { rows, cols }, + ) => *depth as i32 == *rows && *width as i32 == *cols, + ( + SummaryKind::CountSketch, + SummaryParams::CountSketch { width, depth }, + SketchKindHandle::CountSketch, + SketchConfig::CountSketch { rows, cols }, + ) => *depth as i32 == *rows && *width as i32 == *cols, + ( + SummaryKind::CmsWithHeap, + SummaryParams::CmsWithHeap { width, depth, .. }, + SketchKindHandle::CmsWithHeap, + SketchConfig::CountMin { rows, cols }, + ) => *depth as i32 == *rows && *width as i32 == *cols, + ( + SummaryKind::CountSketchWithHeap, + SummaryParams::CountSketchWithHeap { width, depth, .. }, + SketchKindHandle::CountSketchWithHeap, + SketchConfig::CountSketch { rows, cols }, + ) => *depth as i32 == *rows && *width as i32 == *cols, + _ => false, + } +} + +/// `SketchConfig` (data_plane's per-sid stored params) -> `DeltaSketchKind` +/// (`delta_apply`'s decode/merge parameter carrier). +fn to_delta_kind(kind: SketchKindHandle, config: &SketchConfig) -> Option { + // Default heap_size when bootstrapping an empty Heap state for a + // delta-from-empty leading window -- `SketchConfig` carries no + // heap_size (see `summary_params_match`'s doc), so this only matters + // transiently: `CountMinSketchWithHeap::merge` takes `min(self, + // other)`, so it converges to the real decoded value as soon as any + // actual frame merges in. Matches this codebase's existing + // heap_size-absent default (`accuracy.rs`). + const DEFAULT_HEAP_SIZE: usize = 100; + match (kind, config) { + (SketchKindHandle::DDSketch, SketchConfig::DDSketch { relative_accuracy }) => { + Some(DeltaSketchKind::DDSketch { + alpha: *relative_accuracy, + }) + } + (SketchKindHandle::Kll, SketchConfig::Kll { k }) => Some(DeltaSketchKind::Kll { k: *k }), + (SketchKindHandle::Hll, SketchConfig::Hll { precision }) => Some(DeltaSketchKind::Hll { + precision: *precision, + }), + (SketchKindHandle::CountMin, SketchConfig::CountMin { rows, cols }) => { + Some(DeltaSketchKind::Cms { + rows: *rows as usize, + cols: *cols as usize, + }) + } + (SketchKindHandle::CountSketch, SketchConfig::CountSketch { rows, cols }) => { + Some(DeltaSketchKind::CountSketch { + rows: *rows as usize, + cols: *cols as usize, + }) + } + (SketchKindHandle::CmsWithHeap, SketchConfig::CountMin { rows, cols }) => { + Some(DeltaSketchKind::CmsWithHeap { + rows: *rows as usize, + cols: *cols as usize, + heap_size: DEFAULT_HEAP_SIZE, + }) + } + (SketchKindHandle::CountSketchWithHeap, SketchConfig::CountSketch { rows, cols }) => { + Some(DeltaSketchKind::CountSketchWithHeap { + rows: *rows as usize, + cols: *cols as usize, + heap_size: DEFAULT_HEAP_SIZE, + }) + } + _ => None, + } +} + +/// Walk an `L4Node`'s `SummaryAgg`/`SummaryEstimate`/`SummaryMerge` +/// spine down to a `Logical` leaf, then walk that leaf's `QueryExpr` +/// down to a `Scan { source: Source::TimeSeries { metric }, .. }` to +/// recover the metric name -- `SummaryAgg` itself carries no +/// metric/source field (see `SummaryExecutor::find_candidates`'s trait +/// doc). Mirrors `control_plane::asap_tier_implement::collect_aggregate_roots`'s +/// exhaustive-variant recursion style (same `QueryExpr` type), swapping +/// "collect Aggregate roots" for "find the first Scan". +fn find_metric(node: &L4Node) -> Option { + match &node.expr { + SummaryExpr::Logical(qe) => find_metric_in_query_expr(qe), + SummaryExpr::SummaryAgg { child, .. } => find_metric(child), + SummaryExpr::SummaryEstimate { sketch_input, .. } => find_metric(sketch_input), + SummaryExpr::SummaryMerge { children } => children.first().and_then(|c| find_metric(c)), + _ => None, + } +} + +fn find_metric_in_query_expr(qe: &QueryExpr) -> Option { + match qe { + QueryExpr::Scan { + source: Source::TimeSeries { metric }, + .. + } => Some(metric.clone()), + QueryExpr::Window { child, .. } + | QueryExpr::Filter { child, .. } + | QueryExpr::Project { child, .. } + | QueryExpr::Aggregate { child, .. } + | QueryExpr::Distinct { child, .. } + | QueryExpr::Sort { child, .. } + | QueryExpr::Limit { child, .. } + | QueryExpr::Subquery { child, .. } + | QueryExpr::TimeRange { child, .. } + | QueryExpr::TimeShift { child, .. } + | QueryExpr::WindowFunc { child, .. } => find_metric_in_query_expr(child), + QueryExpr::LetBinding { expr, child, .. } => { + find_metric_in_query_expr(expr).or_else(|| find_metric_in_query_expr(child)) + } + QueryExpr::Merge { children } => children.iter().find_map(find_metric_in_query_expr), + QueryExpr::Join { left, .. } | QueryExpr::SetOp { left, .. } => { + find_metric_in_query_expr(left) + } + QueryExpr::BinaryOp { lhs, .. } => find_metric_in_query_expr(lhs), + // The PromQL-surface superset (Scan's siblings: Ref/Scalar/ + // EvalTime/VectorFromScalar/ScalarFromVector/Relabel/InfoJoin/ + // Sample) isn't constructed by this parser today -- mirrors + // `collect_aggregate_roots`'s same no-op default. + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage_engines::sketch_db::index::{ + AccuracyBound, Capability, SketchInstanceMetadata, SketchSampleState, SketchStore, + }; + use asap_ir::intent_algebra::{Column, DataType, Schema}; + use asap_sketch::exec::{execute, ExecOutcome}; + use asap_sketch::schema::{L4DataType, L4Field, L4Schema}; + use std::rc::Rc; + + fn scan_node(metric: &str, group_by_field: Option<&str>) -> Rc { + let qe = QueryExpr::Scan { + source: Source::TimeSeries { + metric: metric.to_string(), + }, + predicates: vec![], + schema: Schema::with_time_index( + vec![ + Column::new("ts", DataType::Timestamp, false), + Column::new("value", DataType::Float64, false), + ], + 0, + vec![], + ), + }; + let mut fields = vec![L4Field { + name: "value".into(), + dtype: L4DataType::Primitive(DataType::Float64), + nullable: false, + }]; + if let Some(name) = group_by_field { + fields.push(L4Field { + name: name.into(), + dtype: L4DataType::Primitive(DataType::Utf8), + nullable: false, + }); + } + Rc::new(L4Node { + expr: SummaryExpr::Logical(Box::new(qe)), + schema: L4Schema { + fields, + time_index: None, + }, + }) + } + + fn kll_agg_node(child: Rc, by: Vec) -> Rc { + Rc::new(L4Node { + expr: SummaryExpr::SummaryAgg { + child, + sketch: SummaryKind::Kll, + params: SummaryParams::Kll { k: 200 }, + col: ColumnRef::SampleValue, + by, + }, + schema: L4Schema { + fields: vec![], + time_index: None, + }, + }) + } + + fn hll_agg_node(child: Rc) -> Rc { + Rc::new(L4Node { + expr: SummaryExpr::SummaryAgg { + child, + sketch: SummaryKind::Hll, + params: SummaryParams::Hll { precision: 10 }, + col: ColumnRef::SampleValue, + by: vec![], + }, + schema: L4Schema { + fields: vec![], + time_index: None, + }, + }) + } + + fn estimate_node(sketch_input: Rc, query: SketchQuery) -> Rc { + Rc::new(L4Node { + expr: SummaryExpr::SummaryEstimate { + sketch_input, + query, + }, + schema: L4Schema { + fields: vec![], + time_index: None, + }, + }) + } + + fn kll_meta(sid: u64, metric: &str, group_by: &[&str]) -> SketchInstanceMetadata { + let cfg = SketchConfig::Kll { k: 200 }; + SketchInstanceMetadata { + sid, + metric_name: metric.to_string(), + group_by_keys: group_by + .iter() + .map(|s| s.to_string()) + .collect::>(), + capability: Some(Capability::QuantileApprox(SketchKindHandle::Kll)), + agg_kind: crate::storage_engines::sketch_db::index::AggKind::Sketch { + kind: SketchKindHandle::Kll, + config: cfg.clone(), + spatial_filter_canonical: String::new(), + }, + accuracy: Some(AccuracyBound::from_config(&cfg)), + first_seen_unix_ms: 0, + retired_at_ms: None, + expires_at_ms: None, + policy_fp: asap_types::PolicyFingerprint::UNSET, + } + } + + fn hll_meta(sid: u64, metric: &str) -> SketchInstanceMetadata { + let cfg = SketchConfig::Hll { precision: 10 }; + SketchInstanceMetadata { + sid, + metric_name: metric.to_string(), + group_by_keys: BTreeSet::new(), + capability: Some(Capability::CardinalityApprox), + agg_kind: crate::storage_engines::sketch_db::index::AggKind::Sketch { + kind: SketchKindHandle::Hll, + config: cfg.clone(), + spatial_filter_canonical: String::new(), + }, + accuracy: Some(AccuracyBound::from_config(&cfg)), + first_seen_unix_ms: 0, + retired_at_ms: None, + expires_at_ms: None, + policy_fp: asap_types::PolicyFingerprint::UNSET, + } + } + + fn encode_kll_items_proto(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() + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(state)), + ..Default::default() + }; + env.encode_to_vec() + } + + fn encode_hll_from_items(precision: u32, items: &[&str]) -> Vec { + use asap_sketchlib::{HllSketch, HllVariant, MessagePackCodec}; + let mut sk = HllSketch::new(HllVariant::Regular, precision); + for item in items { + sk.update(item.as_bytes()); + } + sk.to_msgpack().expect("encode HLL msgpack") + } + + fn cms_meta(sid: u64, metric: &str) -> SketchInstanceMetadata { + let cfg = SketchConfig::CountMin { rows: 4, cols: 256 }; + SketchInstanceMetadata { + sid, + metric_name: metric.to_string(), + group_by_keys: BTreeSet::new(), + capability: Some(Capability::FrequencyEstimate(SketchKindHandle::CountMin)), + agg_kind: crate::storage_engines::sketch_db::index::AggKind::Sketch { + kind: SketchKindHandle::CountMin, + config: cfg.clone(), + spatial_filter_canonical: String::new(), + }, + accuracy: Some(AccuracyBound::from_config(&cfg)), + first_seen_unix_ms: 0, + retired_at_ms: None, + expires_at_ms: None, + policy_fp: asap_types::PolicyFingerprint::UNSET, + } + } + + /// Encode a CMS msgpack frame whose row-0 total is `total_weight` + /// (`update`s a single synthetic key `total_weight` times -- row 0's + /// sum equals the number of insertions regardless of hashing, since + /// every insertion touches every row including row 0). + fn encode_cms_with_total(rows: usize, cols: usize, total_weight: usize) -> Vec { + use asap_sketchlib::{CountMinSketch, MessagePackCodec}; + let mut sk = CountMinSketch::new(rows, cols); + for _ in 0..total_weight { + sk.update("k", 1.0); + } + sk.to_msgpack().expect("encode CountMinSketch msgpack") + } + + fn cms_agg_node(child: Rc) -> Rc { + Rc::new(L4Node { + expr: SummaryExpr::SummaryAgg { + child, + sketch: SummaryKind::Cms, + params: SummaryParams::Cms { + width: 256, + depth: 4, + }, + col: ColumnRef::SampleValue, + by: vec![], + }, + schema: L4Schema { + fields: vec![], + time_index: None, + }, + }) + } + + const T0: u64 = 1_000_000; + const T1: u64 = 2_000_000; + + fn ctx(index: &SketchStore) -> QueryExecutionContext<'_> { + QueryExecutionContext { + index, + t0_ms: T0, + t1_ms: T1, + is_cumulative: true, + } + } + + fn matrix_ctx(index: &SketchStore) -> QueryExecutionContext<'_> { + QueryExecutionContext { + index, + t0_ms: T0, + t1_ms: T1, + is_cumulative: false, + } + } + + #[test] + fn single_kll_sid_quantile_readout() { + let idx = SketchStore::new(); + let sid = 1u64; + idx.register(kll_meta(sid, "latency_ms", &[])); + let items: Vec = (1..=100).map(|i| i as f64).collect(); + idx.append_sample( + sid, + BTreeMap::new(), + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_kll_items_proto(200, &items), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + + let child = scan_node("latency_ms", None); + let tree = estimate_node( + kll_agg_node(child, vec![]), + SketchQuery::Quantile { q: 0.5 }, + ); + + let exec = ctx(&idx); + let ExecOutcome::Value(v) = execute(&tree, &exec).expect("execute should succeed") else { + panic!("expected a value"); + }; + assert_eq!(v.len(), 1, "ungrouped query produces exactly one group"); + let (_group, samples) = &v[0]; + let (_ts, median) = samples[0]; + // Median of 1..=100 is ~50. + assert!( + (45.0..=55.0).contains(&median), + "median {median} out of range" + ); + } + + #[test] + fn two_sids_same_group_actually_merge_not_just_first() { + // Two sids covering the SAME group must be MERGED into one + // combined answer, not silently duplicated or one-of-them-dropped. + let idx = SketchStore::new(); + idx.register(kll_meta(1, "latency_ms", &[])); + idx.register(kll_meta(2, "latency_ms", &[])); + // sid 1: values 1..=50 (median ~25); sid 2: values 51..=100 (median ~75). + // Merged, the combined median should be ~50 -- NOT ~25 (if merge + // silently dropped sid 2) and NOT ~75 (if it dropped sid 1). + let items1: Vec = (1..=50).map(|i| i as f64).collect(); + let items2: Vec = (51..=100).map(|i| i as f64).collect(); + idx.append_sample( + 1, + BTreeMap::new(), + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_kll_items_proto(200, &items1), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + idx.append_sample( + 2, + BTreeMap::new(), + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_kll_items_proto(200, &items2), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + + let child = scan_node("latency_ms", None); + let tree = estimate_node( + kll_agg_node(child, vec![]), + SketchQuery::Quantile { q: 0.5 }, + ); + + let exec = ctx(&idx); + let ExecOutcome::Value(v) = execute(&tree, &exec).expect("execute should succeed") else { + panic!("expected a value"); + }; + assert_eq!(v.len(), 1); + let (_group, samples) = &v[0]; + let (_ts, median) = samples[0]; + assert!( + (40.0..=60.0).contains(&median), + "merged median {median} should be ~50 (both sids' data combined), \ + not ~25 or ~75 (one sid dropped)" + ); + } + + #[test] + fn two_sids_different_groups_produce_two_series_not_one_merged_blob() { + // `quantile by (zone) (...)` must produce one output series per + // zone, not one series merging both zones together. + let idx = SketchStore::new(); + idx.register(kll_meta(1, "latency_ms", &["zone"])); + idx.register(kll_meta(2, "latency_ms", &["zone"])); + let items1: Vec = (1..=50).map(|i| i as f64).collect(); + let items2: Vec = (51..=100).map(|i| i as f64).collect(); + let mut labels_east = BTreeMap::new(); + labels_east.insert("zone".to_string(), "us-east".to_string()); + let mut labels_west = BTreeMap::new(); + labels_west.insert("zone".to_string(), "us-west".to_string()); + idx.append_sample( + 1, + labels_east, + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_kll_items_proto(200, &items1), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + idx.append_sample( + 2, + labels_west, + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_kll_items_proto(200, &items2), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + + let child = scan_node("latency_ms", Some("zone")); + // "zone" is field index 1 in `scan_node`'s schema (0 = value). + let tree = estimate_node( + kll_agg_node(child, vec![1]), + SketchQuery::Quantile { q: 0.5 }, + ); + + let exec = ctx(&idx); + let ExecOutcome::Value(mut v) = execute(&tree, &exec).expect("execute should succeed") + else { + panic!("expected a value"); + }; + assert_eq!( + v.len(), + 2, + "two zones must produce two output series, not one merged blob" + ); + v.sort_by(|a, b| a.0.get("zone").cmp(&b.0.get("zone"))); + let (east_group, east_samples) = &v[0]; + let (west_group, west_samples) = &v[1]; + assert_eq!(east_group.get("zone").map(String::as_str), Some("us-east")); + assert_eq!(west_group.get("zone").map(String::as_str), Some("us-west")); + let east_median = east_samples[0].1; + let west_median = west_samples[0].1; + assert!( + (20.0..=30.0).contains(&east_median), + "us-east median {east_median} should reflect only sid 1's data (~25)" + ); + assert!( + (70.0..=80.0).contains(&west_median), + "us-west median {west_median} should reflect only sid 2's data (~75)" + ); + } + + #[test] + fn hll_cardinality_readout() { + let idx = SketchStore::new(); + let sid = 1u64; + idx.register(hll_meta(sid, "unique_users")); + let items: Vec<&str> = vec!["a", "b", "c", "d", "e"]; + idx.append_sample( + sid, + BTreeMap::new(), + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_hll_from_items(10, &items), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::MsgpackFull, + }, + ); + + let child = scan_node("unique_users", None); + let tree = estimate_node(hll_agg_node(child), SketchQuery::Cardinality); + + let exec = ctx(&idx); + let ExecOutcome::Value(v) = execute(&tree, &exec).expect("execute should succeed") else { + panic!("expected a value"); + }; + let (_group, samples) = &v[0]; + let (_ts, card) = samples[0]; + assert!( + (3.0..=7.0).contains(&card), + "cardinality {card} should be ~5" + ); + } + + #[test] + fn single_cms_sid_total_readout() { + let idx = SketchStore::new(); + let sid = 1u64; + idx.register(cms_meta(sid, "requests_total")); + idx.append_sample( + sid, + BTreeMap::new(), + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_cms_with_total(4, 256, 42), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::MsgpackFull, + }, + ); + + let child = scan_node("requests_total", None); + let tree = estimate_node( + cms_agg_node(child), + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + }, + ); + + let exec = ctx(&idx); + let ExecOutcome::Value(v) = execute(&tree, &exec).expect("execute should succeed") else { + panic!("expected a value"); + }; + let (_group, samples) = &v[0]; + let (_ts, total) = samples[0]; + assert_eq!( + total, 42.0, + "bare total must equal the number of insertions" + ); + } + + #[test] + fn two_cms_sids_same_group_totals_actually_merge() { + // Cross-sid merge for the Frequency family: two sids' totals must + // ADD (matrix merge then row-0 sum), not just report one of them. + let idx = SketchStore::new(); + idx.register(cms_meta(1, "requests_total")); + idx.register(cms_meta(2, "requests_total")); + idx.append_sample( + 1, + BTreeMap::new(), + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_cms_with_total(4, 256, 30), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::MsgpackFull, + }, + ); + idx.append_sample( + 2, + BTreeMap::new(), + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_cms_with_total(4, 256, 12), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::MsgpackFull, + }, + ); + + let child = scan_node("requests_total", None); + let tree = estimate_node( + cms_agg_node(child), + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + }, + ); + + let exec = ctx(&idx); + let ExecOutcome::Value(v) = execute(&tree, &exec).expect("execute should succeed") else { + panic!("expected a value"); + }; + let (_group, samples) = &v[0]; + let (_ts, total) = samples[0]; + assert_eq!( + total, 42.0, + "merged total must be the SUM of both sids (30 + 12)" + ); + } + + #[test] + fn topk_query_is_explicitly_unsupported_not_silently_wrong() { + // TopK's answer shape (K items per timestamp) doesn't fit this + // executor's Value = Vec<(i64, f64)> -- must error, not silently + // return a truncated/wrong scalar. + let idx = SketchStore::new(); + let sid = 1u64; + idx.register(cms_meta(sid, "requests_total")); + idx.append_sample( + sid, + BTreeMap::new(), + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_cms_with_total(4, 256, 1), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::MsgpackFull, + }, + ); + let child = scan_node("requests_total", None); + let tree = estimate_node(cms_agg_node(child), SketchQuery::TopK { k: 5 }); + let exec = ctx(&idx); + match execute(&tree, &exec) { + Err(asap_sketch::exec::ExecError::Executor(SummaryExecutorError::Unsupported(_))) => {} + other => panic!("expected Unsupported, got {}", other.is_ok()), + } + } + + #[test] + fn no_matching_sid_is_no_candidates() { + let idx = SketchStore::new(); + let child = scan_node("nonexistent_metric", None); + let tree = estimate_node( + kll_agg_node(child, vec![]), + SketchQuery::Quantile { q: 0.5 }, + ); + let exec = ctx(&idx); + match execute(&tree, &exec) { + Err(asap_sketch::exec::ExecError::NoCandidates) => {} + other => panic!("expected NoCandidates, got {}", other.is_ok()), + } + } + + #[test] + fn mismatched_params_does_not_match() { + // sid is Kll{k: 200}; query wants Kll{k: 500} -- must NOT match, + // even though both are "Kll". + let idx = SketchStore::new(); + idx.register(kll_meta(1, "latency_ms", &[])); + idx.append_sample( + 1, + BTreeMap::new(), + (T0, T0 + 1000), + SketchSampleState { + bytes: encode_kll_items_proto(200, &[1.0, 2.0, 3.0]), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + let child = scan_node("latency_ms", None); + let mismatched = Rc::new(L4Node { + expr: SummaryExpr::SummaryAgg { + child, + sketch: SummaryKind::Kll, + params: SummaryParams::Kll { k: 500 }, + col: ColumnRef::SampleValue, + by: vec![], + }, + schema: L4Schema { + fields: vec![], + time_index: None, + }, + }); + let tree = estimate_node(mismatched, SketchQuery::Quantile { q: 0.5 }); + let exec = ctx(&idx); + match execute(&tree, &exec) { + Err(asap_sketch::exec::ExecError::NoCandidates) => {} + other => panic!( + "expected NoCandidates (param mismatch), got {}", + other.is_ok() + ), + } + } + + #[test] + fn per_window_matrix_produces_multiple_points_for_one_sid() { + let idx = SketchStore::new(); + let sid = 1u64; + idx.register(kll_meta(sid, "latency_ms", &[])); + let w1_end = T0 + 100_000; + let w2_end = T0 + 200_000; + let items_w1: Vec = (1..=50).map(|i| i as f64).collect(); + let items_w2: Vec = (901..=1000).map(|i| i as f64).collect(); + idx.append_sample( + sid, + BTreeMap::new(), + (T0, w1_end), + SketchSampleState { + bytes: encode_kll_items_proto(200, &items_w1), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + idx.append_sample( + sid, + BTreeMap::new(), + (w1_end, w2_end), + SketchSampleState { + bytes: encode_kll_items_proto(200, &items_w2), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + + let child = scan_node("latency_ms", None); + let tree = estimate_node( + kll_agg_node(child, vec![]), + SketchQuery::Quantile { q: 0.5 }, + ); + + let exec = matrix_ctx(&idx); + let ExecOutcome::Value(v) = execute(&tree, &exec).expect("execute should succeed") else { + panic!("expected a value"); + }; + assert_eq!(v.len(), 1, "ungrouped query produces exactly one group"); + let (_group, mut samples) = v.into_iter().next().unwrap(); + samples.sort_by_key(|(ts, _)| *ts); + assert_eq!( + samples.len(), + 2, + "two distinct windows must produce two output points, not one collapsed answer" + ); + assert_eq!(samples[0].0, w1_end as i64); + assert_eq!(samples[1].0, w2_end as i64); + assert!( + (20.0..=30.0).contains(&samples[0].1), + "window 1 median {} should be ~25 (items 1..=50)", + samples[0].1 + ); + assert!( + (940.0..=960.0).contains(&samples[1].1), + "window 2 median {} should be ~950 (items 901..=1000)", + samples[1].1 + ); + } + + #[test] + fn per_window_matrix_merges_across_sids_per_window() { + // Two sids in the SAME group, both contributing a frame to the + // SAME window_end -- the per-window answer for that window must + // reflect BOTH sids merged, not just one (the cross-sid analog of + // `two_sids_same_group_actually_merge_not_just_first`, but for + // one window instead of the whole cumulative range). + let idx = SketchStore::new(); + idx.register(kll_meta(1, "latency_ms", &[])); + idx.register(kll_meta(2, "latency_ms", &[])); + let w_end = T0 + 100_000; + let items1: Vec = (1..=50).map(|i| i as f64).collect(); + let items2: Vec = (51..=100).map(|i| i as f64).collect(); + idx.append_sample( + 1, + BTreeMap::new(), + (T0, w_end), + SketchSampleState { + bytes: encode_kll_items_proto(200, &items1), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + idx.append_sample( + 2, + BTreeMap::new(), + (T0, w_end), + SketchSampleState { + bytes: encode_kll_items_proto(200, &items2), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + + let child = scan_node("latency_ms", None); + let tree = estimate_node( + kll_agg_node(child, vec![]), + SketchQuery::Quantile { q: 0.5 }, + ); + + let exec = matrix_ctx(&idx); + let ExecOutcome::Value(v) = execute(&tree, &exec).expect("execute should succeed") else { + panic!("expected a value"); + }; + let (_group, samples) = v.into_iter().next().unwrap(); + assert_eq!( + samples.len(), + 1, + "both sids share one window_end -> one output point" + ); + let (ts, median) = samples[0]; + assert_eq!(ts, w_end as i64); + assert!( + (40.0..=60.0).contains(&median), + "merged per-window median {median} should be ~50 (both sids' data combined)" + ); + } + + #[test] + fn per_window_matrix_drops_carry_in_base_before_t0() { + // `SketchStore::query_range` may splice in a carry-in Full ending + // BEFORE `t0_ms` to seed the delta-apply walk. That base must not + // surface as an output point. + let idx = SketchStore::new(); + let sid = 1u64; + idx.register(kll_meta(sid, "latency_ms", &[])); + let carry_in_end = T0 - 50_000; // before t0 + let in_range_end = T0 + 100_000; + idx.append_sample( + sid, + BTreeMap::new(), + (T0 - 100_000, carry_in_end), + SketchSampleState { + bytes: encode_kll_items_proto(200, &[1.0, 2.0, 3.0]), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + idx.append_sample( + sid, + BTreeMap::new(), + (carry_in_end, in_range_end), + SketchSampleState { + bytes: encode_kll_items_proto(200, &[10.0, 20.0, 30.0]), + encoding: crate::storage_engines::sketch_db::index::SketchEncoding::ProtoFull, + }, + ); + + let child = scan_node("latency_ms", None); + let tree = estimate_node( + kll_agg_node(child, vec![]), + SketchQuery::Quantile { q: 0.5 }, + ); + + let exec = matrix_ctx(&idx); + let ExecOutcome::Value(v) = execute(&tree, &exec).expect("execute should succeed") else { + panic!("expected a value"); + }; + let (_group, samples) = v.into_iter().next().unwrap(); + assert_eq!( + samples.len(), + 1, + "the carry-in-base window (ending before t0) must not appear in output" + ); + assert_eq!(samples[0].0, in_range_end as i64); + } +} diff --git a/data_plane/src/storage_engines/sketch_db/query/decoders.rs b/data_plane/src/storage_engines/sketch_db/query/decoders.rs index bb5e8768..f1698c2f 100644 --- a/data_plane/src/storage_engines/sketch_db/query/decoders.rs +++ b/data_plane/src/storage_engines/sketch_db/query/decoders.rs @@ -20,6 +20,8 @@ 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::precompute_engine::operators::count_min_sketch_with_heap_accumulator::CountMinSketchWithHeapAccumulator; @@ -184,6 +186,16 @@ pub fn decode_cms_with_heap_from_msgpack(buffer: &[u8]) -> Result Result { + CountSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketchWithHeap msgpack: {e}")) +} + // --------------------------------------------------------------------------- // Delta decoders. Under the per-window-reset (PWR) contract // (`asap-precompute-go/window.go`: a delta is that window's own state @@ -308,3 +320,63 @@ pub fn decode_cms_with_heap_from_msgpack_delta( .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_accumulator.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)>, + } + + 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 { + return Err(format!( + "CountSketchWithHeap delta frame has zero dims (rows={rows}, cols={cols})" + )); + } + 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; + } + matrix[r][c] += *dc as f64; + } + let heap: Vec = 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, + )) +} diff --git a/data_plane/src/storage_engines/sketch_db/query/delta_apply.rs b/data_plane/src/storage_engines/sketch_db/query/delta_apply.rs index 93177b63..f5de8ce9 100644 --- a/data_plane/src/storage_engines/sketch_db/query/delta_apply.rs +++ b/data_plane/src/storage_engines/sketch_db/query/delta_apply.rs @@ -30,6 +30,10 @@ //! `Full`. For cumulative mode this means the cumulative answer //! starts at the first Full in the range, not at `t0`. +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; @@ -37,38 +41,91 @@ use asap_sketchlib::KllSketch; use asap_sketchlib::MessagePackCodec; use crate::storage_engines::sketch_db::index::{SketchEncoding, SketchSampleState}; +use crate::storage_engines::sketch_db::query::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, +}; -/// Whether a sketch family supports delta-via-merge (DD/KLL) or -/// delta-via-`apply_delta` (HLL). The reducer reads bytes through -/// the appropriate `decode_*_full` path and folds the result into a -/// rolling state. -/// -/// Each variant carries the sketch parameters (`alpha` / `k` / -/// `precision`) needed to *bootstrap an empty rolling 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). DD/KLL deltas embed -/// their own params in the wire fragment, but HLL register deltas are -/// applied onto a pre-sized register array, so the precision must be -/// known up front to allocate it. +/// 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. #[derive(Debug, Clone, Copy)] pub enum DeltaSketchKind { - DDSketch { alpha: f64 }, - Hll { precision: u32 }, - Kll { k: u32 }, + DDSketch { + alpha: f64, + }, + Hll { + precision: u32, + }, + Kll { + k: u32, + }, + Cms { + rows: usize, + cols: usize, + }, + CountSketch { + 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. + CmsWithHeap { + rows: usize, + cols: usize, + heap_size: usize, + }, + CountSketchWithHeap { + rows: usize, + cols: usize, + heap_size: usize, + }, } impl DeltaSketchKind { - /// Construct an EMPTY rolling 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). - fn bootstrap_empty(&self) -> RollingState { + /// 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). + fn bootstrap_empty(&self) -> SummaryState { match self { - DeltaSketchKind::DDSketch { alpha } => RollingState::Dd(DdSketch::new(*alpha)), - DeltaSketchKind::Kll { k } => RollingState::Kll(KllSketch::new(*k as u16)), + DeltaSketchKind::DDSketch { alpha } => SummaryState::Dd(DdSketch::new(*alpha)), + DeltaSketchKind::Kll { k } => SummaryState::Kll(KllSketch::new(*k as u16)), DeltaSketchKind::Hll { precision } => { - RollingState::Hll(HllSketch::new(HllVariant::Regular, *precision)) + SummaryState::Hll(HllSketch::new(HllVariant::Regular, *precision)) + } + DeltaSketchKind::Cms { rows, cols } => { + SummaryState::Cms(CountMinSketch::new(*rows, *cols)) + } + DeltaSketchKind::CountSketch { rows, cols } => { + SummaryState::CountSketch(CountSketch::new(*rows, *cols)) } + DeltaSketchKind::CmsWithHeap { + rows, + cols, + heap_size, + } => SummaryState::CmsWithHeap(CountMinSketchWithHeap::new(*rows, *cols, *heap_size)), + DeltaSketchKind::CountSketchWithHeap { + rows, + cols, + heap_size, + } => SummaryState::CountSketchWithHeap(CountSketchWithHeap::new( + *rows, *cols, *heap_size, + )), } } } @@ -79,47 +136,83 @@ fn decode_full( kind: &DeltaSketchKind, bytes: &[u8], encoding: SketchEncoding, -) -> Result { +) -> Result { match (kind, encoding) { (DeltaSketchKind::DDSketch { .. }, SketchEncoding::ProtoFull) => { let sk = dd_from_proto(bytes)?; - Ok(RollingState::Dd(sk)) + Ok(SummaryState::Dd(sk)) } (DeltaSketchKind::DDSketch { .. }, SketchEncoding::MsgpackFull) => { let sk = DdSketch::from_msgpack(bytes) .map_err(|e| format!("deserialize DDSketch msgpack: {e}"))?; - Ok(RollingState::Dd(sk)) + Ok(SummaryState::Dd(sk)) } (DeltaSketchKind::Hll { .. }, SketchEncoding::ProtoFull) => { let sk = hll_from_proto(bytes)?; - Ok(RollingState::Hll(sk)) + Ok(SummaryState::Hll(sk)) } (DeltaSketchKind::Hll { .. }, SketchEncoding::MsgpackFull) => { let sk = HllSketch::from_msgpack(bytes) .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?; - Ok(RollingState::Hll(sk)) + Ok(SummaryState::Hll(sk)) } (DeltaSketchKind::Kll { .. }, SketchEncoding::ProtoFull) => { let sk = kll_from_proto(bytes)?; - Ok(RollingState::Kll(sk)) + Ok(SummaryState::Kll(sk)) } (DeltaSketchKind::Kll { .. }, SketchEncoding::MsgpackFull) => { let sk = KllSketch::from_msgpack(bytes) .map_err(|e| format!("deserialize KllSketch msgpack: {e}"))?; - Ok(RollingState::Kll(sk)) + Ok(SummaryState::Kll(sk)) + } + (DeltaSketchKind::Cms { .. }, SketchEncoding::ProtoFull) => { + Ok(SummaryState::Cms(decode_cms_from_proto(bytes)?)) + } + (DeltaSketchKind::Cms { .. }, SketchEncoding::MsgpackFull) => { + Ok(SummaryState::Cms(decode_cms_from_msgpack(bytes)?)) + } + (DeltaSketchKind::CountSketch { .. }, SketchEncoding::ProtoFull) => { + Ok(SummaryState::CountSketch(decode_cs_from_proto(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)?, + )), (_, e) => Err(format!("decode_full called with non-Full encoding {e:?}")), } } -/// The rolling state the delta-application loop maintains. -pub enum RollingState { +/// 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. +pub enum SummaryState { 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), } -impl RollingState { +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 @@ -142,7 +235,7 @@ impl RollingState { )); } match self { - RollingState::Dd(sk) => { + SummaryState::Dd(sk) => { match encoding { // PROTO_DELTA: dispatch on the payload SHAPE, mirroring the // edge's own `DDSketchWrapper::apply_delta` @@ -173,7 +266,7 @@ impl RollingState { // 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(RollingState::Dd(other)) = decode_full( + if let Ok(SummaryState::Dd(other)) = decode_full( &DeltaSketchKind::DDSketch { alpha: 0.0 }, bytes, SketchEncoding::ProtoFull, @@ -203,7 +296,7 @@ impl RollingState { bytes, SketchEncoding::MsgpackFull, ) { - Ok(RollingState::Dd(s)) => s, + Ok(SummaryState::Dd(s)) => s, Ok(_) => { return Err( "decode_full(DDSketch) returned non-DDSketch state".to_string() @@ -218,7 +311,7 @@ impl RollingState { _ => unreachable!(), } } - RollingState::Hll(sk) => { + 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). @@ -235,14 +328,14 @@ impl RollingState { Ok(()) } } - RollingState::Kll(sk) => { + 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(RollingState::Kll(s)) => s, + Ok(SummaryState::Kll(s)) => s, Ok(_) => return Err("decode_full(Kll) returned non-Kll state".to_string()), Err(e) => return Err(e), }; @@ -250,24 +343,117 @@ impl RollingState { .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}")) + } + 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}")) + } + 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)? + } else { + decode_cms_with_heap_from_msgpack(bytes)? + }; + sk.merge(&other) + .map_err(|e| format!("merge CmsWithHeap delta: {e}")) + } + SummaryState::CountSketchWithHeap(sk) => { + let other = if encoding == SketchEncoding::MsgpackDelta { + decode_cs_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}")) + } } } pub fn quantile(&self, q: f64) -> f64 { match self { - RollingState::Dd(sk) => sk.quantile(q).unwrap_or(0.0), - RollingState::Kll(sk) => sk.quantile(q), - RollingState::Hll(_) => 0.0, + SummaryState::Dd(sk) => sk.quantile(q).unwrap_or(0.0), + SummaryState::Kll(sk) => sk.quantile(q), + _ => 0.0, } } pub fn cardinality(&self) -> f64 { match self { - RollingState::Hll(sk) => sk.estimate(), + SummaryState::Hll(sk) => sk.estimate(), _ => 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. + 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(), + _ => return 0.0, + }; + matrix + .first() + .map(|row| row.iter().copied().sum::()) + .unwrap_or(0.0) + } + + /// 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. + pub fn topk_items(&self) -> Option> { + match self { + SummaryState::CmsWithHeap(h) => Some( + h.topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect(), + ), + SummaryState::CountSketchWithHeap(h) => Some( + h.topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect(), + ), + _ => None, + } + } + /// Borrow the inner HLL sketch when this rolling state is HLL-backed. /// Used by the GLOBAL cardinality rollup (`count(hll_metric)` with no /// `by`), which must MERGE the per-series HLL registers (register-wise @@ -275,36 +461,85 @@ impl RollingState { /// distinct estimates would double-count items present in multiple series. pub fn as_hll(&self) -> Option<&HllSketch> { match self { - RollingState::Hll(sk) => Some(sk), + SummaryState::Hll(sk) => Some(sk), _ => None, } } + + /// 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`. + pub fn merge_same_family(&mut self, other: &SummaryState) -> Result<(), String> { + match (self, other) { + (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}")), + (a, _) => Err(format!( + "SummaryState family mismatch in merge_same_family (self is {})", + a.family_name() + )), + } + } + + /// Diagnostic family name for error messages — not used for dispatch. + fn family_name(&self) -> &'static str { + match self { + SummaryState::Dd(_) => "DDSketch", + SummaryState::Hll(_) => "Hll", + SummaryState::Kll(_) => "Kll", + SummaryState::Cms(_) => "Cms", + SummaryState::CountSketch(_) => "CountSketch", + SummaryState::CmsWithHeap(_) => "CmsWithHeap", + SummaryState::CountSketchWithHeap(_) => "CountSketchWithHeap", + } + } } -/// Fold every in-range window's frames for ONE series into a single merged -/// `HllSketch` (cumulative over `[t0, t1]`), returning `None` if no Full -/// HLL frame ever landed (every sample was a leading delta). This is the -/// per-series building block for the GLOBAL `count(hll_metric)` rollup: the -/// reducer merges the returned sketches across series (register-wise max) -/// before estimating, so the answer is the distinct UNION cardinality, not -/// the sum of per-series cardinalities. -pub fn cumulative_hll_state( +/// Fold every in-range window's frames for ONE series into a single +/// merged `SummaryState` (cumulative over `[t0, t1]`), returning `None` +/// if no Full frame ever landed (every sample was a leading delta). The +/// per-sid building block for a cross-sid answer: reconstruct each +/// candidate sid's state this way, then merge them (`merge_same_family`) +/// before reading out a quantile/cardinality over the combined data. +pub fn cumulative_summary_state( samples: &[(i64, &SketchSampleState)], - precision: u32, -) -> Result, String> { - let kind = DeltaSketchKind::Hll { precision }; - let mut rolling: Option = None; + kind: DeltaSketchKind, +) -> Result, String> { + let mut rolling: Option = None; for (_window_end, state) in samples { match state.encoding { SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull => { let new_state = decode_full(&kind, &state.bytes, state.encoding)?; - rolling = Some(match (rolling.take(), new_state) { - (None, n) => n, - (Some(RollingState::Hll(mut a)), RollingState::Hll(b)) => { - a.merge(&b).map_err(|e| format!("cum merge HLL: {e}"))?; - RollingState::Hll(a) + rolling = Some(match rolling.take() { + None => new_state, + Some(mut prev) => { + prev.merge_same_family(&new_state)?; + prev } - (Some(prev), _) => prev, }); } SketchEncoding::ProtoDelta | SketchEncoding::MsgpackDelta => { @@ -317,10 +552,7 @@ pub fn cumulative_hll_state( } } } - Ok(rolling.and_then(|rs| match rs { - RollingState::Hll(sk) => Some(sk), - _ => None, - })) + Ok(rolling) } /// Walk a sorted-by-window-end slice of samples in time order and @@ -369,25 +601,47 @@ pub fn per_window_evaluate( eval: E, ) -> Result<(Vec<(i64, f64)>, usize), String> where - E: Fn(&RollingState) -> f64, + E: Fn(&SummaryState) -> f64, { - let mut out: Vec<(i64, f64)> = Vec::new(); + let (states, skipped) = per_window_summary_states(samples, kind)?; + Ok(( + states.into_iter().map(|(w, rs)| (w, eval(&rs))).collect(), + skipped, + )) +} + +/// Walk a sorted-by-window-end slice of samples in time order and +/// reconstruct ONE sid's per-window `SummaryState` (same per-window-reset +/// walk as [`per_window_evaluate`], generalized to return the +/// reconstructed state itself instead of an already-evaluated scalar). +/// The per-sid building block for cross-sid per-window merging (unlike +/// [`cumulative_summary_state`], which folds a whole `[t0, t1]` range +/// into one answer, this keeps each window separate so a caller can +/// merge same-window states across several sids before evaluating -- +/// needed for a matrix/range-query answer, where each output point is +/// itself a cross-sid merge for that one window). +/// +/// Returns `Ok((per_window_states, skipped))`. +pub fn per_window_summary_states( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, +) -> Result<(Vec<(i64, SummaryState)>, usize), String> { + let mut out: Vec<(i64, SummaryState)> = Vec::new(); let mut skipped = 0usize; // Rolling state for the CURRENT window only. Reset to None whenever // `window_end` changes (a new window establishes its own base from // empty). `cur_end` tracks which window `rolling` belongs to. - let mut rolling: Option = None; + let mut rolling: Option = None; let mut cur_end: Option = None; for (window_end, state) in samples { // Window boundary: flush the previous window's final accumulated - // value, then reset the base so this window starts from empty. + // state, then reset the base so this window starts from empty. if cur_end != Some(*window_end) { - if let (Some(prev_end), Some(rs)) = (cur_end, rolling.as_ref()) { - out.push((prev_end, eval(rs))); + if let (Some(prev_end), Some(rs)) = (cur_end, rolling.take()) { + out.push((prev_end, rs)); } - rolling = None; cur_end = Some(*window_end); } @@ -412,8 +666,8 @@ where } // Flush the final window. - if let (Some(prev_end), Some(rs)) = (cur_end, rolling.as_ref()) { - out.push((prev_end, eval(rs))); + if let (Some(prev_end), Some(rs)) = (cur_end, rolling.take()) { + out.push((prev_end, rs)); } Ok((out, skipped)) @@ -433,9 +687,9 @@ pub fn cumulative_evaluate( eval: E, ) -> Result<(Option<(i64, f64)>, usize), String> where - E: Fn(&RollingState) -> f64, + E: Fn(&SummaryState) -> f64, { - let mut rolling: Option = None; + let mut rolling: Option = None; let mut latest_end = i64::MIN; let mut skipped = 0usize; @@ -453,17 +707,17 @@ where // sample inclusion. rolling = Some(match (rolling.take(), new_state) { (None, n) => n, - (Some(RollingState::Dd(mut a)), RollingState::Dd(b)) => { + (Some(SummaryState::Dd(mut a)), SummaryState::Dd(b)) => { a.merge(&b).map_err(|e| format!("cum merge DD: {e}"))?; - RollingState::Dd(a) + SummaryState::Dd(a) } - (Some(RollingState::Hll(mut a)), RollingState::Hll(b)) => { + (Some(SummaryState::Hll(mut a)), SummaryState::Hll(b)) => { a.merge(&b).map_err(|e| format!("cum merge HLL: {e}"))?; - RollingState::Hll(a) + SummaryState::Hll(a) } - (Some(RollingState::Kll(mut a)), RollingState::Kll(b)) => { + (Some(SummaryState::Kll(mut a)), SummaryState::Kll(b)) => { a.merge(&b).map_err(|e| format!("cum merge KLL: {e}"))?; - RollingState::Kll(a) + SummaryState::Kll(a) } (Some(_), _) => { return Err("cumulative merge across sketch family mismatch".to_string()) @@ -883,10 +1137,8 @@ mod tests { let bytes = sk.compute_delta(&empty, 0); frames.push((((w as u64) + 1) * 1000, delta(bytes))); } - let samples: Vec<(i64, &SketchSampleState)> = frames - .iter() - .map(|(t, s)| (*t as i64, s)) - .collect(); + let samples: Vec<(i64, &SketchSampleState)> = + frames.iter().map(|(t, s)| (*t as i64, s)).collect(); let kind = DeltaSketchKind::Hll { precision }; let (out, skipped) = @@ -902,4 +1154,154 @@ mod tests { ); } } + + /// `CmsWithHeap` (min-over-rows estimator, `CountMinSketchWithHeap`) + /// and `CountSketchWithHeap` (median-of-signed-rows estimator, the + /// distinct `CountSketchWithHeap` type) are different sketch + /// algorithms that merely happen to share a storage shape — merging + /// one into the other must be rejected as a family mismatch, the + /// same as merging a `Cms` into a `Kll` would be. Since the two + /// `SummaryState` variants now hold genuinely different Rust types, + /// this is also enforced at compile time — there is no arm in + /// `merge_same_family` that type-checks a mixed pair together. + #[test] + fn cms_with_heap_and_count_sketch_with_heap_are_not_the_same_family() { + use asap_sketchlib::{CountMinSketchWithHeap, CountSketchWithHeap, MessagePackCodec}; + + let mut cms_heap = CountMinSketchWithHeap::new(4, 256, 10); + cms_heap.update("a", 1.0); + let mut cs_heap = CountSketchWithHeap::new(4, 256, 10); + cs_heap.update("b", 1.0); + + let mut a = SummaryState::CmsWithHeap( + CountMinSketchWithHeap::from_msgpack(&cms_heap.to_msgpack().unwrap()).unwrap(), + ); + let b = SummaryState::CountSketchWithHeap( + 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" + ), + } + } + + 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") + } + + /// `SummaryState::CountSketchWithHeap` must decode both FULL and + /// DELTA-HEAP msgpack frames through the genuine + /// `asap_sketchlib::CountSketchWithHeap` (median-of-signed-rows + /// estimator) rather than the CMS-family `CountMinSketchWithHeap` + /// (min-over-rows estimator) it used to alias — the bug this split + /// fixed. Built via real `update()` calls (not a hand-crafted matrix) + /// so the sign-hashed row semantics are genuinely exercised, then + /// checks both decode paths reproduce the same matrix and the same + /// `estimate()` as the in-memory sketch they were encoded from. + #[test] + fn count_sketch_with_heap_full_and_delta_decode_via_new_asap_sketchlib_type() { + use asap_sketchlib::{CountSketchWithHeap, MessagePackCodec}; + + let mut built = CountSketchWithHeap::new(4, 64, 10); + for _ in 0..50 { + built.update("k", 1.0); + } + let expected_matrix = built.sketch_matrix(); + let expected_estimate = built.estimate("k"); + + // FULL path. + let full_bytes = built.to_msgpack().expect("encode full CountSketchWithHeap"); + let full_state = decode_full( + &DeltaSketchKind::CountSketchWithHeap { + rows: 4, + cols: 64, + heap_size: 10, + }, + &full_bytes, + SketchEncoding::MsgpackFull, + ) + .expect("decode_full CountSketchWithHeap"); + match full_state { + SummaryState::CountSketchWithHeap(inner) => { + assert_eq!(inner.sketch_matrix(), expected_matrix); + assert_eq!(inner.estimate("k"), expected_estimate); + } + other => panic!( + "expected CountSketchWithHeap state, got {}", + other.family_name() + ), + } + + // DELTA-HEAP path: same cells + heap against an empty base (PWR + // contract), encoded the way the Go producer does. + let cells: Vec<(u32, u32, i64)> = expected_matrix + .iter() + .enumerate() + .flat_map(|(r, row)| { + row.iter().enumerate().filter_map(move |(c, v)| { + if *v != 0.0 { + Some((r as u32, c as u32, *v as i64)) + } else { + None + } + }) + }) + .collect(); + let heap_pairs: Vec<(String, f64)> = built + .topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect(); + assert!(!heap_pairs.is_empty(), "expected \"k\" in the top-k heap"); + let heap_refs: Vec<(&str, f64)> = + heap_pairs.iter().map(|(k, v)| (k.as_str(), *v)).collect(); + let delta_bytes = encode_delta_heap(4, 64, &cells, &heap_refs, 10); + + let mut rolling = DeltaSketchKind::CountSketchWithHeap { + rows: 4, + cols: 64, + heap_size: 10, + } + .bootstrap_empty(); + rolling + .apply_delta_bytes(&delta_bytes, SketchEncoding::MsgpackDelta) + .expect("apply CountSketchWithHeap delta"); + match rolling { + SummaryState::CountSketchWithHeap(inner) => { + assert_eq!( + inner.sketch_matrix(), + expected_matrix, + "delta path must reconstruct the identical matrix" + ); + assert_eq!(inner.estimate("k"), expected_estimate); + } + other => panic!( + "expected CountSketchWithHeap state, got {}", + other.family_name() + ), + } + } } diff --git a/data_plane/src/storage_engines/sketch_db/query/sketch_reducer.rs b/data_plane/src/storage_engines/sketch_db/query/sketch_reducer.rs index ec974ced..0c3ee3d5 100644 --- a/data_plane/src/storage_engines/sketch_db/query/sketch_reducer.rs +++ b/data_plane/src/storage_engines/sketch_db/query/sketch_reducer.rs @@ -63,6 +63,7 @@ use crate::storage_engines::sketch_db::query::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, }; use crate::storage_engines::sketch_db::query::delta_apply::{ cumulative_evaluate, per_window_evaluate, DeltaSketchKind, @@ -431,11 +432,10 @@ impl<'a> SketchReducer<'a> { t0_ms: u64, t1_ms: u64, ) -> Result { - use super::delta_apply::cumulative_hll_state; + use super::delta_apply::{cumulative_summary_state, DeltaSketchKind, SummaryState}; use crate::storage_engines::sketch_db::data::SketchConfig; - use asap_sketchlib::HllSketch; - let mut merged: Option = None; + let mut merged: Option = None; let mut metric_name_for_err = String::new(); let mut cov_lo: u64 = u64::MAX; let mut cov_hi: u64 = 0; @@ -478,23 +478,24 @@ impl<'a> SketchReducer<'a> { cov_lo = cov_lo.min(w); cov_hi = cov_hi.max(w); } - let series_state = cumulative_hll_state(&samples_vec, precision).map_err(|e| { - ASAPTierError::DeserializeFailure { - sid, - encoding: SketchEncoding::ProtoFull, - reason: e, - } - })?; + let series_state = + cumulative_summary_state(&samples_vec, DeltaSketchKind::Hll { precision }) + .map_err(|e| ASAPTierError::DeserializeFailure { + sid, + encoding: SketchEncoding::ProtoFull, + reason: e, + })?; if let Some(sk) = series_state { merged = Some(match merged.take() { None => sk, Some(mut acc) => { - acc.merge(&sk) - .map_err(|e| ASAPTierError::DeserializeFailure { + acc.merge_same_family(&sk).map_err(|e| { + ASAPTierError::DeserializeFailure { sid, encoding: SketchEncoding::ProtoFull, reason: format!("global HLL merge: {e}"), - })?; + } + })?; acc } }); @@ -508,7 +509,7 @@ impl<'a> SketchReducer<'a> { }); }; let _ = any_window; - let estimate = merged.estimate(); + let estimate = merged.cardinality(); let window_end = if cov_hi > 0 { cov_hi as i64 } else { @@ -667,25 +668,66 @@ impl<'a> SketchReducer<'a> { let mut summed: std::collections::HashMap = std::collections::HashMap::new(); let mut any_frame = false; + // Both heap-bearing kinds share a byte-identical wire + // shape and `topk_heap_items()` just reads back the + // agent-stored `(key, value)` pairs (no re-estimation + // happens here — see the comment above), so decoding + // either through `decode_cms_with_heap_from_msgpack*` + // produces the same items. Still dispatch on the sid's + // actual kind and decode through its own + // `asap_sketchlib` type, matching the other frequency + // paths, so this can't silently paper over a real + // divergence if one is ever introduced here. + let is_count_sketch = matches!( + meta.sketch_kind() + .expect("ASAP-tier reducer only handles sketch-backed sids"), + SketchKindHandle::CountSketchWithHeap + ); for state in frames { // A FULL frame deserializes directly; a MSGPACK_DELTA // frame is reconstructed by applying its sparse matrix // delta + full heap onto an empty base (per-window-reset). - let decoded = match state.encoding { - SketchEncoding::MsgpackDelta => { - decode_cms_with_heap_from_msgpack_delta(&state.bytes) + let items: Vec<(String, f64)> = if is_count_sketch { + let decoded = match state.encoding { + SketchEncoding::MsgpackDelta => { + decode_cs_with_heap_from_msgpack_delta(&state.bytes) + } + _ => decode_cs_with_heap_from_msgpack(&state.bytes), } - _ => decode_cms_with_heap_from_msgpack(&state.bytes), - } - .map_err(|e| { - ASAPTierError::DeserializeFailure { - sid, - encoding: state.encoding, - reason: e, + .map_err(|e| { + ASAPTierError::DeserializeFailure { + sid, + encoding: state.encoding, + reason: e, + } + })?; + decoded + .topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect() + } else { + let decoded = match state.encoding { + SketchEncoding::MsgpackDelta => { + decode_cms_with_heap_from_msgpack_delta(&state.bytes) + } + _ => decode_cms_with_heap_from_msgpack(&state.bytes), } - })?; - for item in decoded.topk_heap_items() { - *summed.entry(item.key).or_insert(0.0) += item.value; + .map_err(|e| { + ASAPTierError::DeserializeFailure { + sid, + encoding: state.encoding, + reason: e, + } + })?; + decoded + .topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect() + }; + for (key, value) in items { + *summed.entry(key).or_insert(0.0) += value; } any_frame = true; } @@ -765,7 +807,7 @@ impl<'a> SketchReducer<'a> { .copied() .filter(|q| (0.0..=1.0).contains(q)) .unwrap_or(0.99); - let evaluator: Box f64> = match family { + let evaluator: Box f64> = match family { QueryFamily::Quantile => Box::new(move |rs| rs.quantile(q)), QueryFamily::Cardinality => Box::new(|rs| rs.cardinality()), _ => unreachable!(), @@ -1528,12 +1570,17 @@ fn decode_frequency_total( }; Ok(row0_sum_cs(&cs)) } - // Heap-bearing variants: decode via the CMS-with-heap envelope - // and read the underlying CMS matrix the same way. A FULL frame + // Heap-bearing variants: the bucket TOTAL is a row-0 sum of the raw + // matrix, which is estimator-agnostic (unlike the per-key point + // estimate below), but each kind is still decoded through its own + // `asap_sketchlib` type — `CmsWithHeap` via `CountMinSketchWithHeap`, + // `CountSketchWithHeap` via the distinct `CountSketchWithHeap` — + // so a future field added to one family's decode/matrix layout + // can't silently leak into the other's arm. A FULL frame // (MSGPACK / PROTO) deserializes directly; a MSGPACK_DELTA frame // (the delta-heap wire form) is reconstructed by applying the // sparse matrix delta + full heap onto an empty base. - SketchKindHandle::CmsWithHeap | SketchKindHandle::CountSketchWithHeap => { + SketchKindHandle::CmsWithHeap => { let heap = match state.encoding { SketchEncoding::MsgpackDelta => { decode_cms_with_heap_from_msgpack_delta(&state.bytes) @@ -1542,8 +1589,18 @@ fn decode_frequency_total( _ => decode_cms_with_heap_from_msgpack(&state.bytes) .map_err(|e| to_err(e, state.encoding))?, }; - let matrix = heap.sketch_matrix(); - Ok(row0_sum_from_matrix(&matrix)) + Ok(row0_sum_from_matrix(&heap.sketch_matrix())) + } + SketchKindHandle::CountSketchWithHeap => { + let heap = match state.encoding { + SketchEncoding::MsgpackDelta => { + decode_cs_with_heap_from_msgpack_delta(&state.bytes) + .map_err(|e| to_err(e, state.encoding))? + } + _ => decode_cs_with_heap_from_msgpack(&state.bytes) + .map_err(|e| to_err(e, state.encoding))?, + }; + Ok(row0_sum_from_matrix(&heap.sketch_matrix())) } // Quantile / cardinality handles can't answer frequency — caller // should have rejected at `require_capability`. Defensive arm. @@ -1614,7 +1671,10 @@ fn decode_frequency_estimate( }; Ok(cs.estimate(key).max(0.0)) } - SketchKindHandle::CmsWithHeap | SketchKindHandle::CountSketchWithHeap => { + // CMS-with-heap: min-over-rows estimator, via `CountMinSketchWithHeap` + // directly — no need to rebuild a bare `CountMinSketch` wrapper, the + // heap type's own `estimate` already does the CMS math. + SketchKindHandle::CmsWithHeap => { let heap = match state.encoding { SketchEncoding::MsgpackDelta => { decode_cms_with_heap_from_msgpack_delta(&state.bytes) @@ -1623,11 +1683,24 @@ fn decode_frequency_estimate( _ => decode_cms_with_heap_from_msgpack(&state.bytes) .map_err(|e| to_err(e, state.encoding))?, }; - let matrix = heap.sketch_matrix(); - let rows = matrix.len(); - let cols = matrix.first().map(|r| r.len()).unwrap_or(0); - let cms = CountMinSketch::from_legacy_matrix(matrix, rows, cols); - Ok(cms.estimate(key).max(0.0)) + Ok(heap.estimate(key).max(0.0)) + } + // CountSketch-with-heap: median-of-signed-rows estimator, via the + // distinct `asap_sketchlib::CountSketchWithHeap`. Previously this + // arm was collapsed with `CmsWithHeap` above and always rebuilt a + // `CountMinSketch` (min-over-rows) regardless of family — silently + // wrong for any CountSketchWithHeap sid. Fixed the same way as + // `SummaryState::CountSketchWithHeap` in `delta_apply.rs`. + SketchKindHandle::CountSketchWithHeap => { + let heap = match state.encoding { + SketchEncoding::MsgpackDelta => { + decode_cs_with_heap_from_msgpack_delta(&state.bytes) + .map_err(|e| to_err(e, state.encoding))? + } + _ => decode_cs_with_heap_from_msgpack(&state.bytes) + .map_err(|e| to_err(e, state.encoding))?, + }; + Ok(heap.estimate(key).max(0.0)) } other => Err(ASAPTierError::UnsupportedCapability { function: "frequency_estimate".to_string(), @@ -1652,3 +1725,54 @@ fn row0_sum_from_matrix(matrix: &[Vec]) -> f64 { .map(|row| row.iter().copied().sum::()) .unwrap_or(0.0) } + +#[cfg(test)] +mod frequency_heap_tests { + use super::*; + use asap_sketchlib::{CountMinSketchWithHeap, CountSketchWithHeap, MessagePackCodec}; + + /// `decode_frequency_estimate`'s heap-bearing arm used to always decode + /// through `CountMinSketchWithHeap` (min-over-rows) regardless of + /// whether the sid was actually `CmsWithHeap` or `CountSketchWithHeap`. + /// Built via real `update()` calls (not a hand-crafted matrix, whose + /// per-row sign bits `estimate()` would reinterpret unpredictably), so + /// each kind's own `estimate("k")` is a ground truth captured before + /// encoding. Proves `decode_frequency_estimate` routes each sid kind + /// through its own `asap_sketchlib` type and reproduces that truth — + /// previously the `CountSketchWithHeap` sid would have silently gone + /// through `CountMinSketchWithHeap::estimate` instead. + #[test] + fn decode_frequency_estimate_uses_each_kinds_own_estimator() { + let mut cms_heap = CountMinSketchWithHeap::new(4, 64, 10); + for _ in 0..50 { + cms_heap.update("k", 1.0); + } + let cms_truth = cms_heap.estimate("k"); + let cms_state = SketchSampleState { + bytes: cms_heap.to_msgpack().expect("encode CmsWithHeap"), + encoding: SketchEncoding::MsgpackFull, + }; + let cms_estimate = + decode_frequency_estimate(1, SketchKindHandle::CmsWithHeap, &cms_state, "k") + .expect("decode CmsWithHeap estimate"); + assert_eq!(cms_estimate, cms_truth.max(0.0)); + + let mut cs_heap = CountSketchWithHeap::new(4, 64, 10); + for _ in 0..50 { + cs_heap.update("k", 1.0); + } + let cs_truth = cs_heap.estimate("k"); + let cs_state = SketchSampleState { + bytes: cs_heap.to_msgpack().expect("encode CountSketchWithHeap"), + encoding: SketchEncoding::MsgpackFull, + }; + let cs_estimate = + decode_frequency_estimate(2, SketchKindHandle::CountSketchWithHeap, &cs_state, "k") + .expect("decode CountSketchWithHeap estimate"); + assert_eq!( + cs_estimate, + cs_truth.max(0.0), + "must be CountSketch's own median-of-rows estimate, not CMS's min-over-rows" + ); + } +}