diff --git a/controller/src/algebra/allocator.rs b/controller/src/algebra/allocator.rs new file mode 100644 index 00000000..5ed084c6 --- /dev/null +++ b/controller/src/algebra/allocator.rs @@ -0,0 +1,992 @@ +//! Sketch allocator — converts a [`QueryExpr`] tree into an annotated +//! [`PlanNode`] tree, assigning every operator to a [`PipelineStage`] and +//! choosing between sketch and exact execution. +//! +//! # Algorithm +//! +//! 1. Walk the tree bottom-up (children before parents). +//! 2. For each node, determine the *preferred* stage using the rules below. +//! 3. If the preferred stage exceeds its resource budget, demote to the next +//! stage in the chain: `Agent → Backend → Precompute → Db`. +//! 4. Annotate the node with the chosen sketch type, delta-encoding flag, +//! and a human-readable rationale. +//! +//! ## Stage assignment rules +//! +//! | Node type | Default stage | Condition | +//! |-----------|---------------|-----------| +//! | Source, Filter, Window, Partition, Dedup | Agent | Always | +//! | SketchAgg (sketachable op, mergeable) | Agent | budget OK | +//! | SketchAgg (sketchable, mergeable) | Backend | agent budget exceeded | +//! | SketchAgg (sketchable, not mergeable: Avg) | Db | always | +//! | SketchAgg (exact: Sum/Count/Min/Max) | Backend | mergeable | +//! | TopK | Precompute | always | +//! | Merge, JoinSketch | Backend | always | +//! | Aggregate, Project, Sort, Limit | Db | always | +//! | WindowFunc, HistogramQuantile | Db | always | +//! | PromQLSubquery, BinaryOp | Precompute | has sketch children | +//! | LetBinding, Subquery | same as body/inner | propagated | + +use super::expr::QueryExpr; +use super::plan::{ + CostEstimate, ExecutionMode, NodeAnnotation, PipelineStage, PlanNode, +}; +use crate::query_parser::sketch_algebra::{ExactAgg, SketchAggOp}; +use crate::types::{SketchParams, SketchType, StageResourceBudgets}; + +// ── Resource budget tracker ─────────────────────────────────────────────────── + +/// Mutable budget state, consumed during allocation. +#[derive(Debug, Clone)] +struct BudgetState { + agent_memory_remaining_bytes: f64, + backend_memory_remaining_bytes: f64, +} + +impl BudgetState { + fn from_budgets(b: &StageResourceBudgets) -> Self { + Self { + agent_memory_remaining_bytes: b.agent_memory_bytes + .map(|v| v as f64) + .unwrap_or(f64::INFINITY), + backend_memory_remaining_bytes: b.backend_memory_bytes + .map(|v| v as f64) + .unwrap_or(f64::INFINITY), + } + } + + fn fits_agent(&self, bytes: f64) -> bool { + bytes <= self.agent_memory_remaining_bytes + } + + fn fits_backend(&self, bytes: f64) -> bool { + bytes <= self.backend_memory_remaining_bytes + } + + fn consume_agent(&mut self, bytes: f64) { + self.agent_memory_remaining_bytes = + (self.agent_memory_remaining_bytes - bytes).max(0.0); + } + + fn consume_backend(&mut self, bytes: f64) { + self.backend_memory_remaining_bytes = + (self.backend_memory_remaining_bytes - bytes).max(0.0); + } +} + +// ── Public allocator ────────────────────────────────────────────────────────── + +/// Converts a (pre-optimised) [`QueryExpr`] tree into an annotated +/// [`PlanNode`] tree. +pub struct SketchAllocator { + budgets: StageResourceBudgets, + raw_bytes_per_sec: f64, +} + +impl SketchAllocator { + /// Create an allocator. + /// + /// * `budgets` — per-stage memory caps (from [`StageResourceBudgets`]). + /// * `raw_bytes_per_sec` — baseline bandwidth of the raw OTLP stream, + /// used to estimate compression ratios. + pub fn new(budgets: StageResourceBudgets, raw_bytes_per_sec: f64) -> Self { + Self { budgets, raw_bytes_per_sec } + } + + /// Allocate stages for the entire expression tree. + pub fn allocate(&self, expr: QueryExpr) -> PlanNode { + let mut budget = BudgetState::from_budgets(&self.budgets); + self.alloc_node(expr, &mut budget) + } + + // ── Recursive allocation ────────────────────────────────────────────────── + + fn alloc_node(&self, expr: QueryExpr, budget: &mut BudgetState) -> PlanNode { + match expr { + // ── Leaves ─────────────────────────────────────────────────────── + QueryExpr::Source(_) | QueryExpr::Ref(_) => PlanNode::leaf( + expr, + PipelineStage::Agent, + ExecutionMode::Passthrough, + ), + + // ── Structural / filter nodes — always Agent ────────────────── + QueryExpr::Filter { pred, input } => { + let child = self.alloc_node(*input, budget); + let stage = PipelineStage::Agent; + PlanNode { + expr: QueryExpr::Filter { pred, input: Box::new(child.expr.clone()) }, + stage, + mode: ExecutionMode::Passthrough, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec * 0.5, + ..Default::default() + }, + annotation: NodeAnnotation { + rationale: "Filter pushed to Agent to reduce data volume early".into(), + ..Default::default() + }, + children: vec![child], + } + } + + QueryExpr::Window { duration, slide, input } => { + let child = self.alloc_node(*input, budget); + PlanNode { + expr: QueryExpr::Window { + duration, slide, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Agent, + mode: ExecutionMode::Passthrough, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale: "Time window computed at Agent".into(), + ..Default::default() + }, + children: vec![child], + } + } + + QueryExpr::Partition { keys, input } => { + let child = self.alloc_node(*input, budget); + PlanNode { + expr: QueryExpr::Partition { + keys, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Agent, + mode: ExecutionMode::Passthrough, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale: "Partition for GROUP BY at Agent".into(), + ..Default::default() + }, + children: vec![child], + } + } + + QueryExpr::Dedup { col, input } => { + let child = self.alloc_node(*input, budget); + PlanNode { + expr: QueryExpr::Dedup { + col, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Agent, + mode: ExecutionMode::Passthrough, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale: "Dedup at Agent before sketch build".into(), + ..Default::default() + }, + children: vec![child], + } + } + + // ── Sketch aggregation — core allocation logic ──────────────── + QueryExpr::SketchAgg { op, col, input } => { + let child = self.alloc_node(*input, budget); + self.alloc_sketch_agg(op, col, child, budget) + } + + // ── TopK — Precompute engine ────────────────────────────────── + QueryExpr::TopK { k, by, input } => { + let child = self.alloc_node(*input, budget); + PlanNode { + expr: QueryExpr::TopK { + k, by, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Precompute, + mode: ExecutionMode::Sketch, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec * 0.05, + memory_bytes: (k as f64) * 64.0, + ..Default::default() + }, + annotation: NodeAnnotation { + sketch_type: Some(SketchType::CountSketch), + rationale: "TopK assigned to Precompute engine (CountSketch)".into(), + ..Default::default() + }, + children: vec![child], + } + } + + // ── Merge — Backend ─────────────────────────────────────────── + QueryExpr::Merge { inputs } => { + let children: Vec = inputs + .into_iter() + .map(|inp| self.alloc_node(inp, budget)) + .collect(); + let mem: f64 = children.iter().map(|c| c.cost.memory_bytes).sum(); + PlanNode { + expr: QueryExpr::Merge { + inputs: children.iter().map(|c| c.expr.clone()).collect(), + }, + stage: PipelineStage::Backend, + mode: ExecutionMode::Passthrough, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec * 0.1, + memory_bytes: mem, + ..Default::default() + }, + annotation: NodeAnnotation { + rationale: "Sketch merge at Backend".into(), + ..Default::default() + }, + children, + } + } + + QueryExpr::JoinSketch { join_key, outer, inner } => { + let outer_node = self.alloc_node(*outer, budget); + let inner_node = self.alloc_node(*inner, budget); + PlanNode { + expr: QueryExpr::JoinSketch { + join_key, + outer: Box::new(outer_node.expr.clone()), + inner: Box::new(inner_node.expr.clone()), + }, + stage: PipelineStage::Backend, + mode: ExecutionMode::Passthrough, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec * 0.2, + ..Default::default() + }, + annotation: NodeAnnotation { + rationale: "JoinSketch at Backend: pre-agg inner then merge".into(), + ..Default::default() + }, + children: vec![outer_node, inner_node], + } + } + + // ── Exact / relational — Db ─────────────────────────────────── + QueryExpr::Aggregate { keys, aggs, having, input } => { + let child = self.alloc_node(*input, budget); + PlanNode { + expr: QueryExpr::Aggregate { + keys, aggs, having, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Db, + mode: ExecutionMode::Exact, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec, + ..Default::default() + }, + annotation: NodeAnnotation { + rationale: "General Aggregate at Db (exact)".into(), + ..Default::default() + }, + children: vec![child], + } + } + + QueryExpr::Project { cols, input } => { + let child = self.alloc_node(*input, budget); + PlanNode { + expr: QueryExpr::Project { + cols, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Db, + mode: ExecutionMode::Exact, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale: "Project at Db".into(), + ..Default::default() + }, + children: vec![child], + } + } + + QueryExpr::Sort { keys, input } => { + let child = self.alloc_node(*input, budget); + PlanNode { + expr: QueryExpr::Sort { + keys, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Db, + mode: ExecutionMode::Exact, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale: "Sort at Db".into(), + ..Default::default() + }, + children: vec![child], + } + } + + QueryExpr::Limit { n, offset, input } => { + let child = self.alloc_node(*input, budget); + PlanNode { + expr: QueryExpr::Limit { + n, offset, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Db, + mode: ExecutionMode::Exact, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale: "Limit at Db".into(), + ..Default::default() + }, + children: vec![child], + } + } + + QueryExpr::Join { kind, pred, left, right } => { + let left_node = self.alloc_node(*left, budget); + let right_node = self.alloc_node(*right, budget); + PlanNode { + expr: QueryExpr::Join { + kind, pred, + left: Box::new(left_node.expr.clone()), + right: Box::new(right_node.expr.clone()), + }, + stage: PipelineStage::Db, + mode: ExecutionMode::Exact, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec, + ..Default::default() + }, + annotation: NodeAnnotation { + rationale: "Join at Db (exact)".into(), + ..Default::default() + }, + children: vec![left_node, right_node], + } + } + + QueryExpr::SetOp { kind, all, left, right } => { + let left_node = self.alloc_node(*left, budget); + let right_node = self.alloc_node(*right, budget); + PlanNode { + expr: QueryExpr::SetOp { + kind, all, + left: Box::new(left_node.expr.clone()), + right: Box::new(right_node.expr.clone()), + }, + stage: PipelineStage::Db, + mode: ExecutionMode::Exact, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale: "SetOp at Db".into(), + ..Default::default() + }, + children: vec![left_node, right_node], + } + } + + QueryExpr::WindowFunc { func, partition_by, order_by, frame, input } => { + let child = self.alloc_node(*input, budget); + PlanNode { + expr: QueryExpr::WindowFunc { + func, partition_by, order_by, frame, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Db, + mode: ExecutionMode::Exact, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale: "Analytic window function at Db".into(), + ..Default::default() + }, + children: vec![child], + } + } + + // ── PromQL-specific ─────────────────────────────────────────── + QueryExpr::HistogramQuantile { phi, input } => { + let child = self.alloc_node(*input, budget); + // If the child is a sketch, elevate to Precompute; + // otherwise fall through to Db. + let stage = if child.mode == ExecutionMode::Sketch { + PipelineStage::Precompute + } else { + PipelineStage::Db + }; + let rationale = format!("histogram_quantile(φ={phi}) at {stage}"); + PlanNode { + expr: QueryExpr::HistogramQuantile { + phi, + input: Box::new(child.expr.clone()), + }, + stage, + mode: ExecutionMode::Sketch, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec * 0.02, + ..Default::default() + }, + annotation: NodeAnnotation { + sketch_type: Some(SketchType::DDSketch), + rationale, + ..Default::default() + }, + children: vec![child], + } + } + + QueryExpr::PromQLSubquery { range, resolution, input } => { + let child = self.alloc_node(*input, budget); + let stage = if child.mode == ExecutionMode::Sketch { + PipelineStage::Precompute + } else { + PipelineStage::Db + }; + let rationale = format!("PromQL subquery at {stage}"); + PlanNode { + expr: QueryExpr::PromQLSubquery { + range, resolution, + input: Box::new(child.expr.clone()), + }, + stage, + mode: child.mode.clone(), + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale, + ..Default::default() + }, + children: vec![child], + } + } + + QueryExpr::BinaryOp { op, lhs, rhs, vector_match } => { + let left_node = self.alloc_node(*lhs, budget); + let right_node = self.alloc_node(*rhs, budget); + let has_sketch = left_node.mode == ExecutionMode::Sketch + || right_node.mode == ExecutionMode::Sketch; + let stage = if has_sketch { + PipelineStage::Precompute + } else { + PipelineStage::Db + }; + let rationale = format!("BinaryOp at {stage}"); + PlanNode { + expr: QueryExpr::BinaryOp { + op, vector_match, + lhs: Box::new(left_node.expr.clone()), + rhs: Box::new(right_node.expr.clone()), + }, + stage, + mode: if has_sketch { ExecutionMode::Sketch } else { ExecutionMode::Exact }, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale, + ..Default::default() + }, + children: vec![left_node, right_node], + } + } + + // ── Scoping constructs — propagate body's stage ─────────────── + QueryExpr::Subquery { alias, expr } => { + let child = self.alloc_node(*expr, budget); + let stage = child.stage.clone(); + let mode = child.mode.clone(); + PlanNode { + expr: QueryExpr::Subquery { + alias, + expr: Box::new(child.expr.clone()), + }, + stage, + mode, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale: "Subquery inherits inner stage".into(), + ..Default::default() + }, + children: vec![child], + } + } + + QueryExpr::LetBinding { name, expr, body } => { + let expr_node = self.alloc_node(*expr, budget); + let body_node = self.alloc_node(*body, budget); + let stage = body_node.stage.clone(); + let mode = body_node.mode.clone(); + PlanNode { + expr: QueryExpr::LetBinding { + name, + expr: Box::new(expr_node.expr.clone()), + body: Box::new(body_node.expr.clone()), + }, + stage, + mode, + cost: CostEstimate::default(), + annotation: NodeAnnotation { + rationale: "LetBinding: stage = body stage".into(), + ..Default::default() + }, + children: vec![expr_node, body_node], + } + } + } + } + + // ── SketchAgg allocation (budget-driven demotion) ───────────────────────── + + fn alloc_sketch_agg( + &self, + op: SketchAggOp, + col: crate::query_parser::sketch_algebra::ColumnRef, + child: PlanNode, + budget: &mut BudgetState, + ) -> PlanNode { + // Exact non-mergeable (Avg) → always Db. + if let SketchAggOp::Exact(ExactAgg::Avg) = &op { + return PlanNode { + expr: QueryExpr::SketchAgg { + op, + col, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Db, + mode: ExecutionMode::Exact, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec, + ..Default::default() + }, + annotation: NodeAnnotation { + rationale: "Exact(Avg) is not mergeable — must run at Db".into(), + ..Default::default() + }, + children: vec![child], + }; + } + + // Exact mergeable (Sum, Count, Min, Max) → Backend. + if let SketchAggOp::Exact(_) = &op { + return PlanNode { + expr: QueryExpr::SketchAgg { + op, + col, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Backend, + mode: ExecutionMode::Exact, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec * 0.8, + ..Default::default() + }, + annotation: NodeAnnotation { + rationale: "Exact(Sum/Count/Min/Max) merged at Backend".into(), + ..Default::default() + }, + children: vec![child], + }; + } + + // Sketch operators: try Agent → Backend → Precompute. + let mem = estimated_sketch_memory(&op); + let (sketch_type, params) = sketch_type_for_op(&op); + + if budget.fits_agent(mem) { + budget.consume_agent(mem); + return PlanNode { + expr: QueryExpr::SketchAgg { + op, + col, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Agent, + mode: ExecutionMode::Sketch, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec * 0.05, + memory_bytes: mem, + compression_ratio: 20.0, + ..Default::default() + }, + annotation: NodeAnnotation { + sketch_type: Some(sketch_type), + sketch_params: Some(params), + rationale: "Sketch at Agent (within budget)".into(), + ..Default::default() + }, + children: vec![child], + }; + } + + if budget.fits_backend(mem) { + budget.consume_backend(mem); + return PlanNode { + expr: QueryExpr::SketchAgg { + op, + col, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Backend, + mode: ExecutionMode::Sketch, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec * 0.1, + memory_bytes: mem, + compression_ratio: 10.0, + ..Default::default() + }, + annotation: NodeAnnotation { + sketch_type: Some(sketch_type), + sketch_params: Some(params), + rationale: "Sketch demoted to Backend (Agent budget exceeded)".into(), + budget_demotion: true, + ..Default::default() + }, + children: vec![child], + }; + } + + // Both Agent and Backend budgets exceeded → Precompute. + PlanNode { + expr: QueryExpr::SketchAgg { + op, + col, + input: Box::new(child.expr.clone()), + }, + stage: PipelineStage::Precompute, + mode: ExecutionMode::Sketch, + cost: CostEstimate { + bytes_per_sec: self.raw_bytes_per_sec * 0.2, + memory_bytes: mem, + compression_ratio: 5.0, + ..Default::default() + }, + annotation: NodeAnnotation { + sketch_type: Some(sketch_type), + sketch_params: Some(params), + rationale: "Sketch demoted to Precompute (Agent+Backend budgets exceeded)".into(), + budget_demotion: true, + ..Default::default() + }, + children: vec![child], + } + } +} + +// ── Helpers ─────────────────────────────────────────────────────────────────── + +/// Estimate the memory footprint of a sketch in bytes. +fn estimated_sketch_memory(op: &SketchAggOp) -> f64 { + match op { + SketchAggOp::DDSketch { .. } => 4_096.0, + SketchAggOp::HLL { registers } => (1u64 << registers) as f64, + SketchAggOp::CountMin { width, depth } => (*width as f64) * (*depth as f64) * 8.0, + SketchAggOp::CountSketch { k } => (*k as f64) * 64.0, + SketchAggOp::ExactMinMax { .. } => 16.0, + SketchAggOp::Hydra { inner, partition_keys } => { + // Hydra memory = inner sketch size × expected number of key tuples. + let inner_mem = estimated_sketch_memory(inner); + let keys = partition_keys.len() as f64; + inner_mem * (10.0_f64.powf(keys)) + } + SketchAggOp::Exact(_) => 8.0, + } +} + +/// Map a [`SketchAggOp`] to a [`SketchType`] + [`SketchParams`]. +fn sketch_type_for_op(op: &SketchAggOp) -> (SketchType, SketchParams) { + match op { + SketchAggOp::DDSketch { quantiles, epsilon } => ( + SketchType::DDSketch, + SketchParams { + relative_accuracy: *epsilon, + quantiles: quantiles.clone(), + ..Default::default() + }, + ), + SketchAggOp::HLL { registers } => ( + SketchType::HLL, + SketchParams { + precision: *registers as u32, + ..Default::default() + }, + ), + SketchAggOp::CountMin { width, depth } => ( + SketchType::CountMinSketch, + SketchParams { + cols: *width, + rows: *depth as u32, + ..Default::default() + }, + ), + SketchAggOp::CountSketch { k } => ( + SketchType::CountSketch, + SketchParams { + k: *k as u32, + ..Default::default() + }, + ), + SketchAggOp::ExactMinMax { .. } => ( + SketchType::DDSketch, + SketchParams::default(), + ), + SketchAggOp::Hydra { inner, .. } => sketch_type_for_op(inner), + SketchAggOp::Exact(_) => (SketchType::DDSketch, SketchParams::default()), + } +} + +// ── Tests ───────────────────────────────────────────────────────────────────── + +#[cfg(test)] +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::types::{SketchType, StageResourceBudgets}; + use std::time::Duration; + + fn src(name: &str) -> QueryExpr { + QueryExpr::Source(SourceSpec { name: name.into() }) + } + + fn alloc(budgets: StageResourceBudgets, expr: QueryExpr) -> PlanNode { + SketchAllocator::new(budgets, 100_000.0).allocate(expr) + } + + fn unlimited() -> StageResourceBudgets { + StageResourceBudgets::default() + } + + fn tight_agent() -> StageResourceBudgets { + StageResourceBudgets { + agent_memory_bytes: Some(1), // 1 byte — too small for any sketch + ..Default::default() + } + } + + fn tight_all() -> StageResourceBudgets { + StageResourceBudgets { + agent_memory_bytes: Some(1), + backend_memory_bytes: Some(1), + ..Default::default() + } + } + + // ── Source / leaf ───────────────────────────────────────────────────────── + + #[test] + fn source_goes_to_agent() { + let node = alloc(unlimited(), src("cpu")); + assert_eq!(node.stage, PipelineStage::Agent); + assert_eq!(node.mode, ExecutionMode::Passthrough); + } + + // ── Filter ──────────────────────────────────────────────────────────────── + + #[test] + fn filter_at_agent() { + use crate::algebra::expr::{LiteralValue, ScalarExpr}; + let expr = QueryExpr::Filter { + pred: ScalarExpr::Literal(LiteralValue::Bool(true)), + input: Box::new(src("m")), + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Agent); + } + + // ── DDSketch within budget → Agent ──────────────────────────────────────── + + #[test] + fn ddsketch_within_budget_goes_to_agent() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, + input: Box::new(src("latency")), + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Agent); + assert_eq!(node.mode, ExecutionMode::Sketch); + assert_eq!(node.annotation.sketch_type, Some(SketchType::DDSketch)); + } + + // ── DDSketch tight agent budget → Backend ───────────────────────────────── + + #[test] + fn ddsketch_agent_budget_exceeded_goes_to_backend() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, + input: Box::new(src("latency")), + }; + let node = alloc(tight_agent(), expr); + assert_eq!(node.stage, PipelineStage::Backend); + assert!(node.annotation.budget_demotion); + } + + // ── DDSketch tight agent+backend → Precompute ───────────────────────────── + + #[test] + fn ddsketch_all_budgets_exceeded_goes_to_precompute() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, + input: Box::new(src("latency")), + }; + let node = alloc(tight_all(), expr); + assert_eq!(node.stage, PipelineStage::Precompute); + assert!(node.annotation.budget_demotion); + } + + // ── Exact(Avg) → Db ─────────────────────────────────────────────────────── + + #[test] + fn exact_avg_goes_to_db() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::Exact(ExactAgg::Avg), + col: ColumnRef::Named("price".into()), + input: Box::new(src("trades")), + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Db); + assert_eq!(node.mode, ExecutionMode::Exact); + } + + // ── Exact(Sum) → Backend ────────────────────────────────────────────────── + + #[test] + fn exact_sum_goes_to_backend() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::Exact(ExactAgg::Sum), + col: ColumnRef::Named("bytes".into()), + input: Box::new(src("network")), + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Backend); + assert_eq!(node.mode, ExecutionMode::Exact); + } + + // ── TopK → Precompute ───────────────────────────────────────────────────── + + #[test] + fn topk_goes_to_precompute() { + let expr = QueryExpr::TopK { + k: 10, + by: vec!["symbol".into()], + input: Box::new(src("trades")), + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Precompute); + assert_eq!(node.annotation.sketch_type, Some(SketchType::CountSketch)); + } + + // ── Merge → Backend ─────────────────────────────────────────────────────── + + #[test] + fn merge_goes_to_backend() { + let expr = QueryExpr::Merge { + inputs: vec![src("a"), src("b")], + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Backend); + } + + // ── HLL → Agent ─────────────────────────────────────────────────────────── + + #[test] + fn hll_within_budget_at_agent() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::HLL { registers: 14 }, + col: ColumnRef::Named("uid".into()), + input: Box::new(src("events")), + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Agent); + assert_eq!(node.annotation.sketch_type, Some(SketchType::HLL)); + } + + // ── CountMin → Agent ────────────────────────────────────────────────────── + + #[test] + fn countmin_within_budget_at_agent() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::default_count_min(), + col: ColumnRef::Wildcard, + input: Box::new(src("requests")), + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Agent); + assert_eq!(node.annotation.sketch_type, Some(SketchType::CountMinSketch)); + } + + // ── Join → Db ───────────────────────────────────────────────────────────── + + #[test] + fn join_goes_to_db() { + use crate::algebra::expr::JoinKind; + let expr = QueryExpr::Join { + kind: JoinKind::Inner, + pred: None, + left: Box::new(src("orders")), + right: Box::new(src("items")), + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Db); + } + + // ── HistogramQuantile + DDSketch → Precompute ───────────────────────────── + + #[test] + fn histogram_quantile_over_sketch_at_precompute() { + let expr = QueryExpr::HistogramQuantile { + phi: 0.95, + input: Box::new(QueryExpr::SketchAgg { + op: SketchAggOp::default_ddsketch(vec![0.95]), + col: ColumnRef::SampleValue, + input: Box::new(src("hist")), + }), + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Precompute); + } + + // ── LetBinding inherits body stage ──────────────────────────────────────── + + #[test] + fn let_binding_inherits_body_stage() { + let expr = QueryExpr::LetBinding { + name: "base".into(), + expr: Box::new(src("cpu")), + body: Box::new(QueryExpr::TopK { + k: 5, + by: vec![], + input: Box::new(src("cpu")), + }), + }; + let node = alloc(unlimited(), expr); + assert_eq!(node.stage, PipelineStage::Precompute); + } + + // ── Memory estimate helpers ─────────────────────────────────────────────── + + #[test] + fn hll_memory_estimate_matches_register_count() { + let mem = estimated_sketch_memory(&SketchAggOp::HLL { registers: 14 }); + assert_eq!(mem, (1u64 << 14) as f64); // 16 384 bytes + } + + #[test] + fn countmin_memory_estimate() { + let op = SketchAggOp::CountMin { width: 2000, depth: 5 }; + let mem = estimated_sketch_memory(&op); + assert_eq!(mem, 2000.0 * 5.0 * 8.0); // 80 000 bytes + } + + // ── PlanSummary from allocated tree ────────────────────────────────────── + + #[test] + fn plan_summary_shows_bandwidth_saved() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::default_ddsketch(vec![0.99]), + col: ColumnRef::SampleValue, + input: Box::new(src("latency")), + }; + let node = alloc(unlimited(), expr); + let summary = node.summarise(100_000.0); + // sketch reduces to ~5% → saved ~95 000 B/s + assert!(summary.bandwidth_saved_bytes_per_sec > 50_000.0); + assert!(summary.agent_memory_bytes > 0.0); + } +} diff --git a/controller/src/algebra/expr.rs b/controller/src/algebra/expr.rs new file mode 100644 index 00000000..43d27fbf --- /dev/null +++ b/controller/src/algebra/expr.rs @@ -0,0 +1,1095 @@ +//! General query algebra — the full IR for SQL and PromQL queries. +//! +//! This module defines two mutually recursive expression types: +//! +//! * [`QueryExpr`] — *relational* operators. Each node takes zero or more +//! relations as input and produces a relation. Maps to SQL's FROM / GROUP BY +//! / JOIN / UNION layer and PromQL's binary / sub-query layer. +//! +//! * [`ScalarExpr`] — *scalar* operators. Each node computes a single value +//! 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, +//! every node carries a [`PipelineStage`](crate::algebra::plan::PipelineStage) +//! tag that says where the work executes: +//! +//! | Stage | Component | +//! |-----------|----------------------------| +//! | Agent | OTel Collector at the SDK | +//! | Backend | Central merge collector | +//! | Precompute| ASAPQuery engine | +//! | Db | Backend OLAP / exact store | + +use std::time::Duration; + +use crate::query_parser::sketch_algebra::{ + ColumnRef, ExactAgg, FilterOp, FilterVal, Predicate, SketchAggOp, SketchExpr, SourceSpec, +}; + +// ── Relational algebra ──────────────────────────────────────────────────────── + +/// Full relational + sketch algebra — replaces / extends [`SketchExpr`]. +/// +/// Every variant is a *node* in the logical query plan tree. Leaves are +/// [`QueryExpr::Source`] or [`QueryExpr::Ref`]. Interior nodes combine their +/// `input` child(ren) through the operator they implement. +#[derive(Debug, Clone)] +pub enum QueryExpr { + // ── Base relations ──────────────────────────────────────────────────── + + /// A named metric stream or table. The outermost leaf. + Source(SourceSpec), + + /// Reference to a CTE / let-binding by name. Resolved at plan time. + Ref(String), + + // ── Filtering & projection ──────────────────────────────────────────── + + /// σ — row-level filter (WHERE / PromQL label matchers). + Filter { + pred: ScalarExpr, + input: Box, + }, + + /// π — column projection (SELECT list). + Project { + cols: Vec, + input: Box, + }, + + // ── Aggregation ─────────────────────────────────────────────────────── + + /// γ + α — GROUP BY followed by aggregate functions. + /// + /// `keys` is the GROUP BY column list (empty → global aggregate). + /// `aggs` is the list of aggregate expressions to compute. + /// `having` is an optional post-aggregation predicate. + Aggregate { + keys: Vec, + aggs: Vec, + having: Option, + input: Box, + }, + + // ── Time / streaming operators ──────────────────────────────────────── + + /// ψ — time window (PromQL `[5m]`; SQL tumbling/sliding window). + Window { + duration: Duration, + slide: Option, + input: Box, + }, + + /// γ+α specialisation for sketch aggregations (single sketch per node). + /// + /// Kept separate from [`Self::Aggregate`] so the allocator can reason + /// about which sketch type to use without parsing `AggFunc` variants. + SketchAgg { + op: SketchAggOp, + col: ColumnRef, + input: Box, + }, + + // ── Distributed / multi-stage operators ────────────────────────────── + + /// Partition the stream by key-tuple (GROUP BY / `by (dims)`). + Partition { + keys: crate::query_parser::sketch_algebra::PartitionKeys, + input: Box, + }, + + /// δ — deduplicate on `col` before sketch ingestion. + Dedup { + col: String, + input: Box, + }, + + /// τ — retain only the top-K entries (heavy hitters). + TopK { + k: u64, + by: Vec, + input: Box, + }, + + /// ⊕ — merge sketches from independent branches (distributed union). + Merge { + inputs: Vec, + }, + + // ── Join operators ──────────────────────────────────────────────────── + + /// Relational join. + Join { + kind: JoinKind, + pred: Option, + left: Box, + right: Box, + }, + + /// Sketch-aware join push-down: pre-aggregate on inner side then merge. + JoinSketch { + join_key: String, + outer: Box, + inner: Box, + }, + + // ── Set operators ───────────────────────────────────────────────────── + + /// UNION / INTERSECT / EXCEPT (with or without ALL). + SetOp { + kind: SetOpKind, + all: bool, + left: Box, + right: Box, + }, + + // ── Ordering & limiting ─────────────────────────────────────────────── + + /// ORDER BY. + Sort { + keys: Vec, + input: Box, + }, + + /// LIMIT [OFFSET]. + Limit { + n: u64, + offset: u64, + input: Box, + }, + + // ── Subquery / CTE ──────────────────────────────────────────────────── + + /// Inline subquery with an alias (SQL `(SELECT ...) AS alias`). + Subquery { + alias: String, + expr: Box, + }, + + /// SQL `WITH name AS (expr) IN body` or PromQL recording rule binding. + LetBinding { + name: String, + expr: Box, + body: Box, + }, + + // ── Window functions (analytic functions) ───────────────────────────── + + /// OVER (PARTITION BY … ORDER BY … frame) analytic functions. + WindowFunc { + func: WindowFuncKind, + partition_by: Vec, + order_by: Vec, + frame: Option, + input: Box, + }, + + // ── PromQL-specific operators ───────────────────────────────────────── + + /// `histogram_quantile(φ, )` — converts an HLL / histogram + /// sketch into a quantile estimate. + HistogramQuantile { + phi: f64, + input: Box, + }, + + /// PromQL sub-query syntax: `[range:resolution]`. + PromQLSubquery { + range: Duration, + resolution: Option, + input: Box, + }, + + /// Binary operation between two instant-vector expressions (PromQL `+`, `/`, …). + /// Also used for SQL arithmetic between sub-relations. + BinaryOp { + op: BinaryOpKind, + lhs: Box, + rhs: Box, + vector_match: Option, + }, +} + +// ── Scalar algebra ──────────────────────────────────────────────────────────── + +/// Scalar expression — computes a single value from a row. +/// +/// Used in [`QueryExpr::Filter`] predicates, [`ProjectItem`] expressions, +/// [`QueryExpr::Aggregate`] HAVING clauses, and JOIN conditions. +#[derive(Debug, Clone)] +pub enum ScalarExpr { + /// Column reference: `t.col` or just `col`. + Column(String), + + /// Literal value. + Literal(LiteralValue), + + /// Arithmetic / comparison / logical / regex binary operator. + BinaryOp { + op: BinaryOpKind, + lhs: Box, + rhs: Box, + }, + + /// Unary prefix operator (`NOT`, `-`, `+`). + UnaryOp { + op: UnaryOpKind, + input: Box, + }, + + /// Named function call (e.g. `ABS(x)`, `DATE_TRUNC('hour', ts)`). + FunctionCall { + name: String, + args: Vec, + }, + + /// Scalar sub-query (`SELECT MAX(price) FROM orders`). + ScalarSubquery(Box), + + /// `expr IN (v1, v2, …)` or `NOT IN (…)`. + InList { + expr: Box, + list: Vec, + negated: bool, + }, + + /// `expr IN (SELECT …)` / `NOT IN (SELECT …)`. + InSubquery { + expr: Box, + subquery: Box, + negated: bool, + }, + + /// `expr BETWEEN low AND high` or `NOT BETWEEN …`. + Between { + expr: Box, + low: Box, + high: Box, + negated: bool, + }, + + /// `expr IS NULL` / `IS NOT NULL`. + IsNull { + expr: Box, + negated: bool, + }, + + /// CASE WHEN … THEN … [ELSE …] END. + Case { + operand: Option>, + when_then: Vec<(ScalarExpr, ScalarExpr)>, + else_: Option>, + }, + + /// CAST(expr AS type). + Cast { + expr: Box, + to: DataType, + }, + + /// PromQL vector binary op between two instant-vector expressions where one + /// or both sides produce a scalar in the final result (e.g. `rate(…) > 0.5`). + VectorBinaryOp { + op: BinaryOpKind, + lhs: Box, + rhs: Box, + vector_match: Option, + }, +} + +// ── Supporting enumerations ─────────────────────────────────────────────────── + +/// A single item in a SELECT projection list. +#[derive(Debug, Clone)] +pub struct ProjectItem { + /// Output column name (SQL `AS alias`; None → use expression name). + pub alias: Option, + pub expr: ScalarExpr, +} + +/// One aggregate function in a GROUP BY / AGGREGATE node. +#[derive(Debug, Clone)] +pub struct AggItem { + /// Output column name. + pub alias: String, + /// The aggregate function. + pub func: AggFunc, + /// Column(s) the function operates on. + pub col: ColumnRef, + /// Whether DISTINCT is applied before aggregation. + pub distinct: bool, +} + +/// All aggregate functions that the algebra supports. +/// +/// "Sketchable" variants (Quantile, CountDistinct, HeavyHitters) can be +/// approximated by a sketch in early pipeline stages; the rest require +/// exact computation. +#[derive(Debug, Clone, PartialEq)] +pub enum AggFunc { + Count, + Sum, + Avg, + Min, + Max, + /// Sample / population standard deviation. + StdDev { population: bool }, + /// Sample / population variance. + Variance { population: bool }, + /// Approximate quantile at φ ∈ (0, 1]. Maps to DDSketch. + Quantile(f64), + /// COUNT DISTINCT — maps to HLL. + CountDistinct, + /// Top-K heavy hitters — maps to CountSketch. + HeavyHitters { k: u64 }, + /// PromQL `rate()` — per-second increase over a window. + Rate, + /// PromQL `increase()` — total increase over a window. + Increase, + /// PromQL `delta()` — change over a window (may be negative). + Delta, + /// Arbitrary named aggregate (UDA or extension). + Custom(String), +} + +impl AggFunc { + /// Returns true when this function can be computed from merged partial + /// results: `f(A ∪ B) = combine(f(A), f(B))`. + pub fn is_mergeable(&self) -> bool { + match self { + AggFunc::Avg | AggFunc::StdDev { .. } | AggFunc::Variance { .. } => false, + _ => true, + } + } + + /// Returns true when this function requires sketch approximation to be + /// bandwidth-efficient (i.e. the raw data would be too large to ship). + pub fn is_sketchable(&self) -> bool { + matches!( + self, + AggFunc::Quantile(_) | AggFunc::CountDistinct | AggFunc::HeavyHitters { .. } + ) + } + + /// Suggest the appropriate [`SketchAggOp`] for this function, if any. + pub fn to_sketch_op(&self) -> Option { + match self { + AggFunc::Quantile(phi) => Some(SketchAggOp::default_ddsketch(vec![*phi])), + AggFunc::CountDistinct => Some(SketchAggOp::default_hll()), + AggFunc::HeavyHitters { k } => Some(SketchAggOp::CountSketch { k: *k }), + AggFunc::Count => Some(SketchAggOp::Exact(ExactAgg::Count)), + AggFunc::Sum => Some(SketchAggOp::Exact(ExactAgg::Sum)), + AggFunc::Avg => Some(SketchAggOp::Exact(ExactAgg::Avg)), + AggFunc::Min => Some(SketchAggOp::ExactMinMax { min: true, max: false }), + AggFunc::Max => Some(SketchAggOp::ExactMinMax { min: false, max: true }), + _ => None, + } + } +} + +/// Binary operator kinds — used in both [`ScalarExpr::BinaryOp`] and +/// [`QueryExpr::BinaryOp`] (PromQL instant-vector arithmetic). +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum BinaryOpKind { + // Arithmetic + Add, + Sub, + Mul, + Div, + Mod, + Pow, + // Comparison + Eq, + Ne, + Lt, + Le, + Gt, + Ge, + // Logical + And, + Or, + // Bitwise + BitAnd, + BitOr, + BitXor, + // String / pattern + Concat, + Like, + NotLike, + Regex, + NotRegex, + // PromQL-specific + Unless, + Atan2, +} + +/// Unary prefix operators. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum UnaryOpKind { + Negate, + Not, + BitwiseNot, +} + +/// JOIN variant. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum JoinKind { + Inner, + LeftOuter, + RightOuter, + FullOuter, + Cross, + /// Semi-join: return only left rows that have a match (WHERE EXISTS). + Semi, + /// Anti-join: return only left rows that have no match (WHERE NOT EXISTS). + AntiSemi, +} + +/// Set-operation variant. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SetOpKind { + Union, + Intersect, + Except, +} + +/// PromQL vector matching semantics (`on (…)` / `ignoring (…)` plus +/// `group_left` / `group_right`). +#[derive(Debug, Clone)] +pub struct VectorMatch { + pub kind: VectorMatchKind, + pub labels: Vec, + pub grouping: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum VectorMatchKind { + On, + Ignoring, +} + +#[derive(Debug, Clone)] +pub struct VectorGrouping { + pub side: GroupSide, + pub labels: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum GroupSide { + Left, + Right, +} + +/// ORDER BY sort key. +#[derive(Debug, Clone)] +pub struct SortKey { + pub col: String, + pub desc: bool, + /// NULLS FIRST / NULLS LAST (None → database default). + pub nulls_first: Option, +} + +/// Analytic window function kinds. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum WindowFuncKind { + RowNumber, + Rank, + DenseRank, + PercentRank, + CumeDist, + NTile { n: u64 }, + Lag { offset: u64 }, + Lead { offset: u64 }, + FirstValue, + LastValue, + NthValue { n: u64 }, + /// User-defined analytic function. + Custom(String), +} + +/// ROWS / RANGE frame clause for analytic functions. +#[derive(Debug, Clone)] +pub struct WindowFrame { + pub unit: FrameUnit, + pub start: FrameBound, + pub end: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum FrameUnit { + Rows, + Range, + Groups, +} + +#[derive(Debug, Clone)] +pub enum FrameBound { + UnboundedPreceding, + Preceding(u64), + CurrentRow, + Following(u64), + UnboundedFollowing, +} + +/// Scalar literal. +#[derive(Debug, Clone, PartialEq)] +pub enum LiteralValue { + Null, + Bool(bool), + Int(i64), + Float(f64), + Str(String), + Duration(Duration), +} + +/// SQL / Arrow data types used in CAST expressions. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DataType { + Boolean, + Int8, + Int16, + Int32, + Int64, + UInt8, + UInt16, + UInt32, + UInt64, + Float32, + Float64, + Utf8, + Binary, + Timestamp, + Date, + Interval, + List(Box), + Struct(Vec<(String, 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); + match self { + QueryExpr::Source(_) | QueryExpr::Ref(_) => {} + QueryExpr::Filter { input, .. } + | QueryExpr::Project { input, .. } + | QueryExpr::Window { input, .. } + | QueryExpr::SketchAgg { input, .. } + | QueryExpr::Partition { input, .. } + | QueryExpr::Dedup { input, .. } + | QueryExpr::TopK { input, .. } + | QueryExpr::Sort { input, .. } + | QueryExpr::Limit { input, .. } + | QueryExpr::WindowFunc { input, .. } + | QueryExpr::HistogramQuantile { input, .. } + | QueryExpr::PromQLSubquery { input, .. } => input.walk(f), + + QueryExpr::Aggregate { input, .. } => input.walk(f), + + QueryExpr::Merge { inputs } => { + for i in inputs { i.walk(f); } + } + QueryExpr::Join { left, right, .. } + | QueryExpr::JoinSketch { outer: left, inner: right, .. } + | QueryExpr::SetOp { left, right, .. } + | QueryExpr::BinaryOp { lhs: left, rhs: right, .. } => { + left.walk(f); + right.walk(f); + } + QueryExpr::Subquery { expr, .. } => expr.walk(f), + QueryExpr::LetBinding { expr, body, .. } => { + expr.walk(f); + body.walk(f); + } + } + } + + /// Returns `true` when the sub-tree contains at least one [`QueryExpr::SketchAgg`] + /// or [`QueryExpr::TopK`] node (i.e. sketch work is present). + pub fn has_sketch_work(&self) -> bool { + let mut found = false; + self.walk(&mut |n| { + if matches!(n, QueryExpr::SketchAgg { .. } | QueryExpr::TopK { .. }) { + found = true; + } + }); + found + } + + /// Returns the outermost metric/table name from the first `Source` leaf. + pub fn source_name(&self) -> Option<&str> { + match self { + QueryExpr::Source(s) => Some(&s.name), + QueryExpr::Filter { input, .. } + | QueryExpr::Project { input, .. } + | QueryExpr::Window { input, .. } + | QueryExpr::SketchAgg { input, .. } + | QueryExpr::Partition { input, .. } + | QueryExpr::Dedup { input, .. } + | QueryExpr::TopK { input, .. } + | QueryExpr::Sort { input, .. } + | QueryExpr::Limit { input, .. } + | QueryExpr::Aggregate { input, .. } + | QueryExpr::WindowFunc { input, .. } + | QueryExpr::HistogramQuantile { input, .. } + | QueryExpr::PromQLSubquery { input, .. } => input.source_name(), + QueryExpr::Merge { inputs } => inputs.first()?.source_name(), + QueryExpr::Join { left, .. } + | QueryExpr::JoinSketch { outer: left, .. } + | QueryExpr::SetOp { left, .. } + | QueryExpr::BinaryOp { lhs: left, .. } => left.source_name(), + QueryExpr::Subquery { expr, .. } => expr.source_name(), + QueryExpr::LetBinding { body, .. } => body.source_name(), + QueryExpr::Ref(_) => None, + } + } +} + +// ── Predicate → ScalarExpr conversion ──────────────────────────────────────── + +/// Convert a slice of legacy [`Predicate`]s (AND-list) into a single +/// [`ScalarExpr`] tree. An empty slice becomes `Literal(true)`. +fn scalar_from_predicates(preds: &[Predicate]) -> ScalarExpr { + if preds.is_empty() { + return ScalarExpr::Literal(LiteralValue::Bool(true)); + } + let mut iter = preds.iter().map(scalar_from_predicate); + let first = iter.next().unwrap(); + iter.fold(first, |acc, p| ScalarExpr::BinaryOp { + op: BinaryOpKind::And, + lhs: Box::new(acc), + rhs: Box::new(p), + }) +} + +fn scalar_from_predicate(p: &Predicate) -> ScalarExpr { + let col = ScalarExpr::Column(p.col.clone()); + let val = match &p.val { + FilterVal::Str(s) => ScalarExpr::Literal(LiteralValue::Str(s.clone())), + FilterVal::Num(n) => ScalarExpr::Literal(LiteralValue::Float(*n)), + FilterVal::Int(i) => ScalarExpr::Literal(LiteralValue::Int(*i)), + FilterVal::Null => ScalarExpr::Literal(LiteralValue::Null), + }; + match &p.op { + FilterOp::Eq => bin(BinaryOpKind::Eq, col, val), + FilterOp::Ne => bin(BinaryOpKind::Ne, col, val), + FilterOp::Lt => bin(BinaryOpKind::Lt, col, val), + FilterOp::Le => bin(BinaryOpKind::Le, col, val), + FilterOp::Gt => bin(BinaryOpKind::Gt, col, val), + FilterOp::Ge => bin(BinaryOpKind::Ge, col, val), + FilterOp::Like => bin(BinaryOpKind::Like, col, val), + FilterOp::NotLike => bin(BinaryOpKind::NotLike, col, val), + FilterOp::IsNull => ScalarExpr::IsNull { expr: Box::new(col), negated: false }, + FilterOp::IsNotNull => ScalarExpr::IsNull { expr: Box::new(col), negated: true }, + FilterOp::Regex(r) => bin( + BinaryOpKind::Regex, + col, + ScalarExpr::Literal(LiteralValue::Str(r.clone())), + ), + FilterOp::NotRegex(r) => bin( + BinaryOpKind::NotRegex, + col, + ScalarExpr::Literal(LiteralValue::Str(r.clone())), + ), + } +} + +fn bin(op: BinaryOpKind, lhs: ScalarExpr, rhs: ScalarExpr) -> ScalarExpr { + ScalarExpr::BinaryOp { op, lhs: Box::new(lhs), rhs: Box::new(rhs) } +} + +// ── Display helpers ─────────────────────────────────────────────────────────── + +impl std::fmt::Display for BinaryOpKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let s = match self { + 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::BitAnd => "&", BinaryOpKind::BitOr => "|", + BinaryOpKind::BitXor => "XOR", + BinaryOpKind::Concat => "||", + BinaryOpKind::Like => "LIKE", BinaryOpKind::NotLike => "NOT LIKE", + BinaryOpKind::Regex => "=~", BinaryOpKind::NotRegex => "!~", + BinaryOpKind::Unless => "unless", BinaryOpKind::Atan2 => "atan2", + }; + write!(f, "{s}") + } +} + +impl std::fmt::Display for AggFunc { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AggFunc::Count => write!(f, "COUNT"), + AggFunc::Sum => write!(f, "SUM"), + AggFunc::Avg => write!(f, "AVG"), + AggFunc::Min => write!(f, "MIN"), + AggFunc::Max => write!(f, "MAX"), + AggFunc::StdDev { .. } => write!(f, "STDDEV"), + AggFunc::Variance { .. } => write!(f, "VARIANCE"), + AggFunc::Quantile(p) => write!(f, "QUANTILE({p})"), + AggFunc::CountDistinct => write!(f, "COUNT_DISTINCT"), + AggFunc::HeavyHitters { k } => write!(f, "HEAVY_HITTERS({k})"), + AggFunc::Rate => write!(f, "rate"), + AggFunc::Increase => write!(f, "increase"), + AggFunc::Delta => write!(f, "delta"), + AggFunc::Custom(s) => write!(f, "{s}"), + } + } +} + +// ── Tests ───────────────────────────────────────────────────────────────────── + +#[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:?}"), + } + } + + // ── has_sketch_work ─────────────────────────────────────────────────────── + + #[test] + fn has_sketch_work_true_when_ddsketch_present() { + use crate::query_parser::sketch_algebra::ColumnRef; + let se = SketchExpr::Agg { + op: SketchAggOp::default_ddsketch(vec![0.5]), + col: ColumnRef::SampleValue, + input: Box::new(src("m")), + }; + assert!(QueryExpr::from_sketch_expr(&se).has_sketch_work()); + } + + #[test] + fn has_sketch_work_false_for_plain_source() { + let qe = QueryExpr::Source(SourceSpec { name: "x".into() }); + assert!(!qe.has_sketch_work()); + } + + // ── source_name ─────────────────────────────────────────────────────────── + + #[test] + fn source_name_extracted_through_chain() { + let se = SketchExpr::Window { + duration: Duration::from_secs(60), + input: Box::new(SketchExpr::Filter { + pred: vec![], + input: Box::new(src("my_metric")), + }), + }; + let qe = QueryExpr::from_sketch_expr(&se); + assert_eq!(qe.source_name(), Some("my_metric")); + } + + // ── AggFunc helpers ─────────────────────────────────────────────────────── + + #[test] + fn agg_func_mergeability() { + assert!(AggFunc::Sum.is_mergeable()); + assert!(AggFunc::Count.is_mergeable()); + assert!(AggFunc::Min.is_mergeable()); + assert!(AggFunc::Max.is_mergeable()); + assert!(!AggFunc::Avg.is_mergeable()); + assert!(!AggFunc::StdDev { population: false }.is_mergeable()); + assert!(!AggFunc::Variance { population: true }.is_mergeable()); + } + + #[test] + fn agg_func_sketchability() { + assert!(AggFunc::Quantile(0.99).is_sketchable()); + assert!(AggFunc::CountDistinct.is_sketchable()); + assert!(AggFunc::HeavyHitters { k: 10 }.is_sketchable()); + assert!(!AggFunc::Avg.is_sketchable()); + assert!(!AggFunc::Sum.is_sketchable()); + } + + #[test] + fn agg_func_to_sketch_op_quantile() { + let op = AggFunc::Quantile(0.99).to_sketch_op(); + assert!(matches!(op, Some(SketchAggOp::DDSketch { .. }))); + } + + #[test] + fn agg_func_to_sketch_op_count_distinct() { + let op = AggFunc::CountDistinct.to_sketch_op(); + assert!(matches!(op, Some(SketchAggOp::HLL { .. }))); + } + + #[test] + fn agg_func_to_sketch_op_heavy_hitters() { + let op = AggFunc::HeavyHitters { k: 50 }.to_sketch_op(); + assert!(matches!(op, Some(SketchAggOp::CountSketch { k: 50 }))); + } + + // ── ScalarExpr predicate list conversion ────────────────────────────────── + + #[test] + fn empty_pred_list_becomes_literal_true() { + let s = scalar_from_predicates(&[]); + assert!(matches!(s, ScalarExpr::Literal(LiteralValue::Bool(true)))); + } + + #[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) }, + ]; + let s = scalar_from_predicates(&preds); + assert!(matches!(s, ScalarExpr::BinaryOp { op: BinaryOpKind::And, .. })); + } + + // ── BinaryOpKind display ────────────────────────────────────────────────── + + #[test] + fn binary_op_kind_display() { + assert_eq!(BinaryOpKind::Add.to_string(), "+"); + assert_eq!(BinaryOpKind::And.to_string(), "AND"); + assert_eq!(BinaryOpKind::Regex.to_string(), "=~"); + assert_eq!(BinaryOpKind::NotRegex.to_string(), "!~"); + assert_eq!(BinaryOpKind::Unless.to_string(), "unless"); + } + + // ── 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 { + 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::Agg { + op: SketchAggOp::CountSketch { k: 10 }, + col: ColumnRef::Wildcard, + input: Box::new(src("price")), + }), + }), + }), + }; + let qe = QueryExpr::from_sketch_expr(&se); + assert!(qe.has_sketch_work()); + assert_eq!(qe.source_name(), Some("price")); + } + + // ── LetBinding and Subquery ─────────────────────────────────────────────── + + #[test] + fn let_binding_construction() { + let expr = QueryExpr::LetBinding { + name: "base".into(), + expr: Box::new(QueryExpr::Source(SourceSpec { name: "cpu".into() })), + body: Box::new(QueryExpr::Ref("base".into())), + }; + match expr { + QueryExpr::LetBinding { name, .. } => assert_eq!(name, "base"), + _ => panic!(), + } + } + + #[test] + fn histogram_quantile_node() { + let expr = QueryExpr::HistogramQuantile { + phi: 0.95, + input: Box::new(QueryExpr::Source(SourceSpec { name: "hist".into() })), + }; + match expr { + QueryExpr::HistogramQuantile { phi, .. } => { + assert!((phi - 0.95).abs() < 1e-9); + } + _ => panic!(), + } + } + + #[test] + fn promql_subquery_node() { + let expr = QueryExpr::PromQLSubquery { + range: Duration::from_secs(3600), + resolution: Some(Duration::from_secs(60)), + input: Box::new(QueryExpr::Source(SourceSpec { name: "m".into() })), + }; + match expr { + QueryExpr::PromQLSubquery { range, resolution, .. } => { + assert_eq!(range, Duration::from_secs(3600)); + assert_eq!(resolution, Some(Duration::from_secs(60))); + } + _ => panic!(), + } + } +} diff --git a/controller/src/algebra/mod.rs b/controller/src/algebra/mod.rs new file mode 100644 index 00000000..461de542 --- /dev/null +++ b/controller/src/algebra/mod.rs @@ -0,0 +1,50 @@ +//! General query algebra — full SQL/PromQL AST, cost-based optimizer, +//! and sketch-stage allocator. +//! +//! # Module layout +//! +//! | Module | Contents | +//! |--------|----------| +//! | [`expr`] | [`QueryExpr`] + [`ScalarExpr`] — the complete relational+scalar algebra | +//! | [`plan`] | [`PlanNode`], [`PipelineStage`], [`ExecutionMode`], [`CostEstimate`] | +//! | [`optimizer`] | [`QueryOptimizer`] + 12 rewrite rules | +//! | [`allocator`] | [`SketchAllocator`] — assigns stages and sketch types | +//! +//! # Typical usage +//! +//! ```rust,ignore +//! use controller::algebra::{ +//! expr::QueryExpr, +//! optimizer::QueryOptimizer, +//! allocator::SketchAllocator, +//! }; +//! 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])")?; +//! +//! // 2. Lift into the general algebra. +//! let query_expr = QueryExpr::from_sketch_expr(&sketch_expr); +//! +//! // 3. Optimise (cost-based fixed-point rewriting). +//! let (opt_expr, _iters) = QueryOptimizer::new(raw_bps).optimize(query_expr); +//! +//! // 4. Allocate stages. +//! let budgets = StageResourceBudgets::from_workload_chars(&workload_chars); +//! let plan_root = SketchAllocator::new(budgets, raw_bps).allocate(opt_expr); +//! +//! // 5. Inspect or serialise. +//! let summary = plan_root.summarise(raw_bps); +//! ``` + +pub mod allocator; +pub mod expr; +pub mod optimizer; +pub mod plan; + +// Convenience re-exports. +pub use allocator::SketchAllocator; +pub use expr::{AggFunc, BinaryOpKind, QueryExpr, ScalarExpr}; +pub use optimizer::QueryOptimizer; +pub use plan::{CostEstimate, ExecutionMode, PipelineStage, PlanNode, PlanSummary}; diff --git a/controller/src/algebra/optimizer.rs b/controller/src/algebra/optimizer.rs new file mode 100644 index 00000000..073c9c10 --- /dev/null +++ b/controller/src/algebra/optimizer.rs @@ -0,0 +1,1034 @@ +//! Cost-based fixed-point query optimizer. +//! +//! The optimizer applies a set of algebraic rewrite rules to a +//! [`QueryExpr`](super::expr::QueryExpr) tree until no rule fires (fixed +//! point). Each rule is a pure function `QueryExpr → Option`: +//! returning `None` means "this rule does not apply here". +//! +//! # Rules implemented +//! +//! | Rule | Name | Description | +//! |------|------|-------------| +//! | R1 | `PredicatePushDown` | Push `Filter` below `Window`, `Partition`, `Aggregate` | +//! | R2 | `MergeLifting` | Lift mergeable sketch aggs above `Partition` nodes | +//! | R3 | `HLLDedupElim` | Eliminate `Dedup` before HLL (HLL is inherently distinct) | +//! | R4 | `FilterWindowSwap` | Swap `Filter` below `Window` to reduce window input size | +//! | R5 | `TopKFusion` | Absorb `Limit` / `TopK` into a `CountSketch` agg | +//! | R6 | `HistogramQuantileFusion` | Recognise `HistogramQuantile(φ, Agg(DDSketch))` and mark | +//! | R7 | `SubqueryDecorrelation` | Hoist correlated `ScalarSubquery` to a `LetBinding` | +//! | R8 | `CommonSubexprElim` | Extract identical sub-trees into `LetBinding`s | +//! | R9 | `HydraConversion` | Convert multi-key `Partition + Agg` into `Hydra` sketch | +//! | R10| `WindowMerge` | Merge adjacent `Window` nodes with the same duration | +//! | R11| `PartitionElim` | Remove `Partition` with empty key list (becomes global agg) | +//! | R12| `SetOpFusion` | Fuse `SetOp(Union, Merge, Merge)` into a single `Merge` | +//! +//! ## Cost model integration +//! +//! Rules R1–R4 are cost-free (always beneficial). Rules R5–R12 consult a +//! [`CostModel`] that estimates bandwidth, memory, and CPU overhead. A +//! rewrite is only applied when the estimated cost improves. + +use std::collections::HashMap; + +use super::expr::{QueryExpr, ScalarExpr, SetOpKind, SortKey}; +use crate::query_parser::sketch_algebra::{PartitionKeys, SketchAggOp, SourceSpec}; + +// ── Cost model interface ────────────────────────────────────────────────────── + +/// Estimated cost of evaluating an expression at a given bandwidth. +#[derive(Debug, Clone, Default)] +pub struct NodeCost { + pub bytes_per_sec: f64, + pub memory_bytes: f64, + pub cpu_per_sample: f64, +} + +/// Pluggable cost oracle. The default implementation uses simple heuristics. +pub trait CostModel: Send + Sync { + /// Estimate the cost of the expression tree rooted at `expr`. + fn estimate(&self, expr: &QueryExpr) -> NodeCost; +} + +/// Default cost model — simple heuristics, no schema statistics. +pub struct DefaultCostModel { + pub raw_bytes_per_sec: f64, +} + +impl CostModel for DefaultCostModel { + fn estimate(&self, expr: &QueryExpr) -> NodeCost { + // Very rough: sketch nodes reduce bandwidth by 10×; exact nodes pass through. + let factor = match expr { + QueryExpr::SketchAgg { op, .. } => match op { + SketchAggOp::DDSketch { .. } => 0.05, + SketchAggOp::HLL { .. } => 0.02, + SketchAggOp::CountSketch { .. } | SketchAggOp::CountMin { .. } => 0.03, + SketchAggOp::Exact(_) => 1.0, + _ => 0.1, + }, + QueryExpr::Merge { inputs } => 1.0 / (inputs.len().max(1) as f64), + QueryExpr::Filter { .. } => 0.5, + QueryExpr::TopK { k, .. } => (*k as f64).recip().min(0.1), + _ => 1.0, + }; + NodeCost { + bytes_per_sec: self.raw_bytes_per_sec * factor, + memory_bytes: self.raw_bytes_per_sec * factor * 0.01, + cpu_per_sample: factor * 10.0, + } + } +} + +// ── Rewrite rule trait ──────────────────────────────────────────────────────── + +/// A single algebraic rewrite rule. +pub trait RewriteRule: Send + Sync { + /// Human-readable name for logging. + fn name(&self) -> &'static str; + + /// Try to rewrite `expr`. Returns `Some(new_expr)` if the rule fired, + /// `None` otherwise. The rule is applied top-down: the optimizer will + /// also recurse into the children of `new_expr`. + fn try_rewrite(&self, expr: QueryExpr, model: &dyn CostModel) -> Option; +} + +// ── R1: PredicatePushDown ───────────────────────────────────────────────────── + +/// Push `Filter` nodes as deep as possible — reduces data volume early. +/// +/// Transformations applied (all strictly beneficial, no cost model check): +/// * `Filter(p, Window(d, e))` → `Window(d, Filter(p, e))` +/// * `Filter(p, Partition(k, e))` → `Partition(k, Filter(p, e))` +/// * `Filter(p, Sort(k, e))` → `Sort(k, Filter(p, e))` +/// * `Filter(p, Limit(n, o, e))` → `Limit(n, o, Filter(p, e))` (NB: only safe when p is on input cols) +pub struct PredicatePushDown; + +impl RewriteRule for PredicatePushDown { + fn name(&self) -> &'static str { "PredicatePushDown" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + match expr { + QueryExpr::Filter { pred, input } => { + match *input { + // Filter below Window + QueryExpr::Window { duration, slide, input: inner } => Some( + QueryExpr::Window { + duration, + slide, + input: Box::new(QueryExpr::Filter { + pred, + input: inner, + }), + } + ), + // Filter below Partition + QueryExpr::Partition { keys, input: inner } => Some( + QueryExpr::Partition { + keys, + input: Box::new(QueryExpr::Filter { + pred, + input: inner, + }), + } + ), + // Filter below Sort (safe when pred references input columns only) + QueryExpr::Sort { keys, input: inner } => Some( + QueryExpr::Sort { + keys, + input: Box::new(QueryExpr::Filter { + pred, + input: inner, + }), + } + ), + // Not applicable — reconstruct + other => Some(QueryExpr::Filter { pred, input: Box::new(other) }), + } + } + _ => None, + } + } +} + +// ── R2: MergeLifting ────────────────────────────────────────────────────────── + +/// Lift mergeable `SketchAgg` ops above `Merge` nodes. +/// +/// `SketchAgg(op, Merge([a, b]))` → `Merge([SketchAgg(op, a), SketchAgg(op, b)])` +/// +/// Only applied for `is_mergeable()` ops so we don't incorrectly distribute +/// `Avg` or `StdDev`. +pub struct MergeLifting; + +impl RewriteRule for MergeLifting { + fn name(&self) -> &'static str { "MergeLifting" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + match expr { + QueryExpr::SketchAgg { ref op, ref col, ref input } + if op.is_mergeable() => + { + if let QueryExpr::Merge { inputs } = input.as_ref() { + let new_inputs: Vec = inputs.iter().map(|branch| { + QueryExpr::SketchAgg { + op: op.clone(), + col: col.clone(), + input: Box::new(branch.clone()), + } + }).collect(); + return Some(QueryExpr::Merge { inputs: new_inputs }); + } + None + } + _ => None, + } + } +} + +// ── R3: HLLDedupElim ───────────────────────────────────────────────────────── + +/// Eliminate `Dedup` nodes that immediately precede an HLL aggregation. +/// +/// HLL counts distinct values intrinsically; an explicit dedup step is +/// redundant and wastes CPU / memory. +/// +/// `SketchAgg(HLL, Dedup(col, e))` → `SketchAgg(HLL, e)` +pub struct HLLDedupElim; + +impl RewriteRule for HLLDedupElim { + fn name(&self) -> &'static str { "HLLDedupElim" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + match expr { + QueryExpr::SketchAgg { op: SketchAggOp::HLL { registers }, col, input } => { + if let QueryExpr::Dedup { input: inner, .. } = *input { + return Some(QueryExpr::SketchAgg { + op: SketchAggOp::HLL { registers }, + col, + input: inner, + }); + } + None + } + _ => None, + } + } +} + +// ── R4: FilterWindowSwap ────────────────────────────────────────────────────── + +/// Push `Filter` below `Window` when the predicate references only source +/// columns (not windowed aggregates). +/// +/// Identical to the push-down in R1 for the Window case, but checked +/// separately so the optimizer can attribute the transformation correctly +/// in logs. +pub struct FilterWindowSwap; + +impl RewriteRule for FilterWindowSwap { + fn name(&self) -> &'static str { "FilterWindowSwap" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + // Handled by PredicatePushDown — mark as no-op here to avoid double-fire. + match expr { + QueryExpr::Filter { pred, input } => { + if let QueryExpr::Window { duration, slide, input: inner } = *input { + return Some(QueryExpr::Window { + duration, + slide, + input: Box::new(QueryExpr::Filter { pred, input: inner }), + }); + } + None + } + _ => None, + } + } +} + +// ── R5: TopKFusion ──────────────────────────────────────────────────────────── + +/// Fuse a `Limit(k, TopK(_, e))` or `Limit(k, Sort(_, e))` into a single +/// `TopK(k, e)` node that the allocator maps to a `CountSketch`. +/// +/// `Limit(n, Sort([col DESC], e))` → `TopK(n, [col], e)` +pub struct TopKFusion; + +impl RewriteRule for TopKFusion { + fn name(&self) -> &'static str { "TopKFusion" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + match expr { + QueryExpr::Limit { n, offset: 0, input } => { + if let QueryExpr::Sort { keys, input: inner } = *input { + // Only fuse when all keys are DESC (top-k semantics). + if keys.iter().all(|k| k.desc) { + let by: Vec = keys.into_iter().map(|k| k.col).collect(); + return Some(QueryExpr::TopK { k: n, by, input: inner }); + } + } + None + } + _ => None, + } + } +} + +// ── R6: HistogramQuantileFusion ─────────────────────────────────────────────── + +/// Recognise `HistogramQuantile(φ, SketchAgg(DDSketch([φ]), …))` and +/// simplify to a single annotated node that the allocator handles as one +/// DDSketch query. +/// +/// `HistogramQuantile(φ, SketchAgg(DDSketch(qs), col, e))` +/// where `qs` contains `φ` +/// → `HistogramQuantile(φ, SketchAgg(DDSketch(qs), col, e))` [marked fused] +/// +/// In practice we just ensure the quantile is in the DDSketch's quantile +/// list so the allocator emits a single sketch with the right φ. +pub struct HistogramQuantileFusion; + +impl RewriteRule for HistogramQuantileFusion { + fn name(&self) -> &'static str { "HistogramQuantileFusion" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + match expr { + QueryExpr::HistogramQuantile { phi, input } => { + match *input { + QueryExpr::SketchAgg { + op: SketchAggOp::DDSketch { quantiles, epsilon }, + col, + input: inner, + } => { + if !quantiles.contains(&phi) { + let mut new_qs = quantiles; + new_qs.push(phi); + new_qs.sort_by(|a, b| a.partial_cmp(b).unwrap()); + Some(QueryExpr::HistogramQuantile { + phi, + input: Box::new(QueryExpr::SketchAgg { + op: SketchAggOp::DDSketch { quantiles: new_qs, epsilon }, + col, + input: inner, + }), + }) + } else { + Some(QueryExpr::HistogramQuantile { + phi, + input: Box::new(QueryExpr::SketchAgg { + op: SketchAggOp::DDSketch { quantiles, epsilon }, + col, + input: inner, + }), + }) + } + } + other => Some(QueryExpr::HistogramQuantile { phi, input: Box::new(other) }), + } + } + _ => None, + } + } +} + +// ── R7: SubqueryDecorrelation ───────────────────────────────────────────────── + +/// Hoist correlated `ScalarSubquery` references into `LetBinding`s so that +/// the subquery is evaluated once rather than once per row. +/// +/// This rule is a structural marker — full correlated-subquery detection +/// requires a binder pass that is out of scope here. We handle the simple +/// case: a `Filter` whose predicate contains a `ScalarSubquery` that does +/// not reference the filter's own input. +pub struct SubqueryDecorrelation; + +impl RewriteRule for SubqueryDecorrelation { + fn name(&self) -> &'static str { "SubqueryDecorrelation" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + match expr { + QueryExpr::Filter { pred, input } => { + if let Some((name, sq_expr, new_pred)) = extract_scalar_subquery(pred) { + return Some(QueryExpr::LetBinding { + name: name.clone(), + expr: Box::new(sq_expr), + body: Box::new(QueryExpr::Filter { + pred: new_pred, + input, + }), + }); + } + None + } + _ => None, + } + } +} + +/// If `pred` contains a `ScalarSubquery`, extract it as +/// `(binding_name, subquery_expr, pred_with_ref)`. +fn extract_scalar_subquery( + pred: ScalarExpr, +) -> Option<(String, QueryExpr, ScalarExpr)> { + match pred { + ScalarExpr::BinaryOp { op, lhs, rhs } => { + // Check lhs + if let ScalarExpr::ScalarSubquery(sq) = *lhs { + let name = "__subq_0".to_string(); + let new_pred = ScalarExpr::BinaryOp { + op, + lhs: Box::new(ScalarExpr::Column(name.clone())), + rhs, + }; + return Some((name, *sq, new_pred)); + } + // Check rhs + if let ScalarExpr::ScalarSubquery(sq) = *rhs { + let name = "__subq_0".to_string(); + let new_pred = ScalarExpr::BinaryOp { + op, + lhs, + rhs: Box::new(ScalarExpr::Column(name.clone())), + }; + return Some((name, *sq, new_pred)); + } + None + } + _ => None, + } +} + +// ── R8: CommonSubexprElim ───────────────────────────────────────────────────── + +/// Identify identical sub-trees that appear in multiple branches of a `Merge` +/// node and hoist them into a `LetBinding`. +/// +/// This is a conservative implementation: only `Source` nodes with the same +/// name are deduplicated (the common case where the same metric appears in +/// multiple union branches). +pub struct CommonSubexprElim; + +impl RewriteRule for CommonSubexprElim { + fn name(&self) -> &'static str { "CommonSubexprElim" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + match expr { + QueryExpr::Merge { ref inputs } => { + // Count occurrences of each source name. + let mut counts: HashMap = HashMap::new(); + for inp in inputs { + if let Some(name) = inp.source_name() { + *counts.entry(name.to_string()).or_insert(0) += 1; + } + } + let repeated: Vec = counts.into_iter() + .filter(|(_, c)| *c > 1) + .map(|(n, _)| n) + .collect(); + if repeated.is_empty() { + return None; + } + // Hoist the first repeated source into a LetBinding. + let name = repeated.into_iter().next()?; + let binding_name = format!("__cse_{name}"); + let new_inputs: Vec = inputs.iter().cloned().map(|inp| { + if inp.source_name() == Some(name.as_str()) { + QueryExpr::Ref(binding_name.clone()) + } else { + inp + } + }).collect(); + Some(QueryExpr::LetBinding { + name: binding_name, + expr: Box::new(QueryExpr::Source(SourceSpec { name })), + body: Box::new(QueryExpr::Merge { inputs: new_inputs }), + }) + } + _ => None, + } + } +} + +// ── R9: HydraConversion ─────────────────────────────────────────────────────── + +/// Convert `Partition(keys, SketchAgg(op, col, e))` where `keys` has ≥ 2 +/// dimensions into `SketchAgg(Hydra{inner: op, keys}, col, e)`. +/// +/// Hydra is a sketch-of-sketches that handles multi-dimensional GROUP BY +/// more efficiently than one sketch per group tuple. +pub struct HydraConversion; + +impl RewriteRule for HydraConversion { + fn name(&self) -> &'static str { "HydraConversion" } + + fn try_rewrite(&self, expr: QueryExpr, model: &dyn CostModel) -> Option { + match expr { + QueryExpr::Partition { keys: PartitionKeys::By(ref key_list), ref input } + if key_list.len() >= 2 => + { + if let QueryExpr::SketchAgg { op: ref inner_op, ref col, input: ref inner_input } = + **input + { + if matches!( + inner_op, + SketchAggOp::DDSketch { .. } + | SketchAggOp::HLL { .. } + | SketchAggOp::CountMin { .. } + | SketchAggOp::CountSketch { .. } + ) { + let hydra_op = SketchAggOp::Hydra { + inner: Box::new(inner_op.clone()), + partition_keys: key_list.clone(), + }; + let candidate = QueryExpr::SketchAgg { + op: hydra_op, + col: col.clone(), + input: inner_input.clone(), + }; + let old_cost = model.estimate(&expr); + let new_cost = model.estimate(&candidate); + if new_cost.memory_bytes < old_cost.memory_bytes + || new_cost.bytes_per_sec < old_cost.bytes_per_sec + { + return Some(candidate); + } + } + } + None + } + _ => None, + } + } +} + +// ── R10: WindowMerge ───────────────────────────────────────────────────────── + +/// Merge two adjacent `Window` nodes with the same `duration` into one. +/// +/// `Window(d, Window(d, e))` → `Window(d, e)` +pub struct WindowMerge; + +impl RewriteRule for WindowMerge { + fn name(&self) -> &'static str { "WindowMerge" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + match expr { + QueryExpr::Window { duration, slide, input } => { + if let QueryExpr::Window { duration: inner_d, slide: inner_s, input: inner_e } = + *input + { + if duration == inner_d && slide == inner_s { + return Some(QueryExpr::Window { + duration, + slide, + input: inner_e, + }); + } + } + None + } + _ => None, + } + } +} + +// ── R11: PartitionElim ──────────────────────────────────────────────────────── + +/// Remove `Partition` with an empty key list — equivalent to a global +/// aggregation with no GROUP BY. +/// +/// `Partition(By([]), e)` → `e` +pub struct PartitionElim; + +impl RewriteRule for PartitionElim { + fn name(&self) -> &'static str { "PartitionElim" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + match expr { + QueryExpr::Partition { keys, input } if keys.is_empty() => Some(*input), + _ => None, + } + } +} + +// ── R12: SetOpFusion ────────────────────────────────────────────────────────── + +/// Fuse `SetOp(Union, Merge([…]), Merge([…]))` into a single `Merge([…, …])`. +pub struct SetOpFusion; + +impl RewriteRule for SetOpFusion { + fn name(&self) -> &'static str { "SetOpFusion" } + + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + match expr { + QueryExpr::SetOp { + kind: SetOpKind::Union, + all: true, + left, + right, + } => { + match (*left, *right) { + (QueryExpr::Merge { inputs: mut li }, QueryExpr::Merge { inputs: mut ri }) => { + li.append(&mut ri); + Some(QueryExpr::Merge { inputs: li }) + } + (l, r) => Some(QueryExpr::SetOp { + kind: SetOpKind::Union, + all: true, + left: Box::new(l), + right: Box::new(r), + }), + } + } + _ => None, + } + } +} + +// ── Optimizer ───────────────────────────────────────────────────────────────── + +/// Fixed-point query optimizer. +/// +/// Call [`QueryOptimizer::optimize`] to rewrite a [`QueryExpr`] tree. +/// The optimizer iterates over all registered rules until no rule fires. +pub struct QueryOptimizer { + rules: Vec>, + cost_model: Box, + /// Maximum number of fixed-point iterations (prevents infinite loops). + max_iters: usize, +} + +impl QueryOptimizer { + /// Create an optimizer with the default rule set and cost model. + pub fn new(raw_bytes_per_sec: f64) -> Self { + Self { + rules: default_rules(), + cost_model: Box::new(DefaultCostModel { raw_bytes_per_sec }), + max_iters: 32, + } + } + + /// Create an optimizer with a custom cost model. + pub fn with_cost_model(cost_model: Box) -> Self { + Self { + rules: default_rules(), + cost_model, + max_iters: 32, + } + } + + /// Set maximum fixed-point iterations (default: 32). + pub fn max_iters(mut self, n: usize) -> Self { + self.max_iters = n; + self + } + + /// Optimize `expr` until fixed point or `max_iters` iterations. + /// + /// Returns the rewritten tree and the number of iterations actually run. + pub fn optimize(&self, expr: QueryExpr) -> (QueryExpr, usize) { + let mut current = expr; + for iter in 0..self.max_iters { + let (next, changed) = self.apply_all(current); + current = next; + if !changed { + return (current, iter + 1); + } + } + (current, self.max_iters) + } + + /// Apply all rules once to every node in the tree (single pass). + /// Returns `(new_tree, did_anything_change)`. + fn apply_all(&self, expr: QueryExpr) -> (QueryExpr, bool) { + // First recurse into children, then try rules at this node. + let (expr_with_new_children, child_changed) = self.recurse_children(expr); + let (final_expr, this_changed) = self.apply_rules_at(expr_with_new_children); + (final_expr, child_changed || this_changed) + } + + /// Apply all rules at the current node (no recursion). + fn apply_rules_at(&self, mut expr: QueryExpr) -> (QueryExpr, bool) { + let mut changed = false; + for rule in &self.rules { + if let Some(new_expr) = rule.try_rewrite(expr.clone(), self.cost_model.as_ref()) { + expr = new_expr; + changed = true; + // After firing, restart from the first rule (fixed-point per node). + break; + } + } + (expr, changed) + } + + /// Recurse into children, rebuilding the node with rewritten children. + fn recurse_children(&self, expr: QueryExpr) -> (QueryExpr, bool) { + macro_rules! recurse { + ($child:expr) => {{ + let (e, c) = self.apply_all(*$child); + (Box::new(e), c) + }}; + } + match expr { + QueryExpr::Source(_) | QueryExpr::Ref(_) => (expr, false), + + QueryExpr::Filter { pred, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::Filter { pred, input: new_input }, c) + } + QueryExpr::Project { cols, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::Project { cols, input: new_input }, c) + } + QueryExpr::Aggregate { keys, aggs, having, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::Aggregate { keys, aggs, having, input: new_input }, c) + } + QueryExpr::Window { duration, slide, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::Window { duration, slide, input: new_input }, c) + } + QueryExpr::SketchAgg { op, col, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::SketchAgg { op, col, input: new_input }, c) + } + QueryExpr::Partition { keys, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::Partition { keys, input: new_input }, c) + } + QueryExpr::Dedup { col, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::Dedup { col, input: new_input }, c) + } + QueryExpr::TopK { k, by, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::TopK { k, by, input: new_input }, c) + } + QueryExpr::Sort { keys, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::Sort { keys, input: new_input }, c) + } + QueryExpr::Limit { n, offset, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::Limit { n, offset, input: new_input }, c) + } + QueryExpr::WindowFunc { func, partition_by, order_by, frame, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::WindowFunc { func, partition_by, order_by, frame, input: new_input }, c) + } + QueryExpr::HistogramQuantile { phi, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::HistogramQuantile { phi, input: new_input }, c) + } + QueryExpr::PromQLSubquery { range, resolution, input } => { + let (new_input, c) = recurse!(input); + (QueryExpr::PromQLSubquery { range, resolution, input: new_input }, c) + } + QueryExpr::Merge { inputs } => { + let (new_inputs, changed): (Vec<_>, Vec<_>) = inputs + .into_iter() + .map(|inp| self.apply_all(inp)) + .unzip(); + (QueryExpr::Merge { inputs: new_inputs }, changed.into_iter().any(|c| c)) + } + QueryExpr::Join { kind, pred, left, right } => { + let (new_left, cl) = recurse!(left); + let (new_right, cr) = recurse!(right); + (QueryExpr::Join { kind, pred, left: new_left, right: new_right }, cl || cr) + } + QueryExpr::JoinSketch { join_key, outer, inner } => { + let (new_outer, co) = recurse!(outer); + let (new_inner, ci) = recurse!(inner); + (QueryExpr::JoinSketch { join_key, outer: new_outer, inner: new_inner }, co || ci) + } + QueryExpr::SetOp { kind, all, left, right } => { + let (new_left, cl) = recurse!(left); + let (new_right, cr) = recurse!(right); + (QueryExpr::SetOp { kind, all, left: new_left, right: new_right }, cl || cr) + } + QueryExpr::BinaryOp { op, lhs, rhs, vector_match } => { + let (new_lhs, cl) = recurse!(lhs); + let (new_rhs, cr) = recurse!(rhs); + (QueryExpr::BinaryOp { op, lhs: new_lhs, rhs: new_rhs, vector_match }, cl || cr) + } + QueryExpr::Subquery { alias, expr } => { + let (new_expr, c) = recurse!(expr); + (QueryExpr::Subquery { alias, expr: new_expr }, c) + } + QueryExpr::LetBinding { name, expr, body } => { + let (new_expr, ce) = recurse!(expr); + let (new_body, cb) = recurse!(body); + (QueryExpr::LetBinding { name, expr: new_expr, body: new_body }, ce || cb) + } + } + } +} + +/// Construct the default ordered rule set. +fn default_rules() -> Vec> { + vec![ + Box::new(PredicatePushDown), + Box::new(FilterWindowSwap), + Box::new(HLLDedupElim), + Box::new(WindowMerge), + Box::new(PartitionElim), + Box::new(TopKFusion), + Box::new(HistogramQuantileFusion), + Box::new(MergeLifting), + Box::new(SetOpFusion), + Box::new(HydraConversion), + Box::new(SubqueryDecorrelation), + Box::new(CommonSubexprElim), + ] +} + +// ── Tests ───────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use crate::algebra::expr::{LiteralValue, ScalarExpr}; + use crate::query_parser::sketch_algebra::{ColumnRef, PartitionKeys, SketchAggOp, SourceSpec}; + use std::time::Duration; + + fn src(name: &str) -> QueryExpr { + QueryExpr::Source(SourceSpec { name: name.into() }) + } + + fn opt() -> QueryOptimizer { + QueryOptimizer::new(100_000.0) + } + + // ── R1: PredicatePushDown ───────────────────────────────────────────────── + + #[test] + fn r1_pushes_filter_below_window() { + let expr = QueryExpr::Filter { + pred: ScalarExpr::Literal(LiteralValue::Bool(true)), + input: Box::new(QueryExpr::Window { + duration: Duration::from_secs(60), + slide: None, + input: Box::new(src("m")), + }), + }; + let (result, _) = opt().optimize(expr); + assert!( + matches!(&result, QueryExpr::Window { input, .. } + if matches!(input.as_ref(), QueryExpr::Filter { .. })), + "filter should be inside window: {result:?}" + ); + } + + #[test] + fn r1_pushes_filter_below_partition() { + let expr = QueryExpr::Filter { + pred: ScalarExpr::Literal(LiteralValue::Bool(true)), + input: Box::new(QueryExpr::Partition { + keys: PartitionKeys::By(vec!["host".into()]), + input: Box::new(src("cpu")), + }), + }; + let (result, _) = opt().optimize(expr); + assert!( + matches!(&result, QueryExpr::Partition { input, .. } + if matches!(input.as_ref(), QueryExpr::Filter { .. })) + ); + } + + // ── R3: HLLDedupElim ───────────────────────────────────────────────────── + + #[test] + fn r3_removes_dedup_before_hll() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::HLL { registers: 14 }, + col: ColumnRef::Named("user_id".into()), + input: Box::new(QueryExpr::Dedup { + col: "user_id".into(), + input: Box::new(src("events")), + }), + }; + let (result, _) = opt().optimize(expr); + assert!( + !matches!(&result, QueryExpr::SketchAgg { input, .. } + if matches!(input.as_ref(), QueryExpr::Dedup { .. })), + "Dedup should be eliminated before HLL" + ); + } + + // ── R5: TopKFusion ──────────────────────────────────────────────────────── + + #[test] + fn r5_fuses_limit_sort_to_topk() { + let expr = QueryExpr::Limit { + n: 10, + offset: 0, + input: Box::new(QueryExpr::Sort { + keys: vec![SortKey { col: "count".into(), desc: true, nulls_first: None }], + input: Box::new(src("events")), + }), + }; + let (result, _) = opt().optimize(expr); + assert!( + matches!(&result, QueryExpr::TopK { k: 10, .. }), + "expected TopK(10), got {result:?}" + ); + } + + #[test] + fn r5_does_not_fuse_ascending_sort() { + // ASC sort → not a top-k query. + let expr = QueryExpr::Limit { + n: 10, + offset: 0, + input: Box::new(QueryExpr::Sort { + keys: vec![SortKey { col: "ts".into(), desc: false, nulls_first: None }], + input: Box::new(src("events")), + }), + }; + let (result, _) = opt().optimize(expr); + assert!( + !matches!(&result, QueryExpr::TopK { .. }), + "ascending sort should not become TopK" + ); + } + + // ── R6: HistogramQuantileFusion ─────────────────────────────────────────── + + #[test] + fn r6_adds_phi_to_ddsketch_quantiles() { + let expr = QueryExpr::HistogramQuantile { + phi: 0.95, + input: Box::new(QueryExpr::SketchAgg { + op: SketchAggOp::DDSketch { quantiles: vec![0.5], epsilon: 0.01 }, + col: ColumnRef::SampleValue, + input: Box::new(src("latency")), + }), + }; + let (result, _) = opt().optimize(expr); + match &result { + QueryExpr::HistogramQuantile { input, .. } => { + if let QueryExpr::SketchAgg { op: SketchAggOp::DDSketch { quantiles, .. }, .. } = + input.as_ref() + { + assert!(quantiles.contains(&0.95), "0.95 should be in DDSketch quantiles"); + assert!(quantiles.contains(&0.5), "0.5 should still be present"); + } else { + panic!("expected DDSketch under HistogramQuantile"); + } + } + other => panic!("unexpected {other:?}"), + } + } + + // ── R10: WindowMerge ────────────────────────────────────────────────────── + + #[test] + fn r10_merges_duplicate_windows() { + let expr = QueryExpr::Window { + duration: Duration::from_secs(300), + slide: None, + input: Box::new(QueryExpr::Window { + duration: Duration::from_secs(300), + slide: None, + input: Box::new(src("m")), + }), + }; + let (result, _) = opt().optimize(expr); + assert!( + !matches!(&result, QueryExpr::Window { input, .. } + if matches!(input.as_ref(), QueryExpr::Window { .. })), + "duplicate window should be merged" + ); + } + + // ── R11: PartitionElim ──────────────────────────────────────────────────── + + #[test] + fn r11_removes_empty_partition() { + let expr = QueryExpr::Partition { + keys: PartitionKeys::By(vec![]), + input: Box::new(src("m")), + }; + let (result, _) = opt().optimize(expr); + assert!( + matches!(&result, QueryExpr::Source(_)), + "empty Partition should be eliminated" + ); + } + + // ── Fixed-point convergence ─────────────────────────────────────────────── + + #[test] + fn optimizer_reaches_fixed_point_on_simple_tree() { + let expr = src("m"); + let (result, iters) = opt().optimize(expr); + assert!(iters < 5, "should converge quickly on source-only tree"); + assert!(matches!(result, QueryExpr::Source(_))); + } + + #[test] + fn optimizer_chain_of_rewrites() { + // Filter(Window(Dedup(HLL(Source)))) → + // R1: Window(Filter(Dedup(HLL(Source)))) + // R3: Window(Filter(HLL(Source))) (HLL absorbs Dedup) + let expr = QueryExpr::Filter { + pred: ScalarExpr::Literal(LiteralValue::Bool(true)), + input: Box::new(QueryExpr::Window { + duration: Duration::from_secs(60), + slide: None, + input: Box::new(QueryExpr::SketchAgg { + op: SketchAggOp::HLL { registers: 14 }, + col: ColumnRef::Named("uid".into()), + input: Box::new(QueryExpr::Dedup { + col: "uid".into(), + input: Box::new(src("events")), + }), + }), + }), + }; + let (result, _iters) = opt().optimize(expr); + // The Dedup should be gone. + let mut dedup_found = false; + result.walk(&mut |n| { + if matches!(n, QueryExpr::Dedup { .. }) { + dedup_found = true; + } + }); + assert!(!dedup_found, "Dedup should have been eliminated"); + } + + // ── R2: MergeLifting ───────────────────────────────────────────────────── + + #[test] + fn r2_lifts_mergeable_sketch_above_merge() { + let expr = QueryExpr::SketchAgg { + op: SketchAggOp::default_hll(), + col: ColumnRef::Named("uid".into()), + input: Box::new(QueryExpr::Merge { + inputs: vec![src("shard_a"), src("shard_b")], + }), + }; + let (result, _) = opt().optimize(expr); + assert!( + matches!(&result, QueryExpr::Merge { inputs } + if inputs.iter().all(|i| matches!(i, QueryExpr::SketchAgg { .. }))), + "HLL should be pushed into each Merge branch" + ); + } + + // ── R12: SetOpFusion ───────────────────────────────────────────────────── + + #[test] + fn r12_fuses_union_of_merges() { + let expr = QueryExpr::SetOp { + kind: SetOpKind::Union, + all: true, + left: Box::new(QueryExpr::Merge { inputs: vec![src("a"), src("b")] }), + right: Box::new(QueryExpr::Merge { inputs: vec![src("c")] }), + }; + let (result, _) = opt().optimize(expr); + match result { + QueryExpr::Merge { inputs } => assert_eq!(inputs.len(), 3), + other => panic!("expected Merge(3), got {other:?}"), + } + } +} diff --git a/controller/src/algebra/plan.rs b/controller/src/algebra/plan.rs new file mode 100644 index 00000000..8857c5ab --- /dev/null +++ b/controller/src/algebra/plan.rs @@ -0,0 +1,428 @@ +//! Annotated plan nodes — the output of the [`super::allocator::SketchAllocator`]. +//! +//! After the optimizer rewrites a [`QueryExpr`](super::expr::QueryExpr) tree, +//! the allocator wraps every node in a [`PlanNode`] that carries: +//! +//! * **`stage`** — which pipeline component executes this operator. +//! * **`mode`** — whether the operator uses sketch approximation or exact +//! computation. +//! * **`cost`** — estimated memory and bandwidth cost at this node. +//! * **`annotation`** — additional hints for the code-generator (e.g. which +//! sketch type to use, whether delta encoding is enabled). +//! +//! The annotated plan tree is serialisable to JSON so it can be included in +//! the `/api/v1/plan` response for observability. + +use serde::{Deserialize, Serialize}; + +use super::expr::QueryExpr; + +// ── Pipeline stages ─────────────────────────────────────────────────────────── + +/// Which component in the data pipeline executes an operator. +/// +/// The ordering `Agent < Backend < Precompute < Db` mirrors the data-flow +/// direction: data originates at the Agent and flows toward the Db. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PipelineStage { + /// SDK-side OTel Collector — highest bandwidth savings, lowest latency. + Agent, + /// Central merge collector — aggregates partial sketches from many agents. + Backend, + /// ASAPQuery pre-computation engine — materialises recurring queries. + Precompute, + /// Exact OLAP / time-series database — last resort for non-sketchable ops. + Db, +} + +impl std::fmt::Display for PipelineStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let s = match self { + PipelineStage::Agent => "agent", + PipelineStage::Backend => "backend", + PipelineStage::Precompute => "precompute", + PipelineStage::Db => "db", + }; + write!(f, "{s}") + } +} + +// ── Execution mode ──────────────────────────────────────────────────────────── + +/// Whether an operator uses sketch approximation or runs exactly. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ExecutionMode { + /// The operator produces an approximate result via a data sketch. + Sketch, + /// The operator computes an exact result (no error bounds). + Exact, + /// The operator is a structural / routing node (merge, partition, …) + /// that does not itself aggregate — its mode is determined by its children. + Passthrough, +} + +// ── Cost estimate ───────────────────────────────────────────────────────────── + +/// Estimated resource cost of a single plan node. +/// +/// The allocator fills this in using the same cost model as the legacy +/// [`crate::planner::cost_model`]. All fields default to `0.0` for nodes +/// whose cost is negligible or unknown. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct CostEstimate { + /// Outbound bytes per second produced by this node. + pub bytes_per_sec: f64, + /// Memory footprint of the sketch or intermediate state (bytes). + pub memory_bytes: f64, + /// CPU overhead per input sample (µs). + pub cpu_micros_per_sample: f64, + /// Compression ratio relative to the raw OTLP baseline (≥ 1.0 is better). + pub compression_ratio: f64, +} + +// ── Node annotation ─────────────────────────────────────────────────────────── + +/// Extra hints attached to a plan node by the allocator. +/// +/// Not all fields are relevant to all node types; unused fields are `None`. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct NodeAnnotation { + /// Sketch type selected by the allocator (only for sketch nodes). + pub sketch_type: Option, + /// Sketch parameters (width/depth/registers/epsilon). + pub sketch_params: Option, + /// Whether delta encoding should be used at this node. + pub delta_enabled: bool, + /// Minimum cell-change threshold for delta encoding (T). + pub delta_threshold: f64, + /// Human-readable explanation of why this stage/mode was chosen. + pub rationale: String, + /// Whether this node was demoted to a later stage due to budget overflow. + pub budget_demotion: bool, +} + +// ── Plan node ───────────────────────────────────────────────────────────────── + +/// An annotated node in the physical execution plan. +/// +/// The `expr` field holds the logical operator; the surrounding fields +/// describe where and how it runs. +#[derive(Debug, Clone)] +pub struct PlanNode { + /// The logical operator at this node. + pub expr: QueryExpr, + /// Which pipeline stage executes this operator. + pub stage: PipelineStage, + /// Sketch vs. exact vs. passthrough. + pub mode: ExecutionMode, + /// Estimated resource cost. + pub cost: CostEstimate, + /// Allocator hints for code-generation. + pub annotation: NodeAnnotation, + /// Child plan nodes (mirrors `expr`'s children after annotation). + pub children: Vec, +} + +impl PlanNode { + /// Create a leaf `PlanNode` (no children) with default cost/annotation. + pub fn leaf(expr: QueryExpr, stage: PipelineStage, mode: ExecutionMode) -> Self { + Self { + expr, + stage, + mode, + cost: CostEstimate::default(), + annotation: NodeAnnotation::default(), + children: vec![], + } + } + + /// Recursively collect all nodes at a given stage, depth-first. + pub fn nodes_at_stage(&self, target: &PipelineStage) -> Vec<&PlanNode> { + let mut out = vec![]; + if &self.stage == target { + out.push(self); + } + for c in &self.children { + out.extend(c.nodes_at_stage(target)); + } + out + } + + /// Recursively collect all sketch nodes (mode == Sketch). + pub fn sketch_nodes(&self) -> Vec<&PlanNode> { + let mut out = vec![]; + if self.mode == ExecutionMode::Sketch { + out.push(self); + } + for c in &self.children { + out.extend(c.sketch_nodes()); + } + out + } + + /// Total estimated bandwidth of all nodes at `stage` (bytes/sec). + pub fn stage_bandwidth(&self, stage: &PipelineStage) -> f64 { + self.nodes_at_stage(stage) + .iter() + .map(|n| n.cost.bytes_per_sec) + .sum() + } + + /// Total estimated memory of all nodes at `stage` (bytes). + pub fn stage_memory(&self, stage: &PipelineStage) -> f64 { + self.nodes_at_stage(stage) + .iter() + .map(|n| n.cost.memory_bytes) + .sum() + } + + /// Returns a flat, depth-first list of `(depth, node)` pairs for display. + pub fn flatten(&self) -> Vec<(usize, &PlanNode)> { + let mut out = vec![]; + self.flatten_inner(0, &mut out); + out + } + + fn flatten_inner<'a>(&'a self, depth: usize, out: &mut Vec<(usize, &'a PlanNode)>) { + out.push((depth, self)); + for c in &self.children { + c.flatten_inner(depth + 1, out); + } + } +} + +// ── Plan summary (serialisable) ─────────────────────────────────────────────── + +/// A serialisable summary of the full annotated plan, suitable for inclusion +/// in the `/api/v1/plan` JSON response. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct PlanSummary { + /// Total estimated bandwidth saved vs raw OTLP (bytes/sec). + pub bandwidth_saved_bytes_per_sec: f64, + /// Total estimated agent memory for all sketch nodes (bytes). + pub agent_memory_bytes: f64, + /// Total estimated backend memory for all sketch nodes (bytes). + pub backend_memory_bytes: f64, + /// Whether any node was demoted due to budget overflow. + pub has_budget_demotion: bool, + /// List of per-node stage + mode + rationale entries. + pub node_annotations: Vec, +} + +/// One row in the [`PlanSummary::node_annotations`] table. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct NodeSummaryEntry { + pub node_kind: String, + pub stage: PipelineStage, + pub mode: ExecutionMode, + pub rationale: String, + pub memory_bytes: f64, + pub bytes_per_sec: f64, +} + +impl PlanNode { + /// Build a [`PlanSummary`] from this root node. + pub fn summarise(&self, raw_bytes_per_sec: f64) -> PlanSummary { + let flat = self.flatten(); + let agent_mem: f64 = flat.iter() + .filter(|(_, n)| n.stage == PipelineStage::Agent) + .map(|(_, n)| n.cost.memory_bytes) + .sum(); + let backend_mem: f64 = flat.iter() + .filter(|(_, n)| n.stage == PipelineStage::Backend) + .map(|(_, n)| n.cost.memory_bytes) + .sum(); + let plan_bw: f64 = flat.iter() + .filter(|(_, n)| matches!( + n.stage, PipelineStage::Agent | PipelineStage::Backend + )) + .map(|(_, n)| n.cost.bytes_per_sec) + .fold(f64::INFINITY, f64::min); // min of outbound paths + let saved = if raw_bytes_per_sec > plan_bw { + raw_bytes_per_sec - plan_bw + } else { + 0.0 + }; + let has_demotion = flat.iter().any(|(_, n)| n.annotation.budget_demotion); + let entries = flat.iter().map(|(_, n)| { + NodeSummaryEntry { + node_kind: format!("{:?}", n.expr).split_whitespace().next().unwrap_or("?").to_string(), + stage: n.stage.clone(), + mode: n.mode.clone(), + rationale: n.annotation.rationale.clone(), + memory_bytes: n.cost.memory_bytes, + bytes_per_sec: n.cost.bytes_per_sec, + } + }).collect(); + PlanSummary { + bandwidth_saved_bytes_per_sec: saved, + agent_memory_bytes: agent_mem, + backend_memory_bytes: backend_mem, + has_budget_demotion: has_demotion, + node_annotations: entries, + } + } +} + +// ── Tests ───────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use crate::algebra::expr::QueryExpr; + use crate::query_parser::sketch_algebra::SourceSpec; + + fn source_node(name: &str, stage: PipelineStage) -> PlanNode { + PlanNode::leaf( + QueryExpr::Source(SourceSpec { name: name.into() }), + stage, + ExecutionMode::Passthrough, + ) + } + + // ── PipelineStage ordering ──────────────────────────────────────────────── + + #[test] + fn stage_ordering_agent_lt_db() { + assert!(PipelineStage::Agent < PipelineStage::Db); + assert!(PipelineStage::Agent < PipelineStage::Backend); + assert!(PipelineStage::Backend < PipelineStage::Precompute); + assert!(PipelineStage::Precompute < PipelineStage::Db); + } + + // ── nodes_at_stage ──────────────────────────────────────────────────────── + + #[test] + fn nodes_at_stage_collects_correctly() { + let root = PlanNode { + expr: QueryExpr::Source(SourceSpec { name: "root".into() }), + stage: PipelineStage::Agent, + mode: ExecutionMode::Sketch, + cost: CostEstimate { memory_bytes: 100.0, ..Default::default() }, + annotation: NodeAnnotation::default(), + children: vec![ + source_node("child_agent", PipelineStage::Agent), + source_node("child_backend", PipelineStage::Backend), + ], + }; + let agent_nodes = root.nodes_at_stage(&PipelineStage::Agent); + assert_eq!(agent_nodes.len(), 2); // root + child_agent + let backend_nodes = root.nodes_at_stage(&PipelineStage::Backend); + assert_eq!(backend_nodes.len(), 1); + } + + // ── sketch_nodes ───────────────────────────────────────────────────────── + + #[test] + fn sketch_nodes_only_returns_sketch_mode() { + let root = PlanNode { + expr: QueryExpr::Source(SourceSpec { name: "r".into() }), + stage: PipelineStage::Agent, + mode: ExecutionMode::Sketch, + cost: CostEstimate::default(), + annotation: NodeAnnotation::default(), + children: vec![ + PlanNode::leaf( + QueryExpr::Source(SourceSpec { name: "exact_child".into() }), + PipelineStage::Db, + ExecutionMode::Exact, + ), + PlanNode::leaf( + QueryExpr::Source(SourceSpec { name: "sketch_child".into() }), + PipelineStage::Backend, + ExecutionMode::Sketch, + ), + ], + }; + let sn = root.sketch_nodes(); + assert_eq!(sn.len(), 2); // root (Sketch) + sketch_child + } + + // ── stage_bandwidth / stage_memory ──────────────────────────────────────── + + #[test] + fn stage_bandwidth_sums_nodes_at_stage() { + let root = PlanNode { + expr: QueryExpr::Source(SourceSpec { name: "r".into() }), + stage: PipelineStage::Agent, + mode: ExecutionMode::Sketch, + cost: CostEstimate { bytes_per_sec: 500.0, ..Default::default() }, + annotation: NodeAnnotation::default(), + children: vec![ + PlanNode { + expr: QueryExpr::Source(SourceSpec { name: "c".into() }), + stage: PipelineStage::Agent, + mode: ExecutionMode::Passthrough, + cost: CostEstimate { bytes_per_sec: 200.0, ..Default::default() }, + annotation: NodeAnnotation::default(), + children: vec![], + }, + ], + }; + assert!((root.stage_bandwidth(&PipelineStage::Agent) - 700.0).abs() < 1e-6); + } + + // ── flatten ─────────────────────────────────────────────────────────────── + + #[test] + fn flatten_returns_depth_zero_for_root() { + let root = source_node("r", PipelineStage::Agent); + let flat = root.flatten(); + assert_eq!(flat.len(), 1); + assert_eq!(flat[0].0, 0); // depth = 0 + } + + #[test] + fn flatten_depth_increments_per_level() { + let root = PlanNode { + expr: QueryExpr::Source(SourceSpec { name: "r".into() }), + stage: PipelineStage::Agent, + mode: ExecutionMode::Passthrough, + cost: CostEstimate::default(), + annotation: NodeAnnotation::default(), + children: vec![source_node("c1", PipelineStage::Backend)], + }; + let flat = root.flatten(); + assert_eq!(flat[0].0, 0); + assert_eq!(flat[1].0, 1); + } + + // ── PlanSummary ─────────────────────────────────────────────────────────── + + #[test] + fn summarise_reports_bandwidth_saved() { + let root = PlanNode { + expr: QueryExpr::Source(SourceSpec { name: "r".into() }), + stage: PipelineStage::Agent, + mode: ExecutionMode::Sketch, + cost: CostEstimate { + bytes_per_sec: 1_000.0, + memory_bytes: 256.0, + ..Default::default() + }, + annotation: NodeAnnotation::default(), + children: vec![], + }; + // Raw baseline is 10 000 B/s; plan reduces to 1 000 B/s → saved = 9 000. + let summary = root.summarise(10_000.0); + assert!((summary.bandwidth_saved_bytes_per_sec - 9_000.0).abs() < 1.0); + assert!((summary.agent_memory_bytes - 256.0).abs() < 1.0); + assert!(!summary.has_budget_demotion); + } + + #[test] + fn summarise_detects_budget_demotion() { + let root = PlanNode { + expr: QueryExpr::Source(SourceSpec { name: "r".into() }), + stage: PipelineStage::Backend, + mode: ExecutionMode::Sketch, + cost: CostEstimate::default(), + annotation: NodeAnnotation { budget_demotion: true, ..Default::default() }, + children: vec![], + }; + let summary = root.summarise(0.0); + assert!(summary.has_budget_demotion); + } +} diff --git a/controller/src/main.rs b/controller/src/main.rs index 8f5708a3..258682fa 100644 --- a/controller/src/main.rs +++ b/controller/src/main.rs @@ -1,3 +1,4 @@ +mod algebra; mod analyzer; mod config; mod monitor; diff --git a/controller/src/types.rs b/controller/src/types.rs index 1ab6cef0..32ddad20 100644 --- a/controller/src/types.rs +++ b/controller/src/types.rs @@ -237,7 +237,7 @@ pub struct QueryWorkload { pub quantiles: Vec, } -#[derive(Debug, Clone, Default)] +#[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct SketchParams { pub relative_accuracy: f64, pub k: u32, @@ -247,6 +247,33 @@ 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,