From cd5e5da8cb58e989dbcb62cd6bc58aacc0e4cc3f Mon Sep 17 00:00:00 2001 From: zz_y Date: Sun, 10 May 2026 14:22:11 -0600 Subject: [PATCH] =?UTF-8?q?feat:=20warm-tier=20follow-ups=20=E2=80=94=20Fr?= =?UTF-8?q?equencyTopk=20+=20Delta=20encoding=20+=20Hybrid=20stitch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three follow-ups from #124 land in one bundle. The warm-tier reducer now covers all four sketch families end-to-end (DD/KLL quantile, HLL cardinality, CMS-with-heap topk), applies delta-encoded windows by folding from the prior Full, and stitches with the archive engine when the warm tier's coverage is narrower than the query range. ## TODO 1 — FrequencyTopk (CMS-with-heap) [IMPLEMENTED] - New `SketchKindHandle::CmsWithHeap` variant in `stores/sketch_db/sketch_index.rs`. - `drivers/ingest/otel.rs::sketch_kind_handle_for` detects CmsWithHeap by attempting `CountMinSketchWithHeap::deserialize_msgpack(&dp.sketch)` on msgpack payloads and checking `topk_heap_items().is_empty()`. The heap is embedded in the msgpack wrapper, not on the OTLP wire (confirmed by inspecting metrics.proto::CountMinSketch). - New `engines/warm_tier/decoders.rs::decode_cms_with_heap_from_msgpack`: calls `asap_sketchlib::sketches::countminsketch_topk:: CountMinSketchWithHeap::deserialize_msgpack`, then `topk_heap_items()` for the heap. - CMS / CountSketch without heap surface as new `WarmTierError::MissingHeap { sid, sketch_kind }` — clearly distinguished from `UnsupportedCapability` so the EngineRouter's failover semantics can tell "no warm-tier handler for this function" from "warm tier has the data but lacks topk metadata". - `topk` output shape: one `(label_values, samples)` per item where `label_values["item"] = key` and `samples = [(window_end, estimated_count)]`. `k` from `function_args[0]`, default 10. ## TODO 2 — Delta encoding stitching [IMPLEMENTED] - New `engines/warm_tier/delta_apply.rs`: `RollingState` (DDSketch / KLL / HLL variants) + `per_window_evaluate` + `cumulative_evaluate`. - Sketch-lib calls used: `DdSketch::merge`, `KllSketch::merge`, `HllSketch::merge`, `HllSketch::apply_delta(&HllSketchDelta)` for proto register deltas, `HllSketch::deserialize_msgpack` for msgpack delta fragments. - Mode pick by function name: quantile_over_time / count_distinct_over_time / topk_over_time → cumulative (one scalar over [t0, t1]) quantile / histogram_quantile / cardinality_estimate → per-window - Leading deltas without a base Full are skipped (don't error); cumulative mode starts at the first Full in range. Tests verify DDSketch + HLL cumulative Full+Delta(s) round-trip within accuracy envelope of fresh-sketch ground truth. ## TODO 3 — Hybrid warm + archive stitch [IMPLEMENTED] - `WarmTierResult.coverage: Option<(u64, u64)>` populated from observed window-end timestamps. - `SimpleEngine.archive_engine: Option>` field + `with_archive_engine` builder. The trait `execute` adapter inspects `result.coverage` post-evaluate; if narrower than `[t0, now]`, it calls `archive.execute(query)` and stitches. - Chose option (b): in-line stitch in `SimpleEngine` rather than a new `EngineError::PartialHit` variant. Reason: the existing `EngineRouter` is variant-agnostic; introducing PartialHit would force every other engine to also handle it. - Stitch logic: index warm series by `Vec` label key, for each archive series merge by timestamp — warm wins on `[cov_lo, cov_hi]`, archive fills prefix/suffix. New private `stitch_warm_and_archive` helper. - Tests: `stitch_fills_archive_prefix_and_suffix` and `stitch_keeps_archive_only_series_in_full` verify both contracts. ## New API surface - `SketchKindHandle::CmsWithHeap` variant - `WarmTierError::MissingHeap { sid, sketch_kind }` variant - `WarmTierResult.coverage: Option<(u64, u64)>` field - `SimpleEngine::with_archive_engine(archive: Arc)` builder - `engines::warm_tier::decoders` (new module) — 5 decode helpers - `engines::warm_tier::delta_apply` (new module) — `DeltaSketchKind`, `RollingState`, `per_window_evaluate`, `cumulative_evaluate` ## Build + test - `cargo build --release -p query_engine_rust` — clean (3 pre-existing warnings) - Targeted run (24/24 pass): engines::warm_tier + the new hybrid_stitch_tests + warm_tier_classify_tests - Full lib run: 814 ok, 4 pre-existing failed (schema_timeline_dispatch + persistence_integration), 1 pre-existing hang (hard_cap_back_pressure) ## Migration note Existing `quantile_over_time` test had to flip its function name to `quantile` (per-window) to match the new cumulative semantics — caller code in production may need similar migration if it depended on per-window emit from `quantile_over_time`. ## Remaining follow-ups - KllSketch proto-delta wire path: today `delta_apply` decodes `Delta` as "Full fragment of matching encoding" then merges. The asap_sketchlib KLL proto wire `KllState.items[]` is sparse-friendly but there's no canonical `KllDelta` proto. If downstream agents start emitting a true sparse KLL delta, revisit. - `decoders::decode_cms_from_{proto,msgpack}`, `decode_cs_from_{proto,msgpack}` are dead-coded for now (kept for the eventual point-query `frequency(key, foo)` path). Co-Authored-By: Claude Opus 4.7 (1M context) --- asap-query-engine/src/drivers/ingest/otel.rs | 32 +- .../src/engines/simple/engine.rs | 243 +++++++++- .../src/engines/warm_tier/decoders.rs | 174 +++++++ .../src/engines/warm_tier/delta_apply.rs | 436 ++++++++++++++++++ .../src/engines/warm_tier/mod.rs | 14 + .../src/engines/warm_tier/sketch_reducer.rs | 228 ++++++++- .../src/engines/warm_tier/tests.rs | 317 ++++++++++++- .../src/stores/sketch_db/sketch_index.rs | 8 + 8 files changed, 1429 insertions(+), 23 deletions(-) create mode 100644 asap-query-engine/src/engines/warm_tier/decoders.rs create mode 100644 asap-query-engine/src/engines/warm_tier/delta_apply.rs diff --git a/asap-query-engine/src/drivers/ingest/otel.rs b/asap-query-engine/src/drivers/ingest/otel.rs index dc7c6b9d..c9003ece 100644 --- a/asap-query-engine/src/drivers/ingest/otel.rs +++ b/asap-query-engine/src/drivers/ingest/otel.rs @@ -893,7 +893,8 @@ async fn route_modified_otlp_sketches_to_precompute( } SketchKindHandle::Hll => Capability::CardinalityApprox, SketchKindHandle::CountSketch - | SketchKindHandle::CountMin => { + | SketchKindHandle::CountMin + | SketchKindHandle::CmsWithHeap => { Capability::FrequencyTopk(kind) } }; @@ -1077,6 +1078,16 @@ async fn route_modified_otlp_sketches_to_precompute( /// Phase 5 helper — map a `ModifiedOtlpSketchDp` to the matching /// `SketchKindHandle` so registration and capability classification /// share one source of truth. +/// +/// CMS-with-heap detection: the OTLP `CountMinSketch` wire struct +/// itself doesn't carry a top-k heap field (see metrics.proto +/// `CountMinSketch`/`CountMinSketchDataPoint`). The heap is embedded +/// inside the msgpack-encoded `CountMinSketchWithHeapSerialized` +/// payload (an outer `{sketch, topk_heap, heap_size}` wrapper). When +/// the encoding is MSGPACK and the bytes round-trip via +/// `CountMinSketchWithHeap::deserialize_msgpack`, we classify the sid +/// as `CmsWithHeap` so the warm-tier reducer can later read the heap +/// directly for `topk` / `topk_over_time` queries. fn sketch_kind_handle_for( dp: &ModifiedOtlpSketchDp, ) -> crate::stores::sketch_db::sketch_index::SketchKindHandle { @@ -1086,7 +1097,24 @@ fn sketch_kind_handle_for( SketchKind::Kll => SketchKindHandle::Kll, SketchKind::Hll => SketchKindHandle::Hll, SketchKind::CountSketch => SketchKindHandle::CountSketch, - SketchKind::CountMin => SketchKindHandle::CountMin, + SketchKind::CountMin => { + // Try a no-cost peek: msgpack-encoded CMS-with-heap payloads + // round-trip through asap_sketchlib's + // `CountMinSketchWithHeap::deserialize_msgpack`. If the + // sketch bytes decode against that wrapper *and* the + // resulting heap is non-empty, treat the sid as + // CmsWithHeap so warm-tier `topk` can read the heap. + // Otherwise stay with vanilla `CountMin`. + if dp.encoding == ENCODING_MSGPACK { + use asap_sketchlib::sketches::countminsketch_topk::CountMinSketchWithHeap; + if let Ok(cms) = CountMinSketchWithHeap::deserialize_msgpack(&dp.sketch) { + if !cms.topk_heap_items().is_empty() { + return SketchKindHandle::CmsWithHeap; + } + } + } + SketchKindHandle::CountMin + } } } diff --git a/asap-query-engine/src/engines/simple/engine.rs b/asap-query-engine/src/engines/simple/engine.rs index 81f6d0a8..7735f837 100644 --- a/asap-query-engine/src/engines/simple/engine.rs +++ b/asap-query-engine/src/engines/simple/engine.rs @@ -290,6 +290,18 @@ pub struct SimpleEngine { /// engine behaves as it did before Phase 5 wire-in (every query /// goes through `handle_query`'s legacy path). sketch_index: Option>, + /// Phase-5 hybrid-stitch hook — set by `with_archive_engine` from + /// `main.rs`'s engine builder. When the warm-tier reducer reports a + /// `WarmTierResult.coverage` narrower than the requested + /// `[t0, t1]`, the engine calls into this archive engine to fetch + /// the missing prefix / suffix and stitches the two answers by + /// `(label_values, timestamp)`. Warm-tier values win on overlap. + /// + /// When `None` (no archive engine wired), the engine returns the + /// warm answer as-is; the existing `EngineRouter` failover handles + /// the rest of the routing matrix. + archive_engine: + Option>, } impl SimpleEngine { @@ -472,9 +484,22 @@ impl SimpleEngine { controller_client: None, schema_registry: Arc::new(crate::stores::sketch_db::SchemaRegistry::empty()), sketch_index: None, + archive_engine: None, } } + /// Phase-5 hybrid-stitch builder — attach an archive engine the + /// `QueryEngine` trait adapter will dispatch to when the warm-tier + /// reducer reports a coverage narrower than the requested range. + /// When `None`, the engine returns whatever the warm tier covers. + pub fn with_archive_engine( + mut self, + archive: Arc, + ) -> Self { + self.archive_engine = Some(archive); + self + } + /// Phase 5 — attach the shared `SketchIndex` so the `QueryEngine` /// trait adapter's classify+failover logic is active. Without this /// call, the engine keeps the pre-Phase-5 behavior (route every @@ -3487,6 +3512,70 @@ impl SimpleEngine { /// `now_ms` is unused for the matrix variant (each sample carries its /// own window-end timestamp); it's plumbed for future extension to /// the instant-vector case (latest-pane projection). +/// Merge a warm-tier `QueryResult::Matrix` with an archive +/// `QueryResult::Matrix` by `(label_values, timestamp)`. Samples whose +/// timestamps fall inside the warm coverage `(cov_lo, cov_hi)` keep +/// the warm value (warm is approximate but more recent); samples +/// outside that window come from the archive answer. For +/// labels-not-present-in-warm series the archive series is taken in +/// full. Used by `SimpleEngine`'s hybrid-stitch path when the +/// warm-tier reducer reports `coverage` narrower than the request. +fn stitch_warm_and_archive( + warm: crate::engines::query_result::QueryResult, + archive: crate::engines::query_result::QueryResult, + cov_lo: u64, + cov_hi: u64, +) -> crate::engines::query_result::QueryResult { + use crate::engines::query_result::{QueryResult, RangeVectorElement, Sample}; + use std::collections::BTreeMap; + + let warm_matrix = match &warm { + QueryResult::Matrix(m) => m.values.clone(), + _ => return archive, + }; + let archive_matrix = match &archive { + QueryResult::Matrix(m) => m.values.clone(), + QueryResult::Vector(_) => return warm, + }; + + // Index warm series by labels for fast lookup. + let mut by_labels: BTreeMap, RangeVectorElement> = BTreeMap::new(); + for el in warm_matrix { + by_labels.insert(el.labels.labels.clone(), el); + } + + // For each archive series, merge into by_labels. + for arch_el in archive_matrix { + let entry = by_labels.entry(arch_el.labels.labels.clone()).or_insert_with( + || RangeVectorElement::new(arch_el.labels.clone()), + ); + // Build a set of warm timestamps inside coverage (kept). + let warm_ts: std::collections::HashSet = entry + .samples + .iter() + .filter(|s| s.timestamp >= cov_lo && s.timestamp <= cov_hi) + .map(|s| s.timestamp) + .collect(); + // Drop any warm samples that ended up outside coverage — + // archive will replace them. + entry + .samples + .retain(|s| s.timestamp >= cov_lo && s.timestamp <= cov_hi); + for s in arch_el.samples { + // Skip archive samples whose timestamps fall inside warm + // coverage AND warm produced a value there (warm wins). + if s.timestamp >= cov_lo && s.timestamp <= cov_hi && warm_ts.contains(&s.timestamp) { + continue; + } + entry.samples.push(Sample::new(s.timestamp, s.value)); + } + entry.samples.sort_by_key(|s| s.timestamp); + } + + let elements: Vec = by_labels.into_values().collect(); + QueryResult::matrix(elements) +} + fn warm_tier_result_to_query_result( result: crate::engines::warm_tier::WarmTierResult, _now_ms: u64, @@ -3639,7 +3728,34 @@ impl crate::routing::engine_router::QueryEngine for SimpleEngine { now_ms, ) { Ok(result) => { - return Ok(warm_tier_result_to_query_result(result, now_ms)); + // Phase-5 hybrid stitch — if the warm tier + // only covers a sub-range of `[t0, t1]` and an + // archive engine is wired, fetch the missing + // prefix / suffix and merge by + // `(label_values, timestamp)`. Warm-tier + // values win on overlap (warm is approximate + // but more recent; archive is the source of + // truth for older data). + let warm_qr = warm_tier_result_to_query_result(result.clone(), now_ms); + if let (Some((cov_lo, cov_hi)), Some(archive)) = + (result.coverage, self.archive_engine.as_ref()) + { + if cov_lo > t0_ms || cov_hi < now_ms { + let archive_qr = archive.execute(query).await; + if let Ok(archive_qr) = archive_qr { + return Ok(stitch_warm_and_archive( + warm_qr, + archive_qr, + cov_lo, + cov_hi, + )); + } + // On archive error, fall back to the + // warm-only answer (router can decide + // on a higher-level retry). + } + } + return Ok(warm_qr); } Err(crate::engines::warm_tier::WarmTierError::UnsupportedFunction( name, @@ -3689,6 +3805,24 @@ impl crate::routing::engine_router::QueryEngine for SimpleEngine { ), )); } + Err(crate::engines::warm_tier::WarmTierError::MissingHeap { + sid, + sketch_kind, + }) => { + // Top-k against a vanilla CMS / CountSketch + // (no embedded heap) — the reducer can't + // enumerate heavy hitters without the + // external item universe. Fall over to + // archive, which can scan raw samples. + return Err(crate::engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchWarmTier.data_source_id(), + format!( + "SketchWarmTier reducer cannot enumerate top-k for sid \ + {sid} (sketch_kind={sketch_kind:?}, no heap) — \ + failing over to archive" + ), + )); + } } } } @@ -6105,3 +6239,110 @@ mod warm_tier_classify_tests { } } } + +// =========================================================================== +// Hybrid warm + archive stitch tests (TODO 3 of the warm-tier follow-ups). +// Exercise `stitch_warm_and_archive` directly with synthetic +// `QueryResult::Matrix` payloads and assert the merged result honors +// the "warm wins on overlap; archive fills gaps" contract. +// =========================================================================== +#[cfg(test)] +mod hybrid_stitch_tests { + use super::stitch_warm_and_archive; + use crate::data_model::KeyByLabelValues; + use crate::engines::query_result::{QueryResult, RangeVectorElement, Sample}; + + fn matrix_with_samples(label: &str, samples: Vec<(u64, f64)>) -> QueryResult { + let labels = KeyByLabelValues::new_with_labels(vec![label.to_string()]); + let mut el = RangeVectorElement::new(labels); + for (t, v) in samples { + el.samples.push(Sample::new(t, v)); + } + QueryResult::matrix(vec![el]) + } + + #[test] + fn stitch_fills_archive_prefix_and_suffix() { + // Warm covers [100, 200] with timestamps 100, 150, 200. + let warm = matrix_with_samples( + "host=a", + vec![(100, 10.0), (150, 11.0), (200, 12.0)], + ); + // Archive covers [50, 250] with timestamps every 50ms. + let archive = matrix_with_samples( + "host=a", + vec![ + (50, 1.0), + (100, 99.0), // overlap: warm wins + (150, 99.0), // overlap: warm wins + (200, 99.0), // overlap: warm wins + (250, 2.0), + ], + ); + let merged = stitch_warm_and_archive(warm, archive, 100, 200); + let m = match merged { + QueryResult::Matrix(m) => m, + _ => panic!("expected matrix"), + }; + assert_eq!(m.values.len(), 1, "one series"); + let samples = &m.values[0].samples; + // Five distinct timestamps in the merged answer. + assert_eq!(samples.len(), 5); + // Warm values preserved on overlap. + let mut by_ts: std::collections::HashMap = + samples.iter().map(|s| (s.timestamp, s.value)).collect(); + assert_eq!(by_ts.remove(&100), Some(10.0)); + assert_eq!(by_ts.remove(&150), Some(11.0)); + assert_eq!(by_ts.remove(&200), Some(12.0)); + // Archive prefix / suffix preserved. + assert_eq!(by_ts.remove(&50), Some(1.0)); + assert_eq!(by_ts.remove(&250), Some(2.0)); + } + + #[test] + fn stitch_keeps_archive_only_series_in_full() { + // Warm has series "a"; archive has series "a" + "b". Both + // need to make it into the merged answer; "b" comes from + // archive in full. + let warm = matrix_with_samples("host=a", vec![(150, 5.0)]); + let archive = { + let a = { + let labels = KeyByLabelValues::new_with_labels(vec!["host=a".to_string()]); + let mut el = RangeVectorElement::new(labels); + el.samples.push(Sample::new(100, 1.0)); + el.samples.push(Sample::new(150, 99.0)); // warm wins + el.samples.push(Sample::new(200, 2.0)); + el + }; + let b = { + let labels = KeyByLabelValues::new_with_labels(vec!["host=b".to_string()]); + let mut el = RangeVectorElement::new(labels); + el.samples.push(Sample::new(100, 7.0)); + el.samples.push(Sample::new(200, 8.0)); + el + }; + QueryResult::matrix(vec![a, b]) + }; + let merged = stitch_warm_and_archive(warm, archive, 150, 150); + let m = match merged { + QueryResult::Matrix(m) => m, + _ => panic!("expected matrix"), + }; + assert_eq!(m.values.len(), 2, "two series after merge"); + let by_label: std::collections::HashMap, &RangeVectorElement> = m + .values + .iter() + .map(|e| (e.labels.labels.clone(), e)) + .collect(); + let a = by_label.get(&vec!["host=a".to_string()]).expect("series a"); + assert_eq!(a.samples.len(), 3); + let a_at_150 = a + .samples + .iter() + .find(|s| s.timestamp == 150) + .expect("warm value at 150 preserved"); + assert_eq!(a_at_150.value, 5.0, "warm wins on overlap"); + let b = by_label.get(&vec!["host=b".to_string()]).expect("series b"); + assert_eq!(b.samples.len(), 2); + } +} diff --git a/asap-query-engine/src/engines/warm_tier/decoders.rs b/asap-query-engine/src/engines/warm_tier/decoders.rs new file mode 100644 index 00000000..0cab96ab --- /dev/null +++ b/asap-query-engine/src/engines/warm_tier/decoders.rs @@ -0,0 +1,174 @@ +//! Per-sketch-kind decoder helpers — out-of-line wrappers around +//! `asap_sketchlib` deserialize / proto-decode paths. +//! +//! Lifted from the inline closures in [`crate::engines::warm_tier::sketch_reducer`] +//! once the reducer started decoding CMS / CountSketch / CMS-with-heap +//! payloads in addition to DDSketch / KLL / HLL. The CMS / CountSketch +//! / CMS-with-heap decoders mirror +//! `precompute_operators::{count_min_sketch, count_sketch, +//! count_min_sketch_with_heap}_accumulator.rs` bit-for-bit so the +//! warm-tier reducer's output matches what the precompute (ingest-side) +//! accumulator would have produced from the same bytes. +//! +//! Each entry point returns the typed sketchlib struct on success or a +//! plain `String` error; the reducer wraps the error into a +//! `WarmTierError::DeserializeFailure` so the engine router falls over +//! to archive cleanly. + +use asap_sketchlib::sketches::countminsketch::CountMinSketch; +use asap_sketchlib::sketches::countminsketch_topk::CountMinSketchWithHeap; +use asap_sketchlib::sketches::countsketch::CountSketch; + +/// 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_accumulator::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 { + return Err(format!("CountMinState has zero dims (rows={rows}, cols={cols})")); + } + 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" + )); + } + }; + 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(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::deserialize_msgpack(buffer) + .map_err(|e| format!("deserialize CountMinSketch msgpack: {e}")) +} + +/// Decode a `CountSketch` from the modified-OTLP proto wire bytes. +/// Mirrors +/// `precompute_operators::count_sketch_accumulator::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 { + return Err(format!("CountSketchState has zero dims (rows={rows}, cols={cols})")); + } + let expected_len = rows * cols; + let counter_type = CounterType::try_from(state.counter_type) + .map_err(|_| format!("CountSketchState 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!( + "CountSketchState 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!( + "CountSketchState counts_float has {} entries, expected {}", + state.counts_float.len(), + expected_len + )); + } + state.counts_float.clone() + } + other => { + return Err(format!( + "CountSketchState counter_type {other:?} not yet supported in reducer" + )); + } + }; + 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(CountSketch::from_legacy_matrix(matrix, rows, cols)) +} + +/// Decode a `CountSketch` from msgpack bytes (sketch-core wire format). +pub fn decode_cs_from_msgpack(buffer: &[u8]) -> Result { + CountSketch::deserialize_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketch msgpack: {e}")) +} + +/// 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). Mirrors +/// `precompute_operators::count_min_sketch_with_heap_accumulator::deserialize_from_bytes_arroyo`. +pub fn decode_cms_with_heap_from_msgpack( + buffer: &[u8], +) -> Result { + CountMinSketchWithHeap::deserialize_msgpack(buffer) + .map_err(|e| format!("deserialize CountMinSketchWithHeap msgpack: {e}")) +} diff --git a/asap-query-engine/src/engines/warm_tier/delta_apply.rs b/asap-query-engine/src/engines/warm_tier/delta_apply.rs new file mode 100644 index 00000000..abf25847 --- /dev/null +++ b/asap-query-engine/src/engines/warm_tier/delta_apply.rs @@ -0,0 +1,436 @@ +//! Per-window delta stitching for the warm-tier sketch reducer. +//! +//! The reducer walks a sid's per-window samples in time order. When a +//! window's payload is a `Full` encoding (PROTO / MSGPACK), it +//! initializes a rolling "current state" sketch. Subsequent `Delta` +//! encodings merge into that rolling state — either via +//! `asap_sketchlib::HllSketch::apply_delta` for HLL register deltas, +//! or via `Sketch::merge(decoded_delta)` for DD / KLL where the wire +//! format ships a sparse-but-mergeable sketch fragment. +//! +//! Two reducer modes, picked by the PromQL function name in +//! [`crate::engines::warm_tier::sketch_reducer`]: +//! +//! * **per-window** (`quantile`, `histogram_quantile`, +//! `cardinality_estimate`): emit one scalar per window. A `Full` +//! resets the rolling state; a `Delta` merges then emits from the +//! merged state. Window-end timestamps come from the index. +//! * **cumulative** (`quantile_over_time`, `count_distinct_over_time`): +//! walk the full `[t0, t1]` range, accumulating Full + all subsequent +//! Deltas into a single rolling state. Emit one scalar at the last +//! window's end_ms (the rolled-up answer covers the whole window). +//! +//! ## Corner case: leading delta +//! +//! If the first sample in a query window is a `Delta`, the base +//! snapshot from the previous (out-of-range) window isn't available +//! to the reducer. The leading delta is dropped (logged via the +//! `DeserializeFailure { reason: "leading delta without base" }` path +//! when no Full follows in the same window) and we wait for the next +//! `Full`. For cumulative mode this means the cumulative answer +//! starts at the first Full in the range, not at `t0`. + +use asap_sketchlib::sketches::ddsketch::DdSketch; +use asap_sketchlib::sketches::hll::HllSketch; +use asap_sketchlib::sketches::kll::KllSketch; + +use crate::stores::sketch_db::sketch_index::{SketchEncoding, SketchSampleState}; + +/// 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. +#[derive(Debug, Clone, Copy)] +pub enum DeltaSketchKind { + DDSketch, + Hll, + Kll, +} + +/// Try to decode a "full" sketch from the bytes (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) { + (DeltaSketchKind::DDSketch, SketchEncoding::ProtoFull) => { + let sk = dd_from_proto(bytes)?; + Ok(RollingState::Dd(sk)) + } + (DeltaSketchKind::DDSketch, SketchEncoding::MsgpackFull) => { + let sk = DdSketch::deserialize_msgpack(bytes) + .map_err(|e| format!("deserialize DDSketch msgpack: {e}"))?; + Ok(RollingState::Dd(sk)) + } + (DeltaSketchKind::Hll, SketchEncoding::ProtoFull) => { + let sk = hll_from_proto(bytes)?; + Ok(RollingState::Hll(sk)) + } + (DeltaSketchKind::Hll, SketchEncoding::MsgpackFull) => { + let sk = HllSketch::deserialize_msgpack(bytes) + .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?; + Ok(RollingState::Hll(sk)) + } + (DeltaSketchKind::Kll, SketchEncoding::ProtoFull) => { + let sk = kll_from_proto(bytes)?; + Ok(RollingState::Kll(sk)) + } + (DeltaSketchKind::Kll, SketchEncoding::MsgpackFull) => { + let sk = KllSketch::deserialize_msgpack(bytes) + .map_err(|e| format!("deserialize KllSketch msgpack: {e}"))?; + Ok(RollingState::Kll(sk)) + } + (_, e) => Err(format!( + "decode_full called with non-Full encoding {e:?}" + )), + } +} + +/// The rolling state the delta-application loop maintains. +pub enum RollingState { + Dd(DdSketch), + Hll(HllSketch), + Kll(KllSketch), +} + +impl RollingState { + /// 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. + pub fn apply_delta_bytes( + &mut self, + bytes: &[u8], + encoding: SketchEncoding, + ) -> Result<(), String> { + if !matches!( + encoding, + SketchEncoding::ProtoDelta | SketchEncoding::MsgpackDelta + ) { + return Err(format!("apply_delta_bytes called with non-Delta encoding {encoding:?}")); + } + match self { + RollingState::Dd(sk) => { + // The wire delta for DD/KLL today is a full-sketch + // fragment (sparse buckets); decode it via the same + // full-state path and merge into `sk`. Treat the bytes + // as a Full payload of the matching encoding family. + let full_enc = match encoding { + SketchEncoding::ProtoDelta => SketchEncoding::ProtoFull, + SketchEncoding::MsgpackDelta => SketchEncoding::MsgpackFull, + _ => unreachable!(), + }; + let other = match decode_full(&DeltaSketchKind::DDSketch, bytes, full_enc) { + Ok(RollingState::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(()) + } + RollingState::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) + } else { + // MsgpackDelta for HLL isn't a sparse encoding; + // it's a serialized HllSketch fragment, mergeable + // via `HllSketch::merge`. + let other = HllSketch::deserialize_msgpack(bytes) + .map_err(|e| format!("deserialize HllSketch (delta-as-msgpack): {e}"))?; + sk.merge(&other) + .map_err(|e| format!("merge HLL delta: {e}"))?; + Ok(()) + } + } + RollingState::Kll(sk) => { + let full_enc = match encoding { + SketchEncoding::ProtoDelta => SketchEncoding::ProtoFull, + SketchEncoding::MsgpackDelta => SketchEncoding::MsgpackFull, + _ => unreachable!(), + }; + let other = match decode_full(&DeltaSketchKind::Kll, bytes, full_enc) { + Ok(RollingState::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(()) + } + } + } + + 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, + } + } + + pub fn cardinality(&self) -> f64 { + match self { + RollingState::Hll(sk) => sk.estimate(), + _ => 0.0, + } + } +} + +/// Walk a sorted-by-window-end slice of samples in time order and +/// produce per-window scalars. On a `Full` payload, replace the +/// rolling state; on a `Delta`, apply it into the rolling state. +/// Each window emits one `(window_end_ms, scalar)`. +/// +/// `eval` reads a scalar from the rolling state (`quantile(q)` / +/// `cardinality()`). Leading deltas (before any Full) are skipped +/// with a debug-grade error returned to the caller for reporting. +/// +/// Returns `Ok(samples, skipped_leading_deltas)`. +pub fn per_window_evaluate( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, + eval: E, +) -> Result<(Vec<(i64, f64)>, usize), String> +where + E: Fn(&RollingState) -> f64, +{ + let mut out: Vec<(i64, f64)> = Vec::with_capacity(samples.len()); + let mut rolling: Option = None; + let mut skipped = 0usize; + + for (window_end, state) in samples { + match state.encoding { + SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull => { + rolling = Some(decode_full(&kind, &state.bytes, state.encoding)?); + if let Some(rs) = &rolling { + out.push((*window_end, eval(rs))); + } + } + SketchEncoding::ProtoDelta | SketchEncoding::MsgpackDelta => { + if let Some(rs) = rolling.as_mut() { + rs.apply_delta_bytes(&state.bytes, state.encoding)?; + out.push((*window_end, eval(rs))); + } else { + skipped += 1; + } + } + } + } + Ok((out, skipped)) +} + +/// Cumulative-mode rollup: fold every window in `[t0, t1]` into a +/// single rolling state and emit one scalar at the latest +/// `window_end_ms` seen (or the largest if all were Deltas that got +/// skipped). Used by `quantile_over_time` / `count_distinct_over_time`. +/// +/// Returns `Ok((window_end, scalar), skipped_leading_deltas)`. Returns +/// `Ok(None, _)` if every sample was a leading delta (no Full ever +/// landed in the range). +pub fn cumulative_evaluate( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, + eval: E, +) -> Result<(Option<(i64, f64)>, usize), String> +where + E: Fn(&RollingState) -> f64, +{ + let mut rolling: Option = None; + let mut latest_end = i64::MIN; + let mut skipped = 0usize; + + for (window_end, state) in samples { + if *window_end > latest_end { + latest_end = *window_end; + } + match state.encoding { + SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull => { + let new_state = decode_full(&kind, &state.bytes, state.encoding)?; + // Merge new_state into any existing rolling state — a + // mid-range Full effectively "restarts" the window in + // the cumulative roll-up if the agent flushed a new + // snapshot. Merging keeps the answer monotonic in + // sample inclusion. + rolling = Some(match (rolling.take(), new_state) { + (None, n) => n, + (Some(RollingState::Dd(mut a)), RollingState::Dd(b)) => { + a.merge(&b).map_err(|e| format!("cum merge DD: {e}"))?; + RollingState::Dd(a) + } + (Some(RollingState::Hll(mut a)), RollingState::Hll(b)) => { + a.merge(&b).map_err(|e| format!("cum merge HLL: {e}"))?; + RollingState::Hll(a) + } + (Some(RollingState::Kll(mut a)), RollingState::Kll(b)) => { + a.merge(&b) + .map_err(|e| format!("cum merge KLL: {e}"))?; + RollingState::Kll(a) + } + (Some(_), _) => { + return Err( + "cumulative merge across sketch family mismatch".to_string() + ) + } + }); + } + SketchEncoding::ProtoDelta | SketchEncoding::MsgpackDelta => { + if let Some(rs) = rolling.as_mut() { + rs.apply_delta_bytes(&state.bytes, state.encoding)?; + } else { + skipped += 1; + } + } + } + } + let out = rolling.map(|rs| (latest_end, eval(&rs))); + Ok((out, skipped)) +} + +// --------------------------------------------------------------------------- +// Proto-envelope decoders — duplicated minimally from the inline forms +// in `sketch_reducer.rs` so this module can decode "delta as full +// fragment" without re-entering the reducer's private functions. +// --------------------------------------------------------------------------- + +fn dd_from_proto(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; + use prost::Message; + let state = match SketchEnvelope::decode(buffer) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::Ddsketch(st)) => st, + Some(_) => return Err("SketchEnvelope contains non-DDSketch sketch".to_string()), + None => DdSketchState::decode(buffer) + .map_err(|e| format!("decode DDSketchState: {e}"))?, + }, + Err(_) => DdSketchState::decode(buffer) + .map_err(|e| format!("decode DDSketchState: {e}"))?, + }; + if !(state.alpha > 0.0 && state.alpha < 1.0) { + return Err(format!( + "DDSketchState alpha {} out of range (expected 0 < alpha < 1)", + state.alpha + )); + } + Ok(DdSketch::from_raw( + state.alpha, + state.store_counts.clone(), + state.store_offset, + state.count, + state.sum, + state.min, + state.max, + )) +} + +fn kll_from_proto(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + let state = match SketchEnvelope::decode(buffer) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::Kll(st)) => st, + Some(_) => return Err("SketchEnvelope contains non-KLL sketch".to_string()), + None => KllState::decode(buffer).map_err(|e| format!("decode KllState: {e}"))?, + }, + Err(_) => KllState::decode(buffer).map_err(|e| format!("decode KllState: {e}"))?, + }; + if state.k < 8 { + return Err(format!("KllState.k must be >= 8 (got {})", state.k)); + } + if state.k > u16::MAX as u32 { + return Err(format!( + "KllState.k does not fit in u16 (got {}, max {})", + state.k, + u16::MAX + )); + } + let k = state.k as u16; + let mut sk = KllSketch::new(k); + for item in &state.items { + sk.update(*item); + } + Ok(sk) +} + +fn hll_from_proto(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use asap_sketchlib::sketches::hll::HllVariant; + use prost::Message; + let state = match SketchEnvelope::decode(buffer) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::Hll(st)) => st, + Some(_) => return Err("SketchEnvelope contains non-HLL sketch".to_string()), + None => HyperLogLogState::decode(buffer) + .map_err(|e| format!("decode HyperLogLogState: {e}"))?, + }, + Err(_) => HyperLogLogState::decode(buffer) + .map_err(|e| format!("decode HyperLogLogState: {e}"))?, + }; + 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; + if state.registers.len() != expected_len { + return Err(format!( + "HyperLogLogState registers has {} bytes, expected 2^precision = {}", + state.registers.len(), + 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, + }; + Ok(HllSketch::from_raw( + variant, + state.precision, + state.registers.clone(), + state.hip_kxq0, + state.hip_kxq1, + state.hip_est, + )) +} + +/// Apply a proto-encoded `HllDelta` frame onto the HLL register vector +/// — mirrors `HllSketchAccumulator::apply_proto_delta_bytes` (sparse +/// `(index, value)` updates, `register = max(register, value)`). +fn apply_hll_proto_delta(sk: &mut HllSketch, buffer: &[u8]) -> Result<(), String> { + use asap_otel_proto::sketchlib::v1::HllDelta as PbDelta; + use asap_sketchlib::sketches::hll::HllSketchDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode HLLDelta: {e}"))?; + let updates = pb + .updates + .into_iter() + .map(|u| (u.index, u.value as u8)) + .collect(); + let delta = HllSketchDelta { updates }; + sk.apply_delta(&delta) + .map_err(|e| format!("apply HLLDelta: {e}"))?; + Ok(()) +} diff --git a/asap-query-engine/src/engines/warm_tier/mod.rs b/asap-query-engine/src/engines/warm_tier/mod.rs index bbac7d2c..858590f8 100644 --- a/asap-query-engine/src/engines/warm_tier/mod.rs +++ b/asap-query-engine/src/engines/warm_tier/mod.rs @@ -44,7 +44,21 @@ //! Phase-5 hybrid stitching (warm `[t0..t1']` + archive //! `[t1'..t1]`) and per-window iteration (rather than today's //! per-sample evaluate-then-merge) remain follow-ups. +//! +//! 2026-05 follow-ups landed here: +//! * **TODO 1**: CMS-with-heap top-k. `Capability::FrequencyTopk(CmsWithHeap)` +//! reads the embedded heap directly; CMS / CountSketch without a heap +//! surface as `WarmTierError::MissingHeap` and fail over to archive. +//! * **TODO 2**: Delta encoding stitching. `ProtoDelta` / `MsgpackDelta` +//! are now applied via [`delta_apply`] — see that module's docs for the +//! per-window vs cumulative modes (selected by function name). +//! * **TODO 3**: Hybrid warm+archive stitch. [`WarmTierResult::coverage`] +//! reports the actual `(min_window_start_ms, max_window_end_ms)` the +//! reducer covered so `SimpleEngine` can stitch the missing prefix / +//! suffix from the archive engine. +pub mod decoders; +pub mod delta_apply; pub mod promql_extract; pub mod sketch_reducer; diff --git a/asap-query-engine/src/engines/warm_tier/sketch_reducer.rs b/asap-query-engine/src/engines/warm_tier/sketch_reducer.rs index 02efd5de..607d02f7 100644 --- a/asap-query-engine/src/engines/warm_tier/sketch_reducer.rs +++ b/asap-query-engine/src/engines/warm_tier/sketch_reducer.rs @@ -57,6 +57,10 @@ use asap_sketchlib::sketches::kll::KllSketch; use asap_sketchlib::sketches::countminsketch::CountMinSketch; use asap_sketchlib::sketches::countsketch::CountSketch; +use crate::engines::warm_tier::decoders::decode_cms_with_heap_from_msgpack; +use crate::engines::warm_tier::delta_apply::{ + cumulative_evaluate, per_window_evaluate, DeltaSketchKind, +}; use crate::stores::sketch_db::sketch_index::{ Capability, SketchEncoding, SketchIndex, SketchInstanceMetadata, SketchKindHandle, SketchSampleState, @@ -73,7 +77,11 @@ pub struct SketchReducer<'a> { /// can't answer this; archive can". `DeserializeFailure` → "the /// warm-tier state didn't decode; defensive fallback". `NoData` → /// "the sketch index has no samples in `[t0, t1]`; archive may have -/// older history". +/// older history". `MissingHeap` → "the sid is FrequencyTopk-classed +/// but the underlying sketch family carries no heap (vanilla +/// CountSketch / CountMinSketch without `CmsWithHeap`), so the +/// reducer can't materialize top-k items without an external item +/// universe". #[derive(Debug)] pub enum WarmTierError { UnsupportedFunction(String), @@ -81,6 +89,15 @@ pub enum WarmTierError { function: String, capability: Capability, }, + /// Top-k requested against a `FrequencyTopk(CountMin)` or + /// `FrequencyTopk(CountSketch)` sid (i.e. the sketch shape + /// supports point queries but not heavy-hitter enumeration). The + /// router falls over to archive — an archive scan can materialize + /// the full item universe and compute the true top-k. + MissingHeap { + sid: u64, + sketch_kind: SketchKindHandle, + }, DeserializeFailure { sid: u64, encoding: SketchEncoding, @@ -104,6 +121,13 @@ impl std::fmt::Display for WarmTierError { f, "warm-tier reducer cannot answer `{function}` against capability {capability:?}" ), + WarmTierError::MissingHeap { sid, sketch_kind } => write!( + f, + "warm-tier reducer cannot enumerate top-k for sid {sid}: \ + sketch kind {sketch_kind:?} carries no top-k heap \ + (CountMin / CountSketch only support point-frequency queries; \ + use CmsWithHeap for top-k)" + ), WarmTierError::DeserializeFailure { sid, encoding, @@ -124,11 +148,23 @@ impl std::fmt::Display for WarmTierError { impl std::error::Error for WarmTierError {} /// Per-series, per-window scalar results. +/// +/// `coverage` is the actual `(min_window_start_ms, max_window_end_ms)` +/// the reducer covered. `None` when the reducer didn't observe any +/// in-range window (defensive default). The caller (`SimpleEngine`) +/// compares `coverage` against the requested `[t0, t1]` and, on a +/// partial hit (`cov_lo > t0 || cov_hi < t1`), falls over to archive +/// for the missing range and stitches the two answers. See TODO 3 in +/// the warm-tier follow-up PR. #[derive(Debug, Clone, Default)] pub struct WarmTierResult { /// `(label_values, samples)` where `samples` is /// `(window_end_unix_ms, value)`. pub series: Vec<(BTreeMap, Vec<(i64, f64)>)>, + /// Effective coverage `(min_window_start_ms, max_window_end_ms)`. + /// Set whenever the reducer observed at least one window; left + /// `None` when `series` is empty. + pub coverage: Option<(u64, u64)>, } impl WarmTierResult { @@ -192,6 +228,26 @@ impl<'a> SketchReducer<'a> { /// verified to classify as `Hit` against `self.index`. We /// re-resolve metadata (via `instance(sid)`) but don't /// re-classify. + /// + /// ## Delta-stitching modes + /// + /// For DD / KLL / HLL the reducer walks per-window samples in + /// time order via [`delta_apply`](super::delta_apply). The function + /// name decides between: + /// - **per-window**: `quantile`, `histogram_quantile`, + /// `cardinality_estimate` — one scalar per window-end. + /// - **cumulative**: `quantile_over_time`, + /// `count_distinct_over_time` — single scalar covering the + /// full `[t0, t1]` range (Full + every subsequent Delta merged). + /// + /// ## Top-k mode + /// + /// For `topk` / `topk_over_time` against a CmsWithHeap sid, the + /// reducer reads the heap directly from the most-recent window's + /// `CountMinSketchWithHeap` state and emits one + /// `(label_values={"item": }, [(window_end, count)])` entry + /// per top-k item, truncated to the user's `k`. CountMin / + /// CountSketch (no heap) surface as `MissingHeap`. pub fn evaluate( &self, sids: &[u64], @@ -201,11 +257,17 @@ impl<'a> SketchReducer<'a> { t1_ms: u64, ) -> Result { let family = Self::function_to_family(function_name)?; + let is_cumulative = matches!( + function_name, + "quantile_over_time" | "count_distinct_over_time" | "topk_over_time" + ); // Per-(sid, label-values) → time-stamped scalar values. let mut out_series: Vec<(BTreeMap, Vec<(i64, f64)>)> = Vec::new(); let mut metric_name_for_err = String::new(); let mut any_window = false; + let mut cov_lo: u64 = u64::MAX; + let mut cov_hi: u64 = 0; for &sid in sids { let meta = match self.index.instance(sid) { @@ -220,21 +282,132 @@ impl<'a> SketchReducer<'a> { continue; } + // Top-k is a different shape — one entry per top-k item. + if family == QueryFamily::FrequencyTopk { + let k = function_args + .first() + .copied() + .filter(|k| *k > 0.0) + .map(|k| k as usize) + .unwrap_or(10); + for ts in series_list { + // Find latest window's CMS-with-heap state. + let Some((window_end, state)) = ts.samples.iter().next_back() else { + continue; + }; + any_window = true; + let w_end_u64 = if *window_end >= 0 { *window_end as u64 } else { 0 }; + if w_end_u64 < cov_lo { + cov_lo = w_end_u64; + } + if w_end_u64 > cov_hi { + cov_hi = w_end_u64; + } + let cms_heap = match meta.sketch_kind { + SketchKindHandle::CmsWithHeap => { + decode_cms_with_heap_from_msgpack(&state.bytes).map_err(|e| { + WarmTierError::DeserializeFailure { + sid, + encoding: state.encoding, + reason: e, + } + })? + } + SketchKindHandle::CountMin | SketchKindHandle::CountSketch => { + return Err(WarmTierError::MissingHeap { + sid, + sketch_kind: meta.sketch_kind, + }); + } + other => { + return Err(WarmTierError::UnsupportedCapability { + function: function_name.to_string(), + capability: Capability::FrequencyTopk(other), + }); + } + }; + let mut items = cms_heap.topk_heap_items(); + // Sort descending by estimated count. + items.sort_by(|a, b| { + b.value + .partial_cmp(&a.value) + .unwrap_or(std::cmp::Ordering::Equal) + }); + for item in items.into_iter().take(k) { + let mut lv = ts.series_label_values.clone(); + lv.insert("item".to_string(), item.key); + out_series.push((lv, vec![(*window_end, item.value)])); + } + } + continue; + } + + // Quantile / Cardinality with delta stitching. + let delta_kind = match (family, meta.sketch_kind) { + (QueryFamily::Quantile, SketchKindHandle::DDSketch) => DeltaSketchKind::DDSketch, + (QueryFamily::Quantile, SketchKindHandle::Kll) => DeltaSketchKind::Kll, + (QueryFamily::Cardinality, SketchKindHandle::Hll) => DeltaSketchKind::Hll, + _ => { + return Err(WarmTierError::UnsupportedCapability { + function: function_name.to_string(), + capability: meta.capability.clone(), + }); + } + }; + let q = function_args + .first() + .copied() + .filter(|q| (0.0..=1.0).contains(q)) + .unwrap_or(0.99); + let evaluator: Box f64> = match family { + QueryFamily::Quantile => Box::new(move |rs| rs.quantile(q)), + QueryFamily::Cardinality => Box::new(|rs| rs.cardinality()), + _ => unreachable!(), + }; + for ts in series_list { - let mut samples: Vec<(i64, f64)> = Vec::with_capacity(ts.samples.len()); - for (window_end, state) in &ts.samples { + // Build sorted-by-window-end slice of refs. + let samples_vec: Vec<(i64, &SketchSampleState)> = + ts.samples.iter().map(|(t, s)| (*t, s)).collect(); + // BTreeMap iteration is already sorted by key; the + // collect preserves order. Track coverage from raw + // window-end timestamps before delta evaluation + // (skipped leading deltas still count toward the + // covered range). + for (w_end, _) in &samples_vec { any_window = true; - let value = self.evaluate_one_state( - sid, - family, - meta.sketch_kind, - function_args, - state, - )?; - samples.push((*window_end, value)); + let w = if *w_end >= 0 { *w_end as u64 } else { 0 }; + if w < cov_lo { + cov_lo = w; + } + if w > cov_hi { + cov_hi = w; + } } - samples.sort_by_key(|(t, _)| *t); - out_series.push((ts.series_label_values, samples)); + + let samples_out: Vec<(i64, f64)> = if is_cumulative { + let (one, _skipped) = + cumulative_evaluate(&samples_vec, delta_kind, &evaluator) + .map_err(|e| WarmTierError::DeserializeFailure { + sid, + encoding: SketchEncoding::ProtoFull, + reason: e, + })?; + match one { + Some(s) => vec![s], + None => Vec::new(), + } + } else { + let (per_win, _skipped) = + per_window_evaluate(&samples_vec, delta_kind, &evaluator) + .map_err(|e| WarmTierError::DeserializeFailure { + sid, + encoding: SketchEncoding::ProtoFull, + reason: e, + })?; + per_win + }; + out_series.push((ts.series_label_values, samples_out)); } } @@ -244,11 +417,27 @@ impl<'a> SketchReducer<'a> { }); } - Ok(WarmTierResult { series: out_series }) + let coverage = if cov_lo <= cov_hi { + Some((cov_lo, cov_hi)) + } else { + None + }; + Ok(WarmTierResult { + series: out_series, + coverage, + }) } /// Decode one window's sketch state and run the family-appropriate /// reduction. + /// + /// Retained as `#[allow(dead_code)]` after the delta-stitching + /// follow-up moved the per-window decode-then-evaluate flow into + /// [`super::delta_apply`]. Callers that want a one-shot evaluate + /// without delta-state plumbing can still reach this entry point; + /// the warm-tier reducer's main loop now goes through + /// [`per_window_evaluate`] / [`cumulative_evaluate`]. + #[allow(dead_code)] fn evaluate_one_state( &self, sid: u64, @@ -287,6 +476,7 @@ impl<'a> SketchReducer<'a> { } } + #[allow(dead_code)] fn evaluate_quantile( &self, sid: u64, @@ -310,6 +500,7 @@ impl<'a> SketchReducer<'a> { } } + #[allow(dead_code)] fn evaluate_cardinality( &self, sid: u64, @@ -336,6 +527,7 @@ impl<'a> SketchReducer<'a> { // surface as decode failure. // --------------------------------------------------------------------------- +#[allow(dead_code)] fn decode_ddsketch( sid: u64, state: &SketchSampleState, @@ -368,6 +560,7 @@ fn decode_ddsketch( } } +#[allow(dead_code)] fn decode_kll( sid: u64, state: &SketchSampleState, @@ -398,6 +591,7 @@ fn decode_kll( } } +#[allow(dead_code)] fn decode_hll( sid: u64, state: &SketchSampleState, @@ -433,7 +627,7 @@ fn decode_hll( // proto envelope wrapping (from DataCollector's `*processor`) is a // product of the OTLP wire layer, not the sketch library. -#[allow(non_snake_case)] +#[allow(non_snake_case, dead_code)] fn DdSketch_from_sketchlib_proto_bytes(buffer: &[u8]) -> Result { use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; use prost::Message; @@ -464,7 +658,7 @@ fn DdSketch_from_sketchlib_proto_bytes(buffer: &[u8]) -> Result Result { use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; use prost::Message; @@ -494,7 +688,7 @@ fn KllSketch_from_sketchlib_proto_bytes(buffer: &[u8]) -> Result Result { use asap_sketchlib::proto::sketchlib::{ sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, diff --git a/asap-query-engine/src/engines/warm_tier/tests.rs b/asap-query-engine/src/engines/warm_tier/tests.rs index 38822eb8..3018f506 100644 --- a/asap-query-engine/src/engines/warm_tier/tests.rs +++ b/asap-query-engine/src/engines/warm_tier/tests.rs @@ -141,13 +141,17 @@ fn hll_meta(sid: u64, precision: u32) -> SketchInstanceMetadata { } // --------------------------------------------------------------------------- -// DDSketch quantile_over_time — three windows, each with a different +// DDSketch per-window `quantile` — three windows, each with a different // data distribution. Verifies (a) per-window evaluation, (b) result // shape, (c) DDSketch's relative-accuracy bound holds. +// +// The cumulative variant `quantile_over_time` is exercised by +// `ddsketch_cumulative_full_plus_two_deltas` (TODO-2 follow-up); this +// test is renamed but otherwise preserves its original assertions. // --------------------------------------------------------------------------- #[test] -fn ddsketch_quantile_over_time_three_windows() { +fn ddsketch_quantile_per_window_three_windows() { let idx = SketchIndex::new(); let sid = 1; idx.register(dd_meta(sid)); @@ -175,9 +179,12 @@ fn ddsketch_quantile_over_time_three_windows() { idx.append_sample(sid, lv, (window_start, window_end), proto_full(bytes)); } + // `quantile` (per-window) emits one scalar per window; the + // cumulative variant `quantile_over_time` is exercised by + // [`quantile_over_time_cumulative_mode`] below. let reducer = SketchReducer::new(&idx); let result = reducer - .evaluate(&[sid], "quantile_over_time", &[0.5], 1000, 1100) + .evaluate(&[sid], "quantile", &[0.5], 1000, 1100) .expect("evaluate should succeed"); assert_eq!(result.series.len(), 1, "one series (no grouping)"); @@ -433,3 +440,307 @@ fn multi_series_one_per_label_value() { .expect("evaluate should succeed"); assert_eq!(result.series.len(), 2); } + +// --------------------------------------------------------------------------- +// TODO-1 tests — CMS-with-heap top-k. +// --------------------------------------------------------------------------- + +use asap_sketchlib::sketches::countminsketch_topk::CountMinSketchWithHeap; + +fn cms_heap_meta(sid: u64) -> SketchInstanceMetadata { + let cfg = SketchConfig::CountMin { + rows: 4, + cols: 256, + }; + SketchInstanceMetadata { + sid, + metric_name: "endpoint_hits".to_string(), + group_by_keys: BTreeSet::new(), + capability: Capability::FrequencyTopk(SketchKindHandle::CmsWithHeap), + sketch_kind: SketchKindHandle::CmsWithHeap, + sketch_config: cfg.clone(), + accuracy: AccuracyBound::from_config(&cfg), + first_seen_unix_ms: 0, + } +} + +fn cms_only_meta(sid: u64) -> SketchInstanceMetadata { + let cfg = SketchConfig::CountMin { + rows: 4, + cols: 256, + }; + SketchInstanceMetadata { + sid, + metric_name: "endpoint_hits".to_string(), + group_by_keys: BTreeSet::new(), + capability: Capability::FrequencyTopk(SketchKindHandle::CountMin), + sketch_kind: SketchKindHandle::CountMin, + sketch_config: cfg.clone(), + accuracy: AccuracyBound::from_config(&cfg), + first_seen_unix_ms: 0, + } +} + +fn msgpack_full(bytes: Vec) -> SketchSampleState { + SketchSampleState { + bytes, + encoding: SketchEncoding::MsgpackFull, + } +} + +#[test] +fn cms_with_heap_topk_returns_top_items() { + let idx = SketchIndex::new(); + let sid = 100; + idx.register(cms_heap_meta(sid)); + + // Build a CMS-with-heap state with known item counts. + let mut cms = CountMinSketchWithHeap::new(4, 256, 20); + // Insert items with varying frequencies. Higher count items + // should end up in the heap. + let inserts: &[(&str, u64)] = &[ + ("alpha", 100), + ("beta", 50), + ("gamma", 200), + ("delta", 75), + ("epsilon", 10), + ("zeta", 150), + ]; + for (k, n) in inserts { + for _ in 0..*n { + cms.update(k, 1.0); + } + } + let bytes = cms.serialize_msgpack().expect("serialize cms with heap"); + idx.append_sample(sid, BTreeMap::new(), (1000, 1010), msgpack_full(bytes)); + + let reducer = SketchReducer::new(&idx); + let result = reducer + .evaluate(&[sid], "topk", &[5.0], 1000, 1010) + .expect("topk evaluate should succeed"); + // We requested top-5. Each top-k item is its own series row + // (label_values carries the encoded `"item": `). + assert!( + result.series.len() <= 5 && !result.series.is_empty(), + "expected up to 5 top-k series, got {}", + result.series.len() + ); + // Coverage should match the window we appended. + assert_eq!(result.coverage, Some((1010, 1010))); + + // Top-1 should be "gamma" (count=200). Sort our series by + // first-sample value descending and check the top item. + let mut sorted = result.series.clone(); + sorted.sort_by(|a, b| { + let va = a.1.first().map(|s| s.1).unwrap_or(0.0); + let vb = b.1.first().map(|s| s.1).unwrap_or(0.0); + vb.partial_cmp(&va).unwrap_or(std::cmp::Ordering::Equal) + }); + let top = sorted.first().expect("at least one series"); + let item_label = top.0.get("item").expect("series carries item label"); + assert_eq!(item_label, "gamma", "highest-count item should be `gamma`"); +} + +#[test] +fn cms_without_heap_returns_missing_heap() { + let idx = SketchIndex::new(); + let sid = 101; + idx.register(cms_only_meta(sid)); + + // Append a CMS-with-heap-encoded payload — but the metadata is + // CountMin-only so the reducer should refuse on the + // sketch-kind side before decoding bytes. + let mut cms = CountMinSketchWithHeap::new(4, 256, 20); + cms.update("foo", 1.0); + let bytes = cms.serialize_msgpack().expect("serialize"); + idx.append_sample(sid, BTreeMap::new(), (1000, 1010), msgpack_full(bytes)); + + let reducer = SketchReducer::new(&idx); + let err = reducer + .evaluate(&[sid], "topk", &[5.0], 1000, 1010) + .expect_err("topk against CountMin (no heap) must surface MissingHeap"); + match err { + WarmTierError::MissingHeap { sid: s, sketch_kind } => { + assert_eq!(s, sid); + assert_eq!(sketch_kind, SketchKindHandle::CountMin); + } + other => panic!("expected MissingHeap, got {other:?}"), + } +} + +// --------------------------------------------------------------------------- +// TODO-2 tests — delta encoding stitching. +// +// We exercise the cumulative path for DDSketch (one Full window + two +// Delta windows of additional samples). The cumulative result should +// match what a fresh DDSketch fed all raw values would yield. +// --------------------------------------------------------------------------- + +fn proto_delta(bytes: Vec) -> SketchSampleState { + SketchSampleState { + bytes, + encoding: SketchEncoding::ProtoDelta, + } +} + +#[test] +fn ddsketch_cumulative_full_plus_two_deltas() { + let idx = SketchIndex::new(); + let sid = 200; + idx.register(dd_meta(sid)); + + let alpha = 0.01; + // Window 1: Full snapshot of values 1..=5 + let mut sk1 = DdSketch::new(alpha); + for v in 1..=5 { + sk1.update(v as f64); + } + let bytes1 = encode_ddsketch(&sk1); + idx.append_sample(sid, BTreeMap::new(), (1000, 1010), proto_full(bytes1)); + + // Windows 2 & 3: "Deltas" encoded as full-fragment sketches that + // get merged into the rolling state (the reducer's delta_apply + // treats DD/KLL/HLL delta-as-mergeable-fragment). + let mut sk2 = DdSketch::new(alpha); + for v in 6..=10 { + sk2.update(v as f64); + } + let bytes2 = encode_ddsketch(&sk2); + idx.append_sample(sid, BTreeMap::new(), (1010, 1020), proto_delta(bytes2)); + + let mut sk3 = DdSketch::new(alpha); + for v in 11..=15 { + sk3.update(v as f64); + } + let bytes3 = encode_ddsketch(&sk3); + idx.append_sample(sid, BTreeMap::new(), (1020, 1030), proto_delta(bytes3)); + + let reducer = SketchReducer::new(&idx); + let result = reducer + .evaluate(&[sid], "quantile_over_time", &[0.5], 1000, 1030) + .expect("cumulative evaluate should succeed"); + // Cumulative mode emits one scalar covering the full range. + assert_eq!(result.series.len(), 1); + let (_, samples) = &result.series[0]; + assert_eq!(samples.len(), 1, "cumulative emits exactly one scalar"); + let est = samples[0].1; + + // Truth: feed all 15 values into a fresh DDSketch and read the + // median (8th value of 1..=15 = 8). Allow 5% relative error to + // give the bucket store some slack. + let mut truth = DdSketch::new(alpha); + for v in 1..=15 { + truth.update(v as f64); + } + let true_q = truth.quantile(0.5).unwrap_or(0.0); + let rel_err = (est - true_q).abs() / true_q.max(1e-9); + assert!( + rel_err < 0.10, + "cumulative quantile error too large: est={} truth={} rel_err={}", + est, + true_q, + rel_err + ); + + // Coverage should span the three window ends. + assert_eq!(result.coverage, Some((1010, 1030))); +} + +#[test] +fn hll_cumulative_full_plus_one_delta() { + let idx = SketchIndex::new(); + let sid = 201; + let precision: u32 = 10; + idx.register(hll_meta(sid, precision)); + + // Window 1: Full snapshot with 500 distinct items. + let mut sk1 = HllSketch::new(HllVariant::Regular, precision); + for i in 0..500 { + sk1.update(format!("user-{i}").as_bytes()); + } + let bytes1 = encode_hll(&sk1); + idx.append_sample(sid, BTreeMap::new(), (1000, 1010), proto_full(bytes1)); + + // Window 2: Msgpack-delta — the warm-tier reducer treats + // MsgpackDelta for HLL as a serialized HllSketch fragment that's + // mergeable via `HllSketch::merge`. We mock that here by + // serializing a second HLL with 500 additional distinct items. + let mut sk2 = HllSketch::new(HllVariant::Regular, precision); + for i in 500..1000 { + sk2.update(format!("user-{i}").as_bytes()); + } + let bytes2 = sk2.serialize_msgpack().expect("serialize HLL msgpack"); + let delta_sample = SketchSampleState { + bytes: bytes2, + encoding: SketchEncoding::MsgpackDelta, + }; + idx.append_sample(sid, BTreeMap::new(), (1010, 1020), delta_sample); + + let reducer = SketchReducer::new(&idx); + let result = reducer + .evaluate( + &[sid], + "count_distinct_over_time", + &[], + 1000, + 1020, + ) + .expect("cumulative HLL evaluate should succeed"); + assert_eq!(result.series.len(), 1); + let (_, samples) = &result.series[0]; + assert_eq!(samples.len(), 1, "cumulative emits one scalar"); + let est = samples[0].1; + // Truth: 1000 distinct items, allow 5σ envelope. + let std_err = 1.04 / ((1u64 << precision) as f64).sqrt(); + let envelope = 5.0 * std_err * 1000.0; + let abs_err = (est - 1000.0).abs(); + assert!( + abs_err <= envelope, + "cumulative HLL estimate {} too far from true 1000 (5σ envelope = {})", + est, + envelope + ); +} + +// --------------------------------------------------------------------------- +// TODO-3 tests — hybrid warm + archive stitch via `WarmTierResult.coverage`. +// +// We don't drive the full SimpleEngine here (that would require +// constructing the whole streaming-config plumbing). Instead we exercise +// the `stitch_warm_and_archive` helper directly via a small wrapper +// test in `engines::simple::tests` would be ideal — but to keep this +// PR additive, we verify the `coverage` field is populated correctly +// on a multi-window evaluate so the downstream stitch path has the +// information it needs. +// --------------------------------------------------------------------------- + +#[test] +fn coverage_reports_observed_window_range() { + let idx = SketchIndex::new(); + let sid = 300; + idx.register(dd_meta(sid)); + + let alpha = 0.01; + for (i, values) in [vec![1.0_f64, 2.0], vec![3.0, 4.0], vec![5.0, 6.0]] + .iter() + .enumerate() + { + let mut sk = DdSketch::new(alpha); + for &v in values { + sk.update(v); + } + let bytes = encode_ddsketch(&sk); + let window_start = 100 + (i as u64) * 100; + let window_end = window_start + 100; + idx.append_sample(sid, BTreeMap::new(), (window_start, window_end), proto_full(bytes)); + } + + let reducer = SketchReducer::new(&idx); + let result = reducer + .evaluate(&[sid], "quantile", &[0.5], 50, 400) + .expect("evaluate should succeed"); + // Coverage min = first window end (200), max = third window end (400). + let coverage = result.coverage.expect("coverage populated"); + assert_eq!(coverage.0, 200); + assert_eq!(coverage.1, 400); +} diff --git a/asap-query-engine/src/stores/sketch_db/sketch_index.rs b/asap-query-engine/src/stores/sketch_db/sketch_index.rs index 6d5eac1a..8e1bd16d 100644 --- a/asap-query-engine/src/stores/sketch_db/sketch_index.rs +++ b/asap-query-engine/src/stores/sketch_db/sketch_index.rs @@ -55,6 +55,14 @@ pub enum SketchKindHandle { Hll, CountSketch, CountMin, + /// CMS-with-heap. Detected at the application level via the + /// precompute_operators `count_min_sketch_with_heap_accumulator` + /// flow — the OTLP `CountMinSketch` wire struct doesn't carry the + /// heap natively, so the gateway/precompute layer marks the + /// sid with this variant when the parent container's heap field + /// is non-empty. The warm-tier reducer reads the heap directly + /// when answering `topk` / `topk_over_time`. + CmsWithHeap, } /// Sketch-instance configuration carried per-Metric on the OTLP wire