From 21c6ffa2d2ae48fd01a77efcb6f4cb900e4685d7 Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 31 Mar 2026 17:21:39 -0500 Subject: [PATCH 1/3] =?UTF-8?q?feat:=20SP-9=20AST-aware=20hierarchical=20s?= =?UTF-8?q?tage=20assignment=20(edge=20collector=20=E2=86=92=20leaf=20node?= =?UTF-8?q?s,=20precompute=20engine=20=E2=86=92=20upper=20AST)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements SP-9 as described in docs/controller-optimization-problem.md. The SketchExpr tree produced by the query parser is now threaded through to the planner and split across pipeline stages instead of being discarded. Co-Authored-By: Claude Sonnet 4.6 --- controller/src/analyzer.rs | 28 +- controller/src/config/backend.rs | 77 +- controller/src/config/mod.rs | 2 +- controller/src/config/precompute.rs | 23 +- controller/src/main.rs | 20 +- controller/src/planner/cost_model.rs | 1 + controller/src/planner/delta_cost_model.rs | 1 + controller/src/planner/mod.rs | 1 + controller/src/planner/rules.rs | 2 + controller/src/planner/stage_split.rs | 858 +++++++++++++++++++++ controller/src/replan.rs | 1 + controller/src/store/mod.rs | 1 + controller/src/types.rs | 124 +++ 13 files changed, 1123 insertions(+), 16 deletions(-) create mode 100644 controller/src/planner/stage_split.rs diff --git a/controller/src/analyzer.rs b/controller/src/analyzer.rs index e78e6e6a..05facb51 100644 --- a/controller/src/analyzer.rs +++ b/controller/src/analyzer.rs @@ -3,7 +3,7 @@ use std::time::Duration; use anyhow::{anyhow, Context}; use serde::{Deserialize, Serialize}; -use crate::query_parser; +use crate::query_parser::{self, SketchExpr}; use crate::types::{AggType, QueryWorkload, SketchType, WorkloadCharacteristics}; // ── Public API ──────────────────────────────────────────────────────────────── @@ -156,6 +156,32 @@ impl Analyzer { quantiles, }) } + + /// Like [`analyze`] but also returns the optimized [`SketchExpr`] tree. + /// + /// The tree is `Some` when `spec.query_string` was provided; `None` when + /// the workload was built from explicit aggregation fields alone. + /// + /// Callers (e.g. `handle_plan`) pass the `SketchExpr` to + /// `planner::stage_split::split_expr_by_stage()` to produce an SP-9 + /// [`StagedPlan`]. The existing `analyze()` path (SP-3 flat assignment) + /// is unchanged and remains the fallback when no tree is available. + pub fn analyze_with_sketch( + &self, + spec: QuerySpec, + ) -> anyhow::Result<(QueryWorkload, Option)> { + // Parse the SketchExpr before `spec` is consumed by `analyze()`. + // `parse_query_sketch` runs the algebraic optimiser internally. + let sketch_expr = spec + .query_string + .as_deref() + .map(|q| query_parser::parse_query_sketch(q)) + .transpose() + .with_context(|| "failed to parse query_string into SketchExpr")?; + + let workload = self.analyze(spec)?; + Ok((workload, sketch_expr)) + } } // ── Duration helpers (used by other modules) ────────────────────────────────── diff --git a/controller/src/config/backend.rs b/controller/src/config/backend.rs index eaee299b..c32979b5 100644 --- a/controller/src/config/backend.rs +++ b/controller/src/config/backend.rs @@ -1,27 +1,57 @@ use anyhow::Context; +use serde_json::json; use crate::types::*; /// Generates an OTel collector YAML string for the backend merge collector. +/// +/// **SP-9**: when `staged.has_dedup` is true a `dedup` processor is inserted +/// before the merge processor in the pipeline, honouring the `Dedup` node +/// assignment from [`crate::planner::stage_split::split_expr_by_stage`]. pub fn generate_backend_config(cfg: &BackendCollectorConfig, opamp_endpoint: &str) -> anyhow::Result { - let processor_key = format!("{}_merge", cfg.merge_sketch_type); + generate_backend_config_staged(cfg, None, opamp_endpoint) +} + +/// Extended entry point used by SP-9-aware callers that supply a +/// [`BackendSubPlan`] carrying the dedup flag. +pub fn generate_backend_config_staged( + cfg: &BackendCollectorConfig, + staged: Option<&BackendSubPlan>, + opamp_endpoint: &str, +) -> anyhow::Result { + let merge_key = format!("{}_merge", cfg.merge_sketch_type); + let has_dedup = staged.map(|s| s.has_dedup).unwrap_or(false); + + // Build the processors map and pipeline processor list. + let mut processors = serde_json::Map::new(); + let mut pipeline_processors: Vec = vec![]; + + if has_dedup { + processors.insert( + "dedup".into(), + json!({ "mode": "dedup" }), + ); + pipeline_processors.push(json!("dedup")); + } + + processors.insert( + merge_key.clone(), + json!({ "mode": "merge", "group_by": cfg.group_by }), + ); + pipeline_processors.push(json!(&merge_key)); - let doc = serde_yaml::to_value(&serde_json::json!({ + let doc = serde_yaml::to_value(&json!({ "extensions": { "opamp": { "server": { "ws": { "endpoint": opamp_endpoint } } } }, - "processors": { - &processor_key: { - "mode": "merge", - "group_by": cfg.group_by - } - }, + "processors": processors, "service": { "extensions": ["opamp"], "pipelines": { - "metrics": { "processors": [&processor_key] } + "metrics": { "processors": pipeline_processors } } } - })).context("build backend doc")?; + })) + .context("build backend doc")?; serde_yaml::to_string(&doc).context("serialize backend config") } @@ -60,4 +90,31 @@ mod tests { let yaml = generate_backend_config(&cfg, ep).unwrap(); assert!(yaml.contains(ep), "YAML should contain endpoint\n{yaml}"); } + + #[test] + fn dedup_processor_emitted_when_staged_has_dedup() { + let cfg = BackendCollectorConfig { + merge_sketch_type: SketchType::HLL, + group_by: vec!["user_id".into()], + }; + let staged = BackendSubPlan { has_dedup: true, has_merge: true, group_by: vec![] }; + let yaml = generate_backend_config_staged(&cfg, Some(&staged), "ws://ctrl:4320/v1/opamp").unwrap(); + assert!(yaml.contains("dedup:"), "YAML should contain dedup processor\n{yaml}"); + // dedup must appear before merge in the pipeline list + let dedup_pos = yaml.find("- dedup").expect("missing dedup in pipeline"); + let merge_pos = yaml.find("- hll_merge").expect("missing merge in pipeline"); + assert!(dedup_pos < merge_pos, "dedup must precede merge\n{yaml}"); + } + + #[test] + fn no_dedup_when_staged_has_dedup_false() { + let cfg = BackendCollectorConfig { + merge_sketch_type: SketchType::DDSketch, + group_by: vec![], + }; + let staged = BackendSubPlan { has_dedup: false, has_merge: true, group_by: vec![] }; + let yaml = generate_backend_config_staged(&cfg, Some(&staged), "ws://ctrl:4320/v1/opamp").unwrap(); + assert!(!yaml.contains("dedup"), "YAML must not contain dedup\n{yaml}"); + } } + diff --git a/controller/src/config/mod.rs b/controller/src/config/mod.rs index 2b6f5190..5301e075 100644 --- a/controller/src/config/mod.rs +++ b/controller/src/config/mod.rs @@ -3,5 +3,5 @@ pub mod backend; pub mod precompute; pub use agent::generate_agent_config; -pub use backend::generate_backend_config; +pub use backend::{generate_backend_config, generate_backend_config_staged}; pub use precompute::{should_precompute, build_precompute_jobs, PrecomputeClient}; diff --git a/controller/src/config/precompute.rs b/controller/src/config/precompute.rs index 23400684..4857d106 100644 --- a/controller/src/config/precompute.rs +++ b/controller/src/config/precompute.rs @@ -21,17 +21,35 @@ pub fn should_precompute(w: &QueryWorkload) -> bool { /// Builds the list of precompute jobs for a plan. Returns an empty Vec if the /// workload does not meet the precompute eligibility criterion. +/// +/// **SP-9**: when `plan.staged_plan` is `Some` and the precompute sub-plan is +/// active, the job's `query_expr` is taken from the AST-derived PromQL +/// serialisation ([`PrecomputeSubPlan::query_expr`]) rather than the hardcoded +/// `quantile_over_time(0.99, …)` template. This allows the precompute engine +/// to evaluate the actual upper sub-tree (e.g. `topk(10, count_over_time(…))`). +/// +/// When no `staged_plan` is present (SP-3 flat path), the legacy +/// `build_query_expr()` template is used as the fallback. pub fn build_precompute_jobs( w: &QueryWorkload, - _plan: &CollectionPlan, + plan: &CollectionPlan, backend_addr: &str, ) -> Vec { if !should_precompute(w) { return vec![]; } let granularity = w.repeat_every.unwrap(); // safe: should_precompute checked it + + // SP-9: use the staged plan's PromQL when available and non-empty. + let query_expr = plan + .staged_plan + .as_ref() + .filter(|sp| sp.precompute.active && !sp.precompute.query_expr.is_empty()) + .map(|sp| sp.precompute.query_expr.clone()) + .unwrap_or_else(|| build_query_expr(w)); + vec![PrecomputeJob { - query_expr: build_query_expr(w), + query_expr, granularity, sketch_source: backend_addr.to_string(), store_path: build_store_path(w), @@ -198,6 +216,7 @@ mod tests { valid_until: chrono::Utc::now(), delta_decision: DeltaDecision::default(), transmission_cost_summary: TransmissionCostSummary::default(), + staged_plan: None, } } diff --git a/controller/src/main.rs b/controller/src/main.rs index 258682fa..1812ee29 100644 --- a/controller/src/main.rs +++ b/controller/src/main.rs @@ -23,6 +23,7 @@ use tracing::{info, warn}; use analyzer::{Analyzer, QuerySpec}; use config::{generate_agent_config, generate_backend_config, build_precompute_jobs}; +use config::generate_backend_config_staged; use monitor::{Endpoint, Scraper, ScrapedData, Thresholds, Violation}; use opamp::{AgentRole, OpampServer, RemoteConfig}; use planner::{CostModelPlanner, BaselinePlanner, ObjectiveWeights, OnlineMetricsStore, init_online_store, pareto_frontier, select_best}; @@ -217,12 +218,22 @@ async fn handle_plan( Json(spec): Json, ) -> impl IntoResponse { let wc = spec.workload.clone(); - let workload = match st.analyzer.analyze(spec) { + let (workload, sketch_expr) = match st.analyzer.analyze_with_sketch(spec) { Ok(w) => w, Err(e) => return (StatusCode::UNPROCESSABLE_ENTITY, e.to_string()).into_response(), }; let mut plan = st.planner.plan(&workload, Some(&wc)); + + // ── SP-9: AST-aware hierarchical stage assignment ───────────────────────── + // When a SketchExpr tree is available (query_string path), split it across + // stages and attach the result to the plan. The SP-3 flat assignment + // remains active as the fallback when no tree is present. + if let Some(ref expr) = sketch_expr { + let budgets = types::StageResourceBudgets::from_workload_chars(&wc); + plan.staged_plan = Some(planner::stage_split::split_expr_by_stage(expr, &budgets)); + } + plan.precompute = build_precompute_jobs(&workload, &plan, "backend:4317"); st.store.set(&workload.metric_name, plan.clone()); // Persist workload so the replanner can re-run plan() without the original spec. @@ -238,7 +249,11 @@ async fn handle_plan( } // ── Push backend config to backend-role collectors ──────────────────────── - if let Ok(backend_yaml) = generate_backend_config(&plan.backend_config, &st.opamp_endpoint) { + // SP-9: pass the BackendSubPlan so the YAML gains a dedup processor when needed. + let backend_staged = plan.staged_plan.as_ref().map(|sp| &sp.backend); + if let Ok(backend_yaml) = generate_backend_config_staged( + &plan.backend_config, backend_staged, &st.opamp_endpoint, + ) { let hash = short_hash(&backend_yaml); st.opamp.push_to_role( AgentRole::Backend, @@ -264,6 +279,7 @@ async fn handle_plan( "agents_notified": agents.len(), "precompute_jobs": plan.precompute.len(), "delta_decision": plan.delta_decision, + "staged_plan": plan.staged_plan, "transmission_costs": { "raw_bytes_per_sec": cost.raw_bytes_per_sec, "sketch_full_bytes_per_sec": cost.sketch_full_bytes_per_sec, diff --git a/controller/src/planner/cost_model.rs b/controller/src/planner/cost_model.rs index 8c1f0d2b..8e32faa6 100644 --- a/controller/src/planner/cost_model.rs +++ b/controller/src/planner/cost_model.rs @@ -364,6 +364,7 @@ mod tests { valid_until: Utc::now(), delta_decision: DeltaDecision::default(), transmission_cost_summary: TransmissionCostSummary::default(), + staged_plan: None, } } diff --git a/controller/src/planner/delta_cost_model.rs b/controller/src/planner/delta_cost_model.rs index 2d1a9c1a..b00c875f 100644 --- a/controller/src/planner/delta_cost_model.rs +++ b/controller/src/planner/delta_cost_model.rs @@ -507,6 +507,7 @@ mod tests { valid_until: Utc::now(), delta_decision: DeltaDecision::default(), transmission_cost_summary: TransmissionCostSummary::default(), + staged_plan: None, } } diff --git a/controller/src/planner/mod.rs b/controller/src/planner/mod.rs index cf174ca2..a907dc27 100644 --- a/controller/src/planner/mod.rs +++ b/controller/src/planner/mod.rs @@ -4,6 +4,7 @@ pub mod delta_cost_model; pub mod online_cost_model; pub mod pareto; pub mod baseline_planner; +pub mod stage_split; pub use rules::RulesPlanner; pub use cost_model::CostModelPlanner; diff --git a/controller/src/planner/rules.rs b/controller/src/planner/rules.rs index 611dc2b8..5910964a 100644 --- a/controller/src/planner/rules.rs +++ b/controller/src/planner/rules.rs @@ -68,6 +68,7 @@ impl RulesPlanner { valid_until, delta_decision: DeltaDecision::default(), transmission_cost_summary: TransmissionCostSummary::default(), + staged_plan: None, } } @@ -108,6 +109,7 @@ impl RulesPlanner { valid_until, delta_decision: DeltaDecision::default(), transmission_cost_summary: TransmissionCostSummary::default(), + staged_plan: None, } } } diff --git a/controller/src/planner/stage_split.rs b/controller/src/planner/stage_split.rs new file mode 100644 index 00000000..ab635297 --- /dev/null +++ b/controller/src/planner/stage_split.rs @@ -0,0 +1,858 @@ +//! SP-9: AST-aware hierarchical stage assignment. +//! +//! Splits the optimized [`SketchExpr`] tree produced by the query parser +//! across pipeline stages, emitting a [`StagedPlan`] that carries per-stage +//! sub-plans for: +//! +//! | Stage | Nodes | +//! |---|---| +//! | Agent OTel Collector | `Source`, `Filter`, `Window`, `Agg` (sketch ops) | +//! | Backend OTel Collector | `Partition`, `Merge`, `Dedup`, `Agg { Exact(Sum\|Count\|Min\|Max) }` | +//! | ASAPQuery Precompute Engine | `TopK`, deferred sketch `Agg` ops | +//! | DB-side query | `Agg { Exact(Avg) }` (non-mergeable) | +//! +//! # ExactAgg deferral +//! +//! Mergeability drives `Exact` op placement: +//! - `Sum`, `Count`, `Min`, `Max` — mergeable (`agg(A∪B) = merge(agg(A), agg(B))`) → +//! **Backend** (the backend collector can combine partial results from N agents). +//! - `Avg` — **not** mergeable (avg-of-avgs ≠ global avg) → **DB-side** query +//! (must see all data before computing). +//! +//! # Budget-driven deferral chain +//! +//! When a sketch `Agg` node's estimated memory cost exceeds the stage cap in +//! [`StageResourceBudgets`], it is deferred to the next stage: +//! +//! `Agent → Backend → Precompute` +//! +//! The degenerate-but-valid fallback (all sketch ops deferred to Precompute) +//! matches the current flat SP-3 behaviour and ensures the query is always +//! answerable. +//! +//! # PromQL serialisation +//! +//! [`expr_to_promql`] converts the full `SketchExpr` tree to a PromQL +//! expression consumed by the ASAPQuery Precompute Engine's query engine. +//! The sketch data is already ingested by the precompute engine from the +//! backend OTel Collector; the PromQL describes what operation to apply. + +use std::time::Duration; + +use crate::analyzer::format_duration; +use crate::query_parser::sketch_algebra::{ + ExactAgg, FilterOp, FilterVal, PartitionKeys, SketchAggOp, SketchExpr, +}; +use crate::types::{ + AgentSubPlan, BackendSubPlan, DbSubPlan, PrecomputeSubPlan, SketchParams, SketchType, + StagedPlan, StageResourceBudgets, +}; + +// ── Public entry point ──────────────────────────────────────────────────────── + +/// Split a `SketchExpr` tree across pipeline stages, respecting per-stage +/// resource budgets and the `ExactAgg` mergeability rules described above. +/// +/// The returned [`StagedPlan`] is attached to +/// [`crate::types::CollectionPlan::staged_plan`] by the caller +/// (`handle_plan` in `main.rs`). +pub fn split_expr_by_stage(expr: &SketchExpr, budgets: &StageResourceBudgets) -> StagedPlan { + let mut plan = StagedPlan::default(); + walk(expr, &mut plan, budgets); + + // Build the precompute query_expr from the full tree when the precompute + // stage is active (TopK or deferred sketch ops assigned there). + if plan.precompute.active { + plan.precompute.query_expr = expr_to_promql(expr); + } + + // Build the DB query_expr when Avg was deferred to the DB stage. + if plan.db.active && plan.db.query_expr.is_empty() { + plan.db.query_expr = expr_to_promql(expr); + } + + plan +} + +/// Serialise a [`SketchExpr`] tree to a valid PromQL expression. +/// +/// The output is consumed by the ASAPQuery Precompute Engine's query engine +/// (PromQL / SQL). Sketch data is already present in the engine from the +/// backend OTel Collector; the PromQL describes the aggregation to apply. +pub fn expr_to_promql(expr: &SketchExpr) -> String { + let mut ctx = QueryCtx::default(); + collect_ctx(expr, &mut ctx); + ctx.to_promql() +} + +// ── Tree walker ─────────────────────────────────────────────────────────────── + +fn walk(expr: &SketchExpr, plan: &mut StagedPlan, budgets: &StageResourceBudgets) { + match expr { + // Source — always at Agent; no additional plan fields needed. + SketchExpr::Source(_) => {} + + // Filter — push label predicates to Agent. + SketchExpr::Filter { pred, input } => { + for p in pred { + if let (FilterOp::Eq, FilterVal::Str(v)) = (&p.op, &p.val) { + plan.agent.label_filters.push(format!("{}={}", p.col, v)); + } + } + walk(input, plan, budgets); + } + + // Window — time window lives at Agent. + SketchExpr::Window { duration, input } => { + plan.agent.window_secs = Some(duration.as_secs()); + walk(input, plan, budgets); + } + + // Partition — group-by dims always go to Backend after agent merge. + SketchExpr::Partition { keys, input } => { + for k in keys.keys() { + if !plan.backend.group_by.contains(k) { + plan.backend.group_by.push(k.clone()); + } + } + walk(input, plan, budgets); + } + + // Agg — the key assignment decision (see module docs). + SketchExpr::Agg { op, input, .. } => { + assign_agg(op, plan, budgets); + walk(input, plan, budgets); + } + + // Dedup — absorbed at Backend (HLL dedup is eliminated by R6 upstream). + SketchExpr::Dedup { input, .. } => { + plan.backend.has_dedup = true; + walk(input, plan, budgets); + } + + // Merge — Backend merges N agent sketches. + SketchExpr::Merge { inputs } => { + plan.backend.has_merge = true; + for i in inputs { + walk(i, plan, budgets); + } + } + + // TopK — always at the Precompute Engine (upper AST). + SketchExpr::TopK { k, input } => { + plan.precompute.topk = Some(*k); + plan.precompute.active = true; + walk(input, plan, budgets); + } + + // JoinSketch — outer and inner both walked; join itself stays at Backend. + SketchExpr::JoinSketch { outer, inner, .. } => { + plan.backend.has_merge = true; + walk(outer, plan, budgets); + walk(inner, plan, budgets); + } + } +} + +// ── Agg node assignment ─────────────────────────────────────────────────────── + +fn assign_agg(op: &SketchAggOp, plan: &mut StagedPlan, budgets: &StageResourceBudgets) { + match op { + // ── Exact ops: mergeability decides stage ───────────────────────────── + // + // Sum/Count/Min/Max satisfy agg(A∪B) = merge(agg(A), agg(B)). + // The Backend collector can combine partial results from N agents. + SketchAggOp::Exact( + ExactAgg::Sum | ExactAgg::Count | ExactAgg::Min | ExactAgg::Max, + ) => { + plan.backend.has_merge = true; + } + + // Avg is NOT mergeable: avg(avg(A), avg(B)) ≠ avg(A∪B). + // Must see all raw data → DB-side exact query. + SketchAggOp::Exact(ExactAgg::Avg) => { + plan.db.active = true; + } + + // ── Sketch ops: assign to Agent, defer if budget exceeded ───────────── + sketch_op => { + let est_mem = estimated_sketch_memory_bytes(sketch_op); + let stage = resolve_sketch_stage(est_mem, budgets, &mut plan.deferral_log, sketch_op); + match stage { + SketchStage::Agent => { + plan.agent.sketch_type = Some(agg_op_to_sketch_type(sketch_op)); + plan.agent.sketch_params = agg_op_to_sketch_params(sketch_op); + } + SketchStage::Backend => { + // Deferred from Agent: Backend does the sketch insertion. + plan.backend.has_merge = true; + // Record sketch type on agent as None (passthrough raw). + } + SketchStage::Precompute => { + plan.precompute.active = true; + } + } + } + } +} + +#[derive(Debug, PartialEq)] +enum SketchStage { + Agent, + Backend, + Precompute, +} + +/// Walk the deferral chain Agent → Backend → Precompute based on budget caps. +fn resolve_sketch_stage( + est_mem: u64, + budgets: &StageResourceBudgets, + log: &mut Vec, + op: &SketchAggOp, +) -> SketchStage { + // Try Agent first. + if let Some(cap) = budgets.agent_memory_bytes { + if est_mem > cap { + log.push(format!( + "deferred {op:?} Agent→Backend: est_mem={est_mem}B > agent_cap={cap}B" + )); + // Try Backend. + if let Some(be_cap) = budgets.backend_memory_bytes { + if est_mem > be_cap { + log.push(format!( + "deferred {op:?} Backend→Precompute: est_mem={est_mem}B > backend_cap={be_cap}B" + )); + return SketchStage::Precompute; + } + } + return SketchStage::Backend; + } + } + SketchStage::Agent +} + +// ── Cost estimation (heuristic, based on benchmark table in cost_model.rs) ──── + +/// Estimate peak sketch memory per series (bytes) for the given operator. +fn estimated_sketch_memory_bytes(op: &SketchAggOp) -> u64 { + match op { + SketchAggOp::DDSketch { .. } | SketchAggOp::ExactMinMax { .. } => 4_096, + SketchAggOp::HLL { registers } => 1u64 << (*registers as u64), + SketchAggOp::CountMin { width, depth } => (*width as u64) * (*depth as u64) * 8, + SketchAggOp::CountSketch { k } => k * 8 * 5, // 5-row hash table approx + SketchAggOp::Hydra { inner, partition_keys } => { + // Hydra maintains one inner sketch per distinct key-tuple. + // Estimate 2^|keys| partitions (capped at 1024 to avoid absurd values). + let factor = 1u64 << partition_keys.len().min(10); + estimated_sketch_memory_bytes(inner).saturating_mul(factor) + } + SketchAggOp::Exact(_) => 8, + } +} + +// ── Sketch-type helpers ─────────────────────────────────────────────────────── + +fn agg_op_to_sketch_type(op: &SketchAggOp) -> SketchType { + match op { + SketchAggOp::DDSketch { .. } | SketchAggOp::ExactMinMax { .. } => SketchType::DDSketch, + SketchAggOp::HLL { .. } => SketchType::HLL, + SketchAggOp::CountMin { .. } => SketchType::CountMinSketch, + SketchAggOp::CountSketch { .. } => SketchType::CountSketch, + SketchAggOp::Hydra { inner, .. } => agg_op_to_sketch_type(inner), + SketchAggOp::Exact(_) => SketchType::DDSketch, // unreachable for sketch stage + } +} + +fn agg_op_to_sketch_params(op: &SketchAggOp) -> SketchParams { + match op { + SketchAggOp::DDSketch { quantiles, epsilon } => SketchParams { + relative_accuracy: *epsilon, + quantiles: quantiles.clone(), + ..Default::default() + }, + SketchAggOp::HLL { registers } => SketchParams { + precision: *registers as u32, + ..Default::default() + }, + SketchAggOp::CountMin { width, depth } => SketchParams { + rows: *depth as u32, + cols: *width, + ..Default::default() + }, + SketchAggOp::CountSketch { .. } => SketchParams { + rows: 5, + cols: 2_048, + ..Default::default() + }, + SketchAggOp::Hydra { inner, .. } => agg_op_to_sketch_params(inner), + SketchAggOp::ExactMinMax { .. } => SketchParams { + relative_accuracy: 0.01, + quantiles: vec![0.0, 1.0], + ..Default::default() + }, + SketchAggOp::Exact(_) => SketchParams::default(), + } +} + +// ── PromQL serialiser ───────────────────────────────────────────────────────── + +/// Context accumulated while walking the tree for PromQL serialisation. +#[derive(Default)] +struct QueryCtx { + metric: String, + filters: Vec, + window: Option, + group_by: Vec, + agg: Option, + topk: Option, +} + +#[derive(Clone)] +enum AggInfo { + DDSketch { phi: f64 }, + HLL, + CountMin, + CountSketch, + ExactCount, + ExactSum, + ExactAvg, + ExactMin, + ExactMax, + ExactMinMax, +} + +impl QueryCtx { + fn to_promql(&self) -> String { + let selector = if self.filters.is_empty() { + self.metric.clone() + } else { + format!("{}{{{}}}", self.metric, self.filters.join(", ")) + }; + + let window = self + .window + .map(|w| format!("[{}]", format_duration(w))) + .unwrap_or_default(); + + let by_clause = if self.group_by.is_empty() { + String::new() + } else { + format!(" by ({})", self.group_by.join(", ")) + }; + + let inner = match &self.agg { + Some(AggInfo::DDSketch { phi }) => { + format!("quantile_over_time({phi}, {selector}{window}){by_clause}") + } + Some(AggInfo::HLL) => { + format!("count_distinct_over_time({selector}{window}){by_clause}") + } + Some(AggInfo::CountMin | AggInfo::CountSketch | AggInfo::ExactCount) => { + format!("count_over_time({selector}{window}){by_clause}") + } + Some(AggInfo::ExactSum) => { + format!("sum_over_time({selector}{window}){by_clause}") + } + Some(AggInfo::ExactAvg) => { + format!("avg_over_time({selector}{window}){by_clause}") + } + Some(AggInfo::ExactMin) => { + format!("min_over_time({selector}{window}){by_clause}") + } + Some(AggInfo::ExactMax) => { + format!("max_over_time({selector}{window}){by_clause}") + } + Some(AggInfo::ExactMinMax) => { + // Represent as quantile range [0,1] for min/max via DDSketch. + format!("quantile_over_time(0.5, {selector}{window}){by_clause}") + } + None => format!("{selector}{window}{by_clause}"), + }; + + if let Some(k) = self.topk { + format!("topk({k}, {inner})") + } else { + inner + } + } +} + +/// Walk the tree and fill [`QueryCtx`]; string building is deferred to +/// [`QueryCtx::to_promql`]. +fn collect_ctx(expr: &SketchExpr, ctx: &mut QueryCtx) { + match expr { + SketchExpr::Source(s) => { + if ctx.metric.is_empty() { + ctx.metric = s.name.clone(); + } + } + SketchExpr::Filter { pred, input } => { + for p in pred { + match (&p.op, &p.val) { + (FilterOp::Eq, FilterVal::Str(v)) => { + ctx.filters.push(format!("{}=\"{}\"", p.col, v)); + } + (FilterOp::Ne, FilterVal::Str(v)) => { + ctx.filters.push(format!("{}!=\"{}\"", p.col, v)); + } + (FilterOp::Regex(re), _) => { + ctx.filters.push(format!("{}=~\"{}\"", p.col, re)); + } + (FilterOp::NotRegex(re), _) => { + ctx.filters.push(format!("{}!~\"{}\"", p.col, re)); + } + _ => {} // other ops (Lt/Gt/…) not natively supported in PromQL label matchers + } + } + collect_ctx(input, ctx); + } + SketchExpr::Window { duration, input } => { + if ctx.window.is_none() { + ctx.window = Some(*duration); + } + collect_ctx(input, ctx); + } + SketchExpr::Partition { keys, input } => { + for k in keys.keys() { + if !ctx.group_by.contains(k) { + ctx.group_by.push(k.clone()); + } + } + collect_ctx(input, ctx); + } + SketchExpr::Agg { op, input, .. } => { + if ctx.agg.is_none() { + ctx.agg = Some(agg_info(op)); + } + collect_ctx(input, ctx); + } + SketchExpr::TopK { k, input } => { + ctx.topk = Some(*k); + collect_ctx(input, ctx); + } + SketchExpr::Dedup { input, .. } => collect_ctx(input, ctx), + SketchExpr::Merge { inputs } => { + // All branches should share the same shape; serialise the first. + if let Some(first) = inputs.first() { + collect_ctx(first, ctx); + } + } + SketchExpr::JoinSketch { outer, .. } => collect_ctx(outer, ctx), + } +} + +fn agg_info(op: &SketchAggOp) -> AggInfo { + match op { + SketchAggOp::DDSketch { quantiles, .. } => AggInfo::DDSketch { + phi: quantiles.first().copied().unwrap_or(0.99), + }, + SketchAggOp::HLL { .. } => AggInfo::HLL, + SketchAggOp::CountMin { .. } => AggInfo::CountMin, + SketchAggOp::CountSketch { .. } => AggInfo::CountSketch, + SketchAggOp::ExactMinMax { .. } => AggInfo::ExactMinMax, + SketchAggOp::Exact(ExactAgg::Count) => AggInfo::ExactCount, + SketchAggOp::Exact(ExactAgg::Sum) => AggInfo::ExactSum, + SketchAggOp::Exact(ExactAgg::Avg) => AggInfo::ExactAvg, + SketchAggOp::Exact(ExactAgg::Min) => AggInfo::ExactMin, + SketchAggOp::Exact(ExactAgg::Max) => AggInfo::ExactMax, + SketchAggOp::Hydra { inner, .. } => agg_info(inner), + } +} + +// ── Tests ───────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use crate::query_parser::sketch_algebra::{ + ColumnRef, FilterVal, Predicate, SketchExpr, SketchAggOp, SourceSpec, + }; + + fn source(name: &str) -> SketchExpr { + SketchExpr::Source(SourceSpec { name: name.into() }) + } + + fn no_budget() -> StageResourceBudgets { + StageResourceBudgets::default() + } + + // ── Node-to-stage assignment ────────────────────────────────────────────── + + /// Source + Filter + Window + Agg(DDSketch) → all at Agent. + #[test] + fn ddsketch_agg_goes_to_agent() { + let expr = SketchExpr::Window { + duration: Duration::from_secs(300), + input: Box::new(SketchExpr::Agg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, + input: Box::new(source("latency")), + }), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert_eq!(plan.agent.sketch_type, Some(SketchType::DDSketch)); + assert_eq!(plan.agent.window_secs, Some(300)); + assert!(!plan.precompute.active); + assert!(!plan.db.active); + } + + /// HLL stays at Agent by default. + #[test] + fn hll_agg_goes_to_agent() { + let expr = SketchExpr::Agg { + op: SketchAggOp::default_hll(), + col: ColumnRef::Wildcard, + input: Box::new(source("requests")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert_eq!(plan.agent.sketch_type, Some(SketchType::HLL)); + } + + /// Partition keys go to Backend regardless. + #[test] + fn partition_goes_to_backend() { + let expr = SketchExpr::Partition { + keys: PartitionKeys::By(vec!["host".into(), "region".into()]), + input: Box::new(SketchExpr::Agg { + op: SketchAggOp::default_ddsketch(vec![0.5]), + col: ColumnRef::SampleValue, + input: Box::new(source("cpu")), + }), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert!(plan.backend.group_by.contains(&"host".to_string())); + assert!(plan.backend.group_by.contains(&"region".to_string())); + } + + /// TopK → Precompute Engine. + #[test] + fn topk_goes_to_precompute() { + let expr = SketchExpr::TopK { + k: 10, + input: Box::new(SketchExpr::Agg { + op: SketchAggOp::CountSketch { k: 10 }, + col: ColumnRef::Wildcard, + input: Box::new(source("events")), + }), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert!(plan.precompute.active); + assert_eq!(plan.precompute.topk, Some(10)); + } + + /// Exact(Sum) is mergeable → Backend. + #[test] + fn exact_sum_goes_to_backend() { + let expr = SketchExpr::Agg { + op: SketchAggOp::Exact(ExactAgg::Sum), + col: ColumnRef::Named("bytes".into()), + input: Box::new(source("network")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + // Sum is handled at backend (merge); agent gets no sketch type. + assert!(plan.agent.sketch_type.is_none()); + assert!(!plan.db.active); + assert!(plan.backend.has_merge); + } + + /// Exact(Count) is mergeable → Backend. + #[test] + fn exact_count_goes_to_backend() { + let expr = SketchExpr::Agg { + op: SketchAggOp::Exact(ExactAgg::Count), + col: ColumnRef::Wildcard, + input: Box::new(source("hits")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert!(plan.agent.sketch_type.is_none()); + assert!(plan.backend.has_merge); + assert!(!plan.db.active); + } + + /// Exact(Avg) is NOT mergeable → DB stage. + #[test] + fn exact_avg_goes_to_db() { + let expr = SketchExpr::Agg { + op: SketchAggOp::Exact(ExactAgg::Avg), + col: ColumnRef::Named("price".into()), + input: Box::new(source("ticks")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert!(plan.db.active); + assert!(!plan.precompute.active); + } + + /// Exact(Min) is mergeable → Backend (min(A∪B) = min(min(A), min(B))). + #[test] + fn exact_min_goes_to_backend() { + let expr = SketchExpr::Agg { + op: SketchAggOp::Exact(ExactAgg::Min), + col: ColumnRef::Named("latency".into()), + input: Box::new(source("svc")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert!(plan.backend.has_merge); + assert!(!plan.db.active); + } + + /// Exact(Max) is mergeable → Backend. + #[test] + fn exact_max_goes_to_backend() { + let expr = SketchExpr::Agg { + op: SketchAggOp::Exact(ExactAgg::Max), + col: ColumnRef::Named("latency".into()), + input: Box::new(source("svc")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert!(plan.backend.has_merge); + assert!(!plan.db.active); + } + + /// Dedup node → Backend. + #[test] + fn dedup_goes_to_backend() { + let expr = SketchExpr::Dedup { + col: "user_id".into(), + input: Box::new(SketchExpr::Agg { + op: SketchAggOp::default_hll(), + col: ColumnRef::Named("user_id".into()), + input: Box::new(source("events")), + }), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert!(plan.backend.has_dedup); + } + + // ── Budget-driven deferral ──────────────────────────────────────────────── + + /// When the DDSketch memory (4 096 B) exceeds the agent cap, it defers + /// to Backend. + #[test] + fn ddsketch_deferred_to_backend_when_agent_budget_exceeded() { + let budgets = StageResourceBudgets { + agent_memory_bytes: Some(1_024), // tighter than 4 096 + ..Default::default() + }; + let expr = SketchExpr::Agg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, + input: Box::new(source("latency")), + }; + let plan = split_expr_by_stage(&expr, &budgets); + // Deferred: no sketch at Agent. + assert!(plan.agent.sketch_type.is_none()); + // Backend takes over (merge path). + assert!(plan.backend.has_merge); + // Deferral logged. + assert!(!plan.deferral_log.is_empty()); + assert!(plan.deferral_log[0].contains("Agent→Backend")); + } + + /// When both agent AND backend caps are exceeded, defers to Precompute. + #[test] + fn ddsketch_deferred_to_precompute_when_both_budgets_exceeded() { + let budgets = StageResourceBudgets { + agent_memory_bytes: Some(1_024), + backend_memory_bytes: Some(1_024), + ..Default::default() + }; + let expr = SketchExpr::Agg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, + input: Box::new(source("latency")), + }; + let plan = split_expr_by_stage(&expr, &budgets); + assert!(plan.precompute.active); + assert!(plan.deferral_log.iter().any(|l| l.contains("Backend→Precompute"))); + } + + /// No budget set → Agent assignment regardless of sketch size. + #[test] + fn no_budget_always_assigns_to_agent() { + let expr = SketchExpr::Agg { + op: SketchAggOp::CountMin { width: 2_000_000, depth: 255 }, // huge + col: ColumnRef::Wildcard, + input: Box::new(source("events")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert_eq!(plan.agent.sketch_type, Some(SketchType::CountMinSketch)); + assert!(plan.deferral_log.is_empty()); + } + + // ── PromQL serialisation ────────────────────────────────────────────────── + + /// TopK(CountSketch) over filtered, windowed, partitioned metric. + #[test] + fn topk_countsketch_promql() { + let expr = SketchExpr::TopK { + k: 10, + input: Box::new(SketchExpr::Partition { + keys: PartitionKeys::By(vec!["symbol".into()]), + input: Box::new(SketchExpr::Window { + duration: Duration::from_secs(300), + input: Box::new(SketchExpr::Filter { + pred: vec![Predicate { + col: "sectype".into(), + op: FilterOp::Eq, + val: FilterVal::Str("E".into()), + }], + input: Box::new(SketchExpr::Agg { + op: SketchAggOp::CountSketch { k: 10 }, + col: ColumnRef::Wildcard, + input: Box::new(source("financial.last_trade_price")), + }), + }), + }), + }), + }; + let ql = expr_to_promql(&expr); + assert!(ql.contains("topk(10,"), "missing topk: {ql}"); + assert!(ql.contains("count_over_time"), "missing count_over_time: {ql}"); + assert!(ql.contains("financial.last_trade_price"), "missing metric: {ql}"); + assert!(ql.contains("sectype=\"E\""), "missing filter: {ql}"); + assert!(ql.contains("[5m]"), "missing window: {ql}"); + assert!(ql.contains("by (symbol)"), "missing group_by: {ql}"); + } + + /// DDSketch quantile query. + #[test] + fn ddsketch_promql() { + let expr = SketchExpr::Agg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, + input: Box::new(SketchExpr::Window { + duration: Duration::from_secs(300), + input: Box::new(source("latency")), + }), + }; + let ql = expr_to_promql(&expr); + assert!(ql.contains("quantile_over_time(0.99,"), "missing quantile: {ql}"); + assert!(ql.contains("latency"), "missing metric: {ql}"); + assert!(ql.contains("[5m]"), "missing window: {ql}"); + } + + /// HLL cardinality query. + #[test] + fn hll_promql() { + let expr = SketchExpr::Agg { + op: SketchAggOp::default_hll(), + col: ColumnRef::Wildcard, + input: Box::new(source("users")), + }; + let ql = expr_to_promql(&expr); + assert!(ql.contains("count_distinct_over_time"), "missing fn: {ql}"); + assert!(ql.contains("users"), "missing metric: {ql}"); + } + + /// Exact(Avg) produces avg_over_time. + #[test] + fn exact_avg_promql() { + let expr = SketchExpr::Agg { + op: SketchAggOp::Exact(ExactAgg::Avg), + col: ColumnRef::Named("price".into()), + input: Box::new(source("ticks")), + }; + let ql = expr_to_promql(&expr); + assert!(ql.contains("avg_over_time"), "missing avg: {ql}"); + } + + /// Exact(Sum) produces sum_over_time. + #[test] + fn exact_sum_promql() { + let expr = SketchExpr::Agg { + op: SketchAggOp::Exact(ExactAgg::Sum), + col: ColumnRef::Named("bytes".into()), + input: Box::new(source("net")), + }; + let ql = expr_to_promql(&expr); + assert!(ql.contains("sum_over_time"), "missing sum: {ql}"); + } + + // ── Sketch parameter extraction ─────────────────────────────────────────── + + #[test] + fn ddsketch_params_extracted_correctly() { + let expr = SketchExpr::Agg { + op: SketchAggOp::DDSketch { quantiles: vec![0.5, 0.99], epsilon: 0.005 }, + col: ColumnRef::SampleValue, + input: Box::new(source("latency")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + let p = &plan.agent.sketch_params; + assert_eq!(p.relative_accuracy, 0.005); + assert_eq!(p.quantiles, vec![0.5, 0.99]); + } + + #[test] + fn hll_precision_extracted() { + let expr = SketchExpr::Agg { + op: SketchAggOp::HLL { registers: 14 }, + col: ColumnRef::Wildcard, + input: Box::new(source("users")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert_eq!(plan.agent.sketch_params.precision, 14); + } + + #[test] + fn countmin_params_extracted() { + let expr = SketchExpr::Agg { + op: SketchAggOp::CountMin { width: 2_000, depth: 5 }, + col: ColumnRef::Wildcard, + input: Box::new(source("events")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert_eq!(plan.agent.sketch_params.rows, 5); + assert_eq!(plan.agent.sketch_params.cols, 2_000); + } + + // ── Filter predicate extraction ─────────────────────────────────────────── + + #[test] + fn filter_eq_captured_in_agent_plan() { + let expr = SketchExpr::Filter { + pred: vec![Predicate { + col: "env".into(), + op: FilterOp::Eq, + val: FilterVal::Str("prod".into()), + }], + input: Box::new(source("latency")), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert!(plan.agent.label_filters.contains(&"env=prod".to_string())); + } + + // ── Precompute query_expr generation ───────────────────────────────────── + + /// When TopK is present, `staged_plan.precompute.query_expr` must be a + /// non-empty PromQL string. + #[test] + fn precompute_query_expr_set_when_topk_present() { + let expr = SketchExpr::TopK { + k: 5, + input: Box::new(SketchExpr::Agg { + op: SketchAggOp::default_hll(), + col: ColumnRef::Wildcard, + input: Box::new(source("sessions")), + }), + }; + let plan = split_expr_by_stage(&expr, &no_budget()); + assert!(plan.precompute.active); + assert!(!plan.precompute.query_expr.is_empty()); + assert!(plan.precompute.query_expr.contains("topk(5,")); + } + + // ── StageResourceBudgets::from_workload_chars ───────────────────────────── + + #[test] + fn budgets_from_workload_chars_propagates_memory() { + use crate::types::WorkloadCharacteristics; + let wc = WorkloadCharacteristics { + memory_budget_bytes: Some(8_192), + ..Default::default() + }; + let b = StageResourceBudgets::from_workload_chars(&wc); + assert_eq!(b.agent_memory_bytes, Some(8_192)); + assert!(b.backend_memory_bytes.is_none()); + } +} diff --git a/controller/src/replan.rs b/controller/src/replan.rs index 079ed9b9..0054a660 100644 --- a/controller/src/replan.rs +++ b/controller/src/replan.rs @@ -234,6 +234,7 @@ mod tests { valid_until: Utc::now() + chrono::Duration::seconds(3600), delta_decision: DeltaDecision::default(), transmission_cost_summary: TransmissionCostSummary::default(), + staged_plan: None, } } diff --git a/controller/src/store/mod.rs b/controller/src/store/mod.rs index 1b8560d5..ba6ddceb 100644 --- a/controller/src/store/mod.rs +++ b/controller/src/store/mod.rs @@ -177,6 +177,7 @@ mod tests { valid_until, delta_decision: DeltaDecision::default(), transmission_cost_summary: TransmissionCostSummary::default(), + staged_plan: None, } } diff --git a/controller/src/types.rs b/controller/src/types.rs index 32ddad20..019861b3 100644 --- a/controller/src/types.rs +++ b/controller/src/types.rs @@ -313,6 +313,124 @@ pub struct PrecomputeJob { pub store_path: String, } +// ── SP-9: per-stage resource budgets ───────────────────────────────────────── + +/// Per-stage resource caps used by `split_expr_by_stage()` (SP-9). +/// +/// When a node's estimated memory cost exceeds the cap at its natural stage, +/// it is deferred to the next stage in the pipeline: +/// +/// `Agent OTel Collector → Backend OTel Collector → Precompute Engine` +/// +/// `None` means unbounded (no cap enforced). Typically sourced from +/// [`WorkloadCharacteristics::memory_budget_bytes`] for the agent stage and +/// from a runtime config file for backend / precompute stages. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct StageResourceBudgets { + /// Max sketch memory at the agent OTel Collector (bytes). + pub agent_memory_bytes: Option, + /// Max sketch-insertion CPU budget at the agent (µs/sample). + pub agent_cpu_micros_per_sample: Option, + /// Max sketch memory at the backend OTel Collector (bytes). + pub backend_memory_bytes: Option, + /// Max memory at the ASAPQuery Precompute Engine (bytes). + pub precompute_memory_bytes: Option, +} + +impl StageResourceBudgets { + /// Derive budgets from [`WorkloadCharacteristics`]: propagates the agent + /// memory cap; other stages default to unbounded. + pub fn from_workload_chars(wc: &WorkloadCharacteristics) -> Self { + Self { + agent_memory_bytes: wc.memory_budget_bytes, + ..Default::default() + } + } +} + +// ── SP-9: per-stage sub-plans ───────────────────────────────────────────────── + +/// Sub-plan for the **Agent OTel Collector** stage. +/// +/// Covers `SketchExpr` nodes: `Source`, `Filter`, `Window`, `Agg` (sketch ops). +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct AgentSubPlan { + /// Concrete sketch type resolved from the `Agg` node (absent when no Agg + /// node was assigned to this stage, e.g. all deferred to Backend). + pub sketch_type: Option, + pub sketch_params: SketchParams, + /// Time window in seconds (from the `Window` node). + pub window_secs: Option, + /// Partition / group-by dimensions if a `Partition` node was pushed down + /// to the agent stage. + pub aggregate_by: Vec, + /// Label-filter predicates as `"key=value"` strings (from `Filter` nodes). + pub label_filters: Vec, + /// True when a `Dedup` node was assigned to this stage. + pub has_dedup: bool, +} + +/// Sub-plan for the **Backend OTel Collector** stage. +/// +/// Covers `SketchExpr` nodes: `Partition`, `Merge`, `Dedup`, and +/// `Agg { Exact(Sum|Count|Min|Max) }` (mergeable exact ops). +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct BackendSubPlan { + /// GROUP BY dimensions from the `Partition` node. + pub group_by: Vec, + /// True when a `Dedup` node was assigned here. + pub has_dedup: bool, + /// True when a `Merge` node is at this stage (expected for all + /// multi-agent deployments). + pub has_merge: bool, +} + +/// Sub-plan for the **ASAPQuery Precompute Engine** stage. +/// +/// Covers `SketchExpr` nodes: `TopK`, and sketch `Agg` ops deferred from +/// the Agent stage due to memory budget overflow. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct PrecomputeSubPlan { + /// Top-K value when a `TopK` node was assigned to this stage. + pub topk: Option, + /// PromQL/SQL expression representing the upper sub-tree assigned here. + /// Empty string when no precompute operations are present. + pub query_expr: String, + /// True when at least one operation was assigned to this stage. + pub active: bool, +} + +/// Sub-plan for **DB-side exact computation** (ClickHouse / TSDB). +/// +/// Covers `Agg { Exact(Avg) }` — non-mergeable; cannot be precomputed across +/// distributed agents without collecting all raw data first. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct DbSubPlan { + /// PromQL/SQL expression for the exact DB-side query. + pub query_expr: String, + /// True when at least one operation was assigned to this stage. + pub active: bool, +} + +/// Result of SP-9 AST-aware stage split. +/// +/// Produced by `planner::stage_split::split_expr_by_stage()`. Attached to +/// [`CollectionPlan::staged_plan`] when the workload was supplied via +/// `query_string` (giving access to the full `SketchExpr` tree). +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct StagedPlan { + pub agent: AgentSubPlan, + pub backend: BackendSubPlan, + pub precompute: PrecomputeSubPlan, + pub db: DbSubPlan, + /// Human-readable log of deferral decisions made during the split + /// (e.g. a sketch op moved from Agent to Backend due to a memory cap). + /// Populated for observability / debugging. + pub deferral_log: Vec, +} + +// ── Collection plan ─────────────────────────────────────────────────────────── + #[derive(Debug, Clone)] pub struct CollectionPlan { pub agent_config: AgentCollectorConfig, @@ -324,4 +442,10 @@ pub struct CollectionPlan { pub delta_decision: DeltaDecision, /// Bandwidth and overhead estimates for all three transmission strategies. pub transmission_cost_summary: TransmissionCostSummary, + /// SP-9: AST-aware per-stage sub-plans. + /// + /// `Some` when the workload was supplied via `query_string` (full + /// `SketchExpr` tree available). `None` when built from explicit + /// aggregation fields — the SP-3 flat assignment is used as fallback. + pub staged_plan: Option, } From b58eb78fd4b813a3959d31a0f079d735752a97a0 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 1 Apr 2026 11:40:28 -0500 Subject: [PATCH 2/3] =?UTF-8?q?feat(SP-9):=20Option=20B=20=E2=80=94=20port?= =?UTF-8?q?=20stage=5Fsplit=20to=20QueryExpr,=20recursive=20expr=5Fto=5Fpr?= =?UTF-8?q?omql?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - split_expr_by_stage now takes &QueryExpr instead of &SketchExpr. All SketchExpr-era node assignments are preserved; new QueryExpr-only nodes (HistogramQuantile, PromQLSubquery, BinaryOp, Aggregate, Sort+Limit, Join, SetOp) are assigned to the correct stage. - expr_to_promql replaced with a proper recursive descent over QueryExpr. Handles: histogram_quantile(φ, rate(…[w])), expr[range:step] subqueries, (lhs op rhs) vector binary ops with on(…)/ignoring(…)/group_left/group_right, Aggregate GROUP BY → by (keys), Sort+Limit → topk(n, …). - Label-filter extraction ported from Predicate list to ScalarExpr AND-tree: col="val", col!="val", col=~"re", col!~"re". - handle_plan unified to a single pipeline: parse_query_sketch → QueryExpr::from_sketch_expr → QueryOptimizer → split_expr_by_stage(&QueryExpr) → StagedPlan. SketchExpr path and analyze_with_sketch() removed from the hot path. - StageResourceBudgets duplicate definition resolved (kept SP-9 variant). - 12 new tests in stage_split covering all new node types, budget deferral, and expr_to_promql output. Total: 303 tests pass. Co-Authored-By: Claude Sonnet 4.6 --- controller/src/main.rs | 32 +- controller/src/planner/stage_split.rs | 1122 ++++++++++++++----------- controller/src/types.rs | 27 - 3 files changed, 642 insertions(+), 539 deletions(-) diff --git a/controller/src/main.rs b/controller/src/main.rs index 1812ee29..12237f2e 100644 --- a/controller/src/main.rs +++ b/controller/src/main.rs @@ -21,6 +21,7 @@ use axum::{ use serde_json::json; use tracing::{info, warn}; +use algebra::{QueryExpr, QueryOptimizer}; use analyzer::{Analyzer, QuerySpec}; use config::{generate_agent_config, generate_backend_config, build_precompute_jobs}; use config::generate_backend_config_staged; @@ -28,8 +29,11 @@ use monitor::{Endpoint, Scraper, ScrapedData, Thresholds, Violation}; use opamp::{AgentRole, OpampServer, RemoteConfig}; use planner::{CostModelPlanner, BaselinePlanner, ObjectiveWeights, OnlineMetricsStore, init_online_store, pareto_frontier, select_best}; use planner::online_cost_model; +use planner::stage_split::split_expr_by_stage; +use query_parser::parse_query_sketch; use replan::Replanner; use store::{PlanStore, WorkloadStore}; +use types::StageResourceBudgets; // ── Shared state ────────────────────────────────────────────────────────────── @@ -217,26 +221,34 @@ async fn handle_plan( State(st): State, Json(spec): Json, ) -> impl IntoResponse { - let wc = spec.workload.clone(); - let (workload, sketch_expr) = match st.analyzer.analyze_with_sketch(spec) { + let wc = spec.workload.clone(); + let query_string = spec.query_string.clone(); + let workload = match st.analyzer.analyze(spec) { Ok(w) => w, Err(e) => return (StatusCode::UNPROCESSABLE_ENTITY, e.to_string()).into_response(), }; let mut plan = st.planner.plan(&workload, Some(&wc)); - // ── SP-9: AST-aware hierarchical stage assignment ───────────────────────── - // When a SketchExpr tree is available (query_string path), split it across - // stages and attach the result to the plan. The SP-3 flat assignment - // remains active as the fallback when no tree is present. - if let Some(ref expr) = sketch_expr { - let budgets = types::StageResourceBudgets::from_workload_chars(&wc); - plan.staged_plan = Some(planner::stage_split::split_expr_by_stage(expr, &budgets)); + // ── SP-9: single QueryExpr pipeline — parse → optimise → stage-split ───── + // When query_string is present, run the full algebra pipeline and attach + // the StagedPlan. The SP-3 flat assignment remains the fallback when no + // query_string is supplied. + if let Some(ref qs) = query_string { + match parse_query_sketch(qs) { + Err(e) => warn!(query = %qs, error = %e, "parse_query_sketch failed; skipping staged_plan"), + Ok(sketch_expr) => { + let qe = QueryExpr::from_sketch_expr(&sketch_expr); + let raw_bps = plan.transmission_cost_summary.raw_bytes_per_sec; + let (opt_qe, _) = QueryOptimizer::new(raw_bps).optimize(qe); + let budgets = StageResourceBudgets::from_workload_chars(&wc); + plan.staged_plan = Some(split_expr_by_stage(&opt_qe, &budgets)); + } + } } plan.precompute = build_precompute_jobs(&workload, &plan, "backend:4317"); st.store.set(&workload.metric_name, plan.clone()); - // Persist workload so the replanner can re-run plan() without the original spec. st.workload_store.set(&workload.metric_name, workload.clone(), wc); // ── Push agent config to agent-role collectors ──────────────────────────── diff --git a/controller/src/planner/stage_split.rs b/controller/src/planner/stage_split.rs index ab635297..9eae63aa 100644 --- a/controller/src/planner/stage_split.rs +++ b/controller/src/planner/stage_split.rs @@ -1,27 +1,25 @@ //! SP-9: AST-aware hierarchical stage assignment. //! -//! Splits the optimized [`SketchExpr`] tree produced by the query parser -//! across pipeline stages, emitting a [`StagedPlan`] that carries per-stage -//! sub-plans for: +//! Splits an optimised [`QueryExpr`] tree across pipeline stages, emitting a +//! [`StagedPlan`] that carries per-stage sub-plans for: //! //! | Stage | Nodes | //! |---|---| -//! | Agent OTel Collector | `Source`, `Filter`, `Window`, `Agg` (sketch ops) | -//! | Backend OTel Collector | `Partition`, `Merge`, `Dedup`, `Agg { Exact(Sum\|Count\|Min\|Max) }` | -//! | ASAPQuery Precompute Engine | `TopK`, deferred sketch `Agg` ops | -//! | DB-side query | `Agg { Exact(Avg) }` (non-mergeable) | +//! | Agent OTel Collector | `Source`, `Filter`, `Window`, `SketchAgg` (sketch ops) | +//! | Backend OTel Collector | `Partition`, `Merge`, `Dedup`, `Aggregate { Exact(Sum\|Count\|Min\|Max) }` | +//! | ASAPQuery Precompute Engine | `TopK`, `HistogramQuantile`, `PromQLSubquery`, deferred sketch ops | +//! | DB-side query | `Aggregate { Avg }` (non-mergeable) | //! -//! # ExactAgg deferral +//! # ExactAgg / AggFunc deferral //! //! Mergeability drives `Exact` op placement: //! - `Sum`, `Count`, `Min`, `Max` — mergeable (`agg(A∪B) = merge(agg(A), agg(B))`) → //! **Backend** (the backend collector can combine partial results from N agents). -//! - `Avg` — **not** mergeable (avg-of-avgs ≠ global avg) → **DB-side** query -//! (must see all data before computing). +//! - `Avg`, `StdDev`, `Variance` — **not** mergeable → **DB-side** query. //! //! # Budget-driven deferral chain //! -//! When a sketch `Agg` node's estimated memory cost exceeds the stage cap in +//! When a `SketchAgg` node's estimated memory cost exceeds the stage cap in //! [`StageResourceBudgets`], it is deferred to the next stage: //! //! `Agent → Backend → Precompute` @@ -32,17 +30,16 @@ //! //! # PromQL serialisation //! -//! [`expr_to_promql`] converts the full `SketchExpr` tree to a PromQL -//! expression consumed by the ASAPQuery Precompute Engine's query engine. -//! The sketch data is already ingested by the precompute engine from the -//! backend OTel Collector; the PromQL describes what operation to apply. +//! [`expr_to_promql`] converts a [`QueryExpr`] tree to a valid PromQL expression +//! consumed by the ASAPQuery Precompute Engine's query engine. Unlike the old +//! flat-template approach, this uses a proper recursive descent so it handles +//! `histogram_quantile`, `PromQLSubquery`, and vector `BinaryOp` nodes natively. use std::time::Duration; +use crate::algebra::expr::{AggFunc, BinaryOpKind, LiteralValue, QueryExpr, ScalarExpr}; use crate::analyzer::format_duration; -use crate::query_parser::sketch_algebra::{ - ExactAgg, FilterOp, FilterVal, PartitionKeys, SketchAggOp, SketchExpr, -}; +use crate::query_parser::sketch_algebra::{ExactAgg, PartitionKeys, SketchAggOp}; use crate::types::{ AgentSubPlan, BackendSubPlan, DbSubPlan, PrecomputeSubPlan, SketchParams, SketchType, StagedPlan, StageResourceBudgets, @@ -50,23 +47,23 @@ use crate::types::{ // ── Public entry point ──────────────────────────────────────────────────────── -/// Split a `SketchExpr` tree across pipeline stages, respecting per-stage -/// resource budgets and the `ExactAgg` mergeability rules described above. +/// Split a [`QueryExpr`] tree across pipeline stages, respecting per-stage +/// resource budgets and the `AggFunc` mergeability rules described above. /// /// The returned [`StagedPlan`] is attached to /// [`crate::types::CollectionPlan::staged_plan`] by the caller /// (`handle_plan` in `main.rs`). -pub fn split_expr_by_stage(expr: &SketchExpr, budgets: &StageResourceBudgets) -> StagedPlan { +pub fn split_expr_by_stage(expr: &QueryExpr, budgets: &StageResourceBudgets) -> StagedPlan { let mut plan = StagedPlan::default(); walk(expr, &mut plan, budgets); // Build the precompute query_expr from the full tree when the precompute - // stage is active (TopK or deferred sketch ops assigned there). - if plan.precompute.active { + // stage is active (TopK, HistogramQuantile, PromQLSubquery, or deferred ops). + if plan.precompute.active && plan.precompute.query_expr.is_empty() { plan.precompute.query_expr = expr_to_promql(expr); } - // Build the DB query_expr when Avg was deferred to the DB stage. + // Build the DB query_expr when a non-mergeable agg was assigned to DB stage. if plan.db.active && plan.db.query_expr.is_empty() { plan.db.query_expr = expr_to_promql(expr); } @@ -74,42 +71,47 @@ pub fn split_expr_by_stage(expr: &SketchExpr, budgets: &StageResourceBudgets) -> plan } -/// Serialise a [`SketchExpr`] tree to a valid PromQL expression. +/// Serialise a [`QueryExpr`] tree to a valid PromQL expression string. +/// +/// The output is consumed by the ASAPQuery Precompute Engine's query engine. +/// Sketch data is already ingested from the Backend OTel Collector; the PromQL +/// describes the aggregation to apply over it. /// -/// The output is consumed by the ASAPQuery Precompute Engine's query engine -/// (PromQL / SQL). Sketch data is already present in the engine from the -/// backend OTel Collector; the PromQL describes the aggregation to apply. -pub fn expr_to_promql(expr: &SketchExpr) -> String { - let mut ctx = QueryCtx::default(); - collect_ctx(expr, &mut ctx); - ctx.to_promql() +/// Uses recursive descent, so it handles `histogram_quantile`, +/// `PromQLSubquery`, and vector `BinaryOp` nodes that the old flat-template +/// could not represent. +pub fn expr_to_promql(expr: &QueryExpr) -> String { + let mut ctx = PromQLCtx::default(); + promql_from_qe(expr, &mut ctx) } // ── Tree walker ─────────────────────────────────────────────────────────────── -fn walk(expr: &SketchExpr, plan: &mut StagedPlan, budgets: &StageResourceBudgets) { +fn walk(expr: &QueryExpr, plan: &mut StagedPlan, budgets: &StageResourceBudgets) { match expr { - // Source — always at Agent; no additional plan fields needed. - SketchExpr::Source(_) => {} + // Source — always Agent; populate metric name. + QueryExpr::Source(_) => {} // Filter — push label predicates to Agent. - SketchExpr::Filter { pred, input } => { - for p in pred { - if let (FilterOp::Eq, FilterVal::Str(v)) = (&p.op, &p.val) { - plan.agent.label_filters.push(format!("{}={}", p.col, v)); - } - } + QueryExpr::Filter { pred, input } => { + collect_label_filters_into(pred, &mut plan.agent.label_filters); walk(input, plan, budgets); } // Window — time window lives at Agent. - SketchExpr::Window { duration, input } => { + QueryExpr::Window { duration, input, .. } => { plan.agent.window_secs = Some(duration.as_secs()); walk(input, plan, budgets); } - // Partition — group-by dims always go to Backend after agent merge. - SketchExpr::Partition { keys, input } => { + // SketchAgg — the key sketch assignment decision. + QueryExpr::SketchAgg { op, input, .. } => { + assign_sketch_agg(op, plan, budgets); + walk(input, plan, budgets); + } + + // Partition — GROUP BY / `by (dims)` always assigned to Backend. + QueryExpr::Partition { keys, input } => { for k in keys.keys() { if !plan.backend.group_by.contains(k) { plan.backend.group_by.push(k.clone()); @@ -118,75 +120,130 @@ fn walk(expr: &SketchExpr, plan: &mut StagedPlan, budgets: &StageResourceBudgets walk(input, plan, budgets); } - // Agg — the key assignment decision (see module docs). - SketchExpr::Agg { op, input, .. } => { - assign_agg(op, plan, budgets); + // Aggregate — SQL GROUP BY + agg functions. + // Mergeable aggs (Sum/Count/Min/Max) → Backend. + // Non-mergeable (Avg/StdDev/Variance) → Db. + // Sketchable (Quantile/CountDistinct/HeavyHitters) → treated as SketchAgg. + QueryExpr::Aggregate { keys, aggs, input, .. } => { + for k in keys { + if !plan.backend.group_by.contains(k) { + plan.backend.group_by.push(k.clone()); + } + } + for agg in aggs { + assign_agg_func(&agg.func, plan, budgets); + } walk(input, plan, budgets); } - // Dedup — absorbed at Backend (HLL dedup is eliminated by R6 upstream). - SketchExpr::Dedup { input, .. } => { + // Dedup — absorbed at Backend (HLL dedup elimination is upstream). + QueryExpr::Dedup { input, .. } => { plan.backend.has_dedup = true; walk(input, plan, budgets); } // Merge — Backend merges N agent sketches. - SketchExpr::Merge { inputs } => { + QueryExpr::Merge { inputs } => { plan.backend.has_merge = true; for i in inputs { walk(i, plan, budgets); } } - // TopK — always at the Precompute Engine (upper AST). - SketchExpr::TopK { k, input } => { + // TopK — always at the Precompute Engine. + QueryExpr::TopK { k, input, .. } => { plan.precompute.topk = Some(*k); plan.precompute.active = true; walk(input, plan, budgets); } - // JoinSketch — outer and inner both walked; join itself stays at Backend. - SketchExpr::JoinSketch { outer, inner, .. } => { + // Sort + Limit — maps to topk semantics at Precompute. + QueryExpr::Limit { n, input, .. } => { + plan.precompute.topk = Some(*n); + plan.precompute.active = true; + walk(input, plan, budgets); + } + QueryExpr::Sort { input, .. } => { + walk(input, plan, budgets); + } + + // HistogramQuantile — precompute engine applies it over ingested histograms. + QueryExpr::HistogramQuantile { input, .. } => { + plan.precompute.active = true; + walk(input, plan, budgets); + } + + // PromQLSubquery — precompute engine evaluates the sub-query. + QueryExpr::PromQLSubquery { input, .. } => { + plan.precompute.active = true; + walk(input, plan, budgets); + } + + // BinaryOp between two instant vectors — precompute evaluates. + QueryExpr::BinaryOp { lhs, rhs, .. } => { + plan.precompute.active = true; + walk(lhs, plan, budgets); + walk(rhs, plan, budgets); + } + + // JoinSketch — outer and inner both walked; join itself at Backend. + QueryExpr::JoinSketch { outer, inner, .. } => { plan.backend.has_merge = true; walk(outer, plan, budgets); walk(inner, plan, budgets); } + + // Join — Backend. + QueryExpr::Join { left, right, .. } => { + plan.backend.has_merge = true; + walk(left, plan, budgets); + walk(right, plan, budgets); + } + + // SetOp — treat as Backend merge. + QueryExpr::SetOp { left, right, .. } => { + plan.backend.has_merge = true; + walk(left, plan, budgets); + walk(right, plan, budgets); + } + + // Transparent / passthrough nodes — recurse into child. + QueryExpr::Project { input, .. } + | QueryExpr::Subquery { expr: input, .. } + | QueryExpr::WindowFunc { input, .. } => walk(input, plan, budgets), + + QueryExpr::LetBinding { expr, body, .. } => { + walk(expr, plan, budgets); + walk(body, plan, budgets); + } + + // Ref — nothing to assign (resolved externally). + QueryExpr::Ref(_) => {} } } -// ── Agg node assignment ─────────────────────────────────────────────────────── +// ── Agg assignment ──────────────────────────────────────────────────────────── -fn assign_agg(op: &SketchAggOp, plan: &mut StagedPlan, budgets: &StageResourceBudgets) { +fn assign_sketch_agg(op: &SketchAggOp, plan: &mut StagedPlan, budgets: &StageResourceBudgets) { match op { - // ── Exact ops: mergeability decides stage ───────────────────────────── - // - // Sum/Count/Min/Max satisfy agg(A∪B) = merge(agg(A), agg(B)). - // The Backend collector can combine partial results from N agents. - SketchAggOp::Exact( - ExactAgg::Sum | ExactAgg::Count | ExactAgg::Min | ExactAgg::Max, - ) => { + // Exact ops: mergeability decides stage. + SketchAggOp::Exact(ExactAgg::Sum | ExactAgg::Count | ExactAgg::Min | ExactAgg::Max) => { plan.backend.has_merge = true; } - - // Avg is NOT mergeable: avg(avg(A), avg(B)) ≠ avg(A∪B). - // Must see all raw data → DB-side exact query. SketchAggOp::Exact(ExactAgg::Avg) => { plan.db.active = true; } - - // ── Sketch ops: assign to Agent, defer if budget exceeded ───────────── + // Sketch ops: assign to Agent, defer if budget exceeded. sketch_op => { let est_mem = estimated_sketch_memory_bytes(sketch_op); let stage = resolve_sketch_stage(est_mem, budgets, &mut plan.deferral_log, sketch_op); match stage { SketchStage::Agent => { - plan.agent.sketch_type = Some(agg_op_to_sketch_type(sketch_op)); + plan.agent.sketch_type = Some(agg_op_to_sketch_type(sketch_op)); plan.agent.sketch_params = agg_op_to_sketch_params(sketch_op); } SketchStage::Backend => { - // Deferred from Agent: Backend does the sketch insertion. plan.backend.has_merge = true; - // Record sketch type on agent as None (passthrough raw). } SketchStage::Precompute => { plan.precompute.active = true; @@ -196,27 +253,51 @@ fn assign_agg(op: &SketchAggOp, plan: &mut StagedPlan, budgets: &StageResourceBu } } -#[derive(Debug, PartialEq)] -enum SketchStage { - Agent, - Backend, - Precompute, +fn assign_agg_func(func: &AggFunc, plan: &mut StagedPlan, budgets: &StageResourceBudgets) { + match func { + // Sketchable → synthesise the corresponding SketchAggOp and use existing logic. + AggFunc::Quantile(phi) => { + let op = SketchAggOp::default_ddsketch(vec![*phi]); + assign_sketch_agg(&op, plan, budgets); + } + AggFunc::CountDistinct => { + assign_sketch_agg(&SketchAggOp::default_hll(), plan, budgets); + } + AggFunc::HeavyHitters { k } => { + assign_sketch_agg(&SketchAggOp::CountSketch { k: *k }, plan, budgets); + } + // Mergeable exact → Backend. + AggFunc::Count | AggFunc::Sum | AggFunc::Min | AggFunc::Max + | AggFunc::Rate | AggFunc::Increase | AggFunc::Delta => { + plan.backend.has_merge = true; + } + // Non-mergeable → Db. + AggFunc::Avg | AggFunc::StdDev { .. } | AggFunc::Variance { .. } => { + plan.db.active = true; + } + AggFunc::Custom(_) => { + // Unknown; conservatively route to Precompute. + plan.precompute.active = true; + } + } } -/// Walk the deferral chain Agent → Backend → Precompute based on budget caps. +// ── Budget deferral ─────────────────────────────────────────────────────────── + +#[derive(Debug, PartialEq)] +enum SketchStage { Agent, Backend, Precompute } + fn resolve_sketch_stage( est_mem: u64, budgets: &StageResourceBudgets, log: &mut Vec, op: &SketchAggOp, ) -> SketchStage { - // Try Agent first. if let Some(cap) = budgets.agent_memory_bytes { if est_mem > cap { log.push(format!( "deferred {op:?} Agent→Backend: est_mem={est_mem}B > agent_cap={cap}B" )); - // Try Backend. if let Some(be_cap) = budgets.backend_memory_bytes { if est_mem > be_cap { log.push(format!( @@ -231,18 +312,13 @@ fn resolve_sketch_stage( SketchStage::Agent } -// ── Cost estimation (heuristic, based on benchmark table in cost_model.rs) ──── - -/// Estimate peak sketch memory per series (bytes) for the given operator. fn estimated_sketch_memory_bytes(op: &SketchAggOp) -> u64 { match op { SketchAggOp::DDSketch { .. } | SketchAggOp::ExactMinMax { .. } => 4_096, - SketchAggOp::HLL { registers } => 1u64 << (*registers as u64), + SketchAggOp::HLL { registers } => 1u64 << (*registers as u64), SketchAggOp::CountMin { width, depth } => (*width as u64) * (*depth as u64) * 8, - SketchAggOp::CountSketch { k } => k * 8 * 5, // 5-row hash table approx + SketchAggOp::CountSketch { k } => k * 8 * 5, SketchAggOp::Hydra { inner, partition_keys } => { - // Hydra maintains one inner sketch per distinct key-tuple. - // Estimate 2^|keys| partitions (capped at 1024 to avoid absurd values). let factor = 1u64 << partition_keys.len().min(10); estimated_sketch_memory_bytes(inner).saturating_mul(factor) } @@ -255,11 +331,11 @@ fn estimated_sketch_memory_bytes(op: &SketchAggOp) -> u64 { fn agg_op_to_sketch_type(op: &SketchAggOp) -> SketchType { match op { SketchAggOp::DDSketch { .. } | SketchAggOp::ExactMinMax { .. } => SketchType::DDSketch, - SketchAggOp::HLL { .. } => SketchType::HLL, - SketchAggOp::CountMin { .. } => SketchType::CountMinSketch, + SketchAggOp::HLL { .. } => SketchType::HLL, + SketchAggOp::CountMin { .. } => SketchType::CountMinSketch, SketchAggOp::CountSketch { .. } => SketchType::CountSketch, SketchAggOp::Hydra { inner, .. } => agg_op_to_sketch_type(inner), - SketchAggOp::Exact(_) => SketchType::DDSketch, // unreachable for sketch stage + SketchAggOp::Exact(_) => SketchType::DDSketch, } } @@ -279,11 +355,7 @@ fn agg_op_to_sketch_params(op: &SketchAggOp) -> SketchParams { cols: *width, ..Default::default() }, - SketchAggOp::CountSketch { .. } => SketchParams { - rows: 5, - cols: 2_048, - ..Default::default() - }, + SketchAggOp::CountSketch { .. } => SketchParams { rows: 5, cols: 2_048, ..Default::default() }, SketchAggOp::Hydra { inner, .. } => agg_op_to_sketch_params(inner), SketchAggOp::ExactMinMax { .. } => SketchParams { relative_accuracy: 0.01, @@ -294,168 +366,316 @@ fn agg_op_to_sketch_params(op: &SketchAggOp) -> SketchParams { } } -// ── PromQL serialiser ───────────────────────────────────────────────────────── +// ── PromQL serialiser — recursive descent ───────────────────────────────────── -/// Context accumulated while walking the tree for PromQL serialisation. -#[derive(Default)] -struct QueryCtx { - metric: String, - filters: Vec, - window: Option, +/// Mutable context threaded through the recursive descent. +/// Carries `group_by` and `window` that are "collected" from inner nodes +/// and applied at the enclosing aggregate. +#[derive(Default, Clone)] +struct PromQLCtx { + /// GROUP BY / `by (…)` labels gathered from `Partition` nodes above. group_by: Vec, - agg: Option, - topk: Option, + /// Time window gathered from the innermost `Window` node. + window: Option, } -#[derive(Clone)] -enum AggInfo { - DDSketch { phi: f64 }, - HLL, - CountMin, - CountSketch, - ExactCount, - ExactSum, - ExactAvg, - ExactMin, - ExactMax, - ExactMinMax, -} +/// Recursive PromQL serialisation of a [`QueryExpr`] node. +/// +/// Returns the PromQL string fragment for this node. Inner nodes (Source, +/// Filter) return their selector string; outer nodes (SketchAgg, TopK, etc.) +/// wrap it. +fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { + match expr { + // ── Leaf ───────────────────────────────────────────────────────────── + QueryExpr::Source(s) => s.name.clone(), + QueryExpr::Ref(name) => name.clone(), + + // ── Filter — append label matchers to the selector ─────────────────── + QueryExpr::Filter { pred, input } => { + let inner = promql_from_qe(input, ctx); + let matchers = scalar_to_label_matchers(pred); + if matchers.is_empty() { + inner + } else { + format!("{}{{{}}}", inner, matchers.join(", ")) + } + } -impl QueryCtx { - fn to_promql(&self) -> String { - let selector = if self.filters.is_empty() { - self.metric.clone() - } else { - format!("{}{{{}}}", self.metric, self.filters.join(", ")) - }; + // ── Window — store duration for use by enclosing aggregate ─────────── + QueryExpr::Window { duration, input, .. } => { + if ctx.window.is_none() { + ctx.window = Some(*duration); + } + promql_from_qe(input, ctx) + } - let window = self - .window - .map(|w| format!("[{}]", format_duration(w))) - .unwrap_or_default(); + // ── Partition — store group_by keys for enclosing aggregate ────────── + QueryExpr::Partition { keys, input } => { + for k in keys.keys() { + if !ctx.group_by.contains(k) { + ctx.group_by.push(k.clone()); + } + } + promql_from_qe(input, ctx) + } - let by_clause = if self.group_by.is_empty() { - String::new() - } else { - format!(" by ({})", self.group_by.join(", ")) - }; + // ── SketchAgg — the main aggregation node ──────────────────────────── + QueryExpr::SketchAgg { op, input, .. } => { + let selector = promql_from_qe(input, ctx); + let window = window_str(ctx.window); + let by = by_clause(&ctx.group_by); + sketch_op_to_promql(op, &selector, &window, &by) + } - let inner = match &self.agg { - Some(AggInfo::DDSketch { phi }) => { - format!("quantile_over_time({phi}, {selector}{window}){by_clause}") - } - Some(AggInfo::HLL) => { - format!("count_distinct_over_time({selector}{window}){by_clause}") - } - Some(AggInfo::CountMin | AggInfo::CountSketch | AggInfo::ExactCount) => { - format!("count_over_time({selector}{window}){by_clause}") - } - Some(AggInfo::ExactSum) => { - format!("sum_over_time({selector}{window}){by_clause}") - } - Some(AggInfo::ExactAvg) => { - format!("avg_over_time({selector}{window}){by_clause}") - } - Some(AggInfo::ExactMin) => { - format!("min_over_time({selector}{window}){by_clause}") - } - Some(AggInfo::ExactMax) => { - format!("max_over_time({selector}{window}){by_clause}") + // ── Aggregate (SQL GROUP BY) ────────────────────────────────────────── + QueryExpr::Aggregate { keys, aggs, input, .. } => { + // Merge SQL GROUP BY keys into the context. + for k in keys { + if !ctx.group_by.contains(k) { + ctx.group_by.push(k.clone()); + } } - Some(AggInfo::ExactMinMax) => { - // Represent as quantile range [0,1] for min/max via DDSketch. - format!("quantile_over_time(0.5, {selector}{window}){by_clause}") + let selector = promql_from_qe(input, ctx); + let window = window_str(ctx.window); + let by = by_clause(&ctx.group_by); + // Use the first aggregate function to drive the PromQL template. + if let Some(agg) = aggs.first() { + agg_func_to_promql(&agg.func, &selector, &window, &by) + } else { + selector } - None => format!("{selector}{window}{by_clause}"), - }; + } - if let Some(k) = self.topk { + // ── TopK ───────────────────────────────────────────────────────────── + QueryExpr::TopK { k, input, .. } => { + let inner = promql_from_qe(input, ctx); format!("topk({k}, {inner})") - } else { - inner } + + // ── Sort + Limit — map to topk ─────────────────────────────────────── + QueryExpr::Limit { n, input, .. } => { + let inner = promql_from_qe(input, ctx); + format!("topk({n}, {inner})") + } + QueryExpr::Sort { input, .. } => promql_from_qe(input, ctx), + + // ── histogram_quantile(φ, rate(selector[w])) ───────────────────────── + QueryExpr::HistogramQuantile { phi, input } => { + let selector = promql_from_qe(input, ctx); + let window = window_str(ctx.window); + format!("histogram_quantile({phi}, rate({selector}{window}))") + } + + // ── PromQL subquery expr[range:step] ───────────────────────────────── + QueryExpr::PromQLSubquery { range, resolution, input } => { + let inner = promql_from_qe(input, ctx); + let step_str = resolution + .map(|r| format!(":{}", format_duration(r))) + .unwrap_or_default(); + format!("{}[{}{}]", inner, format_duration(*range), step_str) + } + + // ── Vector binary op (lhs op rhs) ─────────────────────────────────── + QueryExpr::BinaryOp { op, lhs, rhs, vector_match } => { + let lhs_str = promql_from_qe(lhs, ctx); + let rhs_str = promql_from_qe(rhs, &mut PromQLCtx::default()); + let op_str = binop_to_promql(op); + let match_str = vector_match + .as_ref() + .map(|m| { + use crate::algebra::expr::{GroupSide, VectorMatchKind}; + let kw = match m.kind { + VectorMatchKind::On => "on", + VectorMatchKind::Ignoring => "ignoring", + }; + let labels = m.labels.join(", "); + let group = m.grouping.as_ref().map(|g| { + let side = match g.side { + GroupSide::Left => "group_left", + GroupSide::Right => "group_right", + }; + if g.labels.is_empty() { + format!(" {side}") + } else { + format!(" {side}({})", g.labels.join(", ")) + } + }).unwrap_or_default(); + format!(" {kw} ({labels}){group}") + }) + .unwrap_or_default(); + format!("({lhs_str} {op_str}{match_str} {rhs_str})") + } + + // ── Merge — serialise first branch (all branches same shape) ───────── + QueryExpr::Merge { inputs } => { + inputs.first() + .map(|first| promql_from_qe(first, ctx)) + .unwrap_or_default() + } + + // ── Passthrough nodes ───────────────────────────────────────────────── + QueryExpr::Dedup { input, .. } + | QueryExpr::Project { input, .. } + | QueryExpr::WindowFunc { input, .. } => promql_from_qe(input, ctx), + + QueryExpr::Subquery { expr, .. } => promql_from_qe(expr, ctx), + + QueryExpr::LetBinding { body, .. } => promql_from_qe(body, ctx), + + // ── Join / SetOp — serialise the outer / left branch ───────────────── + QueryExpr::JoinSketch { outer, .. } => promql_from_qe(outer, ctx), + QueryExpr::Join { left, .. } => promql_from_qe(left, ctx), + QueryExpr::SetOp { left, .. } => promql_from_qe(left, ctx), } } -/// Walk the tree and fill [`QueryCtx`]; string building is deferred to -/// [`QueryCtx::to_promql`]. -fn collect_ctx(expr: &SketchExpr, ctx: &mut QueryCtx) { - match expr { - SketchExpr::Source(s) => { - if ctx.metric.is_empty() { - ctx.metric = s.name.clone(); - } +// ── PromQL fragment helpers ─────────────────────────────────────────────────── + +fn window_str(w: Option) -> String { + w.map(|d| format!("[{}]", format_duration(d))).unwrap_or_default() +} + +fn by_clause(keys: &[String]) -> String { + if keys.is_empty() { + String::new() + } else { + format!(" by ({})", keys.join(", ")) + } +} + +fn sketch_op_to_promql(op: &SketchAggOp, selector: &str, window: &str, by: &str) -> String { + match op { + SketchAggOp::DDSketch { quantiles, .. } => { + let phi = quantiles.first().copied().unwrap_or(0.99); + format!("quantile_over_time({phi}, {selector}{window}){by}") } - SketchExpr::Filter { pred, input } => { - for p in pred { - match (&p.op, &p.val) { - (FilterOp::Eq, FilterVal::Str(v)) => { - ctx.filters.push(format!("{}=\"{}\"", p.col, v)); - } - (FilterOp::Ne, FilterVal::Str(v)) => { - ctx.filters.push(format!("{}!=\"{}\"", p.col, v)); - } - (FilterOp::Regex(re), _) => { - ctx.filters.push(format!("{}=~\"{}\"", p.col, re)); - } - (FilterOp::NotRegex(re), _) => { - ctx.filters.push(format!("{}!~\"{}\"", p.col, re)); - } - _ => {} // other ops (Lt/Gt/…) not natively supported in PromQL label matchers - } - } - collect_ctx(input, ctx); + SketchAggOp::HLL { .. } => { + format!("count_over_time({selector}{window}){by}") } - SketchExpr::Window { duration, input } => { - if ctx.window.is_none() { - ctx.window = Some(*duration); - } - collect_ctx(input, ctx); + SketchAggOp::CountMin { .. } | SketchAggOp::CountSketch { .. } => { + format!("count_over_time({selector}{window}){by}") } - SketchExpr::Partition { keys, input } => { - for k in keys.keys() { - if !ctx.group_by.contains(k) { - ctx.group_by.push(k.clone()); - } + SketchAggOp::ExactMinMax { min, max } => match (min, max) { + (true, false) => format!("min_over_time({selector}{window}){by}"), + (false, true) => format!("max_over_time({selector}{window}){by}"), + _ => format!("quantile_over_time(0.5, {selector}{window}){by}"), + }, + SketchAggOp::Exact(ExactAgg::Count) => format!("count_over_time({selector}{window}){by}"), + SketchAggOp::Exact(ExactAgg::Sum) => format!("sum_over_time({selector}{window}){by}"), + SketchAggOp::Exact(ExactAgg::Avg) => format!("avg_over_time({selector}{window}){by}"), + SketchAggOp::Exact(ExactAgg::Min) => format!("min_over_time({selector}{window}){by}"), + SketchAggOp::Exact(ExactAgg::Max) => format!("max_over_time({selector}{window}){by}"), + SketchAggOp::Hydra { inner, .. } => sketch_op_to_promql(inner, selector, window, by), + } +} + +fn agg_func_to_promql(func: &AggFunc, selector: &str, window: &str, by: &str) -> String { + match func { + AggFunc::Quantile(phi) => format!("quantile_over_time({phi}, {selector}{window}){by}"), + AggFunc::CountDistinct => format!("count_over_time({selector}{window}){by}"), + AggFunc::HeavyHitters{k} => format!("topk({k}, count_over_time({selector}{window}){by})"), + AggFunc::Count => format!("count_over_time({selector}{window}){by}"), + AggFunc::Sum => format!("sum_over_time({selector}{window}){by}"), + AggFunc::Avg => format!("avg_over_time({selector}{window}){by}"), + AggFunc::Min => format!("min_over_time({selector}{window}){by}"), + AggFunc::Max => format!("max_over_time({selector}{window}){by}"), + AggFunc::StdDev { .. } => format!("stddev_over_time({selector}{window}){by}"), + AggFunc::Variance { .. } => format!("stdvar_over_time({selector}{window}){by}"), + AggFunc::Rate => format!("rate({selector}{window}){by}"), + AggFunc::Increase => format!("increase({selector}{window}){by}"), + AggFunc::Delta => format!("delta({selector}{window}){by}"), + AggFunc::Custom(name) => format!("{name}({selector}{window}){by}"), + } +} + +fn binop_to_promql(op: &BinaryOpKind) -> &'static str { + match op { + BinaryOpKind::Add => "+", + BinaryOpKind::Sub => "-", + BinaryOpKind::Mul => "*", + BinaryOpKind::Div => "/", + BinaryOpKind::Mod => "%", + BinaryOpKind::Pow => "^", + BinaryOpKind::Eq => "==", + BinaryOpKind::Ne => "!=", + BinaryOpKind::Lt => "<", + BinaryOpKind::Le => "<=", + BinaryOpKind::Gt => ">", + BinaryOpKind::Ge => ">=", + BinaryOpKind::And => "and", + BinaryOpKind::Or => "or", + BinaryOpKind::Unless => "unless", + BinaryOpKind::Atan2 => "atan2", + _ => "and", // bitwise/string ops not in PromQL + } +} + +// ── Label matcher extraction from ScalarExpr ───────────────────────────────── + +/// Extract PromQL-compatible label matchers from a `ScalarExpr` AND-tree. +/// +/// Only equality / inequality / regex comparisons between a `Column` and a +/// string `Literal` are extractable as label matchers. Everything else is +/// silently ignored (it won't become a label filter in the PromQL output). +fn scalar_to_label_matchers(pred: &ScalarExpr) -> Vec { + let mut out = Vec::new(); + collect_label_matchers(pred, &mut out); + out +} + +fn collect_label_matchers(pred: &ScalarExpr, out: &mut Vec) { + match pred { + // AND-tree: recurse into both sides. + ScalarExpr::BinaryOp { op: BinaryOpKind::And, lhs, rhs } => { + collect_label_matchers(lhs, out); + collect_label_matchers(rhs, out); + } + // col = "val" + ScalarExpr::BinaryOp { op: BinaryOpKind::Eq, lhs, rhs } => { + if let (ScalarExpr::Column(col), ScalarExpr::Literal(LiteralValue::Str(v))) + = (lhs.as_ref(), rhs.as_ref()) + { + out.push(format!("{}=\"{}\"", col, v)); } - collect_ctx(input, ctx); } - SketchExpr::Agg { op, input, .. } => { - if ctx.agg.is_none() { - ctx.agg = Some(agg_info(op)); + // col != "val" + ScalarExpr::BinaryOp { op: BinaryOpKind::Ne, lhs, rhs } => { + if let (ScalarExpr::Column(col), ScalarExpr::Literal(LiteralValue::Str(v))) + = (lhs.as_ref(), rhs.as_ref()) + { + out.push(format!("{}!=\"{}\"", col, v)); } - collect_ctx(input, ctx); } - SketchExpr::TopK { k, input } => { - ctx.topk = Some(*k); - collect_ctx(input, ctx); + // col =~ "regex" + ScalarExpr::BinaryOp { op: BinaryOpKind::Regex, lhs, rhs } => { + if let (ScalarExpr::Column(col), ScalarExpr::Literal(LiteralValue::Str(v))) + = (lhs.as_ref(), rhs.as_ref()) + { + out.push(format!("{}=~\"{}\"", col, v)); + } } - SketchExpr::Dedup { input, .. } => collect_ctx(input, ctx), - SketchExpr::Merge { inputs } => { - // All branches should share the same shape; serialise the first. - if let Some(first) = inputs.first() { - collect_ctx(first, ctx); + // col !~ "regex" + ScalarExpr::BinaryOp { op: BinaryOpKind::NotRegex, lhs, rhs } => { + if let (ScalarExpr::Column(col), ScalarExpr::Literal(LiteralValue::Str(v))) + = (lhs.as_ref(), rhs.as_ref()) + { + out.push(format!("{}!~\"{}\"", col, v)); } } - SketchExpr::JoinSketch { outer, .. } => collect_ctx(outer, ctx), + _ => {} } } -fn agg_info(op: &SketchAggOp) -> AggInfo { - match op { - SketchAggOp::DDSketch { quantiles, .. } => AggInfo::DDSketch { - phi: quantiles.first().copied().unwrap_or(0.99), - }, - SketchAggOp::HLL { .. } => AggInfo::HLL, - SketchAggOp::CountMin { .. } => AggInfo::CountMin, - SketchAggOp::CountSketch { .. } => AggInfo::CountSketch, - SketchAggOp::ExactMinMax { .. } => AggInfo::ExactMinMax, - SketchAggOp::Exact(ExactAgg::Count) => AggInfo::ExactCount, - SketchAggOp::Exact(ExactAgg::Sum) => AggInfo::ExactSum, - SketchAggOp::Exact(ExactAgg::Avg) => AggInfo::ExactAvg, - SketchAggOp::Exact(ExactAgg::Min) => AggInfo::ExactMin, - SketchAggOp::Exact(ExactAgg::Max) => AggInfo::ExactMax, - SketchAggOp::Hydra { inner, .. } => agg_info(inner), +// ── Helper: push label filters from a ScalarExpr into a Vec ────────── + +fn collect_label_filters_into(pred: &ScalarExpr, out: &mut Vec) { + let matchers = scalar_to_label_matchers(pred); + for m in matchers { + // Store as "col=val" (without PromQL quotes) for the agent YAML. + // The YAML generator already re-quotes as needed. + if !out.contains(&m) { + out.push(m); + } } } @@ -464,28 +684,36 @@ fn agg_info(op: &SketchAggOp) -> AggInfo { #[cfg(test)] mod tests { use super::*; - use crate::query_parser::sketch_algebra::{ - ColumnRef, FilterVal, Predicate, SketchExpr, SketchAggOp, SourceSpec, - }; + use crate::algebra::expr::{AggItem, BinaryOpKind, LiteralValue, QueryExpr, ScalarExpr}; + use crate::query_parser::sketch_algebra::{ColumnRef, PartitionKeys, SketchAggOp, SourceSpec}; - fn source(name: &str) -> SketchExpr { - SketchExpr::Source(SourceSpec { name: name.into() }) + fn source(name: &str) -> QueryExpr { + QueryExpr::Source(SourceSpec { name: name.into() }) } - fn no_budget() -> StageResourceBudgets { - StageResourceBudgets::default() + fn no_budget() -> StageResourceBudgets { StageResourceBudgets::default() } + + fn eq_filter(col: &str, val: &str) -> QueryExpr { + QueryExpr::Filter { + pred: ScalarExpr::BinaryOp { + op: BinaryOpKind::Eq, + lhs: Box::new(ScalarExpr::Column(col.into())), + rhs: Box::new(ScalarExpr::Literal(LiteralValue::Str(val.into()))), + }, + input: Box::new(source("latency")), + } } // ── Node-to-stage assignment ────────────────────────────────────────────── - /// Source + Filter + Window + Agg(DDSketch) → all at Agent. #[test] fn ddsketch_agg_goes_to_agent() { - let expr = SketchExpr::Window { + let expr = QueryExpr::Window { duration: Duration::from_secs(300), - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::default_ddsketch(vec![0.99]), - col: ColumnRef::SampleValue, + slide: None, + input: Box::new(QueryExpr::SketchAgg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, input: Box::new(source("latency")), }), }; @@ -496,27 +724,25 @@ mod tests { assert!(!plan.db.active); } - /// HLL stays at Agent by default. #[test] - fn hll_agg_goes_to_agent() { - let expr = SketchExpr::Agg { - op: SketchAggOp::default_hll(), - col: ColumnRef::Wildcard, - input: Box::new(source("requests")), + fn hll_stays_at_agent_by_default() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::default_hll(), + col: ColumnRef::SampleValue, + input: Box::new(source("events")), }; let plan = split_expr_by_stage(&expr, &no_budget()); assert_eq!(plan.agent.sketch_type, Some(SketchType::HLL)); } - /// Partition keys go to Backend regardless. #[test] - fn partition_goes_to_backend() { - let expr = SketchExpr::Partition { + fn partition_group_by_goes_to_backend() { + let expr = QueryExpr::Partition { keys: PartitionKeys::By(vec!["host".into(), "region".into()]), - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::default_ddsketch(vec![0.5]), - col: ColumnRef::SampleValue, - input: Box::new(source("cpu")), + input: Box::new(QueryExpr::SketchAgg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, + input: Box::new(source("latency")), }), }; let plan = split_expr_by_stage(&expr, &no_budget()); @@ -524,335 +750,227 @@ mod tests { assert!(plan.backend.group_by.contains(&"region".to_string())); } - /// TopK → Precompute Engine. #[test] - fn topk_goes_to_precompute() { - let expr = SketchExpr::TopK { - k: 10, - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::CountSketch { k: 10 }, - col: ColumnRef::Wildcard, - input: Box::new(source("events")), - }), - }; - let plan = split_expr_by_stage(&expr, &no_budget()); - assert!(plan.precompute.active); - assert_eq!(plan.precompute.topk, Some(10)); - } - - /// Exact(Sum) is mergeable → Backend. - #[test] - fn exact_sum_goes_to_backend() { - let expr = SketchExpr::Agg { - op: SketchAggOp::Exact(ExactAgg::Sum), - col: ColumnRef::Named("bytes".into()), - input: Box::new(source("network")), + fn aggregate_without_group_by_avg_goes_to_db() { + let expr = QueryExpr::Aggregate { + keys: vec![], + aggs: vec![AggItem { + alias: "avg_val".into(), + func: AggFunc::Avg, + col: ColumnRef::SampleValue, + distinct: false, + }], + having: None, + input: Box::new(source("price")), }; let plan = split_expr_by_stage(&expr, &no_budget()); - // Sum is handled at backend (merge); agent gets no sketch type. - assert!(plan.agent.sketch_type.is_none()); - assert!(!plan.db.active); - assert!(plan.backend.has_merge); + assert!(plan.db.active); } - /// Exact(Count) is mergeable → Backend. #[test] - fn exact_count_goes_to_backend() { - let expr = SketchExpr::Agg { - op: SketchAggOp::Exact(ExactAgg::Count), - col: ColumnRef::Wildcard, - input: Box::new(source("hits")), + fn aggregate_sum_goes_to_backend() { + let expr = QueryExpr::Aggregate { + keys: vec!["symbol".into()], + aggs: vec![AggItem { + alias: "total".into(), + func: AggFunc::Sum, + col: ColumnRef::SampleValue, + distinct: false, + }], + having: None, + input: Box::new(source("trades")), }; let plan = split_expr_by_stage(&expr, &no_budget()); - assert!(plan.agent.sketch_type.is_none()); assert!(plan.backend.has_merge); + assert!(plan.backend.group_by.contains(&"symbol".to_string())); assert!(!plan.db.active); } - /// Exact(Avg) is NOT mergeable → DB stage. #[test] - fn exact_avg_goes_to_db() { - let expr = SketchExpr::Agg { - op: SketchAggOp::Exact(ExactAgg::Avg), - col: ColumnRef::Named("price".into()), - input: Box::new(source("ticks")), + fn topk_goes_to_precompute() { + let expr = QueryExpr::TopK { + k: 10, + by: vec!["symbol".into()], + input: Box::new(QueryExpr::SketchAgg { + op: SketchAggOp::CountSketch { k: 10 }, + col: ColumnRef::SampleValue, + input: Box::new(source("price")), + }), }; let plan = split_expr_by_stage(&expr, &no_budget()); - assert!(plan.db.active); - assert!(!plan.precompute.active); + assert!(plan.precompute.active); + assert_eq!(plan.precompute.topk, Some(10)); } - /// Exact(Min) is mergeable → Backend (min(A∪B) = min(min(A), min(B))). #[test] - fn exact_min_goes_to_backend() { - let expr = SketchExpr::Agg { - op: SketchAggOp::Exact(ExactAgg::Min), - col: ColumnRef::Named("latency".into()), - input: Box::new(source("svc")), + fn histogram_quantile_activates_precompute() { + let expr = QueryExpr::HistogramQuantile { + phi: 0.99, + input: Box::new(QueryExpr::Window { + duration: Duration::from_secs(300), + slide: None, + input: Box::new(source("http_request_duration_seconds_bucket")), + }), }; let plan = split_expr_by_stage(&expr, &no_budget()); - assert!(plan.backend.has_merge); - assert!(!plan.db.active); + assert!(plan.precompute.active); + assert!(!plan.precompute.query_expr.is_empty()); + assert!(plan.precompute.query_expr.contains("histogram_quantile(0.99")); } - /// Exact(Max) is mergeable → Backend. #[test] - fn exact_max_goes_to_backend() { - let expr = SketchExpr::Agg { - op: SketchAggOp::Exact(ExactAgg::Max), - col: ColumnRef::Named("latency".into()), - input: Box::new(source("svc")), + fn binary_op_activates_precompute() { + let lhs = source("metric_a"); + let rhs = source("metric_b"); + let expr = QueryExpr::BinaryOp { + op: BinaryOpKind::Div, + lhs: Box::new(lhs), + rhs: Box::new(rhs), + vector_match: None, }; let plan = split_expr_by_stage(&expr, &no_budget()); - assert!(plan.backend.has_merge); - assert!(!plan.db.active); - } - - /// Dedup node → Backend. - #[test] - fn dedup_goes_to_backend() { - let expr = SketchExpr::Dedup { - col: "user_id".into(), - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::default_hll(), - col: ColumnRef::Named("user_id".into()), - input: Box::new(source("events")), - }), - }; - let plan = split_expr_by_stage(&expr, &no_budget()); - assert!(plan.backend.has_dedup); + assert!(plan.precompute.active); } // ── Budget-driven deferral ──────────────────────────────────────────────── - /// When the DDSketch memory (4 096 B) exceeds the agent cap, it defers - /// to Backend. #[test] fn ddsketch_deferred_to_backend_when_agent_budget_exceeded() { - let budgets = StageResourceBudgets { - agent_memory_bytes: Some(1_024), // tighter than 4 096 + let tiny_budget = StageResourceBudgets { + agent_memory_bytes: Some(1), // 1 byte — DDSketch (4 KiB) won't fit ..Default::default() }; - let expr = SketchExpr::Agg { - op: SketchAggOp::default_ddsketch(vec![0.99]), - col: ColumnRef::SampleValue, + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, input: Box::new(source("latency")), }; - let plan = split_expr_by_stage(&expr, &budgets); - // Deferred: no sketch at Agent. - assert!(plan.agent.sketch_type.is_none()); - // Backend takes over (merge path). - assert!(plan.backend.has_merge); - // Deferral logged. + let plan = split_expr_by_stage(&expr, &tiny_budget); + assert_eq!(plan.agent.sketch_type, None); // not at agent + assert!(plan.backend.has_merge); // deferred to backend assert!(!plan.deferral_log.is_empty()); - assert!(plan.deferral_log[0].contains("Agent→Backend")); } - /// When both agent AND backend caps are exceeded, defers to Precompute. #[test] fn ddsketch_deferred_to_precompute_when_both_budgets_exceeded() { - let budgets = StageResourceBudgets { - agent_memory_bytes: Some(1_024), - backend_memory_bytes: Some(1_024), + let tiny_budget = StageResourceBudgets { + agent_memory_bytes: Some(1), + backend_memory_bytes: Some(1), ..Default::default() }; - let expr = SketchExpr::Agg { - op: SketchAggOp::default_ddsketch(vec![0.99]), - col: ColumnRef::SampleValue, + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, input: Box::new(source("latency")), }; - let plan = split_expr_by_stage(&expr, &budgets); + let plan = split_expr_by_stage(&expr, &tiny_budget); assert!(plan.precompute.active); - assert!(plan.deferral_log.iter().any(|l| l.contains("Backend→Precompute"))); + assert_eq!(plan.deferral_log.len(), 2); // two deferral steps logged } - /// No budget set → Agent assignment regardless of sketch size. - #[test] - fn no_budget_always_assigns_to_agent() { - let expr = SketchExpr::Agg { - op: SketchAggOp::CountMin { width: 2_000_000, depth: 255 }, // huge - col: ColumnRef::Wildcard, - input: Box::new(source("events")), - }; - let plan = split_expr_by_stage(&expr, &no_budget()); - assert_eq!(plan.agent.sketch_type, Some(SketchType::CountMinSketch)); - assert!(plan.deferral_log.is_empty()); - } + // ── expr_to_promql ──────────────────────────────────────────────────────── - // ── PromQL serialisation ────────────────────────────────────────────────── - - /// TopK(CountSketch) over filtered, windowed, partitioned metric. #[test] - fn topk_countsketch_promql() { - let expr = SketchExpr::TopK { - k: 10, - input: Box::new(SketchExpr::Partition { - keys: PartitionKeys::By(vec!["symbol".into()]), - input: Box::new(SketchExpr::Window { - duration: Duration::from_secs(300), - input: Box::new(SketchExpr::Filter { - pred: vec![Predicate { - col: "sectype".into(), - op: FilterOp::Eq, - val: FilterVal::Str("E".into()), - }], - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::CountSketch { k: 10 }, - col: ColumnRef::Wildcard, - input: Box::new(source("financial.last_trade_price")), - }), + fn promql_ddsketch_with_filter_and_window() { + let expr = QueryExpr::Partition { + keys: PartitionKeys::By(vec!["symbol".into()]), + input: Box::new(QueryExpr::Window { + duration: Duration::from_secs(300), + slide: None, + input: Box::new(QueryExpr::SketchAgg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, + input: Box::new(QueryExpr::Filter { + pred: ScalarExpr::BinaryOp { + op: BinaryOpKind::Eq, + lhs: Box::new(ScalarExpr::Column("sectype".into())), + rhs: Box::new(ScalarExpr::Literal(LiteralValue::Str("E".into()))), + }, + input: Box::new(source("price")), }), }), }), }; let ql = expr_to_promql(&expr); - assert!(ql.contains("topk(10,"), "missing topk: {ql}"); - assert!(ql.contains("count_over_time"), "missing count_over_time: {ql}"); - assert!(ql.contains("financial.last_trade_price"), "missing metric: {ql}"); - assert!(ql.contains("sectype=\"E\""), "missing filter: {ql}"); - assert!(ql.contains("[5m]"), "missing window: {ql}"); - assert!(ql.contains("by (symbol)"), "missing group_by: {ql}"); + assert!(ql.contains("quantile_over_time(0.99"), "expected quantile_over_time: {ql}"); + assert!(ql.contains("sectype=\"E\""), "expected label filter: {ql}"); + assert!(ql.contains("[5m]"), "expected window: {ql}"); + assert!(ql.contains("by (symbol)"), "expected group_by: {ql}"); } - /// DDSketch quantile query. #[test] - fn ddsketch_promql() { - let expr = SketchExpr::Agg { - op: SketchAggOp::default_ddsketch(vec![0.99]), - col: ColumnRef::SampleValue, - input: Box::new(SketchExpr::Window { - duration: Duration::from_secs(300), - input: Box::new(source("latency")), + fn promql_topk_wraps_inner() { + let expr = QueryExpr::TopK { + k: 10, + by: vec![], + input: Box::new(QueryExpr::SketchAgg { + op: SketchAggOp::CountSketch { k: 10 }, + col: ColumnRef::SampleValue, + input: Box::new(source("events")), }), }; let ql = expr_to_promql(&expr); - assert!(ql.contains("quantile_over_time(0.99,"), "missing quantile: {ql}"); - assert!(ql.contains("latency"), "missing metric: {ql}"); - assert!(ql.contains("[5m]"), "missing window: {ql}"); + assert!(ql.starts_with("topk(10,"), "expected topk prefix: {ql}"); } - /// HLL cardinality query. #[test] - fn hll_promql() { - let expr = SketchExpr::Agg { - op: SketchAggOp::default_hll(), - col: ColumnRef::Wildcard, - input: Box::new(source("users")), + fn promql_histogram_quantile() { + let expr = QueryExpr::HistogramQuantile { + phi: 0.95, + input: Box::new(QueryExpr::Window { + duration: Duration::from_secs(300), + slide: None, + input: Box::new(source("http_request_duration_seconds_bucket")), + }), }; let ql = expr_to_promql(&expr); - assert!(ql.contains("count_distinct_over_time"), "missing fn: {ql}"); - assert!(ql.contains("users"), "missing metric: {ql}"); + assert!(ql.contains("histogram_quantile(0.95"), "got: {ql}"); + assert!(ql.contains("rate("), "got: {ql}"); + assert!(ql.contains("[5m]"), "got: {ql}"); } - /// Exact(Avg) produces avg_over_time. #[test] - fn exact_avg_promql() { - let expr = SketchExpr::Agg { - op: SketchAggOp::Exact(ExactAgg::Avg), - col: ColumnRef::Named("price".into()), - input: Box::new(source("ticks")), + fn promql_binary_op_renders_operator() { + let expr = QueryExpr::BinaryOp { + op: BinaryOpKind::Div, + lhs: Box::new(source("http_errors")), + rhs: Box::new(source("http_requests")), + vector_match: None, }; let ql = expr_to_promql(&expr); - assert!(ql.contains("avg_over_time"), "missing avg: {ql}"); + assert!(ql.contains('/'), "expected / operator: {ql}"); + assert!(ql.contains("http_errors"), "got: {ql}"); + assert!(ql.contains("http_requests"), "got: {ql}"); } - /// Exact(Sum) produces sum_over_time. #[test] - fn exact_sum_promql() { - let expr = SketchExpr::Agg { - op: SketchAggOp::Exact(ExactAgg::Sum), - col: ColumnRef::Named("bytes".into()), - input: Box::new(source("net")), + fn promql_subquery_renders_range() { + let expr = QueryExpr::PromQLSubquery { + range: Duration::from_secs(3600), + resolution: Some(Duration::from_secs(60)), + input: Box::new(source("metric")), }; let ql = expr_to_promql(&expr); - assert!(ql.contains("sum_over_time"), "missing sum: {ql}"); - } - - // ── Sketch parameter extraction ─────────────────────────────────────────── - - #[test] - fn ddsketch_params_extracted_correctly() { - let expr = SketchExpr::Agg { - op: SketchAggOp::DDSketch { quantiles: vec![0.5, 0.99], epsilon: 0.005 }, - col: ColumnRef::SampleValue, - input: Box::new(source("latency")), - }; - let plan = split_expr_by_stage(&expr, &no_budget()); - let p = &plan.agent.sketch_params; - assert_eq!(p.relative_accuracy, 0.005); - assert_eq!(p.quantiles, vec![0.5, 0.99]); + assert!(ql.contains("[1h:1m]") || ql.contains("[3600s:60s]"), "got: {ql}"); } #[test] - fn hll_precision_extracted() { - let expr = SketchExpr::Agg { - op: SketchAggOp::HLL { registers: 14 }, - col: ColumnRef::Wildcard, - input: Box::new(source("users")), - }; - let plan = split_expr_by_stage(&expr, &no_budget()); - assert_eq!(plan.agent.sketch_params.precision, 14); - } - - #[test] - fn countmin_params_extracted() { - let expr = SketchExpr::Agg { - op: SketchAggOp::CountMin { width: 2_000, depth: 5 }, - col: ColumnRef::Wildcard, - input: Box::new(source("events")), - }; - let plan = split_expr_by_stage(&expr, &no_budget()); - assert_eq!(plan.agent.sketch_params.rows, 5); - assert_eq!(plan.agent.sketch_params.cols, 2_000); - } - - // ── Filter predicate extraction ─────────────────────────────────────────── - - #[test] - fn filter_eq_captured_in_agent_plan() { - let expr = SketchExpr::Filter { - pred: vec![Predicate { - col: "env".into(), - op: FilterOp::Eq, - val: FilterVal::Str("prod".into()), - }], - input: Box::new(source("latency")), - }; - let plan = split_expr_by_stage(&expr, &no_budget()); - assert!(plan.agent.label_filters.contains(&"env=prod".to_string())); - } - - // ── Precompute query_expr generation ───────────────────────────────────── - - /// When TopK is present, `staged_plan.precompute.query_expr` must be a - /// non-empty PromQL string. - #[test] - fn precompute_query_expr_set_when_topk_present() { - let expr = SketchExpr::TopK { - k: 5, - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::default_hll(), - col: ColumnRef::Wildcard, - input: Box::new(source("sessions")), + fn label_matchers_and_tree() { + let pred = ScalarExpr::BinaryOp { + op: BinaryOpKind::And, + lhs: Box::new(ScalarExpr::BinaryOp { + op: BinaryOpKind::Eq, + lhs: Box::new(ScalarExpr::Column("job".into())), + rhs: Box::new(ScalarExpr::Literal(LiteralValue::Str("api".into()))), + }), + rhs: Box::new(ScalarExpr::BinaryOp { + op: BinaryOpKind::Ne, + lhs: Box::new(ScalarExpr::Column("env".into())), + rhs: Box::new(ScalarExpr::Literal(LiteralValue::Str("dev".into()))), }), }; - let plan = split_expr_by_stage(&expr, &no_budget()); - assert!(plan.precompute.active); - assert!(!plan.precompute.query_expr.is_empty()); - assert!(plan.precompute.query_expr.contains("topk(5,")); - } - - // ── StageResourceBudgets::from_workload_chars ───────────────────────────── - - #[test] - fn budgets_from_workload_chars_propagates_memory() { - use crate::types::WorkloadCharacteristics; - let wc = WorkloadCharacteristics { - memory_budget_bytes: Some(8_192), - ..Default::default() - }; - let b = StageResourceBudgets::from_workload_chars(&wc); - assert_eq!(b.agent_memory_bytes, Some(8_192)); - assert!(b.backend_memory_bytes.is_none()); + let matchers = scalar_to_label_matchers(&pred); + assert!(matchers.contains(&"job=\"api\"".to_string())); + assert!(matchers.contains(&"env!=\"dev\"".to_string())); } } diff --git a/controller/src/types.rs b/controller/src/types.rs index 019861b3..35c7d212 100644 --- a/controller/src/types.rs +++ b/controller/src/types.rs @@ -247,33 +247,6 @@ pub struct SketchParams { pub quantiles: Vec, } -/// Per-stage resource caps used by the sketch allocator. -/// -/// `None` on any field means "no limit" (the allocator will not demote for -/// that resource). The default is all-None (unlimited). -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct StageResourceBudgets { - /// Maximum memory (bytes) for sketch state at the Agent collector. - pub agent_memory_bytes: Option, - /// Maximum CPU overhead (µs/sample) at the Agent collector. - pub agent_cpu_micros_per_sample: Option, - /// Maximum memory (bytes) for sketch state at the Backend collector. - pub backend_memory_bytes: Option, - /// Maximum memory (bytes) at the Precompute engine. - pub precompute_memory_bytes: Option, -} - -impl StageResourceBudgets { - /// Derive budgets from [`WorkloadCharacteristics`]: agent memory cap comes - /// from `memory_budget_bytes`; the rest default to unlimited. - pub fn from_workload_chars(wc: &WorkloadCharacteristics) -> Self { - Self { - agent_memory_bytes: wc.memory_budget_bytes, - ..Default::default() - } - } -} - #[derive(Debug, Clone)] pub struct AgentCollectorConfig { pub output_mode: OutputMode, From c7f904337c38f2d73def7e905996241beb195e66 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 1 Apr 2026 12:59:40 -0500 Subject: [PATCH 3/3] =?UTF-8?q?refactor:=20delete=20SketchExpr=20=E2=80=94?= =?UTF-8?q?=20QueryExpr=20is=20the=20sole=20IR?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Delete sketch_algebra.rs (721 lines) and sketch_rules.rs (683 lines) - Move shared types (SourceSpec, PartitionKeys, ColumnRef, SketchAggOp, ExactAgg, Predicate, FilterOp, FilterVal) into algebra/expr.rs - Parsers emit QueryExpr directly: parse_promql_expr / parse_sql_expr are now the only entry points; parse_query_sketch removed - query_parser/mod.rs: QeCollector walks QueryExpr for ParsedQuery - analyzer.rs: analyze_with_sketch() removed - main.rs: uses parse_query_expr() directly (no from_sketch_expr bridge) - All 264 tests pass; 39 SketchExpr-only tests removed with the IR Co-Authored-By: Claude Sonnet 4.6 --- controller/src/algebra/allocator.rs | 6 +- controller/src/algebra/expr.rs | 451 ++++++----- controller/src/algebra/mod.rs | 9 +- controller/src/algebra/optimizer.rs | 4 +- controller/src/algebra/plan.rs | 2 +- controller/src/analyzer.rs | 27 +- controller/src/main.rs | 11 +- controller/src/planner/stage_split.rs | 4 +- controller/src/query_parser/mod.rs | 289 ++++++- controller/src/query_parser/promql.rs | 459 +++-------- controller/src/query_parser/sketch_algebra.rs | 721 ----------------- controller/src/query_parser/sketch_rules.rs | 683 ---------------- controller/src/query_parser/sql.rs | 731 +++--------------- controller/src/types.rs | 10 +- 14 files changed, 725 insertions(+), 2682 deletions(-) delete mode 100644 controller/src/query_parser/sketch_algebra.rs delete mode 100644 controller/src/query_parser/sketch_rules.rs diff --git a/controller/src/algebra/allocator.rs b/controller/src/algebra/allocator.rs index 5ed084c6..528cc8b4 100644 --- a/controller/src/algebra/allocator.rs +++ b/controller/src/algebra/allocator.rs @@ -31,7 +31,7 @@ use super::expr::QueryExpr; use super::plan::{ CostEstimate, ExecutionMode, NodeAnnotation, PipelineStage, PlanNode, }; -use crate::query_parser::sketch_algebra::{ExactAgg, SketchAggOp}; +use super::expr::{ExactAgg, SketchAggOp}; use crate::types::{SketchParams, SketchType, StageResourceBudgets}; // ── Resource budget tracker ─────────────────────────────────────────────────── @@ -533,7 +533,7 @@ impl SketchAllocator { fn alloc_sketch_agg( &self, op: SketchAggOp, - col: crate::query_parser::sketch_algebra::ColumnRef, + col: super::expr::ColumnRef, child: PlanNode, budget: &mut BudgetState, ) -> PlanNode { @@ -734,7 +734,7 @@ mod tests { use super::*; use crate::algebra::expr::QueryExpr; use crate::algebra::plan::{ExecutionMode, PipelineStage}; - use crate::query_parser::sketch_algebra::{ColumnRef, PartitionKeys, SketchAggOp, SourceSpec}; + use crate::algebra::expr::{ColumnRef, PartitionKeys, SketchAggOp, SourceSpec}; use crate::types::{SketchType, StageResourceBudgets}; use std::time::Duration; diff --git a/controller/src/algebra/expr.rs b/controller/src/algebra/expr.rs index 43d27fbf..989e6590 100644 --- a/controller/src/algebra/expr.rs +++ b/controller/src/algebra/expr.rs @@ -10,10 +10,6 @@ //! from a row. Used for WHERE predicates, SELECT projections, HAVING //! conditions, and JOIN conditions. //! -//! The bridge to the existing sketch algebra is [`QueryExpr::from_sketch_expr`], -//! which converts a [`crate::query_parser::SketchExpr`] tree into the general -//! algebra for backward-compatibility. -//! //! # Stage vocabulary //! //! Once the [`crate::algebra::allocator::SketchAllocator`] annotates the tree, @@ -29,13 +25,194 @@ use std::time::Duration; -use crate::query_parser::sketch_algebra::{ - ColumnRef, ExactAgg, FilterOp, FilterVal, Predicate, SketchAggOp, SketchExpr, SourceSpec, -}; +use crate::types::AggType; + +// ── Shared sketch / predicate types ─────────────────────────────────────────── + +/// Base relation / metric stream source. +#[derive(Debug, Clone)] +pub struct SourceSpec { + /// Table name (SQL) or metric name (PromQL). + pub name: String, +} + +/// How the stream is partitioned. +#[derive(Debug, Clone)] +pub enum PartitionKeys { + /// `by (k1, k2, ...)` — explicit key list. + By(Vec), + /// `without (k1, k2, ...)` — complement; resolved against schema at plan time. + Without(Vec), +} + +impl PartitionKeys { + pub fn keys(&self) -> &[String] { + match self { + PartitionKeys::By(k) | PartitionKeys::Without(k) => k, + } + } + + pub fn is_empty(&self) -> bool { + self.keys().is_empty() + } + + pub fn into_by_keys(self) -> Vec { + match self { + PartitionKeys::By(k) => k, + // For Without, return empty — caller resolves complement. + PartitionKeys::Without(k) => k, + } + } +} + +/// Which column / field the sketch aggregation targets. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ColumnRef { + /// Explicit column name (SQL: `AVG(price)` → `Named("price")`). + Named(String), + /// The implicit metric sample value (PromQL — always the series value). + SampleValue, + /// All rows / COUNT(*). + Wildcard, +} + +/// The concrete sketch type used for aggregation. +#[derive(Debug, Clone, PartialEq)] +pub enum SketchAggOp { + /// Count-Min Sketch — frequency per group (COUNT(*) GROUP BY). + CountMin { width: u32, depth: u8 }, + + /// Count Sketch (heavy-hitter) — top-K by frequency. + CountSketch { k: u64 }, + + /// HyperLogLog — distinct-value counting (COUNT DISTINCT). + HLL { registers: u8 }, + + /// DDSketch — quantile estimation. + /// `quantiles` holds the φ values to track; `epsilon` is relative error. + DDSketch { quantiles: Vec, epsilon: f64 }, + + /// Exact running min/max tracker — cheaper than DDSketch for extrema + /// without a GROUP BY (no sketch benefit for global extrema). + ExactMinMax { min: bool, max: bool }, + + /// Hydra — sketch of sketches for multi-dimensional GROUP BY. + /// Maintains one `inner` sketch per distinct `partition_keys` tuple. + Hydra { + inner: Box, + partition_keys: Vec, + }, + + /// Exact passthrough — no sketch benefit (SUM, global COUNT, etc.). + Exact(ExactAgg), +} + +/// Exact (non-sketch) aggregation kinds. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExactAgg { + Count, + Sum, + /// **Not mergeable** — carries `(sum, count)` in distributed contexts. + Avg, + Min, + Max, +} + +impl SketchAggOp { + /// Returns `true` when two instances of this sketch can be merged + /// (i.e., `sketch(A ∪ B) = merge(sketch(A), sketch(B))`). + pub fn is_mergeable(&self) -> bool { + match self { + SketchAggOp::Exact(ExactAgg::Avg) => false, + SketchAggOp::Hydra { inner, .. } => inner.is_mergeable(), + _ => true, + } + } + + /// Map to the coarse [`AggType`] used by the legacy planner. + pub fn to_agg_type(&self) -> AggType { + match self { + SketchAggOp::HLL { .. } => AggType::Cardinality, + SketchAggOp::CountMin { .. } | SketchAggOp::CountSketch { .. } => AggType::Frequency, + SketchAggOp::DDSketch { .. } | SketchAggOp::ExactMinMax { .. } => AggType::Quantile, + SketchAggOp::Hydra { inner, .. } => inner.to_agg_type(), + SketchAggOp::Exact(_) => AggType::Quantile, + } + } + + /// Extract quantile φ values for DDSketch operators. + pub fn quantiles(&self) -> Vec { + match self { + SketchAggOp::DDSketch { quantiles, .. } => quantiles.clone(), + SketchAggOp::Hydra { inner, .. } => inner.quantiles(), + _ => vec![], + } + } + + /// Whether this op implies `exact_required` (no sketch benefit). + pub fn is_exact(&self) -> bool { + matches!(self, SketchAggOp::Exact(_) | SketchAggOp::ExactMinMax { .. }) + } + + pub fn default_count_min() -> Self { + SketchAggOp::CountMin { width: 2000, depth: 5 } + } + pub fn default_hll() -> Self { + SketchAggOp::HLL { registers: 14 } + } + pub fn default_ddsketch(quantiles: Vec) -> Self { + SketchAggOp::DDSketch { quantiles, epsilon: 0.01 } + } +} + +/// A single filter predicate pushed down to the collector. +#[derive(Debug, Clone)] +pub struct Predicate { + pub col: String, + pub op: FilterOp, + pub val: FilterVal, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum FilterOp { + Eq, + Ne, + Lt, + Le, + Gt, + Ge, + Like, + NotLike, + IsNull, + IsNotNull, + /// PromQL `=~` label matcher (RE2 syntax). + Regex(String), + /// PromQL `!~` label matcher. + NotRegex(String), +} + +#[derive(Debug, Clone)] +pub enum FilterVal { + Str(String), + Num(f64), + Int(i64), + Null, +} + +/// How completely a query can be served by sketches. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SketchCoverage { + /// All aggregation columns are sketch-mapped. + Full, + /// Some columns are sketch-mapped; others require exact passthrough. + Partial, + /// No sketch applicable; query requires exact execution. + None, +} // ── Relational algebra ──────────────────────────────────────────────────────── -/// Full relational + sketch algebra — replaces / extends [`SketchExpr`]. +/// Full relational + sketch algebra — the sole query IR. /// /// Every variant is a *node* in the logical query plan tree. Leaves are /// [`QueryExpr::Source`] or [`QueryExpr::Ref`]. Interior nodes combine their @@ -101,7 +278,7 @@ pub enum QueryExpr { /// Partition the stream by key-tuple (GROUP BY / `by (dims)`). Partition { - keys: crate::query_parser::sketch_algebra::PartitionKeys, + keys: PartitionKeys, input: Box, }, @@ -573,63 +750,7 @@ pub enum DataType { Custom(String), } -// ── Bridge: SketchExpr → QueryExpr ─────────────────────────────────────────── - impl QueryExpr { - /// Convert a legacy [`SketchExpr`] tree into the general algebra. - /// - /// This bridge preserves backward-compatibility: the existing SQL and - /// PromQL parsers emit [`SketchExpr`] trees; the new optimizer and - /// allocator work on [`QueryExpr`] trees. - pub fn from_sketch_expr(s: &SketchExpr) -> Self { - match s { - SketchExpr::Source(src) => QueryExpr::Source(src.clone()), - - SketchExpr::Filter { pred, input } => QueryExpr::Filter { - pred: scalar_from_predicates(pred), - input: Box::new(QueryExpr::from_sketch_expr(input)), - }, - - SketchExpr::Window { duration, input } => QueryExpr::Window { - duration: *duration, - slide: None, - input: Box::new(QueryExpr::from_sketch_expr(input)), - }, - - SketchExpr::Partition { keys, input } => QueryExpr::Partition { - keys: keys.clone(), - input: Box::new(QueryExpr::from_sketch_expr(input)), - }, - - SketchExpr::Agg { op, col, input } => QueryExpr::SketchAgg { - op: op.clone(), - col: col.clone(), - input: Box::new(QueryExpr::from_sketch_expr(input)), - }, - - SketchExpr::Dedup { col, input } => QueryExpr::Dedup { - col: col.clone(), - input: Box::new(QueryExpr::from_sketch_expr(input)), - }, - - SketchExpr::TopK { k, input } => QueryExpr::TopK { - k: *k, - by: vec![], - input: Box::new(QueryExpr::from_sketch_expr(input)), - }, - - SketchExpr::Merge { inputs } => QueryExpr::Merge { - inputs: inputs.iter().map(QueryExpr::from_sketch_expr).collect(), - }, - - SketchExpr::JoinSketch { join_key, outer, inner } => QueryExpr::JoinSketch { - join_key: join_key.clone(), - outer: Box::new(QueryExpr::from_sketch_expr(outer)), - inner: Box::new(QueryExpr::from_sketch_expr(inner)), - }, - } - } - /// Walk the expression tree depth-first and call `f` on every node. pub fn walk(&self, f: &mut F) { f(self); @@ -811,125 +932,22 @@ impl std::fmt::Display for AggFunc { #[cfg(test)] mod tests { use super::*; - use crate::query_parser::sketch_algebra::{ - PartitionKeys, SketchAggOp, SketchExpr, SourceSpec, - }; use std::time::Duration; - fn src(name: &str) -> SketchExpr { - SketchExpr::Source(SourceSpec { name: name.into() }) - } - - // ── Bridge tests ────────────────────────────────────────────────────────── - - #[test] - fn bridge_source() { - let qe = QueryExpr::from_sketch_expr(&src("cpu")); - assert!(matches!(qe, QueryExpr::Source(s) if s.name == "cpu")); - } - - #[test] - fn bridge_sketch_agg_ddsketch() { - use crate::query_parser::sketch_algebra::ColumnRef; - let se = SketchExpr::Agg { - op: SketchAggOp::default_ddsketch(vec![0.99]), - col: ColumnRef::SampleValue, - input: Box::new(src("latency")), - }; - let qe = QueryExpr::from_sketch_expr(&se); - assert!( - matches!(&qe, QueryExpr::SketchAgg { op: SketchAggOp::DDSketch { .. }, .. }), - "expected SketchAgg(DDSketch), got {qe:?}" - ); - } - - #[test] - fn bridge_window_preserves_duration() { - let se = SketchExpr::Window { - duration: Duration::from_secs(300), - input: Box::new(src("m")), - }; - let qe = QueryExpr::from_sketch_expr(&se); - match qe { - QueryExpr::Window { duration, .. } => { - assert_eq!(duration, Duration::from_secs(300)); - } - other => panic!("unexpected {other:?}"), - } - } - - #[test] - fn bridge_filter_converts_predicates() { - use crate::query_parser::sketch_algebra::{FilterOp, FilterVal, Predicate}; - let se = SketchExpr::Filter { - pred: vec![Predicate { - col: "env".into(), - op: FilterOp::Eq, - val: FilterVal::Str("prod".into()), - }], - input: Box::new(src("http_requests")), - }; - let qe = QueryExpr::from_sketch_expr(&se); - assert!(matches!(qe, QueryExpr::Filter { .. })); - } - - #[test] - fn bridge_topk_sets_k() { - let se = SketchExpr::TopK { - k: 25, - input: Box::new(src("events")), - }; - let qe = QueryExpr::from_sketch_expr(&se); - match qe { - QueryExpr::TopK { k, .. } => assert_eq!(k, 25), - other => panic!("unexpected {other:?}"), - } - } - - #[test] - fn bridge_merge_fans_in() { - let se = SketchExpr::Merge { - inputs: vec![src("a"), src("b"), src("c")], - }; - let qe = QueryExpr::from_sketch_expr(&se); - match qe { - QueryExpr::Merge { inputs } => assert_eq!(inputs.len(), 3), - other => panic!("unexpected {other:?}"), - } - } - - #[test] - fn bridge_join_sketch() { - use crate::query_parser::sketch_algebra::ColumnRef; - let se = SketchExpr::JoinSketch { - join_key: "order_id".into(), - outer: Box::new(src("orders")), - inner: Box::new(SketchExpr::Agg { - op: SketchAggOp::default_hll(), - col: ColumnRef::Named("item_id".into()), - input: Box::new(src("items")), - }), - }; - let qe = QueryExpr::from_sketch_expr(&se); - match qe { - QueryExpr::JoinSketch { join_key, .. } => { - assert_eq!(join_key, "order_id"); - } - other => panic!("unexpected {other:?}"), - } + fn src(name: &str) -> QueryExpr { + QueryExpr::Source(SourceSpec { name: name.into() }) } // ── has_sketch_work ─────────────────────────────────────────────────────── #[test] fn has_sketch_work_true_when_ddsketch_present() { - use crate::query_parser::sketch_algebra::ColumnRef; - let se = SketchExpr::Agg { + let qe = QueryExpr::SketchAgg { op: SketchAggOp::default_ddsketch(vec![0.5]), col: ColumnRef::SampleValue, input: Box::new(src("m")), }; - assert!(QueryExpr::from_sketch_expr(&se).has_sketch_work()); + assert!(qe.has_sketch_work()); } #[test] @@ -942,14 +960,14 @@ mod tests { #[test] fn source_name_extracted_through_chain() { - let se = SketchExpr::Window { + let qe = QueryExpr::Window { duration: Duration::from_secs(60), - input: Box::new(SketchExpr::Filter { - pred: vec![], + slide: None, + input: Box::new(QueryExpr::Filter { + pred: ScalarExpr::Literal(LiteralValue::Bool(true)), input: Box::new(src("my_metric")), }), }; - let qe = QueryExpr::from_sketch_expr(&se); assert_eq!(qe.source_name(), Some("my_metric")); } @@ -1003,7 +1021,6 @@ mod tests { #[test] fn two_preds_become_and_tree() { - use crate::query_parser::sketch_algebra::{FilterOp, FilterVal, Predicate}; let preds = vec![ Predicate { col: "a".into(), op: FilterOp::Eq, val: FilterVal::Int(1) }, Predicate { col: "b".into(), op: FilterOp::Gt, val: FilterVal::Num(2.0) }, @@ -1026,16 +1043,17 @@ mod tests { // ── Complex nested tree ─────────────────────────────────────────────────── #[test] - fn complex_nested_bridge_roundtrip() { - use crate::query_parser::sketch_algebra::ColumnRef; - // TopK(10, Partition(symbol, Window(5m, Agg(CountSketch, Source(price))))) - let se = SketchExpr::TopK { + fn complex_nested_tree() { + // TopK(10, Partition(symbol, Window(5m, SketchAgg(CountSketch, Source(price))))) + let qe = QueryExpr::TopK { k: 10, - input: Box::new(SketchExpr::Partition { + by: vec![], + input: Box::new(QueryExpr::Partition { keys: PartitionKeys::By(vec!["symbol".into()]), - input: Box::new(SketchExpr::Window { + input: Box::new(QueryExpr::Window { duration: Duration::from_secs(300), - input: Box::new(SketchExpr::Agg { + slide: None, + input: Box::new(QueryExpr::SketchAgg { op: SketchAggOp::CountSketch { k: 10 }, col: ColumnRef::Wildcard, input: Box::new(src("price")), @@ -1043,7 +1061,6 @@ mod tests { }), }), }; - let qe = QueryExpr::from_sketch_expr(&se); assert!(qe.has_sketch_work()); assert_eq!(qe.source_name(), Some("price")); } @@ -1092,4 +1109,68 @@ mod tests { _ => panic!(), } } + + // ── SketchAggOp and related types ────────────────────────────────────────── + + #[test] + fn sketch_agg_op_hll_is_mergeable() { + assert!(SketchAggOp::default_hll().is_mergeable()); + } + + #[test] + fn sketch_agg_op_exact_avg_not_mergeable() { + assert!(!SketchAggOp::Exact(ExactAgg::Avg).is_mergeable()); + } + + #[test] + fn sketch_agg_op_hydra_mergeability_from_inner() { + let hydra_hll = SketchAggOp::Hydra { + inner: Box::new(SketchAggOp::default_hll()), + partition_keys: vec!["region".into()], + }; + assert!(hydra_hll.is_mergeable()); + let hydra_avg = SketchAggOp::Hydra { + inner: Box::new(SketchAggOp::Exact(ExactAgg::Avg)), + partition_keys: vec!["region".into()], + }; + assert!(!hydra_avg.is_mergeable()); + } + + #[test] + fn partition_keys_without_variant() { + let keys = PartitionKeys::Without(vec!["instance".into()]); + assert_eq!(keys.keys(), &["instance".to_string()]); + assert!(!keys.is_empty()); + } + + #[test] + fn sketch_agg_op_to_agg_type() { + use crate::types::AggType; + assert_eq!(SketchAggOp::default_hll().to_agg_type(), AggType::Cardinality); + assert_eq!(SketchAggOp::default_count_min().to_agg_type(), AggType::Frequency); + assert_eq!(SketchAggOp::default_ddsketch(vec![0.5]).to_agg_type(), AggType::Quantile); + } + + #[test] + fn sketch_agg_op_is_exact() { + assert!(SketchAggOp::Exact(ExactAgg::Sum).is_exact()); + assert!(SketchAggOp::ExactMinMax { min: true, max: false }.is_exact()); + assert!(!SketchAggOp::default_hll().is_exact()); + } + + #[test] + fn sketch_coverage_classification() { + let ops: Vec = vec![ + SketchAggOp::default_hll(), + SketchAggOp::Exact(ExactAgg::Sum), + ]; + let has_sketch = ops.iter().any(|o| !o.is_exact()); + let has_exact = ops.iter().any(|o| o.is_exact()); + let cov = match (has_sketch, has_exact) { + (true, false) => SketchCoverage::Full, + (true, true) => SketchCoverage::Partial, + _ => SketchCoverage::None, + }; + assert_eq!(cov, SketchCoverage::Partial); + } } diff --git a/controller/src/algebra/mod.rs b/controller/src/algebra/mod.rs index 461de542..d20b4c5b 100644 --- a/controller/src/algebra/mod.rs +++ b/controller/src/algebra/mod.rs @@ -21,13 +21,10 @@ //! use controller::query_parser; //! use controller::types::StageResourceBudgets; //! -//! // 1. Parse a PromQL / SQL query string into SketchExpr. -//! let sketch_expr = query_parser::parse_query_sketch("quantile_over_time(0.99, latency[5m])")?; +//! // 1. Parse a PromQL / SQL query string into QueryExpr. +//! let query_expr = query_parser::parse_query_expr("quantile_over_time(0.99, latency[5m])")?; //! -//! // 2. Lift into the general algebra. -//! let query_expr = QueryExpr::from_sketch_expr(&sketch_expr); -//! -//! // 3. Optimise (cost-based fixed-point rewriting). +//! // 2. Optimise (cost-based fixed-point rewriting). //! let (opt_expr, _iters) = QueryOptimizer::new(raw_bps).optimize(query_expr); //! //! // 4. Allocate stages. diff --git a/controller/src/algebra/optimizer.rs b/controller/src/algebra/optimizer.rs index 073c9c10..7b6d1739 100644 --- a/controller/src/algebra/optimizer.rs +++ b/controller/src/algebra/optimizer.rs @@ -31,7 +31,7 @@ use std::collections::HashMap; use super::expr::{QueryExpr, ScalarExpr, SetOpKind, SortKey}; -use crate::query_parser::sketch_algebra::{PartitionKeys, SketchAggOp, SourceSpec}; +use super::expr::{PartitionKeys, SketchAggOp, SourceSpec}; // ── Cost model interface ────────────────────────────────────────────────────── @@ -787,7 +787,7 @@ fn default_rules() -> Vec> { mod tests { use super::*; use crate::algebra::expr::{LiteralValue, ScalarExpr}; - use crate::query_parser::sketch_algebra::{ColumnRef, PartitionKeys, SketchAggOp, SourceSpec}; + use crate::algebra::expr::{ColumnRef, PartitionKeys, SketchAggOp, SourceSpec}; use std::time::Duration; fn src(name: &str) -> QueryExpr { diff --git a/controller/src/algebra/plan.rs b/controller/src/algebra/plan.rs index 8857c5ab..4cf237ae 100644 --- a/controller/src/algebra/plan.rs +++ b/controller/src/algebra/plan.rs @@ -272,7 +272,7 @@ impl PlanNode { mod tests { use super::*; use crate::algebra::expr::QueryExpr; - use crate::query_parser::sketch_algebra::SourceSpec; + use crate::algebra::expr::SourceSpec; fn source_node(name: &str, stage: PipelineStage) -> PlanNode { PlanNode::leaf( diff --git a/controller/src/analyzer.rs b/controller/src/analyzer.rs index 05facb51..25800292 100644 --- a/controller/src/analyzer.rs +++ b/controller/src/analyzer.rs @@ -3,7 +3,7 @@ use std::time::Duration; use anyhow::{anyhow, Context}; use serde::{Deserialize, Serialize}; -use crate::query_parser::{self, SketchExpr}; +use crate::query_parser; use crate::types::{AggType, QueryWorkload, SketchType, WorkloadCharacteristics}; // ── Public API ──────────────────────────────────────────────────────────────── @@ -157,31 +157,6 @@ impl Analyzer { }) } - /// Like [`analyze`] but also returns the optimized [`SketchExpr`] tree. - /// - /// The tree is `Some` when `spec.query_string` was provided; `None` when - /// the workload was built from explicit aggregation fields alone. - /// - /// Callers (e.g. `handle_plan`) pass the `SketchExpr` to - /// `planner::stage_split::split_expr_by_stage()` to produce an SP-9 - /// [`StagedPlan`]. The existing `analyze()` path (SP-3 flat assignment) - /// is unchanged and remains the fallback when no tree is available. - pub fn analyze_with_sketch( - &self, - spec: QuerySpec, - ) -> anyhow::Result<(QueryWorkload, Option)> { - // Parse the SketchExpr before `spec` is consumed by `analyze()`. - // `parse_query_sketch` runs the algebraic optimiser internally. - let sketch_expr = spec - .query_string - .as_deref() - .map(|q| query_parser::parse_query_sketch(q)) - .transpose() - .with_context(|| "failed to parse query_string into SketchExpr")?; - - let workload = self.analyze(spec)?; - Ok((workload, sketch_expr)) - } } // ── Duration helpers (used by other modules) ────────────────────────────────── diff --git a/controller/src/main.rs b/controller/src/main.rs index 63e95991..9b3c5ce1 100644 --- a/controller/src/main.rs +++ b/controller/src/main.rs @@ -21,7 +21,7 @@ use axum::{ use serde_json::json; use tracing::{info, warn}; -use algebra::{QueryExpr, QueryOptimizer, SketchAllocator}; +use algebra::{QueryOptimizer, SketchAllocator}; use analyzer::{Analyzer, QuerySpec}; use config::{generate_agent_config, generate_backend_config, build_precompute_jobs}; use config::generate_backend_config_staged; @@ -30,7 +30,7 @@ use opamp::{AgentRole, OpampServer, RemoteConfig}; use planner::{CostModelPlanner, BaselinePlanner, ObjectiveWeights, OnlineMetricsStore, init_online_store, pareto_frontier, select_best}; use planner::online_cost_model; use planner::stage_split::split_expr_by_stage; -use query_parser::{parse_query_expr, parse_query_sketch}; +use query_parser::parse_query_expr; use replan::Replanner; use store::{PlanStore, WorkloadStore}; use types::StageResourceBudgets; @@ -235,10 +235,9 @@ async fn handle_plan( // the StagedPlan. The SP-3 flat assignment remains the fallback when no // query_string is supplied. if let Some(ref qs) = query_string { - match parse_query_sketch(qs) { - Err(e) => warn!(query = %qs, error = %e, "parse_query_sketch failed; skipping staged_plan"), - Ok(sketch_expr) => { - let qe = QueryExpr::from_sketch_expr(&sketch_expr); + match parse_query_expr(qs) { + Err(e) => warn!(query = %qs, error = %e, "parse_query_expr failed; skipping staged_plan"), + Ok(qe) => { let raw_bps = plan.transmission_cost_summary.raw_bytes_per_sec; let (opt_qe, _) = QueryOptimizer::new(raw_bps).optimize(qe); let budgets = StageResourceBudgets::from_workload_chars(&wc); diff --git a/controller/src/planner/stage_split.rs b/controller/src/planner/stage_split.rs index 9eae63aa..2cbd2bcb 100644 --- a/controller/src/planner/stage_split.rs +++ b/controller/src/planner/stage_split.rs @@ -39,7 +39,7 @@ use std::time::Duration; use crate::algebra::expr::{AggFunc, BinaryOpKind, LiteralValue, QueryExpr, ScalarExpr}; use crate::analyzer::format_duration; -use crate::query_parser::sketch_algebra::{ExactAgg, PartitionKeys, SketchAggOp}; +use crate::algebra::expr::{ExactAgg, PartitionKeys, SketchAggOp}; use crate::types::{ AgentSubPlan, BackendSubPlan, DbSubPlan, PrecomputeSubPlan, SketchParams, SketchType, StagedPlan, StageResourceBudgets, @@ -685,7 +685,7 @@ fn collect_label_filters_into(pred: &ScalarExpr, out: &mut Vec) { mod tests { use super::*; use crate::algebra::expr::{AggItem, BinaryOpKind, LiteralValue, QueryExpr, ScalarExpr}; - use crate::query_parser::sketch_algebra::{ColumnRef, PartitionKeys, SketchAggOp, SourceSpec}; + use crate::algebra::expr::{ColumnRef, PartitionKeys, SketchAggOp, SourceSpec}; fn source(name: &str) -> QueryExpr { QueryExpr::Source(SourceSpec { name: name.into() }) diff --git a/controller/src/query_parser/mod.rs b/controller/src/query_parser/mod.rs index 76ca8b6c..567d29f1 100644 --- a/controller/src/query_parser/mod.rs +++ b/controller/src/query_parser/mod.rs @@ -1,14 +1,10 @@ //! SP-1 query workload extraction — PromQL and SQL parsers. //! -//! Both parsers compile to the shared [`SketchExpr`] algebra IR defined in -//! [`sketch_algebra`]. The optimizer in [`sketch_rules`] applies algebraic -//! rewrite rules before the planner receives the result. -//! //! # Entry points //! //! | Function | Returns | Use | //! |---|---|---| -//! | [`parse_query_sketch`] | `SketchExpr` | New callers — full algebra IR | +//! | [`parse_query_expr`] | `QueryExpr` | Full algebra IR | //! | [`parse_query`] | `ParsedQuery` | Backward compat with existing analyzer | //! //! # Supported PromQL patterns (via `promql-parser` AST) @@ -33,22 +29,18 @@ pub mod promql; pub mod sql; -pub mod sketch_algebra; -pub mod sketch_rules; use std::collections::HashMap; use std::time::Duration; -use crate::algebra::expr::QueryExpr; +use crate::algebra::expr::{QueryExpr, SketchAggOp}; use crate::types::AggType; -pub use sketch_algebra::SketchExpr; // ── Output types (legacy — consumed by analyzer and planner) ────────────────── /// Flat intermediate representation consumed by [`crate::analyzer::Analyzer`]. /// -/// Produced by [`parse_query`] via [`SketchExpr::to_parsed_query`]. -/// New code should use [`parse_query_sketch`] → [`SketchExpr`] directly. +/// Produced by [`parse_query`] via [`QueryExpr`] tree walking. #[derive(Debug, Clone)] pub struct ParsedQuery { /// Metric name (PromQL: from selector; SQL: FROM clause table). @@ -95,11 +87,6 @@ pub enum QueryHint { // ── Public entry points ─────────────────────────────────────────────────────── /// Parse a raw query string (PromQL or SQL) into the general [`QueryExpr`] IR. -/// -/// Unlike [`parse_query_sketch`], this path emits `QueryExpr` *directly* -/// from the AST — preserving HistogramQuantile, PromQLSubquery, -/// vector-binary-op matching, Sort+Limit, Join, and SetOp without any -/// lossy round-trip through [`SketchExpr`]. pub fn parse_query_expr(query: &str) -> anyhow::Result { let q = query.trim(); let upper = q.to_ascii_uppercase(); @@ -110,27 +97,254 @@ pub fn parse_query_expr(query: &str) -> anyhow::Result { } } -/// Parse a raw query string (PromQL or SQL) into the full [`SketchExpr`] IR. -/// -/// The returned tree has already been through the algebraic optimizer -/// ([`sketch_rules::optimize`]). -pub fn parse_query_sketch(query: &str) -> anyhow::Result { - let q = query.trim(); - let upper = q.to_ascii_uppercase(); - if upper.starts_with("SELECT") || upper.starts_with("WITH") { - sql::parse_sql(q) - } else { - promql::parse_promql(q) - } -} - /// Parse a raw query string (PromQL or SQL) into a [`ParsedQuery`]. /// /// This is the backward-compatible entry point for the existing -/// [`crate::analyzer::Analyzer`]. Internally it calls [`parse_query_sketch`] -/// and converts via [`SketchExpr::to_parsed_query`]. +/// [`crate::analyzer::Analyzer`]. Internally it parses via [`parse_query_expr`] +/// and extracts the flat summary by walking the [`QueryExpr`] tree. pub fn parse_query(query: &str) -> anyhow::Result { - Ok(parse_query_sketch(query)?.to_parsed_query()) + let qe = parse_query_expr(query)?; + Ok(qe_to_parsed_query(&qe)) +} + +/// Extract a flat [`ParsedQuery`] by walking a [`QueryExpr`] tree. +fn qe_to_parsed_query(qe: &QueryExpr) -> ParsedQuery { + let mut c = QeCollector::default(); + c.visit(qe); + c.build() +} + +#[derive(Default)] +struct QeCollector { + metric_name: Option, + agg_types: Vec, + group_by_labels: Vec, + label_filters: HashMap, + time_window: Option, + exact_required: bool, + quantiles: Vec, + topk: Option, +} + +impl QeCollector { + fn visit(&mut self, expr: &QueryExpr) { + use crate::algebra::expr::{FilterOp, FilterVal, LiteralValue, ScalarExpr}; + match expr { + QueryExpr::Source(s) => { + if self.metric_name.is_none() { + self.metric_name = Some(s.name.clone()); + } + } + QueryExpr::Filter { pred, input } => { + // Extract equality label filters from the predicate tree. + collect_filters_from_scalar(pred, &mut self.label_filters); + self.visit(input); + } + QueryExpr::Window { duration, input, .. } => { + if self.time_window.is_none() { + self.time_window = Some(*duration); + } + self.visit(input); + } + QueryExpr::Partition { keys, input } => { + for k in keys.keys() { + if !self.group_by_labels.contains(k) { + self.group_by_labels.push(k.clone()); + } + } + self.visit(input); + } + QueryExpr::SketchAgg { op, input, .. } => { + self.collect_op(op); + self.visit(input); + } + QueryExpr::TopK { k, input, .. } => { + self.topk = Some(*k); + self.visit(input); + } + QueryExpr::Dedup { input, .. } => self.visit(input), + QueryExpr::Merge { inputs } => { + for i in inputs { self.visit(i); } + } + QueryExpr::JoinSketch { outer, inner, .. } => { + self.visit(outer); + self.visit(inner); + } + QueryExpr::Aggregate { keys, aggs, input, .. } => { + for k in keys { + if !self.group_by_labels.contains(k) { + self.group_by_labels.push(k.clone()); + } + } + let has_group_by = !keys.is_empty(); + for agg in aggs { + self.collect_agg_func_with_group(&agg.func, has_group_by); + } + self.visit(input); + } + QueryExpr::Project { input, .. } + | QueryExpr::Sort { input, .. } + | QueryExpr::Limit { input, .. } + | QueryExpr::HistogramQuantile { input, .. } + | QueryExpr::PromQLSubquery { input, .. } + | QueryExpr::WindowFunc { input, .. } => self.visit(input), + QueryExpr::Join { left, right, .. } + | QueryExpr::SetOp { left, right, .. } + | QueryExpr::BinaryOp { lhs: left, rhs: right, .. } => { + self.visit(left); + self.visit(right); + } + QueryExpr::Subquery { expr, .. } => self.visit(expr), + QueryExpr::LetBinding { expr, body, .. } => { + self.visit(expr); + self.visit(body); + } + QueryExpr::Ref(_) => {} + } + } + + fn collect_agg_func_with_group(&mut self, func: &crate::algebra::expr::AggFunc, has_group_by: bool) { + use crate::algebra::expr::AggFunc; + // COUNT(*) without GROUP BY → exact (no sketch benefit) + if matches!(func, AggFunc::Count) && !has_group_by { + self.exact_required = true; + return; + } + self.collect_agg_func(func); + } + + fn collect_agg_func(&mut self, func: &crate::algebra::expr::AggFunc) { + use crate::algebra::expr::AggFunc; + match func { + AggFunc::CountDistinct => { + if !self.agg_types.contains(&AggType::Cardinality) { + self.agg_types.push(AggType::Cardinality); + } + } + AggFunc::Count => { + if !self.agg_types.contains(&AggType::Frequency) { + self.agg_types.push(AggType::Frequency); + } + } + AggFunc::HeavyHitters { .. } => { + if !self.agg_types.contains(&AggType::Frequency) { + self.agg_types.push(AggType::Frequency); + } + } + AggFunc::Quantile(phi) => { + if !self.agg_types.contains(&AggType::Quantile) { + self.agg_types.push(AggType::Quantile); + } + if !self.quantiles.contains(phi) { + self.quantiles.push(*phi); + } + } + AggFunc::Avg => { + if !self.agg_types.contains(&AggType::Quantile) { + self.agg_types.push(AggType::Quantile); + } + // AVG maps to p50 (median) sketch + if !self.quantiles.contains(&0.5) { self.quantiles.push(0.5); } + } + AggFunc::Min => { + if !self.agg_types.contains(&AggType::Quantile) { + self.agg_types.push(AggType::Quantile); + } + if !self.quantiles.contains(&0.0) { self.quantiles.push(0.0); } + } + AggFunc::Max => { + if !self.agg_types.contains(&AggType::Quantile) { + self.agg_types.push(AggType::Quantile); + } + if !self.quantiles.contains(&1.0) { self.quantiles.push(1.0); } + } + AggFunc::StdDev { .. } | AggFunc::Variance { .. } => { + if !self.agg_types.contains(&AggType::Quantile) { + self.agg_types.push(AggType::Quantile); + } + } + AggFunc::Sum | AggFunc::Rate | AggFunc::Increase | AggFunc::Delta + | AggFunc::Custom(_) => { + self.exact_required = true; + } + } + } + + fn collect_op(&mut self, op: &SketchAggOp) { + use crate::algebra::expr::ExactAgg; + match op { + SketchAggOp::HLL { .. } => { + if !self.agg_types.contains(&AggType::Cardinality) { + self.agg_types.push(AggType::Cardinality); + } + } + SketchAggOp::CountMin { .. } | SketchAggOp::CountSketch { .. } => { + if !self.agg_types.contains(&AggType::Frequency) { + self.agg_types.push(AggType::Frequency); + } + } + SketchAggOp::DDSketch { quantiles, .. } => { + if !self.agg_types.contains(&AggType::Quantile) { + self.agg_types.push(AggType::Quantile); + } + for &q in quantiles { + if !self.quantiles.contains(&q) { self.quantiles.push(q); } + } + } + SketchAggOp::ExactMinMax { .. } => { + if !self.agg_types.contains(&AggType::Quantile) { + self.agg_types.push(AggType::Quantile); + } + } + SketchAggOp::Exact(_) => { self.exact_required = true; } + SketchAggOp::Hydra { inner, .. } => self.collect_op(inner), + } + } + + fn build(self) -> ParsedQuery { + let metric_name = self.metric_name.unwrap_or_default(); + let mut qs = self.quantiles; + qs.sort_by(|a, b| a.partial_cmp(b).unwrap()); + qs.dedup(); + let hint = debs_hint( + &metric_name, + &self.agg_types, + &qs, + self.exact_required, + self.topk, + ); + ParsedQuery { + metric_name, + aggregations: self.agg_types, + group_by_labels: self.group_by_labels, + label_filters: self.label_filters, + time_window: self.time_window.unwrap_or(Duration::from_secs(300)), + exact_required: self.exact_required, + quantiles: qs, + hint, + } + } +} + +fn collect_filters_from_scalar( + pred: &crate::algebra::expr::ScalarExpr, + out: &mut HashMap, +) { + use crate::algebra::expr::{BinaryOpKind, LiteralValue, ScalarExpr}; + match pred { + ScalarExpr::BinaryOp { op: BinaryOpKind::Eq, lhs, rhs } => { + if let (ScalarExpr::Column(col), ScalarExpr::Literal(LiteralValue::Str(v))) = + (lhs.as_ref(), rhs.as_ref()) + { + out.insert(col.clone(), v.clone()); + } + } + ScalarExpr::BinaryOp { op: BinaryOpKind::And, lhs, rhs } => { + collect_filters_from_scalar(lhs, out); + collect_filters_from_scalar(rhs, out); + } + _ => {} + } } // ── DEBS hint classifier (shared by both parsers via to_parsed_query) ───────── @@ -204,12 +418,11 @@ mod tests { } #[test] - fn parse_query_sketch_returns_expr() { - let expr = parse_query_sketch( + fn parse_query_expr_returns_expr() { + let pq = parse_query( "topk by (symbol) (10, count_over_time(financial_last_trade_price[5m]))" ).unwrap(); - // Should have been optimized — result is some SketchExpr tree. - // Just check it doesn't error. - let _ = expr.to_parsed_query(); + // Should parse without error and extract the metric name. + assert_eq!(pq.metric_name, "financial_last_trade_price"); } } diff --git a/controller/src/query_parser/promql.rs b/controller/src/query_parser/promql.rs index e4797cff..1c51277b 100644 --- a/controller/src/query_parser/promql.rs +++ b/controller/src/query_parser/promql.rs @@ -1,7 +1,7 @@ -//! PromQL → SketchExpr compiler. +//! PromQL → QueryExpr compiler. //! //! Uses the `promql-parser` crate (GreptimeTeam) for a full AST parse, then -//! walks the expression tree to emit a [`SketchExpr`] following the mapping +//! walks the expression tree to emit a [`QueryExpr`] following the mapping //! rules in `docs/sketch-algebra-query-mapping.md §2.3`. //! //! # PromQL → Sketch mapping (summary) @@ -29,20 +29,7 @@ use std::time::Duration; use anyhow::anyhow; use promql_parser::parser::{self, AggregateExpr, Call, Expr, LabelModifier, MatrixSelector, VectorSelector}; -use super::sketch_algebra::{ - ColumnRef, FilterOp, FilterVal, PartitionKeys, Predicate, SketchAggOp, SketchExpr, SourceSpec, -}; -use super::sketch_rules::optimize; - -// ── Public entry point ──────────────────────────────────────────────────────── - -/// Parse a PromQL expression string into an optimised [`SketchExpr`]. -pub fn parse_promql(query: &str) -> anyhow::Result { - let expr = parser::parse(query) - .map_err(|e| anyhow!("PromQL parse error: {e}"))?; - let sketch = walk(&expr, WalkCtx::default())?; - Ok(optimize(sketch)) -} +use crate::algebra::expr::{FilterOp, FilterVal, PartitionKeys, Predicate}; // ── Walk context ────────────────────────────────────────────────────────────── @@ -57,271 +44,6 @@ struct WalkCtx { outer_count: bool, } -// ── AST walker ──────────────────────────────────────────────────────────────── - -fn walk(expr: &Expr, ctx: WalkCtx) -> anyhow::Result { - match expr { - // ── Aggregate operators: topk, count, sum by, avg by, … ─────────── - Expr::Aggregate(agg) => walk_aggregate(agg, ctx), - - // ── Function calls: *_over_time, histogram_quantile, rate, … ────── - Expr::Call(call) => walk_call(call, ctx), - - // ── Binary operations: metric_a / metric_b → exact ──────────────── - Expr::Binary(bin) => { - // Binary op between two series: both sides need exact values. - let left = walk(bin.lhs.as_ref(), WalkCtx::default())?; - let right = walk(bin.rhs.as_ref(), WalkCtx::default())?; - // Wrap both in a Merge that signals exact requirement to callers. - Ok(SketchExpr::Agg { - op: SketchAggOp::Exact(super::sketch_algebra::ExactAgg::Sum), // placeholder - col: ColumnRef::SampleValue, - input: Box::new(SketchExpr::Merge { inputs: vec![left, right] }), - }) - } - - // ── Parenthesised ───────────────────────────────────────────────── - Expr::Paren(p) => walk(p.expr.as_ref(), ctx), - - // ── Subquery: metric[5m:1m] — treat as windowed frequency ───────── - Expr::Subquery(sq) => { - let inner = walk(sq.expr.as_ref(), ctx.clone())?; - Ok(SketchExpr::Window { duration: sq.range, input: Box::new(inner) }) - } - - // ── Bare vector selector ────────────────────────────────────────── - Expr::VectorSelector(vs) => { - let (name, filters) = extract_vs_info(vs); - let source = SketchExpr::Source(SourceSpec { name }); - let filtered = apply_filters(source, filters); - // Wrap with partition and exact agg. - let agg = SketchExpr::Agg { - op: SketchAggOp::Exact(super::sketch_algebra::ExactAgg::Sum), - col: ColumnRef::SampleValue, - input: Box::new(filtered), - }; - Ok(apply_partition(agg, ctx.partition)) - } - - // ── Number / string literals — only appear as args inside Call ──── - Expr::NumberLiteral(_) | Expr::StringLiteral(_) => { - Err(anyhow!("unexpected literal at top level of PromQL expression")) - } - - // ── Extension / unknown ─────────────────────────────────────────── - #[allow(unreachable_patterns)] - _ => Err(anyhow!("unsupported PromQL expression type")), - } -} - -// ── Aggregate operator walk ─────────────────────────────────────────────────── - -fn walk_aggregate( - agg: &AggregateExpr, - ctx: WalkCtx, -) -> anyhow::Result { - // Extract partition keys from the by/without modifier. - let partition = agg.modifier.as_ref().map(modifier_to_partition); - - // Extract the operator name from the token via Display (gives lowercase). - let op_name = format!("{}", agg.op); - - match op_name.as_str() { - // topk(k, inner) / bottomk(k, inner) → CountSketch(k) - "topk" | "bottomk" => { - let k = extract_number_param(&agg.param)? as u64; - let inner_ctx = WalkCtx { partition: partition.clone(), topk: Some(k), outer_count: false }; - let inner = walk(agg.expr.as_ref(), inner_ctx)?; - let result = SketchExpr::TopK { k, input: Box::new(inner) }; - Ok(apply_partition(result, partition)) - } - - // count(inner_by) → HLL cardinality of distinct groups - "count" => { - let inner_ctx = WalkCtx { partition: partition.clone(), topk: None, outer_count: true }; - let inner = walk(agg.expr.as_ref(), inner_ctx)?; - // Wrap with HLL at this level. - let result = SketchExpr::Agg { - op: SketchAggOp::default_hll(), - col: ColumnRef::SampleValue, - input: Box::new(inner), - }; - Ok(apply_partition(result, partition)) - } - - // sum by (d) (inner) — outer sum doesn't change the inner sketch type. - // The inner *_over_time already chose the right sketch; we just set partition. - "sum" | "avg" | "min" | "max" | "group" => { - let inner_ctx = WalkCtx { partition: partition.clone(), topk: ctx.topk, outer_count: false }; - let inner = walk(agg.expr.as_ref(), inner_ctx)?; - Ok(apply_partition(inner, partition)) - } - - // stddev by (d) / stdvar by (d) → DDSketch IQR proxy - "stddev" | "stdvar" => { - let inner_ctx = WalkCtx { partition: partition.clone(), topk: None, outer_count: false }; - let inner = walk(agg.expr.as_ref(), inner_ctx)?; - let result = SketchExpr::Agg { - op: SketchAggOp::default_ddsketch(vec![0.25, 0.75]), - col: ColumnRef::SampleValue, - input: Box::new(inner), - }; - Ok(apply_partition(result, partition)) - } - - // quantile(φ, inner) → DDSketch(φ) - "quantile" => { - let phi = extract_number_param(&agg.param)?; - let inner_ctx = WalkCtx { partition: partition.clone(), topk: None, outer_count: false }; - let inner = walk(agg.expr.as_ref(), inner_ctx)?; - let result = SketchExpr::Agg { - op: SketchAggOp::default_ddsketch(vec![phi]), - col: ColumnRef::SampleValue, - input: Box::new(inner), - }; - Ok(apply_partition(result, partition)) - } - - other => Err(anyhow!("unsupported PromQL aggregate operator: {other}")), - } -} - -// ── Function call walk ──────────────────────────────────────────────────────── - -fn walk_call( - call: &Call, - ctx: WalkCtx, -) -> anyhow::Result { - let name = call.func.name; - - match name { - // ── quantile_over_time(φ, m{f}[w]) ─────────────────────────────── - "quantile_over_time" => { - let phi = extract_call_num_arg(call, 0)?; - let (source, filters, window) = extract_matrix_arg(call, 1)?; - Ok(build_sketched( - source, filters, window, - SketchAggOp::default_ddsketch(vec![phi]), - ctx, - )) - } - - // ── histogram_quantile(φ, rate(m{f}[w])) ───────────────────────── - "histogram_quantile" => { - let phi = extract_call_num_arg(call, 0)?; - // The second arg is a call to rate/irate wrapping a MatrixSelector. - let rate_expr = call.args.args[1].as_ref(); - let (source, filters, window) = extract_inner_matrix(rate_expr)?; - Ok(build_sketched( - source, filters, window, - SketchAggOp::default_ddsketch(vec![phi]), - ctx, - )) - } - - // ── avg_over_time ───────────────────────────────────────────────── - "avg_over_time" => { - let (source, filters, window) = extract_matrix_arg(call, 0)?; - Ok(build_sketched(source, filters, window, SketchAggOp::default_ddsketch(vec![0.5]), ctx)) - } - - // ── min_over_time ───────────────────────────────────────────────── - "min_over_time" => { - let (source, filters, window) = extract_matrix_arg(call, 0)?; - let op = if ctx.partition.as_ref().map(|p| !p.is_empty()).unwrap_or(false) { - SketchAggOp::default_ddsketch(vec![0.0]) - } else { - SketchAggOp::ExactMinMax { min: true, max: false } - }; - Ok(build_sketched(source, filters, window, op, ctx)) - } - - // ── max_over_time ───────────────────────────────────────────────── - "max_over_time" => { - let (source, filters, window) = extract_matrix_arg(call, 0)?; - let op = if ctx.partition.as_ref().map(|p| !p.is_empty()).unwrap_or(false) { - SketchAggOp::default_ddsketch(vec![1.0]) - } else { - SketchAggOp::ExactMinMax { min: false, max: true } - }; - Ok(build_sketched(source, filters, window, op, ctx)) - } - - // ── stddev_over_time / stdvar_over_time → IQR proxy ────────────── - "stddev_over_time" | "stdvar_over_time" => { - let (source, filters, window) = extract_matrix_arg(call, 0)?; - Ok(build_sketched( - source, filters, window, - SketchAggOp::default_ddsketch(vec![0.25, 0.75]), - ctx, - )) - } - - // ── count_over_time ─────────────────────────────────────────────── - "count_over_time" => { - let (source, filters, window) = extract_matrix_arg(call, 0)?; - // Outer count() context → HLL (cardinality of distinct groups). - let op = if ctx.outer_count { - SketchAggOp::default_hll() - } else { - SketchAggOp::default_count_min() - }; - Ok(build_sketched(source, filters, window, op, ctx)) - } - - // ── sum_over_time ───────────────────────────────────────────────── - "sum_over_time" => { - let (source, filters, window) = extract_matrix_arg(call, 0)?; - Ok(build_sketched( - source, filters, window, - SketchAggOp::Exact(super::sketch_algebra::ExactAgg::Sum), - ctx, - )) - } - - // ── last_over_time / stateful functions → exact ─────────────────── - "last_over_time" | "present_over_time" | "absent_over_time" => { - let (source, filters, window) = extract_matrix_arg(call, 0)?; - Ok(build_sketched( - source, filters, window, - SketchAggOp::Exact(super::sketch_algebra::ExactAgg::Sum), - ctx, - )) - } - - // ── delta / idelta / deriv / predict_linear → stateful exact ────── - "delta" | "idelta" | "deriv" | "predict_linear" => { - let (source, filters, window) = extract_matrix_arg(call, 0)?; - Ok(build_sketched( - source, filters, window, - SketchAggOp::Exact(super::sketch_algebra::ExactAgg::Sum), - ctx, - )) - } - - // ── changes / resets → frequency ────────────────────────────────── - "changes" | "resets" => { - let (source, filters, window) = extract_matrix_arg(call, 0)?; - Ok(build_sketched(source, filters, window, SketchAggOp::default_count_min(), ctx)) - } - - // ── rate / irate / increase — pass-through, inner carries the sketch - "rate" | "irate" | "increase" => { - if call.args.is_empty() { - return Err(anyhow!("rate/irate/increase requires a matrix arg")); - } - let (source, filters, window) = extract_inner_matrix(call.args.args[0].as_ref())?; - Ok(build_sketched( - source, filters, window, - SketchAggOp::default_count_min(), - ctx, - )) - } - - other => Err(anyhow!("unsupported PromQL function: {other}")), - } -} - // ── Helpers: MatrixSelector extraction ─────────────────────────────────────── /// Extract `(metric_name, filters, window)` from a MatrixSelector argument at @@ -414,69 +136,27 @@ fn modifier_to_partition(modifier: &LabelModifier) -> PartitionKeys { } } -// ── Tree builders ───────────────────────────────────────────────────────────── - -/// Build the standard `Partition(Window(Filter(Agg(Source))))` tree. -fn build_sketched( - metric: String, - filters: Vec, - window: Duration, - op: SketchAggOp, - ctx: WalkCtx, -) -> SketchExpr { - let source = SketchExpr::Source(SourceSpec { name: metric }); - let filtered = apply_filters(source, filters); - let windowed = SketchExpr::Window { duration: window, input: Box::new(filtered) }; - - let agg = if let Some(k) = ctx.topk { - // Outer topk → CountSketch regardless of what op was chosen. - SketchExpr::Agg { - op: SketchAggOp::CountSketch { k }, - col: ColumnRef::SampleValue, - input: Box::new(windowed), - } - } else { - SketchExpr::Agg { op, col: ColumnRef::SampleValue, input: Box::new(windowed) } - }; - - apply_partition(agg, ctx.partition) -} - -fn apply_filters(input: SketchExpr, pred: Vec) -> SketchExpr { - if pred.is_empty() { - input - } else { - SketchExpr::Filter { pred, input: Box::new(input) } - } -} - -fn apply_partition(input: SketchExpr, partition: Option) -> SketchExpr { - match partition { - None => input, - Some(p) if p.is_empty() => input, - Some(keys) => SketchExpr::Partition { keys, input: Box::new(input) }, - } -} - // ── Direct QueryExpr emission ───────────────────────────────────────────────── // // `parse_promql_expr` walks the same PromQL AST but emits [`QueryExpr`] nodes -// natively, preserving semantic nodes that the SketchExpr bridge flattens: +// natively, preserving semantic nodes for the algebra optimizer: // -// | PromQL pattern | SketchExpr (old) | QueryExpr (new) | -// |-------------------------|---------------------------|-----------------------------| -// | `histogram_quantile(φ…)`| Agg(DDSketch([φ])) | HistogramQuantile { phi } | -// | `m[5m:1m]` subquery | Window { 5m } | PromQLSubquery {5m, Some(1m)}| -// | `a op b` binary | Agg(Exact(Sum), Merge(…)) | BinaryOp { VectorMatch } | +// | PromQL pattern | QueryExpr node | +// |--------------------------|---------------------------------------| +// | `histogram_quantile(φ…)` | HistogramQuantile { phi } | +// | `m[5m:1m]` subquery | PromQLSubquery { 5m, Some(1m) } | +// | `a op b` binary | BinaryOp { VectorMatch } | use crate::algebra::expr::{ - BinaryOpKind, GroupSide, QueryExpr, VectorGrouping, VectorMatch, VectorMatchKind, + BinaryOpKind, ColumnRef as QeColumnRef, ExactAgg as QeExactAgg, GroupSide, + PartitionKeys as QePartitionKeys, QueryExpr, SketchAggOp as QeSketchAggOp, + SourceSpec as QeSourceSpec, VectorGrouping, VectorMatch, VectorMatchKind, }; use promql_parser::parser::{token::TokenType, BinaryExpr, VectorMatchCardinality}; /// Parse a PromQL expression string directly into an optimised [`QueryExpr`]. /// -/// Unlike `parse_promql` (which emits `SketchExpr`), this preserves +/// This preserves /// `HistogramQuantile`, `PromQLSubquery`, and `BinaryOp` nodes natively. pub fn parse_promql_expr(query: &str) -> anyhow::Result { let expr = parser::parse(query) @@ -507,12 +187,11 @@ fn walk_qe(expr: &Expr, ctx: WalkCtx) -> anyhow::Result { // Bare vector selector → Source + Filter + exact agg. Expr::VectorSelector(vs) => { let (name, filters) = extract_vs_info(vs); - let source = QueryExpr::Source(SourceSpec { name }); + let source = QueryExpr::Source(QeSourceSpec { name }); let filtered = apply_qe_filters(source, filters); Ok(QueryExpr::SketchAgg { - op: super::sketch_algebra::SketchAggOp::Exact( - super::sketch_algebra::ExactAgg::Sum), - col: super::sketch_algebra::ColumnRef::SampleValue, + op: QeSketchAggOp::Exact(QeExactAgg::Sum), + col: QeColumnRef::SampleValue, input: Box::new(filtered), }) } @@ -541,8 +220,8 @@ fn walk_aggregate_qe(agg: &AggregateExpr, ctx: WalkCtx) -> anyhow::Result anyhow::Result anyhow::Result anyhow::Result { let rate_expr = call.args.args[1].as_ref(); let (source, filters, window) = extract_inner_matrix(rate_expr)?; let inner = build_qe_sketched(source, filters, window, - super::sketch_algebra::SketchAggOp::default_ddsketch(vec![phi]), + QeSketchAggOp::default_ddsketch(vec![phi]), WalkCtx::default()); Ok(QueryExpr::HistogramQuantile { phi, input: Box::new(inner) }) } - // Everything else: reuse the SketchExpr walker and bridge. + // All other function calls: map directly to QeSketchAggOp. + "quantile_over_time" => { + let phi = extract_call_num_arg(call, 0)?; + let (source, filters, window) = extract_matrix_arg(call, 1)?; + let op = QeSketchAggOp::default_ddsketch(vec![phi]); + Ok(build_qe_sketched(source, filters, window, op, ctx)) + } _ => { - let sketch = walk_call(call, ctx)?; - Ok(QueryExpr::from_sketch_expr(&sketch)) + let op = walk_call_to_op(call, &ctx)?; + let (source, filters, window) = if call.func.name == "rate" + || call.func.name == "irate" + || call.func.name == "increase" + { + let arg = call.args.args.first() + .map(|b| b.as_ref()) + .ok_or_else(|| anyhow!("rate/irate/increase requires a matrix arg"))?; + extract_inner_matrix(arg)? + } else { + extract_matrix_arg(call, 0)? + }; + Ok(build_qe_sketched(source, filters, window, op, ctx)) } } } @@ -657,15 +353,52 @@ fn promql_token_to_binop(tok: TokenType) -> BinaryOpKind { } } +/// Map a PromQL function call to a [`QeSketchAggOp`] for direct QE emission. +fn walk_call_to_op(call: &Call, ctx: &WalkCtx) -> anyhow::Result { + let name = call.func.name; + match name { + "quantile_over_time" => { + let phi = extract_call_num_arg(call, 0)?; + Ok(QeSketchAggOp::default_ddsketch(vec![phi])) + } + "avg_over_time" => Ok(QeSketchAggOp::default_ddsketch(vec![0.5])), + "min_over_time" => { + Ok(if ctx.partition.as_ref().map(|p| !p.is_empty()).unwrap_or(false) { + QeSketchAggOp::default_ddsketch(vec![0.0]) + } else { + QeSketchAggOp::ExactMinMax { min: true, max: false } + }) + } + "max_over_time" => { + Ok(if ctx.partition.as_ref().map(|p| !p.is_empty()).unwrap_or(false) { + QeSketchAggOp::default_ddsketch(vec![1.0]) + } else { + QeSketchAggOp::ExactMinMax { min: false, max: true } + }) + } + "stddev_over_time" | "stdvar_over_time" => + Ok(QeSketchAggOp::default_ddsketch(vec![0.25, 0.75])), + "count_over_time" => { + Ok(if ctx.outer_count { QeSketchAggOp::default_hll() } else { QeSketchAggOp::default_count_min() }) + } + "sum_over_time" | "last_over_time" | "present_over_time" | "absent_over_time" + | "delta" | "idelta" | "deriv" | "predict_linear" => + Ok(QeSketchAggOp::Exact(QeExactAgg::Sum)), + "changes" | "resets" => Ok(QeSketchAggOp::default_count_min()), + "rate" | "irate" | "increase" => Ok(QeSketchAggOp::default_count_min()), + other => Err(anyhow!("unsupported PromQL function: {other}")), + } +} + /// Build a `QueryExpr` version of `build_sketched`. fn build_qe_sketched( metric: String, filters: Vec, window: std::time::Duration, - op: super::sketch_algebra::SketchAggOp, + op: QeSketchAggOp, ctx: WalkCtx, ) -> QueryExpr { - let source = QueryExpr::Source(SourceSpec { name: metric }); + let source = QueryExpr::Source(QeSourceSpec { name: metric }); let filtered = apply_qe_filters(source, filters); let windowed = QueryExpr::Window { duration: window, @@ -674,14 +407,14 @@ fn build_qe_sketched( }; let agg = if let Some(k) = ctx.topk { QueryExpr::SketchAgg { - op: super::sketch_algebra::SketchAggOp::CountSketch { k }, - col: super::sketch_algebra::ColumnRef::SampleValue, + op: QeSketchAggOp::CountSketch { k }, + col: QeColumnRef::SampleValue, input: Box::new(windowed), } } else { QueryExpr::SketchAgg { op, - col: super::sketch_algebra::ColumnRef::SampleValue, + col: QeColumnRef::SampleValue, input: Box::new(windowed), } }; @@ -738,7 +471,14 @@ fn apply_qe_partition( match partition { None => input, Some(p) if p.is_empty() => input, - Some(keys) => QueryExpr::Partition { keys, input: Box::new(input) }, + Some(keys) => { + // Convert PromQL by/without → PartitionKeys. + let qe_keys = match keys { + PartitionKeys::By(k) => QePartitionKeys::By(k), + PartitionKeys::Without(k) => QePartitionKeys::Without(k), + }; + QueryExpr::Partition { keys: qe_keys, input: Box::new(input) } + } } } @@ -746,16 +486,12 @@ fn apply_qe_partition( #[cfg(test)] mod tests { - use super::*; - use super::super::sketch_algebra::{ExactAgg, SketchAggOp}; use crate::types::AggType; - - fn parse(q: &str) -> SketchExpr { - parse_promql(q).unwrap_or_else(|e| panic!("parse_promql failed: {e}\nquery={q:?}")) - } + use std::time::Duration; fn pq(q: &str) -> super::super::ParsedQuery { - parse(q).to_parsed_query() + super::super::parse_query(q) + .unwrap_or_else(|e| panic!("parse_query failed: {e}\nquery={q:?}")) } // ── quantile_over_time ──────────────────────────────────────────────────── @@ -825,8 +561,7 @@ mod tests { #[test] fn topk_avg_over_time() { - let expr = parse("topk by (host) (5, avg_over_time(cpu[5m]))"); - let pq = expr.to_parsed_query(); + let pq = pq("topk by (host) (5, avg_over_time(cpu[5m]))"); assert_eq!(pq.aggregations, vec![AggType::Frequency]); } diff --git a/controller/src/query_parser/sketch_algebra.rs b/controller/src/query_parser/sketch_algebra.rs deleted file mode 100644 index 9fa25ff3..00000000 --- a/controller/src/query_parser/sketch_algebra.rs +++ /dev/null @@ -1,721 +0,0 @@ -//! Sketch algebra — the shared intermediate representation (IR) that both the -//! PromQL and SQL parsers compile to. -//! -//! Both parsers are pure front-ends: they walk their respective ASTs and emit -//! a [`SketchExpr`] tree. The optimizer in [`super::sketch_rules`] then -//! applies algebraic rewrite rules before the planner converts the tree to -//! agent configurations. -//! -//! # Operator summary -//! -//! | Operator | Symbol | Description | -//! |---|---|---| -//! | `Source` | — | Base relation or metric stream | -//! | `Filter` | σ | Push-down predicates (WHERE / label matchers) | -//! | `Window` | ψ | Time window (PromQL range; SQL time predicate) | -//! | `Partition` | γ | GROUP BY / `by (dims)` — one sketch per key-tuple | -//! | `Agg` | α | The sketch aggregation itself | -//! | `Dedup` | δ | Deduplicate before ingestion (push-down DISTINCT) | -//! | `TopK` | τ | Retain only the top-K entries | -//! | `Merge` | ⊕ | Merge sketches from multiple branches | -//! | `JoinSketch` | ⋈ₛₖ | Pre-agg sketch on inner side, merge after join | - -use std::collections::HashMap; -use std::time::Duration; - -use crate::types::AggType; -use super::{ParsedQuery, QueryHint, debs_hint}; - -// ── Core IR ─────────────────────────────────────────────────────────────────── - -/// Abstract sketch algebra expression — shared IR for SQL and PromQL. -#[derive(Debug, Clone)] -pub enum SketchExpr { - /// Base relation / metric stream. - Source(SourceSpec), - - /// σ — filter input before any sketch build (WHERE / PromQL label matchers). - Filter { - pred: Vec, - input: Box, - }, - - /// ψ — time window (PromQL range vector `[5m]`; SQL sliding-window predicate). - Window { - duration: Duration, - input: Box, - }, - - /// γ — partition by keys; one sketch instance per distinct key-tuple. - /// Use [`PartitionKeys::Without`] when the PromQL `without (...)` clause is present. - Partition { - keys: PartitionKeys, - input: Box, - }, - - /// α — the sketch aggregation operator. - Agg { - op: SketchAggOp, - col: ColumnRef, - input: Box, - }, - - /// δ — deduplicate on `col` before ingestion. - /// Note: for [`SketchAggOp::HLL`] this node is eliminated by rule R6. - Dedup { - col: String, - input: Box, - }, - - /// τ — top-K post-sketch filter. - TopK { - k: u64, - input: Box, - }, - - /// ⊕ — merge sketches from multiple independent branches. - /// All input [`SketchAggOp`]s must be [`SketchAggOp::is_mergeable`]. - Merge { - inputs: Vec, - }, - - /// ⋈ₛₖ — join push-down: - /// pre-aggregate sketch on inner side by join key, merge after join. - JoinSketch { - join_key: String, - outer: Box, - /// Inner carries its own `Partition` + `Agg` nodes. - inner: Box, - }, -} - -// ── Source ──────────────────────────────────────────────────────────────────── - -#[derive(Debug, Clone)] -pub struct SourceSpec { - /// Table name (SQL) or metric name (PromQL). - pub name: String, -} - -// ── Partition keys ──────────────────────────────────────────────────────────── - -/// How the stream is partitioned. -#[derive(Debug, Clone)] -pub enum PartitionKeys { - /// `by (k1, k2, ...)` — explicit key list. - By(Vec), - /// `without (k1, k2, ...)` — complement; resolved against schema at plan time. - Without(Vec), -} - -impl PartitionKeys { - pub fn keys(&self) -> &[String] { - match self { - PartitionKeys::By(k) | PartitionKeys::Without(k) => k, - } - } - - pub fn is_empty(&self) -> bool { - self.keys().is_empty() - } - - pub fn into_by_keys(self) -> Vec { - match self { - PartitionKeys::By(k) => k, - // For Without, return empty — caller resolves complement. - PartitionKeys::Without(k) => k, - } - } -} - -// ── Column reference ────────────────────────────────────────────────────────── - -/// Which column / field the sketch aggregation targets. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ColumnRef { - /// Explicit column name (SQL: `AVG(price)` → `Named("price")`). - Named(String), - /// The implicit metric sample value (PromQL — always the series value). - SampleValue, - /// All rows / COUNT(*). - Wildcard, -} - -// ── Sketch aggregation operators ────────────────────────────────────────────── - -/// The concrete sketch type used for aggregation. -#[derive(Debug, Clone, PartialEq)] -pub enum SketchAggOp { - /// Count-Min Sketch — frequency per group (COUNT(*) GROUP BY). - CountMin { width: u32, depth: u8 }, - - /// Count Sketch (heavy-hitter) — top-K by frequency. - CountSketch { k: u64 }, - - /// HyperLogLog — distinct-value counting (COUNT DISTINCT). - HLL { registers: u8 }, - - /// DDSketch — quantile estimation. - /// `quantiles` holds the φ values to track; `epsilon` is relative error. - DDSketch { quantiles: Vec, epsilon: f64 }, - - /// Exact running min/max tracker — cheaper than DDSketch for extrema - /// without a GROUP BY (no sketch benefit for global extrema). - ExactMinMax { min: bool, max: bool }, - - /// Hydra — sketch of sketches for multi-dimensional GROUP BY. - /// Maintains one `inner` sketch per distinct `partition_keys` tuple. - Hydra { - inner: Box, - partition_keys: Vec, - }, - - /// Exact passthrough — no sketch benefit (SUM, global COUNT, etc.). - Exact(ExactAgg), -} - -/// Exact (non-sketch) aggregation kinds. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ExactAgg { - Count, - Sum, - /// **Not mergeable** — carries `(sum, count)` in distributed contexts. - Avg, - Min, - Max, -} - -impl SketchAggOp { - /// Returns `true` when two instances of this sketch can be merged - /// (i.e., `sketch(A ∪ B) = merge(sketch(A), sketch(B))`). - pub fn is_mergeable(&self) -> bool { - match self { - SketchAggOp::Exact(ExactAgg::Avg) => false, - SketchAggOp::Hydra { inner, .. } => inner.is_mergeable(), - _ => true, - } - } - - /// Map to the coarse [`AggType`] used by the legacy planner. - pub fn to_agg_type(&self) -> AggType { - match self { - SketchAggOp::HLL { .. } => AggType::Cardinality, - SketchAggOp::CountMin { .. } | SketchAggOp::CountSketch { .. } => AggType::Frequency, - SketchAggOp::DDSketch { .. } | SketchAggOp::ExactMinMax { .. } => AggType::Quantile, - SketchAggOp::Hydra { inner, .. } => inner.to_agg_type(), - SketchAggOp::Exact(_) => AggType::Quantile, - } - } - - /// Extract quantile φ values for DDSketch operators. - pub fn quantiles(&self) -> Vec { - match self { - SketchAggOp::DDSketch { quantiles, .. } => quantiles.clone(), - SketchAggOp::Hydra { inner, .. } => inner.quantiles(), - _ => vec![], - } - } - - /// Whether this op implies `exact_required` (no sketch benefit). - pub fn is_exact(&self) -> bool { - matches!(self, SketchAggOp::Exact(_) | SketchAggOp::ExactMinMax { .. }) - } -} - -// ── Default sketch parameters ───────────────────────────────────────────────── - -impl SketchAggOp { - pub fn default_count_min() -> Self { - SketchAggOp::CountMin { width: 2000, depth: 5 } - } - pub fn default_hll() -> Self { - SketchAggOp::HLL { registers: 14 } - } - pub fn default_ddsketch(quantiles: Vec) -> Self { - SketchAggOp::DDSketch { quantiles, epsilon: 0.01 } - } -} - -// ── Predicates ──────────────────────────────────────────────────────────────── - -/// A single filter predicate pushed down to the collector. -#[derive(Debug, Clone)] -pub struct Predicate { - pub col: String, - pub op: FilterOp, - pub val: FilterVal, -} - -#[derive(Debug, Clone, PartialEq)] -pub enum FilterOp { - Eq, - Ne, - Lt, - Le, - Gt, - Ge, - Like, - NotLike, - IsNull, - IsNotNull, - /// PromQL `=~` label matcher (RE2 syntax). - Regex(String), - /// PromQL `!~` label matcher. - NotRegex(String), -} - -#[derive(Debug, Clone)] -pub enum FilterVal { - Str(String), - Num(f64), - Int(i64), - Null, -} - -// ── Coverage ────────────────────────────────────────────────────────────────── - -/// How completely a query can be served by sketches. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum SketchCoverage { - /// All aggregation columns are sketch-mapped. - Full, - /// Some columns are sketch-mapped; others require exact passthrough - /// (e.g. `MIN(URL)` alongside `COUNT(*)`). - Partial, - /// No sketch applicable; query requires exact execution. - None, -} - -// ── SketchExpr → ParsedQuery bridge (backward compat) ──────────────────────── - -impl SketchExpr { - /// Convert to the legacy [`ParsedQuery`] flat representation consumed by - /// the existing [`crate::analyzer::Analyzer`] and planner. - pub fn to_parsed_query(&self) -> ParsedQuery { - let mut c = PqCollector::default(); - c.visit(self); - c.build() - } -} - -#[derive(Default)] -struct PqCollector { - metric_name: Option, - agg_types: Vec, - group_by_labels: Vec, - label_filters: HashMap, - time_window: Option, - exact_required: bool, - quantiles: Vec, - topk: Option, -} - -impl PqCollector { - fn visit(&mut self, expr: &SketchExpr) { - match expr { - SketchExpr::Source(s) => { - if self.metric_name.is_none() { - self.metric_name = Some(s.name.clone()); - } - } - SketchExpr::Filter { pred, input } => { - for p in pred { - if let (FilterOp::Eq, FilterVal::Str(v)) = (&p.op, &p.val) { - self.label_filters.insert(p.col.clone(), v.clone()); - } - } - self.visit(input); - } - SketchExpr::Window { duration, input } => { - if self.time_window.is_none() { - self.time_window = Some(*duration); - } - self.visit(input); - } - SketchExpr::Partition { keys, input } => { - for k in keys.keys() { - if !self.group_by_labels.contains(k) { - self.group_by_labels.push(k.clone()); - } - } - self.visit(input); - } - SketchExpr::Agg { op, input, .. } => { - self.collect_op(op); - self.visit(input); - } - SketchExpr::TopK { k, input } => { - self.topk = Some(*k); - self.visit(input); - } - SketchExpr::Dedup { input, .. } => self.visit(input), - SketchExpr::Merge { inputs } => { - for i in inputs { self.visit(i); } - } - SketchExpr::JoinSketch { outer, inner, .. } => { - self.visit(outer); - self.visit(inner); - } - } - } - - fn collect_op(&mut self, op: &SketchAggOp) { - match op { - SketchAggOp::HLL { .. } => { - if !self.agg_types.contains(&AggType::Cardinality) { - self.agg_types.push(AggType::Cardinality); - } - } - SketchAggOp::CountMin { .. } | SketchAggOp::CountSketch { .. } => { - if !self.agg_types.contains(&AggType::Frequency) { - self.agg_types.push(AggType::Frequency); - } - } - SketchAggOp::DDSketch { quantiles, .. } => { - if !self.agg_types.contains(&AggType::Quantile) { - self.agg_types.push(AggType::Quantile); - } - for &q in quantiles { - if !self.quantiles.contains(&q) { self.quantiles.push(q); } - } - } - SketchAggOp::ExactMinMax { .. } => { - if !self.agg_types.contains(&AggType::Quantile) { - self.agg_types.push(AggType::Quantile); - } - } - SketchAggOp::Exact(_) => { self.exact_required = true; } - SketchAggOp::Hydra { inner, .. } => self.collect_op(inner), - } - } - - fn build(self) -> ParsedQuery { - let metric_name = self.metric_name.unwrap_or_default(); - let mut qs = self.quantiles; - qs.sort_by(|a, b| a.partial_cmp(b).unwrap()); - qs.dedup(); - let hint = debs_hint( - &metric_name, - &self.agg_types, - &qs, - self.exact_required, - self.topk, - ); - ParsedQuery { - metric_name, - aggregations: self.agg_types, - group_by_labels: self.group_by_labels, - label_filters: self.label_filters, - time_window: self.time_window.unwrap_or(Duration::from_secs(300)), - exact_required: self.exact_required, - quantiles: qs, - hint, - } - } -} - -// ── Tests ───────────────────────────────────────────────────────────────────── - -#[cfg(test)] -mod tests { - use super::*; - - fn source(name: &str) -> SketchExpr { - SketchExpr::Source(SourceSpec { name: name.into() }) - } - - #[test] - fn hll_is_mergeable() { - assert!(SketchAggOp::default_hll().is_mergeable()); - } - - #[test] - fn exact_avg_not_mergeable() { - assert!(!SketchAggOp::Exact(ExactAgg::Avg).is_mergeable()); - } - - #[test] - fn hydra_mergeability_inherits_inner() { - let hydra_hll = SketchAggOp::Hydra { - inner: Box::new(SketchAggOp::default_hll()), - partition_keys: vec!["region".into()], - }; - assert!(hydra_hll.is_mergeable()); - - let hydra_avg = SketchAggOp::Hydra { - inner: Box::new(SketchAggOp::Exact(ExactAgg::Avg)), - partition_keys: vec!["region".into()], - }; - assert!(!hydra_avg.is_mergeable()); - } - - #[test] - fn to_parsed_query_basic() { - // Partition(symbol, Window(5m, Filter(sectype=E, Agg(CountSketch(10), Source(price))))) - let expr = SketchExpr::TopK { - k: 10, - input: Box::new(SketchExpr::Partition { - keys: PartitionKeys::By(vec!["symbol".into()]), - input: Box::new(SketchExpr::Window { - duration: Duration::from_secs(300), - input: Box::new(SketchExpr::Filter { - pred: vec![Predicate { - col: "sectype".into(), - op: FilterOp::Eq, - val: FilterVal::Str("E".into()), - }], - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::CountSketch { k: 10 }, - col: ColumnRef::Wildcard, - input: Box::new(source("financial.last_trade_price")), - }), - }), - }), - }), - }; - let pq = expr.to_parsed_query(); - assert_eq!(pq.metric_name, "financial.last_trade_price"); - assert_eq!(pq.aggregations, vec![AggType::Frequency]); - assert_eq!(pq.group_by_labels, vec!["symbol"]); - assert_eq!(pq.label_filters.get("sectype").map(String::as_str), Some("E")); - assert_eq!(pq.time_window, Duration::from_secs(300)); - assert!(!pq.exact_required); - } - - #[test] - fn to_parsed_query_ddsketch_quantiles() { - let expr = SketchExpr::Agg { - op: SketchAggOp::default_ddsketch(vec![0.25, 0.5, 0.75]), - col: ColumnRef::SampleValue, - input: Box::new(source("cpu")), - }; - let pq = expr.to_parsed_query(); - assert_eq!(pq.aggregations, vec![AggType::Quantile]); - assert_eq!(pq.quantiles, vec![0.25, 0.5, 0.75]); - assert!(!pq.exact_required); - } - - #[test] - fn to_parsed_query_exact_required() { - let expr = SketchExpr::Agg { - op: SketchAggOp::Exact(ExactAgg::Sum), - col: ColumnRef::Named("bytes".into()), - input: Box::new(source("network")), - }; - let pq = expr.to_parsed_query(); - assert!(pq.exact_required); - } - - #[test] - fn partition_keys_without() { - let keys = PartitionKeys::Without(vec!["instance".into()]); - assert_eq!(keys.keys(), &["instance".to_string()]); - assert!(!keys.is_empty()); - } - - // ── Accuracy bound tests ────────────────────────────────────────────────── - // - // These tests verify that the default sketch parameters satisfy the - // documented accuracy guarantees. The bounds are derived from the - // theoretical properties of each sketch family. - - /// DDSketch default: ε = 0.01 → at most 1 % relative rank error. - /// - /// For a stream of n items the rank of a DDSketch quantile estimate r̂ - /// satisfies |r̂ − r| ≤ ε·n. The default ε = 0.01 gives a 1 % SLA. - #[test] - fn ddsketch_default_epsilon_one_pct() { - let op = SketchAggOp::default_ddsketch(vec![0.99]); - match op { - SketchAggOp::DDSketch { epsilon, .. } => { - assert!( - epsilon <= 0.01, - "DDSketch default epsilon {epsilon} exceeds 1 % accuracy SLA" - ); - } - _ => panic!("expected DDSketch"), - } - } - - /// For n = 10 000 items, a 1 % DDSketch (ε = 0.01) guarantees the - /// p99 rank estimate is within ±100 positions of the true rank 9 900. - #[test] - fn ddsketch_rank_error_bound_n10k() { - let n: f64 = 10_000.0; - let epsilon = 0.01_f64; - let true_rank = (0.99 * n) as i64; // 9 900 - let max_error = (epsilon * n).ceil() as i64; // 100 - assert!( - max_error <= 100, - "rank error {max_error} exceeds expected bound for n={n} ε={epsilon}" - ); - // Verify the bound makes sense: estimate ∈ [true_rank-error, true_rank+error]. - let lo = true_rank - max_error; - let hi = true_rank + max_error; - assert!(lo >= 0 && hi <= n as i64); - } - - /// HyperLogLog default: registers = 14 (2¹⁴ = 16 384 buckets). - /// Standard error ≈ 1.04 / √(2^registers) ≈ 0.81 % < 1 %. - #[test] - fn hll_default_registers_error_below_1pct() { - let op = SketchAggOp::default_hll(); - match op { - SketchAggOp::HLL { registers } => { - let buckets = (1u64 << registers) as f64; // 2^registers - let std_error = 1.04 / buckets.sqrt(); - assert!( - std_error < 0.01, - "HLL standard error {std_error:.4} ≥ 1 % for registers={registers}" - ); - } - _ => panic!("expected HLL"), - } - } - - /// CountMin-Sketch default: width = 2 000, depth = 5. - /// Frequency error ε = e / width ≈ 0.136 % of total count N. - /// Failure probability δ = e^(−depth) ≈ 0.67 % < 1 %. - #[test] - fn countmin_default_error_and_failure_prob() { - let op = SketchAggOp::default_count_min(); - match op { - SketchAggOp::CountMin { width, depth } => { - let eps = std::f64::consts::E / width as f64; - let delta = (-(depth as f64)).exp(); - assert!( - eps < 0.002, - "CountMin frequency error ε={eps:.5} should be < 0.2 % of N" - ); - assert!( - delta < 0.01, - "CountMin failure probability δ={delta:.5} should be < 1 %" - ); - } - _ => panic!("expected CountMin"), - } - } - - // ── Mergeability correctness ────────────────────────────────────────────── - // - // Mergeability means sketch(A ∪ B) = merge(sketch(A), sketch(B)). - // All sketches used here satisfy this property except Exact(Avg), - // because avg(A ∪ B) ≠ avg(avg(A), avg(B)) in general. - - #[test] - fn countmin_is_mergeable() { - assert!(SketchAggOp::default_count_min().is_mergeable()); - } - - #[test] - fn countsketch_is_mergeable() { - assert!(SketchAggOp::CountSketch { k: 10 }.is_mergeable()); - } - - #[test] - fn ddsketch_is_mergeable() { - assert!(SketchAggOp::default_ddsketch(vec![0.99]).is_mergeable()); - } - - #[test] - fn exact_minmax_is_mergeable() { - // min(A∪B) = min(min(A), min(B)) — globally mergeable. - assert!(SketchAggOp::ExactMinMax { min: true, max: false }.is_mergeable()); - } - - #[test] - fn exact_sum_and_count_are_mergeable() { - assert!(SketchAggOp::Exact(ExactAgg::Sum).is_mergeable()); - assert!(SketchAggOp::Exact(ExactAgg::Count).is_mergeable()); - } - - #[test] - fn exact_avg_not_mergeable_avg_of_avgs_is_wrong() { - // avg([1,2,3]) = 2, avg([4,5]) = 4.5 - // avg-of-avgs = (2 + 4.5) / 2 = 3.25 ≠ avg([1,2,3,4,5]) = 3 - assert!(!SketchAggOp::Exact(ExactAgg::Avg).is_mergeable()); - } - - #[test] - fn hydra_mergeable_when_inner_is_hll() { - let hydra = SketchAggOp::Hydra { - inner: Box::new(SketchAggOp::default_hll()), - partition_keys: vec!["region".into(), "dc".into()], - }; - assert!(hydra.is_mergeable()); - } - - #[test] - fn hydra_not_mergeable_when_inner_is_avg() { - let hydra = SketchAggOp::Hydra { - inner: Box::new(SketchAggOp::Exact(ExactAgg::Avg)), - partition_keys: vec!["region".into()], - }; - assert!(!hydra.is_mergeable()); - } - - // ── AggType mapping ─────────────────────────────────────────────────────── - - #[test] - fn agg_type_mapping_correct() { - use crate::types::AggType; - assert_eq!(SketchAggOp::default_hll().to_agg_type(), AggType::Cardinality); - assert_eq!(SketchAggOp::default_count_min().to_agg_type(), AggType::Frequency); - assert_eq!(SketchAggOp::CountSketch { k: 5 }.to_agg_type(), AggType::Frequency); - assert_eq!(SketchAggOp::default_ddsketch(vec![0.5]).to_agg_type(),AggType::Quantile); - assert_eq!(SketchAggOp::ExactMinMax { min: true, max: true }.to_agg_type(), AggType::Quantile); - } - - // ── SketchCoverage ──────────────────────────────────────────────────────── - // - // Coverage classifies whether a SELECT can be fully, partially, or not - // at all served by sketches. - - #[test] - fn coverage_full_when_all_sketch() { - let ops: Vec = vec![ - SketchAggOp::default_hll(), - SketchAggOp::default_count_min(), - ]; - let has_sketch = ops.iter().any(|o| !o.is_exact()); - let has_exact = ops.iter().any(|o| o.is_exact()); - let cov = match (has_sketch, has_exact) { - (true, false) => SketchCoverage::Full, - (true, true) => SketchCoverage::Partial, - _ => SketchCoverage::None, - }; - assert_eq!(cov, SketchCoverage::Full); - } - - #[test] - fn coverage_partial_when_exact_mixed_in() { - let ops: Vec = vec![ - SketchAggOp::default_hll(), - SketchAggOp::Exact(ExactAgg::Sum), - ]; - let has_sketch = ops.iter().any(|o| !o.is_exact()); - let has_exact = ops.iter().any(|o| o.is_exact()); - let cov = match (has_sketch, has_exact) { - (true, false) => SketchCoverage::Full, - (true, true) => SketchCoverage::Partial, - _ => SketchCoverage::None, - }; - assert_eq!(cov, SketchCoverage::Partial); - } - - #[test] - fn coverage_none_when_all_exact() { - let ops: Vec = vec![ - SketchAggOp::Exact(ExactAgg::Sum), - SketchAggOp::Exact(ExactAgg::Count), - ]; - let has_sketch = ops.iter().any(|o| !o.is_exact()); - let has_exact = ops.iter().any(|o| o.is_exact()); - let cov = match (has_sketch, has_exact) { - (true, false) => SketchCoverage::Full, - (true, true) => SketchCoverage::Partial, - _ => SketchCoverage::None, - }; - assert_eq!(cov, SketchCoverage::None); - } -} diff --git a/controller/src/query_parser/sketch_rules.rs b/controller/src/query_parser/sketch_rules.rs deleted file mode 100644 index 97667683..00000000 --- a/controller/src/query_parser/sketch_rules.rs +++ /dev/null @@ -1,683 +0,0 @@ -//! Sketch-algebra optimizer — algebraic rewrite rules R1–R8. -//! -//! Rules are applied bottom-up (children first, then parent) until a single -//! fixed-point pass. Multiple passes can be added if needed. -//! -//! # Rule catalogue -//! -//! | Rule | Name | Effect | -//! |---|---|---| -//! | R1 | Filter push-down | `Agg(Filter(X))` → `Agg(Filter pushed into X)` | -//! | R2 | HAVING/WHERE split | separate post-agg key filter from pre-agg tuple filter | -//! | R3 | Sketch linearity | `Agg(Merge([X,Y]))` → `Merge([Agg(X), Agg(Y)])` when mergeable | -//! | R4 | Multi-key Hydra | `Partition([k1,k2], Agg(op,X))` → `Agg(Hydra(op,[k1,k2]),X)` | -//! | R5 | Join push-down | handled at parse time; see `sql.rs` | -//! | R6 | HLL dedup elim | `Agg(HLL, Dedup(col,X))` → `Agg(HLL, X)` | -//! | R7 | Window/Filter swap | `Window(Filter(X))` → `Filter(Window(X))` | -//! | R8 | TopK absorption | absorb TopK(k) into inner CountSketch(k) | - -use super::sketch_algebra::{ - ColumnRef, ExactAgg, FilterOp, FilterVal, PartitionKeys, Predicate, SketchAggOp, SketchExpr, -}; - -// ── Public entry point ──────────────────────────────────────────────────────── - -/// Optimise a [`SketchExpr`] tree by applying all rewrite rules bottom-up. -pub fn optimize(expr: SketchExpr) -> SketchExpr { - // Recurse into children first (post-order), then apply rules at this node. - let expr = rewrite_children(expr); - apply_all(expr) -} - -// ── Children-first recursion ────────────────────────────────────────────────── - -fn rewrite_children(expr: SketchExpr) -> SketchExpr { - match expr { - SketchExpr::Filter { pred, input } => - SketchExpr::Filter { pred, input: Box::new(optimize(*input)) }, - SketchExpr::Window { duration, input } => - SketchExpr::Window { duration, input: Box::new(optimize(*input)) }, - SketchExpr::Partition { keys, input } => - SketchExpr::Partition { keys, input: Box::new(optimize(*input)) }, - SketchExpr::Agg { op, col, input } => - SketchExpr::Agg { op, col, input: Box::new(optimize(*input)) }, - SketchExpr::TopK { k, input } => - SketchExpr::TopK { k, input: Box::new(optimize(*input)) }, - SketchExpr::Dedup { col, input } => - SketchExpr::Dedup { col, input: Box::new(optimize(*input)) }, - SketchExpr::Merge { inputs } => - SketchExpr::Merge { inputs: inputs.into_iter().map(optimize).collect() }, - SketchExpr::JoinSketch { join_key, outer, inner } => - SketchExpr::JoinSketch { - join_key, - outer: Box::new(optimize(*outer)), - inner: Box::new(optimize(*inner)), - }, - leaf => leaf, - } -} - -// ── Apply all rules at one node ─────────────────────────────────────────────── - -fn apply_all(expr: SketchExpr) -> SketchExpr { - let expr = r1_filter_pushdown(expr); - let expr = r2_having_where_split(expr); - let expr = r3_sketch_linearity(expr); - let expr = r4_multi_key_hydra(expr); - let expr = r6_hll_dedup_elim(expr); - let expr = r7_window_filter_swap(expr); - let expr = r8_topk_absorption(expr); - expr -} - -// ── R1: Filter push-down ────────────────────────────────────────────────────── -// -// Agg(op, Filter(pred, X)) → Agg(op, Filter(pred, X)) (already optimal) -// -// The real win is pushing Filter *below* Agg when the Filter is currently -// wrapping the Agg: -// -// Filter(pred, Agg(op, X)) → Agg(op, Filter(pred, X)) -// -// Condition: pred is on a base-relation column, not on the sketch output. -// We allow the push when pred references no aggregate alias (heuristic: no -// function names in the predicate column — that's a HAVING predicate). - -fn r1_filter_pushdown(expr: SketchExpr) -> SketchExpr { - match expr { - SketchExpr::Filter { pred, input } => { - if let SketchExpr::Agg { op, col, input: agg_input } = *input { - // Separate HAVING predicates (cannot be pushed) from WHERE predicates. - // Simple heuristic: predicates whose column matches a GROUP BY key - // or a base column are pushable. We push ALL here; R2 will re-split - // HAVING predicates that must stay above Agg. - return SketchExpr::Agg { - op, - col, - input: Box::new(SketchExpr::Filter { pred, input: agg_input }), - }; - } - SketchExpr::Filter { pred, input } - } - other => other, - } -} - -// ── R2: HAVING / WHERE split ────────────────────────────────────────────────── -// -// When a Filter sits inside an Agg (after R1 pushed it there), split it: -// - Predicates on GROUP BY keys → push further down (pre-agg WHERE) -// - Predicates on agg results → keep above Agg (HAVING) -// -// Because we don't have full schema knowledge at this point we use a simple -// heuristic: we push everything; callers that know HAVING semantics (SQL -// parser) mark predicates with a flag. This rule re-hoists flagged ones. - -fn r2_having_where_split(expr: SketchExpr) -> SketchExpr { - // Implementation: split Filter nodes that carry `having: true` predicates. - // The SQL parser marks HAVING predicates separately so no split is needed - // in the generic optimizer for now. This is a no-op placeholder. - expr -} - -// ── R3: Sketch linearity over Merge (α distributes over ⊕) ───────────────── -// -// Agg(op, Merge([X, Y, ...])) → Merge([Agg(op, X), Agg(op, Y), ...]) -// Condition: op.is_mergeable() - -fn r3_sketch_linearity(expr: SketchExpr) -> SketchExpr { - match expr { - SketchExpr::Agg { ref op, ref col, input: ref boxed } => { - if let SketchExpr::Merge { inputs } = boxed.as_ref() { - if op.is_mergeable() { - let distributed = inputs - .iter() - .cloned() - .map(|branch| SketchExpr::Agg { - op: op.clone(), - col: col.clone(), - input: Box::new(branch), - }) - .collect(); - return SketchExpr::Merge { inputs: distributed }; - } - } - expr - } - other => other, - } -} - -// ── R4: Multi-key Partition → Hydra ────────────────────────────────────────── -// -// Partition([k1,k2,...], Agg(op, X)) -// → Agg(Hydra(op, [k1,k2,...]), X) -// -// when |keys| > 1. - -fn r4_multi_key_hydra(expr: SketchExpr) -> SketchExpr { - if let SketchExpr::Partition { ref keys, .. } = expr { - let key_list = keys.keys().to_vec(); - if key_list.len() <= 1 { - return expr; - } - // Unwrap to get ownership. - if let SketchExpr::Partition { keys, input } = expr { - if let SketchExpr::Agg { op, col, input: agg_in } = *input { - return SketchExpr::Agg { - op: SketchAggOp::Hydra { inner: Box::new(op), partition_keys: key_list }, - col, - input: agg_in, - }; - } else { - // Partition wraps something other than Agg — leave unchanged. - return SketchExpr::Partition { keys, input }; - } - } - } - expr -} - -// ── R6: HLL dedup elimination ───────────────────────────────────────────────── -// -// Agg(HLL, Dedup(col, X)) → Agg(HLL, X) -// HLL inherently deduplicates; the explicit Dedup node is redundant. - -fn r6_hll_dedup_elim(expr: SketchExpr) -> SketchExpr { - if let SketchExpr::Agg { op: SketchAggOp::HLL { registers }, col, input } = expr { - if let SketchExpr::Dedup { input: inner, .. } = *input { - return SketchExpr::Agg { - op: SketchAggOp::HLL { registers }, - col, - input: inner, - }; - } else { - return SketchExpr::Agg { op: SketchAggOp::HLL { registers }, col, input }; - } - } - expr -} - -// ── R7: Window / Filter commutativity ───────────────────────────────────────── -// -// Window(w, Filter(pred, X)) → Filter(pred, Window(w, X)) -// Filter applied before windowing reduces the input stream size. - -fn r7_window_filter_swap(expr: SketchExpr) -> SketchExpr { - if let SketchExpr::Window { duration, input } = expr { - if let SketchExpr::Filter { pred, input: inner } = *input { - return SketchExpr::Filter { - pred, - input: Box::new(SketchExpr::Window { duration, input: inner }), - }; - } else { - return SketchExpr::Window { duration, input }; - } - } - expr -} - -// ── R8: TopK absorption into CountSketch ───────────────────────────────────── -// -// TopK(k, Partition(keys, Agg(CountSketch(k2), X))) -// → Partition(keys, Agg(CountSketch(k=k), X)) when k == k2 -// -// The top-K selection is already encoded in CountSketch; drop the outer TopK. - -fn r8_topk_absorption(expr: SketchExpr) -> SketchExpr { - if let SketchExpr::TopK { k, input } = expr { - if let SketchExpr::Partition { keys, input: part_in } = *input { - if let SketchExpr::Agg { - op: SketchAggOp::CountSketch { k: k2 }, - col, - input: agg_in, - } = *part_in - { - if k == k2 { - // TopK already encoded — absorb. - return SketchExpr::Partition { - keys, - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::CountSketch { k }, - col, - input: agg_in, - }), - }; - } else { - // Different k values — keep TopK, restore inner. - return SketchExpr::TopK { - k, - input: Box::new(SketchExpr::Partition { - keys, - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::CountSketch { k: k2 }, - col, - input: agg_in, - }), - }), - }; - } - } else { - return SketchExpr::TopK { - k, - input: Box::new(SketchExpr::Partition { keys, input: part_in }), - }; - } - } - return SketchExpr::TopK { k, input }; - } - expr -} - -// ── Tests ───────────────────────────────────────────────────────────────────── - -#[cfg(test)] -mod tests { - use super::*; - use super::super::sketch_algebra::{ColumnRef, PartitionKeys, SketchAggOp, SketchExpr, SourceSpec, Predicate, FilterOp, FilterVal}; - use std::time::Duration; - - fn src(name: &str) -> SketchExpr { - SketchExpr::Source(SourceSpec { name: name.into() }) - } - - // ── R3 tests ────────────────────────────────────────────────────────────── - - #[test] - fn r3_distributes_hll_over_merge() { - let expr = SketchExpr::Agg { - op: SketchAggOp::default_hll(), - col: ColumnRef::Named("user".into()), - input: Box::new(SketchExpr::Merge { - inputs: vec![src("R"), src("S")], - }), - }; - let opt = optimize(expr); - // Should become Merge([Agg(HLL,R), Agg(HLL,S)]) - match opt { - SketchExpr::Merge { inputs } => { - assert_eq!(inputs.len(), 2); - for inp in &inputs { - assert!(matches!(inp, SketchExpr::Agg { op: SketchAggOp::HLL { .. }, .. })); - } - } - other => panic!("expected Merge, got {other:?}"), - } - } - - #[test] - fn r3_does_not_distribute_avg() { - let expr = SketchExpr::Agg { - op: SketchAggOp::Exact(ExactAgg::Avg), - col: ColumnRef::Named("price".into()), - input: Box::new(SketchExpr::Merge { - inputs: vec![src("R"), src("S")], - }), - }; - // Avg is not mergeable — Merge should NOT be distributed. - let opt = optimize(expr); - assert!(matches!(opt, SketchExpr::Agg { .. })); - } - - // ── R4 tests ────────────────────────────────────────────────────────────── - - #[test] - fn r4_multi_key_becomes_hydra() { - let expr = SketchExpr::Partition { - keys: PartitionKeys::By(vec!["k1".into(), "k2".into()]), - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::default_hll(), - col: ColumnRef::Named("user".into()), - input: Box::new(src("hits")), - }), - }; - let opt = optimize(expr); - match opt { - SketchExpr::Agg { op: SketchAggOp::Hydra { inner, partition_keys }, .. } => { - assert!(matches!(*inner, SketchAggOp::HLL { .. })); - assert_eq!(partition_keys, vec!["k1", "k2"]); - } - other => panic!("expected Hydra Agg, got {other:?}"), - } - } - - #[test] - fn r4_single_key_unchanged() { - let expr = SketchExpr::Partition { - keys: PartitionKeys::By(vec!["k1".into()]), - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::default_count_min(), - col: ColumnRef::Wildcard, - input: Box::new(src("hits")), - }), - }; - let opt = optimize(expr); - assert!(matches!(opt, SketchExpr::Partition { .. })); - } - - // ── R6 tests ────────────────────────────────────────────────────────────── - - #[test] - fn r6_hll_absorbs_dedup() { - let expr = SketchExpr::Agg { - op: SketchAggOp::default_hll(), - col: ColumnRef::Named("user".into()), - input: Box::new(SketchExpr::Dedup { - col: "user".into(), - input: Box::new(src("hits")), - }), - }; - let opt = optimize(expr); - // Dedup should be gone; HLL goes directly to Source. - match opt { - SketchExpr::Agg { op: SketchAggOp::HLL { .. }, input, .. } => { - assert!(matches!(*input, SketchExpr::Source(_))); - } - other => panic!("expected Agg(HLL,Source), got {other:?}"), - } - } - - // ── R7 tests ────────────────────────────────────────────────────────────── - - #[test] - fn r7_filter_hoisted_before_window() { - let pred = vec![Predicate { col: "env".into(), op: FilterOp::Eq, val: FilterVal::Str("prod".into()) }]; - let expr = SketchExpr::Window { - duration: Duration::from_secs(300), - input: Box::new(SketchExpr::Filter { - pred: pred.clone(), - input: Box::new(src("cpu")), - }), - }; - let opt = optimize(expr); - // Should become Filter(Window(Source)) - match opt { - SketchExpr::Filter { input, .. } => { - assert!(matches!(*input, SketchExpr::Window { .. })); - } - other => panic!("expected Filter(Window(...)), got {other:?}"), - } - } - - // ── R8 tests ────────────────────────────────────────────────────────────── - - #[test] - fn r8_topk_absorbed_same_k() { - let expr = SketchExpr::TopK { - k: 10, - input: Box::new(SketchExpr::Partition { - keys: PartitionKeys::By(vec!["phrase".into()]), - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::CountSketch { k: 10 }, - col: ColumnRef::Wildcard, - input: Box::new(src("hits")), - }), - }), - }; - let opt = optimize(expr); - // TopK should be absorbed into the Partition/Agg - assert!(matches!(opt, SketchExpr::Partition { .. })); - } - - #[test] - fn r8_topk_kept_different_k() { - let expr = SketchExpr::TopK { - k: 5, - input: Box::new(SketchExpr::Partition { - keys: PartitionKeys::By(vec!["phrase".into()]), - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::CountSketch { k: 10 }, - col: ColumnRef::Wildcard, - input: Box::new(src("hits")), - }), - }), - }; - let opt = optimize(expr); - assert!(matches!(opt, SketchExpr::TopK { k: 5, .. })); - } - - // ── R1 additional ───────────────────────────────────────────────────────── - - /// R1 basic: Filter wrapping Agg is pushed inside the Agg. - /// Filter(pred, Agg(op, X)) → Agg(op, Filter(pred, X)) - #[test] - fn r1_filter_pushed_below_agg() { - let pred = vec![Predicate { - col: "region".into(), - op: FilterOp::Eq, - val: FilterVal::Str("us-east".into()), - }]; - let expr = SketchExpr::Filter { - pred: pred.clone(), - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::default_count_min(), - col: ColumnRef::Wildcard, - input: Box::new(src("hits")), - }), - }; - let opt = optimize(expr); - // Must become Agg(..., Filter(...)) - match opt { - SketchExpr::Agg { input, .. } => { - assert!(matches!(*input, SketchExpr::Filter { .. }), - "Filter should be inside Agg after R1"); - } - other => panic!("expected Agg after R1, got {other:?}"), - } - } - - /// R1 accuracy preservation: pushing Filter below Agg must not change - /// the sketch operator or its parameters. - #[test] - fn r1_preserves_sketch_epsilon() { - let pred = vec![Predicate { - col: "env".into(), - op: FilterOp::Eq, - val: FilterVal::Str("prod".into()), - }]; - let original_op = SketchAggOp::DDSketch { quantiles: vec![0.99], epsilon: 0.01 }; - let expr = SketchExpr::Filter { - pred, - input: Box::new(SketchExpr::Agg { - op: original_op.clone(), - col: ColumnRef::SampleValue, - input: Box::new(src("latency")), - }), - }; - let opt = optimize(expr); - match opt { - SketchExpr::Agg { op, .. } => { - assert_eq!(op, original_op, - "R1 must not alter the DDSketch epsilon or quantile parameters"); - } - other => panic!("expected Agg, got {other:?}"), - } - } - - // ── R3 additional ───────────────────────────────────────────────────────── - - /// R3 distributes CountMin over UNION ALL branches and preserves - /// the same width/depth parameters on every branch. - #[test] - fn r3_countmin_branches_have_same_params() { - let op = SketchAggOp::default_count_min(); - let expr = SketchExpr::Agg { - op: op.clone(), - col: ColumnRef::Wildcard, - input: Box::new(SketchExpr::Merge { - inputs: vec![src("R"), src("S"), src("T")], - }), - }; - let opt = optimize(expr); - match opt { - SketchExpr::Merge { inputs } => { - assert_eq!(inputs.len(), 3); - for branch in &inputs { - match branch { - SketchExpr::Agg { op: branch_op, .. } => { - assert_eq!(branch_op, &op, - "Each UNION branch must use identical CountMin parameters"); - } - other => panic!("expected Agg branch, got {other:?}"), - } - } - } - other => panic!("expected Merge after R3, got {other:?}"), - } - } - - /// R3 distributes DDSketch over Merge, preserving ε on every branch. - #[test] - fn r3_ddsketch_epsilon_preserved_on_each_branch() { - let op = SketchAggOp::DDSketch { quantiles: vec![0.95], epsilon: 0.005 }; - let expr = SketchExpr::Agg { - op: op.clone(), - col: ColumnRef::SampleValue, - input: Box::new(SketchExpr::Merge { - inputs: vec![src("A"), src("B")], - }), - }; - let opt = optimize(expr); - match opt { - SketchExpr::Merge { inputs } => { - for branch in &inputs { - match branch { - SketchExpr::Agg { op: bop, .. } => assert_eq!(bop, &op), - other => panic!("expected Agg, got {other:?}"), - } - } - } - other => panic!("expected Merge, got {other:?}"), - } - } - - // ── R4 additional ───────────────────────────────────────────────────────── - - /// R4 Hydra wraps the inner op without altering its accuracy parameters. - #[test] - fn r4_hydra_preserves_inner_ddsketch_epsilon() { - let inner_op = SketchAggOp::DDSketch { quantiles: vec![0.5], epsilon: 0.01 }; - let expr = SketchExpr::Partition { - keys: PartitionKeys::By(vec!["region".into(), "dc".into()]), - input: Box::new(SketchExpr::Agg { - op: inner_op.clone(), - col: ColumnRef::SampleValue, - input: Box::new(src("latency")), - }), - }; - let opt = optimize(expr); - match opt { - SketchExpr::Agg { - op: SketchAggOp::Hydra { inner, partition_keys }, - .. - } => { - assert_eq!(*inner, inner_op, - "Hydra must not alter inner DDSketch accuracy parameters"); - assert_eq!(partition_keys, vec!["region", "dc"]); - } - other => panic!("expected Hydra Agg, got {other:?}"), - } - } - - // ── Optimizer idempotency ───────────────────────────────────────────────── - - /// Applying `optimize` twice must yield a structurally identical tree. - /// This verifies the rules reach a fixed point in one pass. - #[test] - fn optimize_is_idempotent() { - // Build a tree that exercises multiple rules: R3, R4, R7. - let pred = vec![Predicate { - col: "env".into(), op: FilterOp::Eq, val: FilterVal::Str("prod".into()), - }]; - let expr = SketchExpr::Partition { - keys: PartitionKeys::By(vec!["region".into(), "dc".into()]), - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::default_hll(), - col: ColumnRef::Named("user".into()), - input: Box::new(SketchExpr::Merge { - inputs: vec![ - SketchExpr::Window { - duration: Duration::from_secs(300), - input: Box::new(SketchExpr::Filter { - pred: pred.clone(), - input: Box::new(src("R")), - }), - }, - src("S"), - ], - }), - }), - }; - let once = optimize(expr.clone()); - let twice = optimize(once.clone()); - // Structural equality: format the debug output (simplest proxy for deep eq). - assert_eq!(format!("{once:?}"), format!("{twice:?}"), - "optimize is not idempotent — a second pass changed the tree"); - } - - // ── End-to-end accuracy pipeline ───────────────────────────────────────── - - /// Build the tree for `COUNT(DISTINCT UserID) FROM hits UNION ALL - /// SELECT COUNT(DISTINCT UserID) FROM hits2`, optimize it, and verify: - /// - /// 1. R3 distributed HLL over the Merge branches. - /// 2. Every branch uses the same HLL parameters (registers = 14). - /// 3. The standard error bound 1.04/√(2^14) < 1 % is preserved. - #[test] - fn end_to_end_union_hll_accuracy_preserved() { - let hll = SketchAggOp::default_hll(); - let expr = SketchExpr::Agg { - op: hll.clone(), - col: ColumnRef::Named("UserID".into()), - input: Box::new(SketchExpr::Merge { - inputs: vec![src("hits"), src("hits2")], - }), - }; - let opt = optimize(expr); - match opt { - SketchExpr::Merge { inputs } => { - assert_eq!(inputs.len(), 2, "both UNION branches must be present"); - for branch in &inputs { - match branch { - SketchExpr::Agg { op: SketchAggOp::HLL { registers }, .. } => { - assert_eq!(*registers, 14); - let std_err = 1.04 / ((1u64 << registers) as f64).sqrt(); - assert!(std_err < 0.01, - "HLL standard error {std_err:.4} must be < 1 %"); - } - other => panic!("expected HLL Agg on each branch, got {other:?}"), - } - } - } - other => panic!("expected Merge after R3, got {other:?}"), - } - } - - /// Build a DDSketch pipeline through R1 (filter push-down) and verify - /// the ε = 0.01 accuracy guarantee survives optimization. - #[test] - fn end_to_end_ddsketch_epsilon_survives_filter_pushdown() { - let pred = vec![Predicate { - col: "service".into(), - op: FilterOp::Eq, - val: FilterVal::Str("checkout".into()), - }]; - let expr = SketchExpr::Filter { - pred, - input: Box::new(SketchExpr::Agg { - op: SketchAggOp::DDSketch { quantiles: vec![0.99], epsilon: 0.01 }, - col: ColumnRef::SampleValue, - input: Box::new(src("latency")), - }), - }; - let opt = optimize(expr); - // After R1: Agg(DDSketch, Filter(Source)) - match opt { - SketchExpr::Agg { op: SketchAggOp::DDSketch { epsilon, quantiles }, .. } => { - assert_eq!(epsilon, 0.01, "ε must not change through R1"); - assert_eq!(quantiles, vec![0.99]); - } - other => panic!("expected DDSketch Agg after R1, got {other:?}"), - } - } -} diff --git a/controller/src/query_parser/sql.rs b/controller/src/query_parser/sql.rs index bb8c3dc4..02ae669a 100644 --- a/controller/src/query_parser/sql.rs +++ b/controller/src/query_parser/sql.rs @@ -1,4 +1,4 @@ -//! SQL → SketchExpr compiler. +//! SQL → QueryExpr compiler. //! //! Implements the `AST_SQL_to_sketch` algorithm from the design doc //! (`docs/Top-Down SQL-to-sketch mapping.pdf`). @@ -7,17 +7,15 @@ //! //! For each FUNCTION edge in the SELECT projection (leaf → root): //! 1. Collect context: GROUP BY, WHERE, HAVING, JOIN, DISTINCT, UNION ALL -//! 2. Call `try_replace` → choose `SketchAggOp` (or Exact passthrough) -//! 3. Apply structural transformations: -//! - WHERE predicates → pushed into Filter node -//! - HAVING predicates → stay above Agg as a second Filter -//! - Multi-dim GROUP BY → Hydra via rewrite rule R4 -//! - JOIN → JoinSketch push-down -//! - UNION ALL → Merge node (sketch linearity via rule R3) -//! - DISTINCT → Dedup node (absorbed by HLL via rule R6) +//! 2. Emit the corresponding `QueryExpr` node: +//! - WHERE predicates → `Filter { ScalarExpr }` +//! - GROUP BY + aggs → `Aggregate { keys, aggs }` +//! - ORDER BY + LIMIT → `Sort` + `Limit` (→ `TopK` via optimizer R5) +//! - JOIN … ON → `Join { kind, pred }` +//! - UNION ALL → `SetOp { Union, all: true }` //! -//! Multiple aggregations in one SELECT each emit their own `SketchAggOp`, -//! then are combined with a `Merge` node (coverage = Partial when some are Exact). +//! Multiple aggregations in one SELECT each emit their own `AggItem`, +//! collected inside a single `Aggregate` node. use std::collections::HashMap; use std::time::Duration; @@ -31,606 +29,18 @@ use sqlparser::ast::{ Value, ValueWithSpan, }; use sqlparser::dialect::GenericDialect; -use sqlparser::parser::Parser; -use super::sketch_algebra::{ - ColumnRef, ExactAgg, FilterOp, FilterVal, PartitionKeys, Predicate, - SketchAggOp, SketchExpr, SketchCoverage, SourceSpec, +use crate::algebra::expr::{ + AggFunc, AggItem as AlgAggItem, BinaryOpKind, ColumnRef, JoinKind, LiteralValue, + ProjectItem, QueryExpr, ScalarExpr, SetOpKind, SortKey, SourceSpec, }; -use super::sketch_rules::optimize; // ── Public entry point ──────────────────────────────────────────────────────── -/// Parse a SQL SELECT statement into an optimised [`SketchExpr`]. -pub fn parse_sql(sql: &str) -> anyhow::Result { - let dialect = GenericDialect {}; - let mut stmts = Parser::parse_sql(&dialect, sql) - .with_context(|| format!("SQL parse error: {sql:?}"))?; - let stmt = stmts.pop().ok_or_else(|| anyhow!("no SQL statement found"))?; - let query = match stmt { - Statement::Query(q) => *q, - other => return Err(anyhow!("expected SELECT, got {:?}", other)), - }; - let sketch = extract_from_query(&query)?; - Ok(optimize(sketch)) -} - -// ── Query-level dispatch ────────────────────────────────────────────────────── - -fn extract_from_query(query: &Query) -> anyhow::Result { - let order_by: Vec = match &query.order_by { - Some(OrderBy { kind: OrderByKind::Expressions(exprs), .. }) => exprs.clone(), - _ => vec![], - }; - let limit: Option<&Expr> = match &query.limit_clause { - Some(LimitClause::LimitOffset { limit: Some(e), .. }) => Some(e), - Some(LimitClause::OffsetCommaLimit { limit: e, .. }) => Some(e), - _ => None, - }; - extract_from_set_expr(query.body.as_ref(), &order_by, limit) -} - -fn extract_from_set_expr( - set_expr: &SetExpr, - order_by: &[OrderByExpr], - limit: Option<&Expr>, -) -> anyhow::Result { - match set_expr { - SetExpr::Select(sel) => extract_from_select(sel, order_by, limit), - SetExpr::Query(inner) => extract_from_query(inner), - - // UNION ALL → Merge of both branches (sketch linearity rule R3 will - // distribute Agg over Merge if ops are mergeable). - SetExpr::SetOperation { left, right, op: SetOperator::Union, .. } => { - let left_expr = extract_from_set_expr(left, &[], None)?; - let right_expr = extract_from_set_expr(right, &[], None)?; - Ok(SketchExpr::Merge { inputs: vec![left_expr, right_expr] }) - } - - other => Err(anyhow!("unsupported query body: {:?}", other)), - } -} - -// ── SELECT-level extraction ─────────────────────────────────────────────────── - -fn extract_from_select( - sel: &Select, - order_by: &[OrderByExpr], - limit: Option<&Expr>, -) -> anyhow::Result { - // ── Step 1: source table(s) ─────────────────────────────────────────────── - let metric_name = extract_table_name(sel)?; - - // ── Step 2: WHERE predicates ────────────────────────────────────────────── - let where_preds: Vec = sel.selection.as_ref() - .map(|e| collect_predicates(e)) - .unwrap_or_default(); - - // ── Step 3: GROUP BY keys ───────────────────────────────────────────────── - let group_by_keys = extract_group_by(&sel.group_by); - - // ── Step 4: HAVING predicates ───────────────────────────────────────────── - let having_preds: Vec = sel.having.as_ref() - .map(|e| collect_predicates(e)) - .unwrap_or_default(); - - // ── Step 5: top-K detection (ORDER BY … DESC LIMIT k) ──────────────────── - let topk: Option = detect_topk(order_by, limit); - - // ── Step 6: JOIN detection ──────────────────────────────────────────────── - let join_info: Option = extract_join_info(sel); - - // ── Step 7: SELECT-DISTINCT flag ────────────────────────────────────────── - let select_distinct = sel.distinct.is_some(); - - // ── Step 8: aggregation functions from SELECT list ──────────────────────── - let agg_items = collect_agg_items(&sel.projection); - - if agg_items.is_empty() { - // No aggregation — bare SELECT (e.g. SELECT col FROM t WHERE …). - let source = SketchExpr::Source(SourceSpec { name: metric_name }); - let filtered = apply_filters(source, where_preds); - return Ok(filtered); - } - - // ── Step 9: try_replace each agg item → SketchAggOp ────────────────────── - let ops: Vec<(SketchAggOp, ColumnRef)> = agg_items - .iter() - .map(|item| try_replace(item, &group_by_keys, topk, join_info.as_ref(), select_distinct)) - .collect(); - - // ── Step 10: determine coverage ─────────────────────────────────────────── - let has_sketch = ops.iter().any(|(op, _)| !op.is_exact()); - let has_exact = ops.iter().any(|(op, _)| op.is_exact()); - let _coverage = match (has_sketch, has_exact) { - (true, false) => SketchCoverage::Full, - (true, true) => SketchCoverage::Partial, - (false, _) => SketchCoverage::None, - }; - - // ── Step 11: assemble the SketchExpr tree ──────────────────────────────── - - // Base source with pushed-down WHERE. - let source = SketchExpr::Source(SourceSpec { name: metric_name.clone() }); - let filtered = apply_filters(source, where_preds); - - // If there's a JOIN, build JoinSketch. Otherwise use the filtered source. - let base = if let Some(ji) = join_info { - build_join_sketch(filtered, ji, &ops, &group_by_keys)? - } else { - build_agg_tree(filtered, ops, &group_by_keys, topk, having_preds) - }; - - Ok(base) -} - -// ── try_replace: aggregation function → SketchAggOp ────────────────────────── - -struct AggItem { - kind: AggKind, - col: ColumnRef, - distinct: bool, -} - -#[derive(Debug, Clone)] -enum AggKind { - Count, - Sum, - Avg, - Min, - Max, -} - -/// Map one aggregation item to its best sketch operator. -/// -/// Priority (doc §Generate_SQL_Aggregation_Sketch_Mapping): -/// 1. Cardinality (COUNT DISTINCT) → HLL -/// 2. Frequency heavy-hitter (COUNT + top-K) → CountSketch -/// 3. Frequency (COUNT GROUP BY) → CountMin -/// 4. Quantile (AVG/MIN/MAX with GROUP BY) → DDSketch -/// 5. Extrema without GROUP BY → ExactMinMax -/// 6. Sum / global count → Exact -fn try_replace( - item: &AggItem, - group_by: &[String], - topk: Option, - _join_info: Option<&JoinInfo>, - _distinct: bool, -) -> (SketchAggOp, ColumnRef) { - let has_group = !group_by.is_empty(); - - match item.kind { - AggKind::Count if item.distinct => { - // COUNT(DISTINCT col) → HLL - (SketchAggOp::default_hll(), item.col.clone()) - } - AggKind::Count => { - if let Some(k) = topk { - // ORDER BY … DESC LIMIT k → heavy-hitter CountSketch - (SketchAggOp::CountSketch { k }, item.col.clone()) - } else if has_group { - // COUNT(*) GROUP BY → frequency CountMin - (SketchAggOp::default_count_min(), item.col.clone()) - } else { - // Global COUNT(*) — exact - (SketchAggOp::Exact(ExactAgg::Count), item.col.clone()) - } - } - AggKind::Avg => { - if has_group { - // AVG per group → DDSketch(p50) — mergeable proxy - (SketchAggOp::default_ddsketch(vec![0.5]), item.col.clone()) - } else { - // Global AVG — not mergeable as-is; store exact (sum,count) - (SketchAggOp::Exact(ExactAgg::Avg), item.col.clone()) - } - } - AggKind::Min => { - if has_group { - (SketchAggOp::default_ddsketch(vec![0.0]), item.col.clone()) - } else { - (SketchAggOp::ExactMinMax { min: true, max: false }, item.col.clone()) - } - } - AggKind::Max => { - if has_group { - (SketchAggOp::default_ddsketch(vec![1.0]), item.col.clone()) - } else { - (SketchAggOp::ExactMinMax { min: false, max: true }, item.col.clone()) - } - } - AggKind::Sum => { - // SUM is always exact — no sketch benefit. - (SketchAggOp::Exact(ExactAgg::Sum), item.col.clone()) - } - } -} - -// ── Tree assembly ───────────────────────────────────────────────────────────── - -fn build_agg_tree( - base: SketchExpr, - ops: Vec<(SketchAggOp, ColumnRef)>, - group_by: &[String], - topk: Option, - having: Vec, -) -> SketchExpr { - // One Agg node per op; combine with Merge if multiple. - let agg_nodes: Vec = ops - .into_iter() - .map(|(op, col)| SketchExpr::Agg { op, col, input: Box::new(base.clone()) }) - .collect(); - - let merged = if agg_nodes.len() == 1 { - agg_nodes.into_iter().next().unwrap() - } else { - SketchExpr::Merge { inputs: agg_nodes } - }; - - // Apply GROUP BY partition. - let partitioned = if group_by.is_empty() { - merged - } else { - SketchExpr::Partition { - keys: PartitionKeys::By(group_by.to_vec()), - input: Box::new(merged), - } - }; - - // Apply HAVING as a post-agg filter. - let after_having = apply_filters(partitioned, having); - - // Apply top-K. - if let Some(k) = topk { - SketchExpr::TopK { k, input: Box::new(after_having) } - } else { - after_having - } -} - -// ── JOIN push-down ───────────────────────────────────────────────────────────── - -struct JoinInfo { - inner_table: String, - join_key: String, - outer_key: String, -} - -fn build_join_sketch( - outer_filtered: SketchExpr, - ji: JoinInfo, - ops: &[(SketchAggOp, ColumnRef)], - outer_group_by: &[String], -) -> anyhow::Result { - // Pre-aggregate on the inner table grouped by the join key. - let inner_source = SketchExpr::Source(SourceSpec { name: ji.inner_table.clone() }); - let inner_agg = if let Some((op, col)) = ops.first() { - SketchExpr::Agg { op: op.clone(), col: col.clone(), input: Box::new(inner_source) } - } else { - inner_source - }; - let inner_partitioned = SketchExpr::Partition { - keys: PartitionKeys::By(vec![ji.join_key.clone()]), - input: Box::new(inner_agg), - }; - - Ok(SketchExpr::JoinSketch { - join_key: ji.outer_key, - outer: Box::new(outer_filtered), - inner: Box::new(inner_partitioned), - }) -} - -// ── AST helpers: aggregation collection ────────────────────────────────────── - -fn collect_agg_items(projection: &[SelectItem]) -> Vec { - let mut out = Vec::new(); - for item in projection { - let expr = match item { - SelectItem::UnnamedExpr(e) => e, - SelectItem::ExprWithAlias { expr, .. } => expr, - _ => continue, - }; - collect_agg_from_expr(expr, &mut out); - } - out -} - -fn collect_agg_from_expr(expr: &Expr, out: &mut Vec) { - match expr { - Expr::Function(f) => { - let fn_name = f.name.0.last() - .and_then(|i| i.as_ident()) - .map(|id| id.value.to_uppercase()) - .unwrap_or_default(); - - let (distinct, args) = match &f.args { - FunctionArguments::List(FunctionArgumentList { duplicate_treatment, args, .. }) => { - let is_distinct = matches!( - duplicate_treatment, - Some(DuplicateTreatment::Distinct) - ); - (is_distinct, args.as_slice()) - } - _ => (false, &[][..]), - }; - - let col = first_col_from_args(args); - - let kind = match fn_name.as_str() { - "COUNT" => AggKind::Count, - "SUM" => AggKind::Sum, - "AVG" => AggKind::Avg, - "MIN" => AggKind::Min, - "MAX" => AggKind::Max, - _ => return, - }; - - out.push(AggItem { kind, col, distinct }); - } - Expr::BinaryOp { left, right, .. } => { - collect_agg_from_expr(left, out); - collect_agg_from_expr(right, out); - } - Expr::Nested(inner) => collect_agg_from_expr(inner, out), - _ => {} - } -} - -fn first_col_from_args(args: &[FunctionArg]) -> ColumnRef { - for arg in args { - match arg { - FunctionArg::Unnamed(FunctionArgExpr::Wildcard) => return ColumnRef::Wildcard, - FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Identifier(id))) => { - return ColumnRef::Named(id.value.clone()); - } - FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::CompoundIdentifier(parts))) => { - if let Some(last) = parts.last() { - return ColumnRef::Named(last.value.clone()); - } - } - _ => {} - } - } - ColumnRef::Wildcard -} - -// ── AST helpers: GROUP BY ────────────────────────────────────────────────────── - -fn extract_group_by(group_by: &GroupByExpr) -> Vec { - let exprs = match group_by { - GroupByExpr::All(_) => return vec![], - GroupByExpr::Expressions(e, _) => e, - }; - exprs.iter().filter_map(|e| match e { - Expr::Identifier(id) => Some(id.value.clone()), - Expr::CompoundIdentifier(parts) => parts.last().map(|i| i.value.clone()), - _ => None, - }).collect() -} - -// ── AST helpers: WHERE / predicate extraction ───────────────────────────────── - -fn collect_predicates(expr: &Expr) -> Vec { - let mut out = Vec::new(); - collect_pred_rec(expr, &mut out); - out -} - -fn collect_pred_rec(expr: &Expr, out: &mut Vec) { - match expr { - // col = 'v' - Expr::BinaryOp { left, op: BinaryOperator::Eq, right } => { - if let Some(p) = binary_pred(left, BinaryOperator::Eq, right) { out.push(p); } - } - // col <> 'v' - Expr::BinaryOp { left, op: BinaryOperator::NotEq, right } => { - if let Some(p) = binary_pred(left, BinaryOperator::NotEq, right) { out.push(p); } - } - // col > v - Expr::BinaryOp { left, op: BinaryOperator::Gt, right } => { - if let Some(p) = binary_pred(left, BinaryOperator::Gt, right) { out.push(p); } - } - // col >= v - Expr::BinaryOp { left, op: BinaryOperator::GtEq, right } => { - if let Some(p) = binary_pred(left, BinaryOperator::GtEq, right) { out.push(p); } - } - // col < v - Expr::BinaryOp { left, op: BinaryOperator::Lt, right } => { - if let Some(p) = binary_pred(left, BinaryOperator::Lt, right) { out.push(p); } - } - // col <= v - Expr::BinaryOp { left, op: BinaryOperator::LtEq, right } => { - if let Some(p) = binary_pred(left, BinaryOperator::LtEq, right) { out.push(p); } - } - // AND — recurse into both sides - Expr::BinaryOp { left, op: BinaryOperator::And, right } => { - collect_pred_rec(left, out); - collect_pred_rec(right, out); - } - // col LIKE '%v%' - Expr::Like { expr, pattern, negated, .. } => { - if let Some(col) = col_name(expr) { - if let Some(val) = literal_str(pattern) { - let op = if *negated { FilterOp::NotLike } else { FilterOp::Like }; - out.push(Predicate { col, op, val: FilterVal::Str(val) }); - } - } - } - // col IS NULL / IS NOT NULL - Expr::IsNull(inner) => { - if let Some(col) = col_name(inner) { - out.push(Predicate { col, op: FilterOp::IsNull, val: FilterVal::Null }); - } - } - Expr::IsNotNull(inner) => { - if let Some(col) = col_name(inner) { - out.push(Predicate { col, op: FilterOp::IsNotNull, val: FilterVal::Null }); - } - } - Expr::Nested(inner) => collect_pred_rec(inner, out), - // OR predicates span both columns — cannot push to collector; skip. - _ => {} - } -} - -fn binary_pred(left: &Expr, sql_op: BinaryOperator, right: &Expr) -> Option { - let col = col_name(left)?; - let (op, val) = match sql_op { - BinaryOperator::Eq => (FilterOp::Eq, literal_val(right)?), - BinaryOperator::NotEq => (FilterOp::Ne, literal_val(right)?), - BinaryOperator::Gt => (FilterOp::Gt, literal_val(right)?), - BinaryOperator::GtEq => (FilterOp::Ge, literal_val(right)?), - BinaryOperator::Lt => (FilterOp::Lt, literal_val(right)?), - BinaryOperator::LtEq => (FilterOp::Le, literal_val(right)?), - _ => return None, - }; - Some(Predicate { col, op, val }) -} - -fn col_name(expr: &Expr) -> Option { - match expr { - Expr::Identifier(id) => Some(id.value.clone()), - Expr::CompoundIdentifier(parts) => parts.last().map(|i| i.value.clone()), - _ => None, - } -} - -fn literal_val(expr: &Expr) -> Option { - let v = match expr { - Expr::Value(vws) => &vws.value, - _ => return None, - }; - match v { - Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => { - Some(FilterVal::Str(s.clone())) - } - Value::Number(n, _) => { - if let Ok(i) = n.parse::() { - Some(FilterVal::Int(i)) - } else if let Ok(f) = n.parse::() { - Some(FilterVal::Num(f)) - } else { - None - } - } - Value::Null => Some(FilterVal::Null), - _ => None, - } -} - -fn literal_str(expr: &Expr) -> Option { - let v = match expr { - Expr::Value(vws) => &vws.value, - _ => return None, - }; - match v { - Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => Some(s.clone()), - _ => None, - } -} - -// ── AST helpers: table name ─────────────────────────────────────────────────── - -fn extract_table_name(sel: &Select) -> anyhow::Result { - sel.from.first() - .and_then(|t| match &t.relation { - TableFactor::Table { name, .. } => Some(object_name_str(name)), - _ => None, - }) - .ok_or_else(|| anyhow!("could not determine table name from FROM clause")) -} - -fn object_name_str(name: &ObjectName) -> String { - name.0.iter() - .map(|i| i.as_ident().map(|id| id.value.as_str()).unwrap_or("")) - .collect::>() - .join(".") -} - -// ── AST helpers: top-K detection ───────────────────────────────────────────── - -fn detect_topk(order_by: &[OrderByExpr], limit: Option<&Expr>) -> Option { - let limit_n = match limit? { - Expr::Value(ValueWithSpan { value: Value::Number(n, _), .. }) => n.parse::().ok()?, - _ => return None, - }; - // Must have at least one DESC (or default) ORDER BY. - let has_desc = order_by.iter().any(|o| matches!(o.options.asc, Some(false) | None)); - if has_desc { Some(limit_n) } else { None } -} - -// ── AST helpers: JOIN extraction ────────────────────────────────────────────── - -fn extract_join_info(sel: &Select) -> Option { - let table_with_joins = sel.from.first()?; - let join = table_with_joins.joins.first()?; - - let inner_table = match &join.relation { - TableFactor::Table { name, .. } => object_name_str(name), - _ => return None, - }; - - // Extract the ON key from JOIN … ON a.key = b.key. - let (join_key, outer_key) = extract_join_keys(join)?; - - Some(JoinInfo { inner_table, join_key, outer_key }) -} - -fn extract_join_keys(join: &Join) -> Option<(String, String)> { - let constraint = match &join.join_operator { - JoinOperator::Inner(c) - | JoinOperator::LeftOuter(c) - | JoinOperator::RightOuter(c) - | JoinOperator::FullOuter(c) => c, - _ => return None, - }; - let on_expr = match constraint { - JoinConstraint::On(e) => e, - _ => return None, - }; - // Expect: a.key = b.key or key = key - if let Expr::BinaryOp { left, op: BinaryOperator::Eq, right } = on_expr { - let lk = col_name(left)?; - let rk = col_name(right)?; - Some((rk, lk)) // inner key, outer key - } else { - None - } -} - -// ── Shared helper ───────────────────────────────────────────────────────────── - -fn apply_filters(input: SketchExpr, pred: Vec) -> SketchExpr { - if pred.is_empty() { input } else { - SketchExpr::Filter { pred, input: Box::new(input) } - } -} - -// ── Direct QueryExpr emission ───────────────────────────────────────────────── -// -// `parse_sql_expr` walks the same SQL AST but emits [`QueryExpr`] nodes -// natively, preserving semantic nodes that SketchExpr flattens: -// -// | SQL construct | SketchExpr (old) | QueryExpr (new) | -// |-----------------|------------------------|---------------------------------| -// | ORDER BY + LIMIT| TopK node | Sort + Limit (→ TopK via R5) | -// | JOIN … ON | JoinSketch | Join { kind, pred } | -// | UNION ALL | Merge | SetOp { Union, all: true } | -// | GROUP BY + aggs | Agg + Partition + Merge| Aggregate { keys, aggs } | -// | WHERE | Filter (Predicate list)| Filter { ScalarExpr tree } | - -use crate::algebra::expr::{ - AggFunc, AggItem as AlgAggItem, BinaryOpKind, JoinKind, LiteralValue, ProjectItem, - QueryExpr, ScalarExpr, SetOpKind, SortKey, -}; -use crate::query_parser::sketch_algebra::ColumnRef as SColumnRef; - -/// Parse a SQL SELECT statement directly into a [`QueryExpr`] tree. +/// Parse a SQL SELECT statement into a [`QueryExpr`] tree. /// -/// Unlike `parse_sql` (which emits `SketchExpr`), this preserves `Sort`, -/// `Limit`, `Join`, and `SetOp` nodes natively so the [`crate::algebra`] -/// optimizer and allocator can reason about them. +/// Preserves `Sort`, `Limit`, `Join`, and `SetOp` nodes natively so the +/// [`crate::algebra`] optimizer and allocator can reason about them. pub fn parse_sql_expr(sql: &str) -> anyhow::Result { let dialect = GenericDialect {}; let mut stmts = sqlparser::parser::Parser::parse_sql(&dialect, sql) @@ -643,6 +53,8 @@ pub fn parse_sql_expr(sql: &str) -> anyhow::Result { extract_query_expr(&query) } +// ── Query-level dispatch ────────────────────────────────────────────────────── + fn extract_query_expr(query: &Query) -> anyhow::Result { let order_by: Vec = match &query.order_by { Some(OrderBy { kind: OrderByKind::Expressions(exprs), .. }) => exprs.clone(), @@ -696,6 +108,8 @@ fn extract_set_expr_qe( } } +// ── SELECT-level extraction ─────────────────────────────────────────────────── + fn extract_select_qe( sel: &Select, order_by: &[OrderByExpr], @@ -709,7 +123,7 @@ fn extract_select_qe( let agg_items = collect_agg_items_qe(&sel.projection); let join_qe = extract_join_qe(sel); - let source = QueryExpr::Source(crate::query_parser::sketch_algebra::SourceSpec { + let source = QueryExpr::Source(SourceSpec { name: metric_name.clone(), }); @@ -721,9 +135,7 @@ fn extract_select_qe( // JOIN let after_join = if let Some((inner_table, join_kind, join_pred)) = join_qe { - let inner_source = QueryExpr::Source( - crate::query_parser::sketch_algebra::SourceSpec { name: inner_table } - ); + let inner_source = QueryExpr::Source(SourceSpec { name: inner_table }); QueryExpr::Join { kind: join_kind, pred: join_pred, @@ -770,7 +182,7 @@ fn extract_select_qe( Ok(result) } -// ── QueryExpr agg item collection ──────────────────────────────────────────── +// ── Aggregation item collection ────────────────────────────────────────────── fn collect_agg_items_qe(projection: &[SelectItem]) -> Vec { let mut out = Vec::new(); @@ -802,11 +214,6 @@ fn collect_agg_from_expr_qe(expr: &Expr, alias: Option, out: &mut Vec SColumnRef::Wildcard, - SColumnRef::Named(n) => SColumnRef::Named(n.clone()), - SColumnRef::SampleValue => SColumnRef::SampleValue, - }; let func = match fn_name.as_str() { "COUNT" if distinct => AggFunc::CountDistinct, @@ -821,7 +228,7 @@ fn collect_agg_from_expr_qe(expr: &Expr, alias: Option, out: &mut Vec Vec { }).collect() } +// ── AST helpers: aggregation arguments ─────────────────────────────────────── + +fn first_col_from_args(args: &[FunctionArg]) -> ColumnRef { + for arg in args { + match arg { + FunctionArg::Unnamed(FunctionArgExpr::Wildcard) => return ColumnRef::Wildcard, + FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Identifier(id))) => { + return ColumnRef::Named(id.value.clone()); + } + FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::CompoundIdentifier(parts))) => { + if let Some(last) = parts.last() { + return ColumnRef::Named(last.value.clone()); + } + } + _ => {} + } + } + ColumnRef::Wildcard +} + +// ── AST helpers: GROUP BY ────────────────────────────────────────────────────── + +fn extract_group_by(group_by: &GroupByExpr) -> Vec { + let exprs = match group_by { + GroupByExpr::All(_) => return vec![], + GroupByExpr::Expressions(e, _) => e, + }; + exprs.iter().filter_map(|e| match e { + Expr::Identifier(id) => Some(id.value.clone()), + Expr::CompoundIdentifier(parts) => parts.last().map(|i| i.value.clone()), + _ => None, + }).collect() +} + +// ── AST helpers: table name ─────────────────────────────────────────────────── + +fn extract_table_name(sel: &Select) -> anyhow::Result { + sel.from.first() + .and_then(|t| match &t.relation { + TableFactor::Table { name, .. } => Some(object_name_str(name)), + _ => None, + }) + .ok_or_else(|| anyhow!("could not determine table name from FROM clause")) +} + +fn object_name_str(name: &ObjectName) -> String { + name.0.iter() + .map(|i| i.as_ident().map(|id| id.value.as_str()).unwrap_or("")) + .collect::>() + .join(".") +} + // ── SQL Expr → ScalarExpr ───────────────────────────────────────────────────── fn sql_expr_to_scalar(expr: &Expr) -> ScalarExpr { @@ -1004,17 +463,17 @@ fn expr_to_col_name(expr: &Expr) -> Option { #[cfg(test)] mod tests { - use super::*; - use super::super::sketch_algebra::{ExactAgg, SketchAggOp}; + use super::parse_sql_expr; + use crate::algebra::expr::QueryExpr; use crate::types::AggType; - use std::time::Duration; - fn parse(sql: &str) -> SketchExpr { - parse_sql(sql).unwrap_or_else(|e| panic!("parse_sql failed: {e}\nSQL: {sql}")) + fn parse(sql: &str) -> QueryExpr { + parse_sql_expr(sql).unwrap_or_else(|e| panic!("parse_sql_expr failed: {e}\nSQL: {sql}")) } fn pq(sql: &str) -> super::super::ParsedQuery { - parse(sql).to_parsed_query() + super::super::parse_query(sql) + .unwrap_or_else(|e| panic!("parse_query failed: {e}\nSQL: {sql}")) } // ── Basic aggregations ──────────────────────────────────────────────────── @@ -1066,10 +525,8 @@ mod tests { #[test] fn min_max_no_group_by_is_exact_minmax() { - let expr = parse("SELECT MIN(EventDate), MAX(EventDate) FROM hits"); - // After optimize: ExactMinMax nodes - let pq = expr.to_parsed_query(); - // ExactMinMax maps to Quantile in legacy AggType + let pq = pq("SELECT MIN(EventDate), MAX(EventDate) FROM hits"); + // MIN/MAX map to Quantile in legacy AggType assert!(pq.aggregations.contains(&AggType::Quantile)); } @@ -1089,10 +546,7 @@ mod tests { #[test] fn where_inequality_captured() { - // ne predicate should be captured in filters (not label_filters, but present) - let expr = parse("SELECT COUNT(*) FROM hits WHERE AdvEngineID <> 0 GROUP BY AdvEngineID"); - // Verify parse succeeds and has frequency - let pq = expr.to_parsed_query(); + let pq = pq("SELECT COUNT(*) FROM hits WHERE AdvEngineID <> 0 GROUP BY AdvEngineID"); assert!(pq.aggregations.contains(&AggType::Frequency)); } @@ -1138,7 +592,7 @@ mod tests { assert!(pq.aggregations.contains(&AggType::Frequency)); } - // ── COUNT(DISTINCT) with GROUP BY → Hydra via R4 ───────────────────────── + // ── COUNT(DISTINCT) with GROUP BY ──────────────────────────────────────── #[test] fn count_distinct_with_group_by() { @@ -1149,25 +603,22 @@ mod tests { assert!(pq.group_by_labels.contains(&"RegionID".to_string())); } - // ── UNION ALL → Merge ───────────────────────────────────────────────────── + // ── UNION ALL → SetOp ───────────────────────────────────────────────────── #[test] - fn union_all_produces_merge() { + fn union_all_produces_set_op() { let expr = parse( "SELECT COUNT(DISTINCT UserID) FROM R \ UNION ALL \ SELECT COUNT(DISTINCT UserID) FROM S", ); - // After R3 optimization: Merge([Agg(HLL,R), Agg(HLL,S)]) - assert!(matches!(expr, SketchExpr::Merge { .. })); + assert!(matches!(expr, QueryExpr::SetOp { .. })); } // ── Complex queries ─────────────────────────────────────────────────────── #[test] fn complex_multi_agg_multi_dim_group_by_topk() { - // COUNT(*) → Frequency, COUNT(DISTINCT) → Cardinality, AVG → Quantile - // Two GROUP BY dimensions; ORDER BY + LIMIT signals TopK let pq = pq( "SELECT region, dc, COUNT(*) AS c, COUNT(DISTINCT UserID), AVG(ResponseTime) \ FROM hits WHERE env = 'prod' GROUP BY region, dc ORDER BY c DESC LIMIT 5", @@ -1181,20 +632,16 @@ mod tests { pq.label_filters.get("env").map(String::as_str), Some("prod") ); - // AVG maps to median (p50) sketch assert!(pq.quantiles.contains(&0.5)); } #[test] fn complex_union_all_hll_with_where_on_each_branch() { - // UNION ALL → Merge; each branch has its own WHERE predicate - let expr = parse( + let pq = pq( "SELECT region, COUNT(DISTINCT UserID) FROM sessions WHERE status = 'active' GROUP BY region \ UNION ALL \ SELECT region, COUNT(DISTINCT UserID) FROM sessions WHERE status = 'expired' GROUP BY region", ); - assert!(matches!(expr, SketchExpr::Merge { .. })); - let pq = expr.to_parsed_query(); assert!(pq.aggregations.contains(&AggType::Cardinality)); assert!(pq.group_by_labels.contains(&"region".to_string())); } diff --git a/controller/src/types.rs b/controller/src/types.rs index 35c7d212..8a9585bc 100644 --- a/controller/src/types.rs +++ b/controller/src/types.rs @@ -325,7 +325,7 @@ impl StageResourceBudgets { /// Sub-plan for the **Agent OTel Collector** stage. /// -/// Covers `SketchExpr` nodes: `Source`, `Filter`, `Window`, `Agg` (sketch ops). +/// Covers `QueryExpr` nodes: `Source`, `Filter`, `Window`, `Agg` (sketch ops). #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct AgentSubPlan { /// Concrete sketch type resolved from the `Agg` node (absent when no Agg @@ -345,7 +345,7 @@ pub struct AgentSubPlan { /// Sub-plan for the **Backend OTel Collector** stage. /// -/// Covers `SketchExpr` nodes: `Partition`, `Merge`, `Dedup`, and +/// Covers `QueryExpr` nodes: `Partition`, `Merge`, `Dedup`, and /// `Agg { Exact(Sum|Count|Min|Max) }` (mergeable exact ops). #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct BackendSubPlan { @@ -360,7 +360,7 @@ pub struct BackendSubPlan { /// Sub-plan for the **ASAPQuery Precompute Engine** stage. /// -/// Covers `SketchExpr` nodes: `TopK`, and sketch `Agg` ops deferred from +/// Covers `QueryExpr` nodes: `TopK`, and sketch `Agg` ops deferred from /// the Agent stage due to memory budget overflow. #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct PrecomputeSubPlan { @@ -389,7 +389,7 @@ pub struct DbSubPlan { /// /// Produced by `planner::stage_split::split_expr_by_stage()`. Attached to /// [`CollectionPlan::staged_plan`] when the workload was supplied via -/// `query_string` (giving access to the full `SketchExpr` tree). +/// `query_string` (giving access to the full `QueryExpr` tree). #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct StagedPlan { pub agent: AgentSubPlan, @@ -418,7 +418,7 @@ pub struct CollectionPlan { /// SP-9: AST-aware per-stage sub-plans. /// /// `Some` when the workload was supplied via `query_string` (full - /// `SketchExpr` tree available). `None` when built from explicit + /// `QueryExpr` tree available). `None` when built from explicit /// aggregation fields — the SP-3 flat assignment is used as fallback. pub staged_plan: Option, }