From 545174cfc64d74734413ff3d7b8fed52d96537b7 Mon Sep 17 00:00:00 2001 From: zz_y Date: Sun, 10 May 2026 19:16:25 -0600 Subject: [PATCH] =?UTF-8?q?feat:=20unify=20extract=5Fpromql=5Fcall=20with?= =?UTF-8?q?=20controller=20=E2=80=94=20single=20PromQL=20=E2=86=92=20Capab?= =?UTF-8?q?ility=20owner?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Today the warm-tier reducer's `extract_promql_call` did a naive AST walk to extract `(function_name, args)` from a PromQL query. This duplicated knowledge the controller already encodes (PromQL → Intent → Capability via `query_parser`, `intent_algebra`, `sketch_algebra`, `algebra::lower`) and silently routed the MVP demo's compound queries (e.g. `sum by (zone) (rate(http_requests_total[5m]))`, `histogram_quantile(0.99, sum(rate(bucket[5m])) by (le))`) through the archive engine because the outermost-call heuristic returned `None` on nested shapes. The controller is now the single owner of "is this PromQL warm-tier-answerable" knowledge. ## What landed ### `controller/src/warm_tier_analysis.rs` (775 lines) Public API: ```rust pub fn analyze_promql_for_warm_tier(promql: &str) -> WarmTierAnalysis; pub struct WarmTierAnalysis { pub candidates: Vec, pub unsupported: Option, } pub struct WarmTierCandidate { pub metric_name: String, pub group_by_keys: BTreeSet, pub required_capability: Capability, pub function: String, pub function_args: Vec, pub range_seconds: u64, } pub enum UnsupportedReason { UnsupportedFunction(String), UnsupportedComposition(String), NoCallNodeFound, UnparseablePromql(String), } ``` Walks the `promql_parser` AST and identifies sub-expressions that can be served from sketches. Mapping table: | PromQL shape | Capability | |---|---| | `quantile_over_time(q, m[r])` | `QuantileApprox(Any)` | | `histogram_quantile(q, m)` | `QuantileApprox(Any)` | | `count_distinct_over_time(m[r])` | `CardinalityApprox` | | `cardinality_estimate(m)` | `CardinalityApprox` | | `topk(k, m)` | `FrequencyTopk(CmsWithHeap)` | | `topk_over_time(k, m[r])` | `FrequencyTopk(CmsWithHeap)` | Explicitly rejected (now surface as the right `UnsupportedReason`, not silent misroute): - `sum by (label_set) (rate(metric[range]))` → `UnsupportedFunction("rate")` - `histogram_quantile(q, sum(rate(bucket[r])) by (le))` → same - `sum by (label_set) (metric)` → `UnsupportedComposition` (Sum-over-CountSketch is a follow-up) - `increase` / `irate` → `UnsupportedFunction` - `topk(k, rate(metric[r]))` → `UnsupportedComposition` ### `controller/src/lib.rs` `pub mod warm_tier_analysis;` — exposes the new module. ### `asap-query-engine/src/stores/sketch_db/sketch_index.rs` (+85) `From` and `From` adapters at the boundary. New `Capability::is_satisfied_by` helper handles the `SketchKindHandle::Any` wildcard ("any sketch impl in the family is acceptable" — e.g. `QuantileApprox(Any)` is satisfied by both DDSketch and KLL instances). ### `asap-query-engine/src/engines/warm_tier/mod.rs` Doc comments rewritten to reference the new controller analyzer. `pub use` re-exports for `promql_extract::*` removed. ### `asap-query-engine/src/engines/warm_tier/promql_extract.rs` **Deleted** (159 lines). ### `asap-query-engine/src/engines/simple/engine.rs::execute` Warm-tier hook rewritten: 1. `controller::warm_tier_analysis::analyze_promql_for_warm_tier(query)` 2. If `analysis.unsupported.is_some()` or `analysis.candidates.is_empty()` → `EngineError::CapabilityMiss`. 3. For each candidate: `index.instances_matching(metric, group_by)`, verify `Capability::is_satisfied_by`, classify, dispatch reducer per Capability. Cold-tier fallthrough is now explicit: - `UnsupportedFunction` / `UnsupportedComposition` / `UnparseablePromql` → archive (router CapabilityMiss failover) - `NoCallNodeFound` (bare selector) → archive - Candidates populated but `instances_matching` empty → archive - All candidates resolve to Hit → warm-tier ## Build + test - `cargo build --release -p query_engine_rust` clean (only pre-existing warnings) - `cargo build --release -p controller` clean - `cargo test --release -p query_engine_rust --lib -- engines::warm_tier` — 13/13 pass ## Diff ``` asap-query-engine/src/engines/simple/engine.rs | 98 ++++++-- asap-query-engine/src/engines/warm_tier/mod.rs | 20 ++ asap-query-engine/src/engines/warm_tier/promql_extract.rs | 159 --- (deleted) asap-query-engine/src/stores/sketch_db/sketch_index.rs | 85 ++++ controller/src/lib.rs | 5 + controller/src/warm_tier_analysis.rs | 775 +++++ (new) 6 files changed, 983 insertions(+), 166 deletions(-) ``` ## Follow-ups (out of scope) - Per-candidate hybrid stitch (PR #126 stitches at engine level; per-candidate is a follow-up). - `Sum-over-CountSketch` reducer for the `sum by (zone) (metric)` shape. - Wire the controller analyzer's `range_seconds` into the reducer's `query_range(t0, t1)` bounds — today the dispatch uses `[now - 5min, now]`. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/engines/simple/engine.rs | 338 ++++--- .../src/engines/warm_tier/mod.rs | 20 +- .../src/engines/warm_tier/promql_extract.rs | 159 ---- .../src/stores/sketch_db/sketch_index.rs | 85 ++ controller/src/lib.rs | 5 + controller/src/warm_tier_analysis.rs | 868 ++++++++++++++++++ 6 files changed, 1166 insertions(+), 309 deletions(-) delete mode 100644 asap-query-engine/src/engines/warm_tier/promql_extract.rs create mode 100644 controller/src/warm_tier_analysis.rs diff --git a/asap-query-engine/src/engines/simple/engine.rs b/asap-query-engine/src/engines/simple/engine.rs index 7735f837..28c54b29 100644 --- a/asap-query-engine/src/engines/simple/engine.rs +++ b/asap-query-engine/src/engines/simple/engine.rs @@ -92,6 +92,15 @@ fn replace_metric_token(haystack: &str, needle: &str, replacement: &str) -> Stri /// is tolerable — the candidates are subsequently classified, and on /// `Ghost` / `Unknown` outcomes the query falls through to the archive /// engine via the EngineRouter's `CapabilityMiss` failover. +/// Legacy `(metric_name, group_by_keys)` extractor — superseded by +/// `controller::warm_tier_analysis::analyze_promql_for_warm_tier`, +/// which returns the full `WarmTierAnalysis` (capability, function +/// name + args, range). Kept around as `#[allow(dead_code)]` because +/// downstream code (range-query pipeline, range-step planner) still +/// uses bare `(metric, keys)` projections for sid candidate filtering; +/// once those callers also migrate to `WarmTierAnalysis`, this can be +/// deleted in a follow-up. +#[allow(dead_code)] fn extract_metric_and_label_keys( query: &str, ) -> Option<(String, std::collections::BTreeSet)> { @@ -3617,146 +3626,168 @@ impl crate::routing::engine_router::QueryEngine for SimpleEngine { &self, query: &str, ) -> Result { - // Phase 5 wire-in (refactor 2026-05) — classify against the - // sketch-warm-tier index BEFORE handing the query to - // `handle_query`. The classification is intentionally crude - // because the warm-tier sketch reducer is still a stub: as soon - // as ANY sid is `Ghost` / `Unknown`, OR no instance even matches - // the metric / group-by KEY set, we surface - // `EngineError::CapabilityMiss(SketchWarmTier, ...)` so the - // EngineRouter (Phase 6) fails over to the archive engine. + // Phase 9 controller-unification (2026-05) — the warm-tier + // hook is now a thin driver around the controller's + // `analyze_promql_for_warm_tier`. The analyzer is the single + // owner of "is this PromQL warm-tier-answerable" knowledge. + // We drop into one of three branches: // - // When `instances_matching` returns sids that all classify as - // `Hit`, we still fall through to `handle_query`'s legacy code - // path — wiring per-Capability sketch reducers on top of - // `SketchIndex.query_range` is out of scope for this PR and - // tracked as a follow-up. The "hybrid stitch" covering - // `[t0..t1']` from warm + `[t1'..t1]` from archive is also - // deferred (`QueryResult` would need timestamp coverage - // metadata to express it). + // 1. `WarmTierAnalysis::unsupported` is `Some(_)` — the + // PromQL shape isn't warm-tier-servable. Surface as + // `EngineError::CapabilityMiss(SketchWarmTier, …)` with the + // structured `UnsupportedReason` in the detail string. The + // EngineRouter fails over to the archive engine. This + // covers all of: + // * `MissReason::UnsupportedFunction(_)` (rate, irate, + // increase, etc.) → cold tier (archive) + // * `MissReason::UnsupportedComposition(_)` (sum-by, + // topk-over-rate, etc.) → cold tier + // * `MissReason::NoCallNodeFound` (bare selector) → + // cold tier (archive answers raw selectors) + // * `MissReason::UnparseablePromql(_)` → cold tier + // (archive's parser may be more permissive, or it'll + // also reject and the user sees the error) + // + // 2. `WarmTierAnalysis::candidates` is populated, but ANY + // candidate's `instances_matching` returns empty OR a + // sid that classifies as `Ghost`/`Unknown` — surface + // as CapabilityMiss. The EngineRouter falls over. + // + // 3. All candidates resolve to all-`Hit` sids — dispatch + // each to the per-`Capability` sketch reducer. Today's + // semantic: ANY candidate-level reducer error → fall + // over to archive (no per-candidate hybrid stitch yet — + // that's the documented follow-up). if let Some(idx) = self.sketch_index.as_ref() { - if let Some((metric_name, required_keys)) = - extract_metric_and_label_keys(query) - { - let candidates = idx.instances_matching(&metric_name, &required_keys); - if candidates.is_empty() { + let analysis = controller::warm_tier_analysis::analyze_promql_for_warm_tier(query); + + // Branch 1 — the controller analyzer rejects the shape. + if let Some(reason) = &analysis.unsupported { + return Err(crate::engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchWarmTier.data_source_id(), + format!( + "SketchWarmTier analyzer rejected `{query}`: {reason:?} — \ + failing over to archive" + ), + )); + } + if analysis.candidates.is_empty() { + // Defensive — `is_warm_tier_answerable` would have + // caught this; analyzer guarantees `unsupported.is_some()` + // when `candidates.is_empty()` but we keep the + // belt-and-braces miss-path for safety. + return Err(crate::engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchWarmTier.data_source_id(), + format!( + "SketchWarmTier analyzer produced no warm-tier candidates for \ + `{query}` — failing over to archive" + ), + )); + } + + // Branch 2 + 3 — resolve each candidate's sids and + // dispatch the reducer. Today this is single-candidate + // for every supported PromQL shape; the loop is here + // for the per-candidate hybrid-stitch follow-up. + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::SystemTime::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0); + // Time bounds: the trait's `execute(&str)` adapter + // doesn't carry an explicit range today (it's an + // instant-query surface). For each candidate, prefer the + // candidate's `range_seconds` (extracted from `[5m]` / + // `[30s]` selectors); fall back to a 5-minute default + // for instant-vector candidates (range_seconds == 0). + const DEFAULT_LOOKBACK_MS: u64 = 5 * 60 * 1000; + + let reducer = crate::engines::warm_tier::SketchReducer::new(idx); + // Multi-candidate aggregation is deferred (single-result + // shapes today). On the first reducer error we surface + // CapabilityMiss; on Ok we keep the result for the + // hybrid-stitch path below. (When more than one + // candidate is supported, a follow-up will fold + // per-candidate WarmTierResults.) + let mut combined_result: Option = None; + let mut combined_t0: u64 = u64::MAX; + + for candidate in &analysis.candidates { + let sids = idx.instances_matching( + &candidate.metric_name, + &candidate.group_by_keys, + ); + if sids.is_empty() { return Err(crate::engines::EngineError::capability_miss( asap_types::StorageBackend::SketchWarmTier.data_source_id(), format!( - "SketchWarmTier has no instance for metric `{metric_name}` \ - with group_by_keys ⊇ {:?}", - required_keys + "SketchWarmTier has no instance for metric `{}` \ + with group_by_keys ⊇ {:?} (analyzer required \ + {:?}) — failing over to archive", + candidate.metric_name, + candidate.group_by_keys, + candidate.required_capability, ), )); } - let mut all_hit = true; - for sid in &candidates { + + // Verify each sid carries the analyzer's required + // capability (the controller's + // `Capability::is_satisfied_by` honors `Any` semantics). + let required: crate::stores::sketch_db::sketch_index::Capability = + candidate.required_capability.clone().into(); + let mut hit_sids: Vec = Vec::with_capacity(sids.len()); + for sid in &sids { match idx.classify(*sid) { crate::stores::sketch_db::sketch_index::SidLookup::Hit => {} crate::stores::sketch_db::sketch_index::SidLookup::Ghost | crate::stores::sketch_db::sketch_index::SidLookup::Unknown => { - all_hit = false; - break; + return Err(crate::engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchWarmTier.data_source_id(), + format!( + "SketchWarmTier ghost/unknown sid {sid} for metric \ + `{}` — failing over to archive", + candidate.metric_name + ), + )); } } + let meta = match idx.instance(*sid) { + Some(m) => m, + None => continue, + }; + if required.is_satisfied_by(&meta.capability) { + hit_sids.push(*sid); + } } - if !all_hit { + if hit_sids.is_empty() { return Err(crate::engines::EngineError::capability_miss( asap_types::StorageBackend::SketchWarmTier.data_source_id(), format!( - "SketchWarmTier ghost/unknown sid for metric `{metric_name}` \ - — failing over to archive" + "SketchWarmTier has no sid satisfying capability \ + {:?} for metric `{}` — failing over to archive", + candidate.required_capability, candidate.metric_name ), )); } - // All sids `Hit` → dispatch to the per-Capability - // sketch reducer (`feat/sketch-reducer-warm-tier-evaluator`, - // 2026-05). The reducer decodes each window's sketch - // state via `asap_sketchlib`, evaluates the - // canonical query (`quantile`, `estimate`, …), and - // returns per-series timestamped scalars that we - // adapt to `QueryResult::Matrix`. - // - // On any failure mode the reducer surfaces, we - // translate to `CapabilityMiss` so the - // `EngineRouter` falls over to archive — including - // `UnsupportedFunction`/`UnsupportedCapability` - // (the user's PromQL doesn't map onto a warm-tier - // capability), `DeserializeFailure` (defensive — - // the warm-tier state didn't decode; archive can - // answer truthfully), and `NoData` (no samples in - // window). - let call = crate::engines::warm_tier::extract_promql_call(query); - let (function_name, function_args) = match &call { - Some(c) if !c.func.is_empty() => (c.func.clone(), c.args.clone()), - _ => { - return Err(crate::engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchWarmTier.data_source_id(), - format!( - "SketchWarmTier reducer cannot extract a PromQL function \ - from `{query}` — failing over to archive" - ), - )); - } + let lookback_ms = if candidate.range_seconds > 0 { + candidate.range_seconds.saturating_mul(1000) + } else { + DEFAULT_LOOKBACK_MS }; - - // Time bounds: the trait's `execute(&str)` adapter - // doesn't carry an explicit range today (it's an - // instant-query surface). Use `[now - default_range, - // now]` matching how `handle_query` used to pick - // its window. Phase-5 hybrid stitch (warm + archive - // for ranges that exceed warm coverage) is a - // follow-up. - let now_ms = std::time::SystemTime::now() - .duration_since(std::time::SystemTime::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0); - // Default lookback: 5 minutes. Engine code that - // wants a precise range (range-query pipeline) - // should call into `SketchReducer::evaluate` - // directly with its own bounds. - let lookback_ms: u64 = 5 * 60 * 1000; let t0_ms = now_ms.saturating_sub(lookback_ms); + if t0_ms < combined_t0 { + combined_t0 = t0_ms; + } - let reducer = crate::engines::warm_tier::SketchReducer::new(idx); - match reducer.evaluate( - &candidates, - &function_name, - &function_args, + let result = match reducer.evaluate( + &hit_sids, + &candidate.function, + &candidate.function_args, t0_ms, now_ms, ) { - Ok(result) => { - // 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); - } + Ok(r) => r, Err(crate::engines::warm_tier::WarmTierError::UnsupportedFunction( name, )) => { @@ -3809,11 +3840,6 @@ impl crate::routing::engine_router::QueryEngine for SimpleEngine { 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!( @@ -3823,7 +3849,37 @@ impl crate::routing::engine_router::QueryEngine for SimpleEngine { ), )); } + }; + combined_result = Some(result); + } + + // All candidates resolved successfully — adapt to + // QueryResult and run the hybrid-stitch path if archive + // is wired and warm coverage is narrower than request. + if let Some(result) = combined_result { + 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()) + { + let stitch_t0 = if combined_t0 == u64::MAX { + now_ms.saturating_sub(DEFAULT_LOOKBACK_MS) + } else { + combined_t0 + }; + if cov_lo > stitch_t0 || 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 warm-only. + } } + return Ok(warm_qr); } } @@ -6167,13 +6223,17 @@ mod warm_tier_classify_tests { #[tokio::test] async fn execute_returns_capability_miss_when_classify_is_ghost() { // Register instance metadata but never call append_sample → the - // sid classifies as Ghost. Adapter must short-circuit to - // CapabilityMiss so the EngineRouter (Phase 6) fails over. + // sid classifies as Ghost. The Phase-9 controller-unified + // adapter expects a call-shaped query (the analyzer rejects + // bare selectors with NoCallNodeFound BEFORE any sid lookup); + // use `quantile_over_time(...)` so the analyzer accepts the + // shape and the per-candidate sid classification surfaces + // the ghost miss. let idx = Arc::new(SketchIndex::new()); idx.register(dd_meta(1, "http_latency_ms", &["zone"])); let engine = build_engine_with_index(idx); let err = engine - .execute("http_latency_ms{zone=\"z0\"}") + .execute("quantile_over_time(0.99, http_latency_ms{zone=\"z0\"}[5m])") .await .expect_err("ghost classification must yield CapabilityMiss"); match err { @@ -6192,21 +6252,17 @@ mod warm_tier_classify_tests { } #[tokio::test] - async fn execute_proceeds_to_handle_query_when_all_sids_hit() { - // Register an instance AND append a sample so the sid Hits. - // - // Pre-`feat/sketch-reducer-warm-tier-evaluator` (this PR): the - // adapter fell through to `handle_query`, producing the legacy - // "no compatible aggregation" miss for an empty SimpleMapStore. - // - // Post-PR: the adapter dispatches to the warm-tier - // `SketchReducer`. A bare vector selector (no PromQL call) - // surfaces as the warm-tier-specific "cannot extract a PromQL - // function" miss (the reducer can only answer call-shaped - // queries; raw selectors fall over to archive). The contract - // verified here is still "Hit does NOT short-circuit to the - // warm-tier ghost/unknown miss"; only the downstream miss - // detail changes. + async fn execute_rejects_bare_selector_via_analyzer() { + // Phase-9 controller-unified behavior: a bare vector selector + // (no call node) is rejected by + // `controller::warm_tier_analysis::analyze_promql_for_warm_tier` + // with `UnsupportedReason::NoCallNodeFound` BEFORE the sid + // index is even consulted. The archive engine answers raw + // selectors directly, so this is the right place for the + // routing decision. Replaces the legacy + // `execute_proceeds_to_handle_query_when_all_sids_hit` test — + // the new behavior is "bare selectors short-circuit on shape + // rejection, regardless of index state". let idx = Arc::new(SketchIndex::new()); idx.register(dd_meta(2, "http_latency_ms", &["zone"])); idx.append_sample( @@ -6224,17 +6280,13 @@ mod warm_tier_classify_tests { match result { Err(EngineError::CapabilityMiss { detail, .. }) => { assert!( - !detail.contains("ghost") && !detail.contains("Ghost"), - "Hit path must NOT short-circuit to the ghost-miss path: {detail}" - ); - assert!( - detail.contains("cannot extract a PromQL function") - || detail.contains("no compatible aggregation"), - "Hit path produced expected post-Hit miss: {detail}" + detail.contains("NoCallNodeFound") + || detail.contains("analyzer rejected"), + "expected NoCallNodeFound analyzer rejection: {detail}" ); } other => panic!( - "expected post-Hit CapabilityMiss, got {other:?}" + "expected analyzer-rejected CapabilityMiss, got {other:?}" ), } } diff --git a/asap-query-engine/src/engines/warm_tier/mod.rs b/asap-query-engine/src/engines/warm_tier/mod.rs index 858590f8..bc8e7d08 100644 --- a/asap-query-engine/src/engines/warm_tier/mod.rs +++ b/asap-query-engine/src/engines/warm_tier/mod.rs @@ -35,11 +35,19 @@ //! "no data in window" (router falls over). //! * [`WarmTierResult`] — per-series timestamped scalar samples //! matching the shape of [`crate::engines::query_result::QueryResult::Matrix`]. -//! * [`extract_promql_call`] — small AST walker that pulls the -//! outermost call's function name + numeric args. Lives here -//! rather than in `simple/engine.rs` because the existing -//! `extract_metric_and_label_keys` already handles the -//! metric-and-keys side; this is the function-name + args side. +//! +//! ## Controller unification (PromQL-shape recognition) +//! +//! The PromQL → `(function_name, args)` AST walker that used to live +//! here in `promql_extract.rs` has been folded into +//! [`controller::warm_tier_analysis::analyze_promql_for_warm_tier`]. +//! That function is the single owner of "is this PromQL +//! warm-tier-answerable" knowledge — it returns a +//! [`controller::warm_tier_analysis::WarmTierAnalysis`] enumerating +//! the warm-tier-servable sub-expressions and the explicit +//! [`controller::warm_tier_analysis::UnsupportedReason`] for the rest. +//! The reducer keys off the analyzer's `required_capability` rather +//! than re-string-matching the PromQL function name. //! //! Phase-5 hybrid stitching (warm `[t0..t1']` + archive //! `[t1'..t1]`) and per-window iteration (rather than today's @@ -59,11 +67,9 @@ pub mod decoders; pub mod delta_apply; -pub mod promql_extract; pub mod sketch_reducer; #[cfg(test)] pub mod tests; -pub use promql_extract::{extract_promql_call, PromqlCall}; pub use sketch_reducer::{SketchReducer, WarmTierError, WarmTierResult}; diff --git a/asap-query-engine/src/engines/warm_tier/promql_extract.rs b/asap-query-engine/src/engines/warm_tier/promql_extract.rs deleted file mode 100644 index b8f6da96..00000000 --- a/asap-query-engine/src/engines/warm_tier/promql_extract.rs +++ /dev/null @@ -1,159 +0,0 @@ -//! PromQL → (function_name, scalar_args) extraction for the warm-tier -//! sketch reducer. -//! -//! Companion to [`crate::engines::simple::engine::extract_metric_and_label_keys`] -//! (which extracts `(metric_name, label_keys)` for warm-tier candidate -//! selection). This module pulls the outer-most function call — -//! identifying name (`quantile_over_time`, `histogram_quantile`, -//! `topk`, `count_distinct_over_time`, …) and any leading numeric -//! arguments (the quantile rank `q`, the `k` for top-k, …). -//! -//! Intentionally tiny: only handles the call shapes the warm-tier -//! reducer can answer today, and rejects anything more complex -//! (binary ops, aggregates over ranges, math on sketch outputs) -//! by returning `None` so the engine surfaces an -//! `UnsupportedFunction` and falls over to the archive engine. -//! That's the right behavior — the warm tier is a fast path; richer -//! query shapes must go through `handle_query` or archive. -//! -//! # Supported shapes -//! -//! - `quantile_over_time(q, foo[5m])` -//! - `histogram_quantile(q, foo)` -//! - `count_distinct_over_time(foo[5m])` -//! - `cardinality_estimate(foo)` (custom function name) -//! - `topk(k, foo)` (PromQL aggregation; `k` lifted from -//! `AggregateExpr::param`) -//! - `topk_over_time(k, foo[5m])` (custom function name) -//! - bare `foo` / `foo{matchers}` (no call → `func == "" `) -//! -//! Anything more nested (`rate(foo[5m]) > 0.5`, `sum by (a) (foo)`, -//! …) returns `None`; the dispatcher then surfaces -//! `WarmTierError::UnsupportedFunction` and the engine falls back to -//! archive. - -use promql_parser::parser::{self, Expr}; - -/// One extracted call site. -#[derive(Debug, Clone)] -pub struct PromqlCall { - /// Lower-case function name. Empty string for a bare vector - /// selector (no function call). - pub func: String, - /// Already-evaluated leading scalar args. Order matches the - /// PromQL surface (`quantile_over_time(q, foo[5m])` → - /// `args[0] = q`). - pub args: Vec, -} - -/// Walk the PromQL AST and return the outer-most call's name + leading -/// scalar args, or a bare-vector marker (`func.is_empty()`). -/// -/// Returns `None` if parsing fails or the query shape isn't one the -/// warm-tier reducer can answer. -pub fn extract_promql_call(query: &str) -> Option { - let ast = parser::parse(query).ok()?; - extract_from_expr(&ast) -} - -fn extract_from_expr(expr: &Expr) -> Option { - match expr { - // `quantile_over_time(q, foo[5m])`, - // `count_distinct_over_time(foo[5m])`, - // `histogram_quantile(q, …)`, etc. — pull the function name - // and leading numeric literal args. - Expr::Call(call) => { - let mut args = Vec::new(); - for a in &call.args.args { - match a.as_ref() { - Expr::NumberLiteral(nl) => args.push(nl.val), - // First non-scalar marks the end of the leading - // scalar args; the remainder is the vector / - // matrix selector. - _ => break, - } - } - Some(PromqlCall { - func: call.func.name.to_string(), - args, - }) - } - // `topk(5, foo)` / `bottomk(3, foo)` / `quantile(0.99, foo)` — - // PromQL aggregations carrying a single `param`. The - // aggregator op displays as its name (see - // `promql_parser::parser::token::token_display`). - Expr::Aggregate(agg) => { - let func = agg.op.to_string(); - let mut args = Vec::new(); - if let Some(p) = &agg.param { - if let Expr::NumberLiteral(nl) = p.as_ref() { - args.push(nl.val); - } - } - Some(PromqlCall { func, args }) - } - Expr::Paren(p) => extract_from_expr(&p.expr), - Expr::Subquery(sq) => extract_from_expr(&sq.expr), - // A bare vector / matrix selector — no call, so the warm-tier - // reducer treats it as "raw select"; the dispatcher will - // reject it as `UnsupportedFunction` (no scalar reduction - // implied, and the warm tier doesn't materialize raw counter - // values, only sketch-state-reduced scalars). - Expr::VectorSelector(_) | Expr::MatrixSelector(_) => Some(PromqlCall { - func: String::new(), - args: Vec::new(), - }), - // Binary ops, unary ops, extensions — out of scope for the - // warm-tier fast path. - _ => None, - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn extracts_quantile_over_time() { - let c = extract_promql_call("quantile_over_time(0.99, http_latency_ms[5m])").unwrap(); - assert_eq!(c.func, "quantile_over_time"); - assert_eq!(c.args, vec![0.99]); - } - - #[test] - fn extracts_histogram_quantile() { - let c = extract_promql_call("histogram_quantile(0.5, http_latency_ms)").unwrap(); - assert_eq!(c.func, "histogram_quantile"); - assert_eq!(c.args, vec![0.5]); - } - - #[test] - fn extracts_topk_aggregate() { - let c = extract_promql_call("topk(5, requests)").unwrap(); - assert_eq!(c.func, "topk"); - assert_eq!(c.args, vec![5.0]); - } - - #[test] - fn extracts_count_distinct_over_time() { - let c = extract_promql_call("count_distinct_over_time(uniq_users[1h])"); - // count_distinct_over_time isn't a known PromQL function in - // the parser's function table — falls into the catch-all - // `_ => None` branch via parser failure. Test we surface - // None so the dispatcher correctly maps to UnsupportedFunction. - assert!(c.is_none() || c.unwrap().func == "count_distinct_over_time"); - } - - #[test] - fn bare_vector_selector_returns_empty_func() { - let c = extract_promql_call("http_requests_total{zone=\"z0\"}").unwrap(); - assert!(c.func.is_empty()); - assert!(c.args.is_empty()); - } - - #[test] - fn rejects_binary_ops() { - let c = extract_promql_call("rate(foo[5m]) > 0.5"); - assert!(c.is_none()); - } -} 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 8e1bd16d..de622d74 100644 --- a/asap-query-engine/src/stores/sketch_db/sketch_index.rs +++ b/asap-query-engine/src/stores/sketch_db/sketch_index.rs @@ -65,6 +65,91 @@ pub enum SketchKindHandle { CmsWithHeap, } +// ── Controller ↔ backend Capability adapters ───────────────────────────────── +// +// `controller::warm_tier_analysis::Capability` is the canonical +// PromQL-shape-recognition output (the controller is the single owner +// of "is this PromQL warm-tier-answerable" knowledge — see +// `controller/src/warm_tier_analysis.rs`). The backend's `Capability` +// here is the source-of-truth for INDEXED sketch instances (driven by +// the OTLP receive path). The two enums are kept structurally +// identical and the backend adapts at the boundary so the controller +// crate stays free of any backend dep. +// +// `SketchKindHandle::Any` from the controller side maps to ALL +// concrete handles when used in capability matching — the backend's +// `is_compatible_with` helper consumes that semantics so callers don't +// need to enumerate the cross product. + +impl From for SketchKindHandle { + fn from(h: controller::warm_tier_analysis::SketchKindHandle) -> Self { + use controller::warm_tier_analysis::SketchKindHandle as C; + match h { + C::DDSketch => SketchKindHandle::DDSketch, + C::Kll => SketchKindHandle::Kll, + C::Hll => SketchKindHandle::Hll, + C::CountSketch => SketchKindHandle::CountSketch, + C::CountMin => SketchKindHandle::CountMin, + C::CmsWithHeap => SketchKindHandle::CmsWithHeap, + // `Any` has no single concrete handle. Callers that need + // to match against a specific instance should use + // `Capability::is_satisfied_by` instead of `From` for the + // handle directly. Defensive default: return DDSketch (the + // QuantileApprox catalog default) so a stray `Any` doesn't + // panic, though the canonical flow goes through + // `Capability::is_satisfied_by`. + C::Any => SketchKindHandle::DDSketch, + } + } +} + +impl Capability { + /// True when the indexed sketch instance's capability satisfies + /// the controller-side required capability. The controller emits + /// `SketchKindHandle::Any` to mean "any implementation in the + /// family is acceptable" (e.g. QuantileApprox(Any) is satisfied + /// by both DDSketch and KLL); concrete handles must match + /// exactly. + pub fn is_satisfied_by(&self, indexed: &Capability) -> bool { + use controller::warm_tier_analysis::SketchKindHandle as Any; + let _ = Any::Any; // silence unused-import warning when compiled standalone + match (self, indexed) { + (Capability::QuantileApprox(_), Capability::QuantileApprox(_)) => true, + (Capability::CardinalityApprox, Capability::CardinalityApprox) => true, + // FrequencyTopk(CmsWithHeap) is the only sub-variant the + // warm tier can answer top-k against (CountMin / + // CountSketch carry no heap). Match exactly on the inner + // handle so the reducer's MissingHeap path stays + // accessible. + (Capability::FrequencyTopk(req), Capability::FrequencyTopk(have)) + if req == have => + { + true + } + _ => false, + } + } +} + +/// Adapt a controller-side analyzed Capability into the backend +/// `Capability` enum. Used by the engine warm-tier hook to compare +/// the analyzer's required-capability against the index's recorded +/// per-instance capability. `SketchKindHandle::Any` from the +/// controller side is preserved as the catalog default for the +/// outer family (DDSketch for QuantileApprox); call sites that need +/// the "matches any concrete impl" semantic should call +/// [`Capability::is_satisfied_by`] instead. +impl From for Capability { + fn from(c: controller::warm_tier_analysis::Capability) -> Self { + use controller::warm_tier_analysis::Capability as C; + match c { + C::QuantileApprox(h) => Capability::QuantileApprox(h.into()), + C::CardinalityApprox => Capability::CardinalityApprox, + C::FrequencyTopk(h) => Capability::FrequencyTopk(h.into()), + } + } +} + /// Sketch-instance configuration carried per-Metric on the OTLP wire /// (Phase 2 lifted these from per-DP up to the parent sketch container). /// Backend reads the relevant variant at ingest time and stores it in diff --git a/controller/src/lib.rs b/controller/src/lib.rs index 9e072b0c..130a5280 100644 --- a/controller/src/lib.rs +++ b/controller/src/lib.rs @@ -46,3 +46,8 @@ pub mod stage_split; pub mod store; pub mod types; pub mod types_v2; +/// PromQL → warm-tier candidate analyzer. Phase-9 unification of the +/// per-`Capability` dispatch knowledge that previously lived in +/// `asap-query-engine/src/engines/warm_tier/promql_extract.rs`. See +/// the module docs for the full PromQL shape coverage matrix. +pub mod warm_tier_analysis; diff --git a/controller/src/warm_tier_analysis.rs b/controller/src/warm_tier_analysis.rs new file mode 100644 index 00000000..368c7097 --- /dev/null +++ b/controller/src/warm_tier_analysis.rs @@ -0,0 +1,868 @@ +//! PromQL → warm-tier candidate analyzer. +//! +//! Single owner of "is this PromQL warm-tier-answerable" knowledge, +//! pulled out of `asap-query-engine/src/engines/warm_tier/promql_extract.rs` +//! (deleted in the same change). The controller already encodes the +//! PromQL → Intent → Capability pipeline via `query_parser`, +//! `intent_algebra`, `sketch_algebra`, and `algebra::lower`; this module +//! is the **query-time** analog: rather than emitting a full plan, it +//! decides which sub-expressions of a PromQL query CAN be answered from +//! the warm tier and what [`Capability`] each requires. +//! +//! The output is consumed by the warm-tier reducer in +//! `asap-query-engine/src/engines/warm_tier/sketch_reducer.rs` and by +//! the `SimpleEngine::execute` warm-tier hook, replacing the previous +//! per-PromQL string-matched dispatch. +//! +//! # API +//! +//! - [`analyze_promql_for_warm_tier`] — pure function; parses + walks +//! the PromQL AST and returns either a populated [`WarmTierAnalysis`] +//! or an [`UnsupportedReason`]. +//! - [`WarmTierAnalysis`] — a vector of [`WarmTierCandidate`]s (the +//! sub-expressions the warm tier CAN serve) plus an +//! [`UnsupportedReason`] when the query has parts that cannot be +//! served (or cannot be parsed). +//! - [`Capability`] — warm-tier capability tag. Mirrors the +//! `sketch_index::Capability` enum in `asap-query-engine`; this is +//! the controller-side authority for the type. The +//! `asap-query-engine` side type-aliases / converts via small From +//! adapters at the call site. +//! +//! # Supported PromQL shapes (and the Capability each maps to) +//! +//! | Shape | Capability | +//! |---|---| +//! | `quantile_over_time(q, m[r])` | `QuantileApprox(Any)` | +//! | `quantile_over_time(q, m)` | `QuantileApprox(Any)` (instant — no range) | +//! | `histogram_quantile(q, m)` | `QuantileApprox(Any)` | +//! | `count_distinct_over_time(m[r])` | `CardinalityApprox` | +//! | `cardinality_estimate(m)` | `CardinalityApprox` | +//! | `topk(k, m)` | `FrequencyTopk(CmsWithHeap)` | +//! | `topk_over_time(k, m[r])` | `FrequencyTopk(CmsWithHeap)` | +//! | bare `m{filters}` | `UnsupportedReason::NoCallNodeFound` | +//! +//! # Explicitly rejected PromQL shapes +//! +//! The demo's compound queries that today silently route through the +//! archive engine are surfaced explicitly: +//! +//! - `sum by (label_set) (rate(metric[range]))` — `rate` is raw +//! counter math, not a sketch op. Surface as +//! `UnsupportedFunction("rate")`. +//! - `histogram_quantile(q, sum(rate(bucket[r])) by (le))` — same +//! reason: nested `rate`. +//! - `sum by (label_set) (metric)` — `Sum-over-CountSketch` reducer +//! is a future follow-up. Surface as `UnsupportedComposition(...)`. +//! - `increase` / `irate` — raw counter math. Surface as +//! `UnsupportedFunction(...)`. +//! - `topk(k, rate(metric[r]))` — `topk` is only meaningful over an +//! instant vector of items. Surface as `UnsupportedComposition(...)`. +//! +//! # Time range extraction +//! +//! Each candidate carries `range_seconds: u64` — the matrix-vector +//! selector's `[r]` parsed into seconds. A `0` value means "no matrix +//! selector" (instant-vector query). The reducer uses this when +//! deciding the per-window vs cumulative dispatch. + +use std::collections::BTreeSet; +use std::time::Duration; + +use promql_parser::parser::{self, AggregateExpr, Call, Expr, MatrixSelector, VectorSelector}; + +// ── Public types ───────────────────────────────────────────────────────────── + +/// Controller-side warm-tier capability tag. Mirrors the +/// `asap-query-engine`-side `sketch_index::Capability` enum so the +/// controller can emit capability requirements without depending on +/// the backend's sketch_index module. The two enums are kept +/// structurally identical and adapted via a small `From` impl at the +/// call site. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum Capability { + QuantileApprox(SketchKindHandle), + CardinalityApprox, + FrequencyTopk(SketchKindHandle), +} + +/// Compact handle for sketch family choice. Mirrors +/// `sketch_index::SketchKindHandle`. The `Any` variant is the +/// controller's "any implementation that satisfies the family works" +/// signal — e.g. for QuantileApprox the controller doesn't pick +/// DDSketch vs KLL at analysis time; the resolver picks whichever +/// instance the index already carries. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum SketchKindHandle { + DDSketch, + Kll, + Hll, + CountSketch, + CountMin, + CmsWithHeap, + /// "Any implementation that satisfies the family". Used at analysis + /// time when the capability is family-bound but not + /// implementation-bound. + Any, +} + +/// One sub-expression of the input PromQL that CAN be served from +/// the warm tier. The reducer resolves each candidate to a vector of +/// sids via `SketchIndex::instances_matching(metric_name, group_by_keys)` +/// and verifies each sid carries the required capability. +#[derive(Debug, Clone, PartialEq)] +pub struct WarmTierCandidate { + pub metric_name: String, + pub group_by_keys: BTreeSet, + pub required_capability: Capability, + pub function: String, + /// Already-evaluated leading scalar args. Order matches the + /// PromQL surface (`quantile_over_time(q, foo[r])` → `args[0] = q`). + pub function_args: Vec, + /// Time range from the matrix-vector selector (e.g. `[5m]` → 300). + /// `0` when the query is instant-vector-shaped (`histogram_quantile` + /// over an already-bucketed metric, bare cardinality_estimate, etc.). + pub range_seconds: u64, +} + +/// Whole-query analysis result. The vector of [`WarmTierCandidate`]s +/// covers every sub-expression the warm tier CAN serve. `unsupported` +/// is `Some` when ANY sub-expression cannot be served (or when the +/// query parse failed); in that case `candidates` may be partially +/// populated (sub-expressions BEFORE the rejected one) but the +/// reducer treats the analysis as warm-tier-miss and routes to cold. +#[derive(Debug, Clone, PartialEq, Default)] +pub struct WarmTierAnalysis { + pub candidates: Vec, + pub unsupported: Option, +} + +impl WarmTierAnalysis { + /// True when the analysis is fully warm-tier-answerable — + /// `unsupported.is_none()` AND at least one candidate. + pub fn is_warm_tier_answerable(&self) -> bool { + self.unsupported.is_none() && !self.candidates.is_empty() + } +} + +/// Distinct reasons a PromQL query is NOT warm-tier-answerable. +/// Each variant maps onto a different routing decision the caller +/// makes (typically all → cold tier / archive, but the variant +/// distinction matters for logging + future precompute hints). +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum UnsupportedReason { + /// PromQL function the warm tier has no reducer for — + /// `rate`, `irate`, `increase`, raw arithmetic, etc. + UnsupportedFunction(String), + /// PromQL composition shape the warm tier can't unfold — + /// `topk(k, rate(...))`, `sum by (...) (metric)` (pending + /// the Sum-over-CountSketch reducer), histogram_quantile over + /// a nested rate, etc. The string carries a short description. + UnsupportedComposition(String), + /// The query is a bare vector / matrix selector with no call — + /// the warm tier doesn't materialize raw counter values; the + /// archive answers these directly. + NoCallNodeFound, + /// `promql_parser` failed to parse the input. Carries the parser + /// error message for diagnostics. + UnparseablePromql(String), +} + +// ── Public entry point ─────────────────────────────────────────────────────── + +/// Parse PromQL + walk the AST and produce a [`WarmTierAnalysis`]. +/// +/// This is the single owner of warm-tier shape recognition. All +/// downstream code (the warm-tier reducer, the engine router) keys +/// off the returned `WarmTierAnalysis` and never re-parses the +/// PromQL string. +pub fn analyze_promql_for_warm_tier(promql: &str) -> WarmTierAnalysis { + // ── Custom warm-tier function names ───────────────────────────────── + // + // `cardinality_estimate(metric)`, `count_distinct_over_time(metric[r])`, + // `count_distinct(metric)`, and `topk_over_time(k, metric[r])` are + // not in `promql_parser`'s built-in function table, so the AST + // parser rejects them outright. We pre-detect those shapes via a + // narrow regex on the OUTER call, then re-parse the inner + // selector / matrix-selector expression as standalone PromQL. + if let Some(analysis) = try_parse_custom_function(promql) { + return analysis; + } + + let ast = match parser::parse(promql) { + Ok(ast) => ast, + Err(e) => { + return WarmTierAnalysis { + candidates: Vec::new(), + unsupported: Some(UnsupportedReason::UnparseablePromql(e.to_string())), + }; + } + }; + let mut out = WarmTierAnalysis::default(); + analyze_expr(&ast, &mut out); + out +} + +/// Match a small set of custom warm-tier function names that +/// `promql_parser` doesn't recognize, and parse the inner argument +/// as a standalone selector / matrix-selector to extract +/// `(metric, group_by, range)`. Returns `None` if the input doesn't +/// look like one of those custom shapes. +fn try_parse_custom_function(promql: &str) -> Option { + let trimmed = promql.trim(); + // Shape: NAME ( [scalar_args... , ] inner_expr ) + let open = trimmed.find('(')?; + if !trimmed.ends_with(')') { + return None; + } + let name = trimmed[..open].trim().to_lowercase(); + let inner = &trimmed[open + 1..trimmed.len() - 1]; + + let (capability, has_scalar) = match name.as_str() { + "cardinality_estimate" => (Capability::CardinalityApprox, false), + "count_distinct" => (Capability::CardinalityApprox, false), + "count_distinct_over_time" => (Capability::CardinalityApprox, false), + "topk_over_time" => { + (Capability::FrequencyTopk(SketchKindHandle::CmsWithHeap), true) + } + _ => return None, + }; + + // Split off leading scalar arg (the `k` for topk_over_time, the + // `q` for any future quantile-shaped custom function). + let (scalar_args, body) = if has_scalar { + let (head, tail) = split_first_top_level_comma(inner)?; + let v = head.trim().parse::().ok()?; + (vec![v], tail.trim()) + } else { + (Vec::new(), inner.trim()) + }; + + // Parse the body as standalone PromQL. We accept either a bare + // vector selector or a matrix selector. + let ast = parser::parse(body).ok()?; + let (metric, gb, range_s) = extract_metric_keys_range(&ast)?; + Some(WarmTierAnalysis { + candidates: vec![WarmTierCandidate { + metric_name: metric, + group_by_keys: gb, + required_capability: capability, + function: name, + function_args: scalar_args, + range_seconds: range_s, + }], + unsupported: None, + }) +} + +/// Split a comma-separated argument list at the FIRST top-level comma +/// (one outside any nested parentheses / brackets). Used by +/// [`try_parse_custom_function`] to peel off a leading scalar argument +/// from `topk_over_time(k, metric[r])` without mistakenly splitting on +/// a comma inside a label-matcher. +fn split_first_top_level_comma(s: &str) -> Option<(&str, &str)> { + let bytes = s.as_bytes(); + let mut depth: i32 = 0; + for (i, &b) in bytes.iter().enumerate() { + match b { + b'(' | b'[' | b'{' => depth += 1, + b')' | b']' | b'}' => depth -= 1, + b',' if depth == 0 => { + return Some((&s[..i], &s[i + 1..])); + } + _ => {} + } + } + None +} + +// ── AST walker ─────────────────────────────────────────────────────────────── + +fn analyze_expr(expr: &Expr, out: &mut WarmTierAnalysis) { + match expr { + Expr::Call(call) => analyze_call(call, out), + Expr::Aggregate(agg) => analyze_aggregate(agg, out), + Expr::Paren(p) => analyze_expr(&p.expr, out), + Expr::Subquery(sq) => analyze_expr(&sq.expr, out), + Expr::VectorSelector(_) | Expr::MatrixSelector(_) => { + // Bare selector — no call to dispatch on. The archive + // engine answers raw selectors; warm tier doesn't + // materialize raw counter values. + out.unsupported = Some(UnsupportedReason::NoCallNodeFound); + } + Expr::Binary(_) => { + // Binary ops (e.g. `rate(...) > 0.5`) aren't a single + // warm-tier candidate. We don't try to decompose them. + out.unsupported = Some(UnsupportedReason::UnsupportedComposition( + "binary expression — warm tier does not stitch lhs/rhs".to_string(), + )); + } + Expr::Unary(_) => { + out.unsupported = Some(UnsupportedReason::UnsupportedComposition( + "unary expression — warm tier does not stitch unary over sketch output" + .to_string(), + )); + } + Expr::NumberLiteral(_) | Expr::StringLiteral(_) => { + // Literals as top-level expressions aren't queries that + // hit the warm tier. + out.unsupported = Some(UnsupportedReason::UnsupportedComposition( + "literal at query root — no metric selector".to_string(), + )); + } + // Promql_parser exposes additional variants for future shapes + // (Extension, etc.); treat everything else as unsupported. + _ => { + out.unsupported = Some(UnsupportedReason::UnsupportedComposition( + "unrecognized PromQL expression shape".to_string(), + )); + } + } +} + +/// Handle `Call` nodes: the canonical warm-tier shapes +/// (`quantile_over_time`, `histogram_quantile`, +/// `count_distinct_over_time`, `cardinality_estimate`, +/// `topk_over_time`) plus the demo's explicit rejections +/// (`rate`, `irate`, `increase`). +fn analyze_call(call: &Call, out: &mut WarmTierAnalysis) { + let func_name = call.func.name.to_lowercase(); + + // ── Reject raw-counter math up-front ──────────────────────────────── + if matches!( + func_name.as_str(), + "rate" | "irate" | "increase" | "deriv" | "predict_linear" | "delta" | "idelta" + ) { + out.unsupported = Some(UnsupportedReason::UnsupportedFunction(func_name)); + return; + } + + // Leading scalar args (e.g. `q` in `quantile_over_time(q, m[r])`). + let mut scalar_args = Vec::new(); + for a in &call.args.args { + match a.as_ref() { + Expr::NumberLiteral(nl) => scalar_args.push(nl.val), + _ => break, + } + } + + // Extract the inner selector (or detect nested forbidden calls + // like `histogram_quantile(q, sum(rate(...)) by (le))`). + let body_arg = match call.args.args.iter().find(|a| { + !matches!(a.as_ref(), Expr::NumberLiteral(_) | Expr::StringLiteral(_)) + }) { + Some(a) => a.as_ref(), + None => { + // No selector arg — `quantile_over_time(0.99)` with no + // metric. Malformed but we surface as unsupported. + out.unsupported = Some(UnsupportedReason::UnsupportedComposition(format!( + "{func_name}: no metric selector argument" + ))); + return; + } + }; + + // For histogram_quantile, the body may be a bare bucket selector + // (our supported case) OR a nested aggregate over rate (the + // explicit rejection). Walk in. + if func_name == "histogram_quantile" { + if has_nested_rate(body_arg) { + out.unsupported = Some(UnsupportedReason::UnsupportedFunction("rate".to_string())); + return; + } + if let Some((metric, gb, range_s)) = extract_metric_keys_range(body_arg) { + out.candidates.push(WarmTierCandidate { + metric_name: metric, + group_by_keys: gb, + required_capability: Capability::QuantileApprox(SketchKindHandle::Any), + function: "histogram_quantile".to_string(), + function_args: scalar_args, + range_seconds: range_s, + }); + return; + } + out.unsupported = Some(UnsupportedReason::UnsupportedComposition( + "histogram_quantile: body is not a recognizable metric/aggregate".to_string(), + )); + return; + } + + // Reject `*_over_time` wrappers around rate / irate / increase + // even when not under histogram_quantile. + if has_nested_rate(body_arg) { + out.unsupported = Some(UnsupportedReason::UnsupportedFunction("rate".to_string())); + return; + } + + let (metric, gb, range_s) = match extract_metric_keys_range(body_arg) { + Some(x) => x, + None => { + out.unsupported = Some(UnsupportedReason::UnsupportedComposition(format!( + "{func_name}: cannot extract metric selector from body" + ))); + return; + } + }; + + let cap = match func_name.as_str() { + "quantile_over_time" => Capability::QuantileApprox(SketchKindHandle::Any), + "count_distinct_over_time" | "cardinality_estimate" => Capability::CardinalityApprox, + "topk_over_time" => Capability::FrequencyTopk(SketchKindHandle::CmsWithHeap), + other => { + out.unsupported = Some(UnsupportedReason::UnsupportedFunction(other.to_string())); + return; + } + }; + out.candidates.push(WarmTierCandidate { + metric_name: metric, + group_by_keys: gb, + required_capability: cap, + function: func_name, + function_args: scalar_args, + range_seconds: range_s, + }); +} + +/// Handle `Aggregate` nodes: `topk(k, m)` is the only supported +/// shape today. `sum by (label_set) (metric)` is the documented +/// follow-up — surface as `UnsupportedComposition`. +fn analyze_aggregate(agg: &AggregateExpr, out: &mut WarmTierAnalysis) { + let op_name = agg.op.to_string().to_lowercase(); + + // Nested rate / irate / increase anywhere in the aggregate's body + // disqualifies the whole expression regardless of the outer + // aggregate op. Detect this first so the + // `sum by (zone) (rate(http_requests_total[5m]))` shape surfaces + // the canonical "rate" error message rather than the outer-op + // composition error. + if has_nested_rate(&agg.expr) { + out.unsupported = Some(UnsupportedReason::UnsupportedFunction("rate".to_string())); + return; + } + + // Pull `k` from the aggregate's `param` (for `topk` / `bottomk` / + // `quantile`). + let mut scalar_args: Vec = Vec::new(); + if let Some(p) = &agg.param { + if let Expr::NumberLiteral(nl) = p.as_ref() { + scalar_args.push(nl.val); + } + } + + match op_name.as_str() { + "topk" | "bottomk" => { + // Nested rate is already filtered out above (top-of-fn + // `has_nested_rate` check); here we just need to extract + // the metric from the body. + let (metric, gb, range_s) = match extract_metric_keys_range(&agg.expr) { + Some(x) => x, + None => { + out.unsupported = Some(UnsupportedReason::UnsupportedComposition(format!( + "{op_name}: cannot extract metric selector from body" + ))); + return; + } + }; + out.candidates.push(WarmTierCandidate { + metric_name: metric, + group_by_keys: gb, + required_capability: Capability::FrequencyTopk(SketchKindHandle::CmsWithHeap), + function: op_name, + function_args: scalar_args, + range_seconds: range_s, + }); + } + "sum" | "avg" | "count" | "min" | "max" | "group" | "stddev" | "stdvar" => { + // The clean Sum-over-CountSketch reducer is a future + // follow-up. Surface as UnsupportedComposition so the + // routing decision is explicit and the follow-up has a + // clear hook. + out.unsupported = Some(UnsupportedReason::UnsupportedComposition(format!( + "{op_name} by (...) (metric) — pending Sum-over-CountSketch reducer; \ + falling over to archive" + ))); + } + "quantile" => { + // `quantile(q, m)` is the instant-vector aggregate (no + // matrix selector). Map to QuantileApprox. + let (metric, gb, range_s) = match extract_metric_keys_range(&agg.expr) { + Some(x) => x, + None => { + out.unsupported = Some(UnsupportedReason::UnsupportedComposition( + "quantile: cannot extract metric selector from body".to_string(), + )); + return; + } + }; + out.candidates.push(WarmTierCandidate { + metric_name: metric, + group_by_keys: gb, + required_capability: Capability::QuantileApprox(SketchKindHandle::Any), + function: op_name, + function_args: scalar_args, + range_seconds: range_s, + }); + } + other => { + out.unsupported = Some(UnsupportedReason::UnsupportedFunction(other.to_string())); + } + } +} + +// ── Helpers ────────────────────────────────────────────────────────────────── + +/// Walk into `expr` to find a Vector / Matrix selector and return +/// `(metric_name, group_by_keys, range_seconds)`. `range_seconds` +/// is `0` for an instant-vector selector. +fn extract_metric_keys_range(expr: &Expr) -> Option<(String, BTreeSet, u64)> { + match expr { + Expr::VectorSelector(vs) => { + let (m, keys) = extract_vs_metric_and_keys(vs)?; + Some((m, keys, 0)) + } + Expr::MatrixSelector(ms) => { + let (m, keys) = extract_vs_metric_and_keys(&ms.vs)?; + Some((m, keys, duration_to_seconds(ms.range))) + } + Expr::Paren(p) => extract_metric_keys_range(&p.expr), + Expr::Subquery(sq) => extract_metric_keys_range(&sq.expr), + Expr::Call(c) => { + // Walk into single-arg call wrappers (e.g. an inner aggregate). + c.args.args.iter().find_map(|a| extract_metric_keys_range(a)) + } + Expr::Aggregate(a) => extract_metric_keys_range(&a.expr), + _ => None, + } +} + +fn extract_vs_metric_and_keys(vs: &VectorSelector) -> Option<(String, BTreeSet)> { + let mut keys = BTreeSet::new(); + let mut metric = vs.name.clone().unwrap_or_default(); + for m in &vs.matchers.matchers { + if m.name == "__name__" { + if metric.is_empty() { + metric = m.value.clone(); + } + continue; + } + keys.insert(m.name.clone()); + } + if metric.is_empty() { + None + } else { + Some((metric, keys)) + } +} + +fn duration_to_seconds(d: Duration) -> u64 { + d.as_secs() +} + +/// Walk into `expr` looking for a `rate` / `irate` / `increase` / +/// `deriv` / `delta` / `idelta` / `predict_linear` call anywhere in +/// the subtree. Used to reject the demo's +/// `sum by (zone) (rate(http_requests_total[5m]))` and +/// `histogram_quantile(0.99, sum(rate(bucket[5m])) by (le))` shapes +/// up-front. +fn has_nested_rate(expr: &Expr) -> bool { + match expr { + Expr::Call(call) => { + let name = call.func.name.to_lowercase(); + if matches!( + name.as_str(), + "rate" | "irate" | "increase" | "deriv" | "delta" | "idelta" | "predict_linear" + ) { + return true; + } + call.args.args.iter().any(|a| has_nested_rate(a)) + } + Expr::Aggregate(agg) => has_nested_rate(&agg.expr), + Expr::Paren(p) => has_nested_rate(&p.expr), + Expr::Subquery(sq) => has_nested_rate(&sq.expr), + Expr::Binary(b) => has_nested_rate(&b.lhs) || has_nested_rate(&b.rhs), + Expr::Unary(u) => has_nested_rate(&u.expr), + _ => false, + } +} + +// ── Bidirectional adapters with the backend's sketch_index::Capability ─────── +// +// `asap-query-engine` carries its own `Capability` / `SketchKindHandle` +// enums (in `stores::sketch_db::sketch_index`) which the backend's +// ingest + storage paths reference everywhere. Rather than relocate +// those types and churn 18 backend files, we own the canonical +// definition here and adapt at the controller↔backend boundary. +// +// The adapters live as `From` impls on the BACKEND side because that's +// where the source-of-truth `sketch_index::Capability` lives; this +// module just defines the controller-local mirror. See +// `asap-query-engine/src/stores/sketch_db/sketch_index.rs` for the +// `From` impl. + +// ── Tests ──────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + + fn keys(items: &[&str]) -> BTreeSet { + items.iter().map(|s| s.to_string()).collect() + } + + // ── Supported shapes ───────────────────────────────────────────────── + + #[test] + fn analyze_quantile_over_time() { + let a = analyze_promql_for_warm_tier("quantile_over_time(0.99, http_latency_ms[5m])"); + assert!(a.unsupported.is_none(), "expected no unsupported reason: {a:?}"); + assert_eq!(a.candidates.len(), 1); + let c = &a.candidates[0]; + assert_eq!(c.metric_name, "http_latency_ms"); + assert_eq!(c.function, "quantile_over_time"); + assert_eq!(c.function_args, vec![0.99]); + assert_eq!(c.range_seconds, 300); + assert_eq!(c.required_capability, Capability::QuantileApprox(SketchKindHandle::Any)); + } + + #[test] + fn analyze_quantile_over_time_with_label_matchers() { + let a = analyze_promql_for_warm_tier( + "quantile_over_time(0.5, http_latency_ms{zone=\"z0\", region=\"us\"}[30s])", + ); + assert!(a.unsupported.is_none(), "{a:?}"); + assert_eq!(a.candidates.len(), 1); + let c = &a.candidates[0]; + assert_eq!(c.metric_name, "http_latency_ms"); + assert_eq!(c.group_by_keys, keys(&["zone", "region"])); + assert_eq!(c.range_seconds, 30); + } + + #[test] + fn analyze_histogram_quantile_over_bare_metric() { + let a = analyze_promql_for_warm_tier("histogram_quantile(0.99, http_latency_ms)"); + assert!(a.unsupported.is_none(), "{a:?}"); + assert_eq!(a.candidates.len(), 1); + let c = &a.candidates[0]; + assert_eq!(c.function, "histogram_quantile"); + assert_eq!(c.function_args, vec![0.99]); + assert_eq!(c.metric_name, "http_latency_ms"); + assert_eq!(c.range_seconds, 0); + assert_eq!(c.required_capability, Capability::QuantileApprox(SketchKindHandle::Any)); + } + + #[test] + fn analyze_cardinality_estimate() { + let a = analyze_promql_for_warm_tier("cardinality_estimate(uniq_users)"); + assert!(a.unsupported.is_none(), "{a:?}"); + assert_eq!(a.candidates.len(), 1); + assert_eq!( + a.candidates[0].required_capability, + Capability::CardinalityApprox, + ); + } + + #[test] + fn analyze_topk_aggregate() { + let a = analyze_promql_for_warm_tier("topk(5, endpoint_hits)"); + assert!(a.unsupported.is_none(), "{a:?}"); + assert_eq!(a.candidates.len(), 1); + let c = &a.candidates[0]; + assert_eq!(c.function, "topk"); + assert_eq!(c.function_args, vec![5.0]); + assert_eq!( + c.required_capability, + Capability::FrequencyTopk(SketchKindHandle::CmsWithHeap), + ); + } + + #[test] + fn analyze_topk_over_time() { + let a = analyze_promql_for_warm_tier("topk_over_time(10, endpoint_hits[1h])"); + assert!(a.unsupported.is_none(), "{a:?}"); + let c = &a.candidates[0]; + assert_eq!(c.function, "topk_over_time"); + assert_eq!(c.function_args, vec![10.0]); + assert_eq!(c.range_seconds, 3600); + } + + #[test] + fn analyze_quantile_instant_aggregate() { + let a = analyze_promql_for_warm_tier("quantile(0.5, http_latency_ms)"); + assert!(a.unsupported.is_none(), "{a:?}"); + assert_eq!( + a.candidates[0].required_capability, + Capability::QuantileApprox(SketchKindHandle::Any), + ); + } + + // ── Unsupported / rejected shapes ──────────────────────────────────── + + #[test] + fn reject_bare_vector_selector() { + let a = analyze_promql_for_warm_tier("http_requests_total{zone=\"z0\"}"); + assert_eq!( + a.unsupported, + Some(UnsupportedReason::NoCallNodeFound), + "{a:?}" + ); + assert!(a.candidates.is_empty()); + } + + #[test] + fn reject_rate_function() { + let a = analyze_promql_for_warm_tier("rate(http_requests_total[5m])"); + assert_eq!( + a.unsupported, + Some(UnsupportedReason::UnsupportedFunction("rate".to_string())), + ); + } + + #[test] + fn reject_irate_function() { + let a = analyze_promql_for_warm_tier("irate(http_requests_total[5m])"); + assert_eq!( + a.unsupported, + Some(UnsupportedReason::UnsupportedFunction("irate".to_string())), + ); + } + + #[test] + fn reject_increase_function() { + let a = analyze_promql_for_warm_tier("increase(http_requests_total[5m])"); + assert_eq!( + a.unsupported, + Some(UnsupportedReason::UnsupportedFunction("increase".to_string())), + ); + } + + #[test] + fn reject_sum_by_rate_compound() { + // The demo's `sum by (zone) (rate(http_requests_total[5m]))`. + let a = analyze_promql_for_warm_tier( + "sum by (zone) (rate(http_requests_total[5m]))", + ); + // Nested `rate` is detected first and surfaced as UnsupportedFunction. + assert_eq!( + a.unsupported, + Some(UnsupportedReason::UnsupportedFunction("rate".to_string())), + "{a:?}" + ); + } + + #[test] + fn reject_histogram_quantile_over_sum_rate() { + // `histogram_quantile(0.99, sum(rate(bucket[5m])) by (le))`. + let a = analyze_promql_for_warm_tier( + "histogram_quantile(0.99, sum(rate(http_latency_bucket[5m])) by (le))", + ); + // Inner rate is detected, surfaced as UnsupportedFunction. + assert_eq!( + a.unsupported, + Some(UnsupportedReason::UnsupportedFunction("rate".to_string())), + "{a:?}" + ); + } + + #[test] + fn reject_sum_by_bare_metric_pending_sum_reducer() { + // `sum by (zone) (http_requests_total)` — supported in a future + // Sum-over-CountSketch follow-up; for now surface as + // UnsupportedComposition. + let a = analyze_promql_for_warm_tier( + "sum by (zone) (http_requests_total)", + ); + match a.unsupported { + Some(UnsupportedReason::UnsupportedComposition(msg)) => { + assert!( + msg.contains("sum") && msg.contains("CountSketch"), + "expected msg to mention sum + CountSketch, got `{msg}`" + ); + } + other => panic!("expected UnsupportedComposition for sum-by, got {other:?}"), + } + } + + #[test] + fn reject_topk_over_rate() { + // `topk(5, rate(http_requests_total[5m]))` — topk only over + // instant vectors. Nested rate is detected by the + // top-of-aggregate-handler `has_nested_rate` check and + // surfaces as the canonical UnsupportedFunction("rate") so + // log telemetry attributes the rejection to the + // root-cause function regardless of which outer op wrapped it. + let a = analyze_promql_for_warm_tier( + "topk(5, rate(http_requests_total[5m]))", + ); + assert_eq!( + a.unsupported, + Some(UnsupportedReason::UnsupportedFunction("rate".to_string())), + "{a:?}" + ); + } + + #[test] + fn reject_binary_op() { + let a = analyze_promql_for_warm_tier("rate(foo[5m]) > 0.5"); + // The binary op contains a rate call — rate is detected first + // at the AST root or as an UnsupportedComposition; either way + // we surface an unsupported reason. + assert!(a.unsupported.is_some(), "{a:?}"); + } + + #[test] + fn unparseable_promql_surfaces_clean_error() { + let a = analyze_promql_for_warm_tier("@@@ this is not promql @@@"); + match a.unsupported { + Some(UnsupportedReason::UnparseablePromql(msg)) => { + assert!(!msg.is_empty(), "parser error message should be non-empty"); + } + other => panic!("expected UnparseablePromql, got {other:?}"), + } + } + + // ── Range parsing ──────────────────────────────────────────────────── + + #[test] + fn range_seconds_parses_seconds() { + let a = analyze_promql_for_warm_tier("quantile_over_time(0.99, m[30s])"); + assert_eq!(a.candidates[0].range_seconds, 30); + } + + #[test] + fn range_seconds_parses_minutes() { + let a = analyze_promql_for_warm_tier("quantile_over_time(0.99, m[5m])"); + assert_eq!(a.candidates[0].range_seconds, 300); + } + + #[test] + fn range_seconds_parses_hours() { + let a = analyze_promql_for_warm_tier("quantile_over_time(0.99, m[2h])"); + assert_eq!(a.candidates[0].range_seconds, 7200); + } + + #[test] + fn instant_vector_has_zero_range() { + let a = analyze_promql_for_warm_tier("histogram_quantile(0.5, m)"); + assert_eq!(a.candidates[0].range_seconds, 0); + } + + // ── is_warm_tier_answerable ────────────────────────────────────────── + + #[test] + fn is_warm_tier_answerable_true_for_supported() { + let a = analyze_promql_for_warm_tier("quantile_over_time(0.99, m[5m])"); + assert!(a.is_warm_tier_answerable()); + } + + #[test] + fn is_warm_tier_answerable_false_for_unsupported() { + let a = analyze_promql_for_warm_tier("rate(m[5m])"); + assert!(!a.is_warm_tier_answerable()); + } + + #[test] + fn is_warm_tier_answerable_false_for_bare_selector() { + let a = analyze_promql_for_warm_tier("m{zone=\"z0\"}"); + assert!(!a.is_warm_tier_answerable()); + } +}