diff --git a/controller/src/intent_algebra/column_resolution.rs b/controller/src/intent_algebra/column_resolution.rs new file mode 100644 index 00000000..80490dff --- /dev/null +++ b/controller/src/intent_algebra/column_resolution.rs @@ -0,0 +1,256 @@ +//! Schema-driven column resolution for the legacy `QueryExpr` IR +//! (Step β of the legacy_expr migration). +//! +//! Step α (PR #138) replaced the legacy `AggIntent` enum with the canonical +//! [`crate::intent_algebra::agg_intent::AggIntent`]. The legacy IR still +//! uses [`crate::intent_algebra::legacy_expr::ColumnRef::Named(String)`] for +//! column references; the canonical IR uses positional +//! [`crate::intent_algebra::schema::ColumnId`] resolved against a per-node +//! [`crate::intent_algebra::schema::Schema`]. +//! +//! Step β plumbs `Schema` through every consumer that walks the legacy +//! tree so Step γ can migrate variant-by-variant to positional column ids +//! without first having to acquire a Schema everywhere. +//! +//! ## Threading approach +//! +//! Approach **(c)** per the migration spec: each traversal function takes a +//! `parent_schema: &Schema` parameter (the schema in scope at that node). +//! The root walker derives the base schema via [`infer_schema_for_root`] +//! (which walks to the outermost `Source` leaf and synthesises a default +//! schema for it) and threads it downward. +//! +//! Legacy operators here are mostly pass-through with respect to schema — +//! they don't add columns, they just filter / partition / sort / window +//! the same rows. The few that DO transform the schema (`Aggregate`, +//! `Project`, `Distinct`) are flagged as Step γ TODOs; Step γ will +//! migrate them onto the canonical `QueryExpr::output_schema_in` path so +//! the schema is locally derivable from the variant fields. +//! +//! ## Why a synthesized default +//! +//! The legacy_expr migration plan's "Synthesise from metric name: +//! `(ts, value, *labels)`" decision applies here. There is no +//! `SchemaCatalog` in the controller today, so the source leaf has to +//! produce a schema purely from the metric / table name. This module +//! ships the minimal time-series schema that's correct for the DC + PromQL +//! use case (the only consumers that exist in tree today). Tabular +//! sources (`SourceSpec::Table { table_ref }`) inherit the same shape as +//! a placeholder; a real catalog lookup is a follow-up Step γ TODO. + +use thiserror::Error; + +use crate::intent_algebra::legacy_expr::{ColumnRef, QueryExpr, SourceSpec}; +use crate::intent_algebra::schema::{Column, ColumnId, DataType, Schema}; + +/// Errors returned by [`resolve_column_ref`] / [`resolve_column_refs`]. +#[derive(Debug, Error, PartialEq, Eq)] +pub enum ResolveError { + /// `ColumnRef::Named(name)` did not match any column in the supplied + /// schema. Common cause: the schema was synthesized from a metric name + /// but the name being resolved is a free-form label (label sets aren't + /// representable in the canonical `Schema` today — see the module + /// doc-comment). + #[error( + "column `{name}` not found in schema (have: {available:?}) \ + — Step γ TODO: label-set resolution against a real catalog" + )] + NotFound { + name: String, + available: Vec, + }, + /// `ColumnRef::SampleValue` was resolved against a schema that has no + /// `value` column. The synthesized time-series schema always has one, + /// so this only fires for tabular sources without the convention. + #[error("ColumnRef::SampleValue has no `value` column in schema (have: {available:?})")] + NoSampleValue { available: Vec }, + /// `ColumnRef::Wildcard` cannot be resolved to a single + /// [`ColumnId`] — by definition it refers to every row, not a specific + /// column. Callers that hit this branch should special-case `Wildcard` + /// instead of asking for a positional id. + #[error("ColumnRef::Wildcard cannot be resolved to a single ColumnId")] + WildcardNotPositional, +} + +/// Synthesize a default time-series schema for a metric / table source. +/// +/// The shape is the conventional PromQL leaf: `(ts: Timestamp, value: +/// Float64)`. Open-set labels are intentionally omitted — the canonical +/// [`Schema`] model uses a positional `Vec` and has no +/// representation for "any number of label columns whose names are +/// data-dependent." That's tracked as a Step γ TODO at the module level. +/// +/// Per the legacy_expr migration plan's "Synthesise from metric name: +/// `(ts, value, *labels)`" decision — minus the `*labels` part the +/// canonical schema model can't express today. +pub fn infer_source_schema(_metric_or_table_name: &str) -> Schema { + Schema::with_time_index( + vec![ + Column { + name: "ts".into(), + dtype: DataType::Timestamp, + nullable: false, + }, + Column { + name: "value".into(), + dtype: DataType::Float64, + nullable: false, + }, + ], + 0, + // No provable unique key without a catalog — leave empty so the + // CSE gatekeeper (`cse_reuse_is_legal`) conservatively refuses to + // share a producer across queries until a real catalog ships. + Vec::new(), + ) +} + +/// Walk the legacy `QueryExpr` tree to find the outermost (left-most) +/// `Source` leaf and synthesise its default schema via +/// [`infer_source_schema`]. This is the entry-point Schema that consumer +/// walkers thread downward through the tree. +/// +/// Falls back to an empty `Schema` when the root has no `Source` leaf +/// (e.g. a bare `Ref(name)` that resolves outside the supplied tree). +/// Callers that need real Ref resolution should pre-resolve via a +/// `BindingScope` analogue — for Step β the empty fallback is fine +/// because consumers only read the schema in pass-through paths. +pub fn infer_schema_for_root(expr: &QueryExpr) -> Schema { + match expr.source_name() { + Some(name) => infer_source_schema(name), + None => Schema::default(), + } +} + +/// Resolve a single [`ColumnRef`] against a schema, returning the +/// positional [`ColumnId`]. Step γ consumers will call this at the point +/// where they need a `ColumnId` instead of a `String` — for now the +/// helper exists so the migration is mechanical at that point. +/// +/// Resolution rules (mirror canonical +/// [`crate::intent_algebra::query_expr::ColumnRef`] semantics where they +/// already exist, e.g. `Distinct { cols }` in `output_schema_in`): +/// +/// * `ColumnRef::Named(name)` → `schema.column_id(name)`. +/// * `ColumnRef::SampleValue` → `schema.column_id("value")` (PromQL +/// convention; the canonical lowering emits a column literally named +/// `value` for time-series scans). +/// * `ColumnRef::Wildcard` → [`ResolveError::WildcardNotPositional`] +/// (callers must special-case). +pub fn resolve_column_ref(col: &ColumnRef, schema: &Schema) -> Result { + match col { + ColumnRef::Named(name) => schema + .column_id(name) + .ok_or_else(|| ResolveError::NotFound { + name: name.clone(), + available: schema.columns.iter().map(|c| c.name.clone()).collect(), + }), + ColumnRef::SampleValue => { + schema + .column_id("value") + .ok_or_else(|| ResolveError::NoSampleValue { + available: schema.columns.iter().map(|c| c.name.clone()).collect(), + }) + } + ColumnRef::Wildcard => Err(ResolveError::WildcardNotPositional), + } +} + +/// Slice-flavoured [`resolve_column_ref`]: resolves every entry, +/// short-circuiting on the first error. Used by `Distinct { cols }` and +/// the future `Aggregate { by }` migration when keys move from +/// `Vec` to `Vec`. +pub fn resolve_column_refs( + cols: &[ColumnRef], + schema: &Schema, +) -> Result, ResolveError> { + cols.iter().map(|c| resolve_column_ref(c, schema)).collect() +} + +// ── Tests ───────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use crate::intent_algebra::legacy_expr::SourceSpec; + + fn src(name: &str) -> QueryExpr { + QueryExpr::Source(SourceSpec { name: name.into() }) + } + + #[test] + fn source_schema_has_ts_and_value() { + let s = infer_source_schema("http_requests_total"); + assert_eq!(s.columns.len(), 2); + assert_eq!(s.columns[0].name, "ts"); + assert_eq!(s.columns[1].name, "value"); + assert_eq!(s.time_index, Some(0)); + // No provable unique key without a real catalog — conservative. + assert!(!s.has_unique_key()); + } + + #[test] + fn root_schema_via_walk() { + let expr = QueryExpr::Filter { + pred: crate::intent_algebra::legacy_expr::ScalarExpr::Literal( + crate::intent_algebra::legacy_expr::LiteralValue::Bool(true), + ), + input: Box::new(src("cpu_usage")), + }; + let s = infer_schema_for_root(&expr); + assert_eq!(s.columns.len(), 2); + assert_eq!(s.columns[1].name, "value"); + } + + #[test] + fn resolve_named_against_value_column() { + let s = infer_source_schema("m"); + let col = ColumnRef::Named("value".into()); + assert_eq!(resolve_column_ref(&col, &s), Ok(1)); + } + + #[test] + fn resolve_sample_value() { + let s = infer_source_schema("m"); + assert_eq!(resolve_column_ref(&ColumnRef::SampleValue, &s), Ok(1)); + } + + #[test] + fn resolve_wildcard_errors() { + let s = infer_source_schema("m"); + assert_eq!( + resolve_column_ref(&ColumnRef::Wildcard, &s), + Err(ResolveError::WildcardNotPositional) + ); + } + + #[test] + fn resolve_unknown_name_errors() { + let s = infer_source_schema("m"); + let err = resolve_column_ref(&ColumnRef::Named("host".into()), &s).unwrap_err(); + assert!(matches!(err, ResolveError::NotFound { .. })); + } + + #[test] + fn resolve_slice_short_circuits() { + let s = infer_source_schema("m"); + let cols = vec![ + ColumnRef::Named("value".into()), + ColumnRef::Named("not_there".into()), + ]; + let r = resolve_column_refs(&cols, &s); + assert!(matches!(r, Err(ResolveError::NotFound { .. }))); + } + + #[test] + fn root_schema_falls_back_when_no_source() { + let expr = QueryExpr::Ref("dangling".into()); + let s = infer_schema_for_root(&expr); + assert_eq!(s.columns.len(), 0); + } +} + +// Silence "field never read" on `SourceSpec` when this module is the only +// consumer in some configurations. +#[allow(dead_code)] +const _: fn(&SourceSpec) -> &str = |s| s.name.as_str(); diff --git a/controller/src/intent_algebra/legacy_lower.rs b/controller/src/intent_algebra/legacy_lower.rs index ae139549..eea2da58 100644 --- a/controller/src/intent_algebra/legacy_lower.rs +++ b/controller/src/intent_algebra/legacy_lower.rs @@ -14,6 +14,7 @@ //! unchanged — the physical planner handles them. use crate::intent_algebra::legacy_expr::*; +use crate::intent_algebra::{infer_schema_for_root, Schema}; /// Lower a Layer 2 `QueryExpr` (relational operators only) to Layer 3 /// (sketch algebra with `AggIntent`). @@ -21,24 +22,43 @@ use crate::intent_algebra::legacy_expr::*; /// This pass walks the tree and converts `Aggregate { AggFunc }` nodes /// to `SketchAgg { AggIntent }` where the aggregation can benefit from /// sketch-based execution. +/// +/// Step β: derives the root-level [`Schema`] from the outermost `Source` +/// leaf (via [`infer_schema_for_root`]) and threads it through the +/// recursive descent. The lowering pass is structurally schema-agnostic +/// today (it preserves `ColumnRef::Named(_)` verbatim); Step γ will +/// migrate the per-variant column-list rewrites onto positional +/// `ColumnId` once the canonical aggregator surface lands consumer-side. pub fn lower_to_sketch_algebra(expr: QueryExpr) -> QueryExpr { + let schema = infer_schema_for_root(&expr); + lower_to_sketch_algebra_with_schema(expr, &schema) +} + +/// Variant of [`lower_to_sketch_algebra`] that takes an explicit +/// inherited [`Schema`]. The public entry point derives it from the +/// outermost source; callers that already hold a Schema (the optimiser +/// pipeline, for instance) can pass it directly. +pub fn lower_to_sketch_algebra_with_schema( + expr: QueryExpr, + parent_schema: &Schema, +) -> QueryExpr { match expr { QueryExpr::Aggregate { keys, aggs, having, input } => { - let input = lower_to_sketch_algebra(*input); + let input = lower_to_sketch_algebra_with_schema(*input, parent_schema); lower_aggregate(keys, aggs, having, Box::new(input)) } // ── Single-input nodes: recurse ────────────────────────────────── QueryExpr::Filter { pred, input } => QueryExpr::Filter { pred, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, QueryExpr::Project { cols, input } => QueryExpr::Project { cols, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, QueryExpr::Window { duration, slide, input } => { - let lowered_input = lower_to_sketch_algebra(*input); + let lowered_input = lower_to_sketch_algebra_with_schema(*input, parent_schema); // Fuse Window + SketchAgg → WindowedAgg (the window defines sketch lifecycle). if let QueryExpr::SketchAgg { op, col, input: sketch_input } = lowered_input { let window = WindowSpec { @@ -56,74 +76,77 @@ pub fn lower_to_sketch_algebra(expr: QueryExpr) -> QueryExpr { QueryExpr::SketchAgg { op, col, input } => QueryExpr::SketchAgg { op, col, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, QueryExpr::WindowedAgg { agg, window, col, input } => QueryExpr::WindowedAgg { agg, window, col, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, QueryExpr::Partition { keys, input } => QueryExpr::Partition { keys, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, QueryExpr::Distinct { cols, input } => QueryExpr::Distinct { cols, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, QueryExpr::TopK { k, by, input } => QueryExpr::TopK { k, by, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, QueryExpr::Sort { keys, input } => QueryExpr::Sort { keys, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, QueryExpr::Limit { n, offset, input } => QueryExpr::Limit { n, offset, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, QueryExpr::HistogramQuantile { phi, input } => QueryExpr::HistogramQuantile { phi, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, QueryExpr::PromQLSubquery { range, resolution, input } => QueryExpr::PromQLSubquery { range, resolution, - input: Box::new(lower_to_sketch_algebra(*input)), + input: Box::new(lower_to_sketch_algebra_with_schema(*input, parent_schema)), }, - // ── Two-input nodes: recurse into both ───���────────────────────── + // ── Two-input nodes: recurse into both ────────────────────────── QueryExpr::BinaryOp { op, lhs, rhs, vector_match } => QueryExpr::BinaryOp { op, - lhs: Box::new(lower_to_sketch_algebra(*lhs)), - rhs: Box::new(lower_to_sketch_algebra(*rhs)), + lhs: Box::new(lower_to_sketch_algebra_with_schema(*lhs, parent_schema)), + rhs: Box::new(lower_to_sketch_algebra_with_schema(*rhs, parent_schema)), vector_match, }, QueryExpr::Join { kind, pred, left, right } => QueryExpr::Join { kind, pred, - left: Box::new(lower_to_sketch_algebra(*left)), - right: Box::new(lower_to_sketch_algebra(*right)), + left: Box::new(lower_to_sketch_algebra_with_schema(*left, parent_schema)), + right: Box::new(lower_to_sketch_algebra_with_schema(*right, parent_schema)), }, QueryExpr::SetOp { kind, all, left, right } => QueryExpr::SetOp { kind, all, - left: Box::new(lower_to_sketch_algebra(*left)), - right: Box::new(lower_to_sketch_algebra(*right)), + left: Box::new(lower_to_sketch_algebra_with_schema(*left, parent_schema)), + right: Box::new(lower_to_sketch_algebra_with_schema(*right, parent_schema)), }, // ── Multi-input / container nodes ─────��────────────────────────── QueryExpr::Merge { inputs } => QueryExpr::Merge { - inputs: inputs.into_iter().map(lower_to_sketch_algebra).collect(), + inputs: inputs + .into_iter() + .map(|i| lower_to_sketch_algebra_with_schema(i, parent_schema)) + .collect(), }, QueryExpr::LetBinding { name, expr, body } => QueryExpr::LetBinding { name, - expr: Box::new(lower_to_sketch_algebra(*expr)), - body: Box::new(lower_to_sketch_algebra(*body)), + expr: Box::new(lower_to_sketch_algebra_with_schema(*expr, parent_schema)), + body: Box::new(lower_to_sketch_algebra_with_schema(*body, parent_schema)), }, // ── Leaf nodes: pass through ────────���──────────────────────────── diff --git a/controller/src/intent_algebra/mod.rs b/controller/src/intent_algebra/mod.rs index 2f3ce464..a6aa24b2 100644 --- a/controller/src/intent_algebra/mod.rs +++ b/controller/src/intent_algebra/mod.rs @@ -87,6 +87,7 @@ pub mod schema; // `QueryExpr` / `AggIntent` types used by the planner, allocator, // physical planner, query_parser, and language_logical_plan modules // today. +pub mod column_resolution; pub mod legacy_expr; pub mod legacy_lower; @@ -102,3 +103,15 @@ pub use query_expr::{ WindowKind, }; pub use schema::{cse_reuse_is_legal, Column, ColumnId, CseError, DataType, Schema}; + +// Step β plumbing: schema-driven column resolution helpers used by the +// legacy planning stack (`optimizer/engine.rs`, `physical/{allocator, +// planner, stage_split}.rs`, `query_parser/*`, `intent_algebra:: +// legacy_lower`) to carry an inherited `Schema` alongside every legacy +// `QueryExpr` traversal. Consumers call `resolve_column_ref` at the point +// where they need a positional `ColumnId` — Step γ migrates variants +// one at a time onto the canonical positional form. +pub use column_resolution::{ + infer_schema_for_root, infer_source_schema, resolve_column_ref, resolve_column_refs, + ResolveError, +}; diff --git a/controller/src/optimizer/engine.rs b/controller/src/optimizer/engine.rs index 1140cb74..1e6fb348 100644 --- a/controller/src/optimizer/engine.rs +++ b/controller/src/optimizer/engine.rs @@ -32,6 +32,7 @@ use std::collections::HashMap; use crate::intent_algebra::legacy_expr::{QueryExpr, ScalarExpr, SetOpKind, SortKey}; use crate::intent_algebra::legacy_expr::{AggIntent, PartitionKeys, SourceSpec}; +use crate::intent_algebra::{infer_schema_for_root, Schema}; use crate::sketch_algebra::capability::{ default_capability_table, load_capability_overrides, SketchCapability, }; @@ -278,6 +279,18 @@ impl CostModel for DefaultCostModel { // ── Rewrite rule trait ──────────────────────────────────────────────────────── /// A single algebraic rewrite rule. +/// +/// # Schema parameter (Step β of the legacy_expr migration) +/// +/// `parent_schema` is the [`Schema`] in scope at this node — the output +/// schema of `expr`'s parent in the tree (or the root-level schema when +/// `expr` is the root). Step β threads this through every walker so Step +/// γ rules can call +/// [`crate::intent_algebra::resolve_column_ref`] to convert +/// `ColumnRef::Named("host")` → `ColumnId(2)` without changing the +/// legacy `QueryExpr` shape. The default rules in this file are +/// structural rewrites that don't read the schema today — the parameter +/// is plumbing for Step γ. pub trait RewriteRule: Send + Sync { /// Human-readable name for logging. fn name(&self) -> &'static str; @@ -285,7 +298,16 @@ pub trait RewriteRule: Send + Sync { /// 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; + /// + /// `parent_schema` is the schema in scope at `expr` — see the trait + /// docs. Rules that need positional column resolution should consult + /// it; rules that are purely structural can ignore it. + fn try_rewrite( + &self, + expr: QueryExpr, + model: &dyn CostModel, + parent_schema: &Schema, + ) -> Option; } // ── R1: PredicatePushDown ───────────────────────────────────────────────────── @@ -302,7 +324,7 @@ pub struct PredicatePushDown; impl RewriteRule for PredicatePushDown { fn name(&self) -> &'static str { "PredicatePushDown" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { match expr { QueryExpr::Filter { pred, input } => { match *input { @@ -359,7 +381,7 @@ pub struct MergeLifting; impl RewriteRule for MergeLifting { fn name(&self) -> &'static str { "MergeLifting" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { match expr { QueryExpr::SketchAgg { ref op, ref col, ref input } if crate::intent_algebra::legacy_expr::agg_is_mergeable(op) => @@ -394,7 +416,7 @@ pub struct HLLDedupElim; impl RewriteRule for HLLDedupElim { fn name(&self) -> &'static str { "HLLDedupElim" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { match expr { QueryExpr::SketchAgg { op: AggIntent::Cardinality { accuracy }, col, input } => { if let QueryExpr::Distinct { input: inner, .. } = *input { @@ -424,7 +446,7 @@ pub struct FilterWindowSwap; impl RewriteRule for FilterWindowSwap { fn name(&self) -> &'static str { "FilterWindowSwap" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { // Handled by PredicatePushDown — mark as no-op here to avoid double-fire. match expr { QueryExpr::Filter { pred, input } => { @@ -453,7 +475,7 @@ pub struct TopKFusion; impl RewriteRule for TopKFusion { fn name(&self) -> &'static str { "TopKFusion" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { match expr { QueryExpr::Limit { n, offset: 0, input } => { if let QueryExpr::Sort { keys, input: inner } = *input { @@ -487,7 +509,7 @@ pub struct HistogramQuantileFusion; impl RewriteRule for HistogramQuantileFusion { fn name(&self) -> &'static str { "HistogramQuantileFusion" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { // Canonical Quantile is single-φ post Step α — multi-φ fan-out // happens at construction time, so the "merge phi into the // existing quantile list" branch is now a "if phis match, keep @@ -555,7 +577,7 @@ pub struct SubqueryDecorrelation; impl RewriteRule for SubqueryDecorrelation { fn name(&self) -> &'static str { "SubqueryDecorrelation" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { match expr { QueryExpr::Filter { pred, input } => { if let Some((name, sq_expr, new_pred)) = extract_scalar_subquery(pred) { @@ -621,7 +643,7 @@ pub struct CommonSubexprElim; impl RewriteRule for CommonSubexprElim { fn name(&self) -> &'static str { "CommonSubexprElim" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { match expr { QueryExpr::Merge { ref inputs } => { // Count occurrences of each source name. @@ -682,7 +704,7 @@ pub struct HydraConversion; impl RewriteRule for HydraConversion { fn name(&self) -> &'static str { "HydraConversion" } - fn try_rewrite(&self, _expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, _expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { // Step α: disabled — see struct docs. Step γ TODO: re-emit as a // canonical `Aggregate { by, aggs: [inner] }` wrapper around the // unwrapped `SketchAgg.op`. @@ -734,7 +756,7 @@ pub struct WindowMerge; impl RewriteRule for WindowMerge { fn name(&self) -> &'static str { "WindowMerge" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { match expr { QueryExpr::Window { duration, slide, input } => { if let QueryExpr::Window { duration: inner_d, slide: inner_s, input: inner_e } = @@ -766,7 +788,7 @@ pub struct PartitionElim; impl RewriteRule for PartitionElim { fn name(&self) -> &'static str { "PartitionElim" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { match expr { QueryExpr::Partition { keys, input } if keys.is_empty() => Some(*input), _ => None, @@ -782,7 +804,7 @@ pub struct SetOpFusion; impl RewriteRule for SetOpFusion { fn name(&self) -> &'static str { "SetOpFusion" } - fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel) -> Option { + fn try_rewrite(&self, expr: QueryExpr, _model: &dyn CostModel, _parent_schema: &Schema) -> Option { match expr { QueryExpr::SetOp { kind: SetOpKind::Union, @@ -861,10 +883,16 @@ impl QueryOptimizer { /// Optimize `expr` until fixed point or `max_iters` iterations. /// /// Returns the rewritten tree and the number of iterations actually run. + /// + /// Step β: derives the root-level [`Schema`] from `expr`'s outermost + /// `Source` leaf (via [`infer_schema_for_root`]) and threads it + /// through the recursive walk so every [`RewriteRule::try_rewrite`] + /// call sees the schema in scope at its node. pub fn optimize(&self, expr: QueryExpr) -> (QueryExpr, usize) { + let schema = infer_schema_for_root(&expr); let mut current = expr; for iter in 0..self.max_iters { - let (next, changed) = self.apply_all(current); + let (next, changed) = self.apply_all(current, &schema); current = next; if !changed { return (current, iter + 1); @@ -875,18 +903,32 @@ impl QueryOptimizer { /// 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) { + /// + /// `parent_schema` is the schema in scope at `expr` per Step β — + /// threaded through to [`RewriteRule::try_rewrite`] and the child + /// recursion. Legacy variants don't transform schema today; Step γ + /// will refine the per-variant schema propagation (e.g. `Aggregate` + /// strips the time axis, `Project` reshapes columns). + fn apply_all(&self, expr: QueryExpr, parent_schema: &Schema) -> (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); + let (expr_with_new_children, child_changed) = + self.recurse_children(expr, parent_schema); + let (final_expr, this_changed) = + self.apply_rules_at(expr_with_new_children, parent_schema); (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) { + fn apply_rules_at( + &self, + mut expr: QueryExpr, + parent_schema: &Schema, + ) -> (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()) { + if let Some(new_expr) = + rule.try_rewrite(expr.clone(), self.cost_model.as_ref(), parent_schema) + { expr = new_expr; changed = true; // After firing, restart from the first rule (fixed-point per node). @@ -897,10 +939,20 @@ impl QueryOptimizer { } /// Recurse into children, rebuilding the node with rewritten children. - fn recurse_children(&self, expr: QueryExpr) -> (QueryExpr, bool) { + /// + /// The child walk inherits the same `parent_schema` — Step β only + /// threads the schema-in-scope without per-variant transformations. + /// Step γ will replace this with proper canonical + /// [`QueryExpr::output_schema_in`]-style propagation as each variant + /// migrates. + fn recurse_children( + &self, + expr: QueryExpr, + parent_schema: &Schema, + ) -> (QueryExpr, bool) { macro_rules! recurse { ($child:expr) => {{ - let (e, c) = self.apply_all(*$child); + let (e, c) = self.apply_all(*$child, parent_schema); (Box::new(e), c) }}; } @@ -962,7 +1014,7 @@ impl QueryOptimizer { QueryExpr::Merge { inputs } => { let (new_inputs, changed): (Vec<_>, Vec<_>) = inputs .into_iter() - .map(|inp| self.apply_all(inp)) + .map(|inp| self.apply_all(inp, parent_schema)) .unzip(); (QueryExpr::Merge { inputs: new_inputs }, changed.into_iter().any(|c| c)) } @@ -1223,7 +1275,8 @@ mod tests { }; let rule = HistogramQuantileFusion; let model = DefaultCostModel { raw_bytes_per_sec: 100_000.0, deployment: None }; - let result = rule.try_rewrite(expr, &model).expect("R6 should fire"); + let schema = infer_schema_for_root(&expr); + let result = rule.try_rewrite(expr, &model, &schema).expect("R6 should fire"); match &result { QueryExpr::HistogramQuantile { input, .. } => { match input.as_ref() { diff --git a/controller/src/physical/allocator.rs b/controller/src/physical/allocator.rs index b73e5d7d..92d5f186 100644 --- a/controller/src/physical/allocator.rs +++ b/controller/src/physical/allocator.rs @@ -32,6 +32,7 @@ use super::plan::{ CostEstimate, ExecutionMode, NodeAnnotation, PipelineStage, PlanNode, }; use crate::intent_algebra::legacy_expr::{agg_is_exact, AggIntent}; +use crate::intent_algebra::{infer_schema_for_root, Schema}; use crate::types::{SketchType, StageResourceBudgets}; // ── Resource budget tracker ─────────────────────────────────────────────────── @@ -94,14 +95,27 @@ impl SketchAllocator { } /// Allocate stages for the entire expression tree. + /// + /// Step β: derives the root-level [`Schema`] from the outermost + /// `Source` leaf (via [`infer_schema_for_root`]) and threads it + /// through every recursive [`Self::alloc_node`] call. The allocator's + /// stage-assignment logic is purely structural today, but the schema + /// parameter is in place for Step γ when sketch-placement heuristics + /// start consulting column types. pub fn allocate(&self, expr: QueryExpr) -> PlanNode { + let schema = infer_schema_for_root(&expr); let mut budget = BudgetState::from_budgets(&self.budgets); - self.alloc_node(expr, &mut budget) + self.alloc_node(expr, &mut budget, &schema) } // ── Recursive allocation ────────────────────────────────────────────────── - fn alloc_node(&self, expr: QueryExpr, budget: &mut BudgetState) -> PlanNode { + fn alloc_node( + &self, + expr: QueryExpr, + budget: &mut BudgetState, + parent_schema: &Schema, + ) -> PlanNode { match expr { // ── Leaves ─────────────────────────────────────────────────────── QueryExpr::Source(_) | QueryExpr::Ref(_) => PlanNode::leaf( @@ -112,7 +126,7 @@ impl SketchAllocator { // ── Structural / filter nodes — always Agent ────────────────── QueryExpr::Filter { pred, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); let stage = PipelineStage::Agent; PlanNode { expr: QueryExpr::Filter { pred, input: Box::new(child.expr.clone()) }, @@ -131,7 +145,7 @@ impl SketchAllocator { } QueryExpr::Window { duration, slide, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); PlanNode { expr: QueryExpr::Window { duration, slide, @@ -149,7 +163,7 @@ impl SketchAllocator { } QueryExpr::Partition { keys, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); PlanNode { expr: QueryExpr::Partition { keys, @@ -167,7 +181,7 @@ impl SketchAllocator { } QueryExpr::Distinct { cols, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); PlanNode { expr: QueryExpr::Distinct { cols, @@ -186,19 +200,19 @@ impl SketchAllocator { // ── 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) + let child = self.alloc_node(*input, budget, parent_schema); + self.alloc_sketch_agg(op, col, child, budget, parent_schema) } // ── WindowedAgg — treat as SketchAgg (window is informational) ── QueryExpr::WindowedAgg { agg, window: _, col, input } => { - let child = self.alloc_node(*input, budget); - self.alloc_sketch_agg(agg, col, child, budget) + let child = self.alloc_node(*input, budget, parent_schema); + self.alloc_sketch_agg(agg, col, child, budget, parent_schema) } // ── TopK — Precompute engine ────────────────────────────────── QueryExpr::TopK { k, by, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); PlanNode { expr: QueryExpr::TopK { k, by, @@ -224,7 +238,7 @@ impl SketchAllocator { QueryExpr::Merge { inputs } => { let children: Vec = inputs .into_iter() - .map(|inp| self.alloc_node(inp, budget)) + .map(|inp| self.alloc_node(inp, budget, parent_schema)) .collect(); let mem: f64 = children.iter().map(|c| c.cost.memory_bytes).sum(); PlanNode { @@ -248,7 +262,7 @@ impl SketchAllocator { // ── Exact / relational — Db ─────────────────────────────────── QueryExpr::Aggregate { keys, aggs, having, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); PlanNode { expr: QueryExpr::Aggregate { keys, aggs, having, @@ -269,7 +283,7 @@ impl SketchAllocator { } QueryExpr::Project { cols, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); PlanNode { expr: QueryExpr::Project { cols, @@ -287,7 +301,7 @@ impl SketchAllocator { } QueryExpr::Sort { keys, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); PlanNode { expr: QueryExpr::Sort { keys, @@ -305,7 +319,7 @@ impl SketchAllocator { } QueryExpr::Limit { n, offset, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); PlanNode { expr: QueryExpr::Limit { n, offset, @@ -323,8 +337,8 @@ impl SketchAllocator { } QueryExpr::Join { kind, pred, left, right } => { - let left_node = self.alloc_node(*left, budget); - let right_node = self.alloc_node(*right, budget); + let left_node = self.alloc_node(*left, budget, parent_schema); + let right_node = self.alloc_node(*right, budget, parent_schema); PlanNode { expr: QueryExpr::Join { kind, pred, @@ -346,8 +360,8 @@ impl SketchAllocator { } QueryExpr::SetOp { kind, all, left, right } => { - let left_node = self.alloc_node(*left, budget); - let right_node = self.alloc_node(*right, budget); + let left_node = self.alloc_node(*left, budget, parent_schema); + let right_node = self.alloc_node(*right, budget, parent_schema); PlanNode { expr: QueryExpr::SetOp { kind, all, @@ -367,7 +381,7 @@ impl SketchAllocator { // ── PromQL-specific ─────────────────────────────────────────── QueryExpr::HistogramQuantile { phi, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); // If the child is a sketch, elevate to Precompute; // otherwise fall through to Db. let stage = if child.mode == ExecutionMode::Sketch { @@ -397,7 +411,7 @@ impl SketchAllocator { } QueryExpr::PromQLSubquery { range, resolution, input } => { - let child = self.alloc_node(*input, budget); + let child = self.alloc_node(*input, budget, parent_schema); let stage = if child.mode == ExecutionMode::Sketch { PipelineStage::Precompute } else { @@ -421,8 +435,8 @@ impl SketchAllocator { } QueryExpr::BinaryOp { op, lhs, rhs, vector_match } => { - let left_node = self.alloc_node(*lhs, budget); - let right_node = self.alloc_node(*rhs, budget); + let left_node = self.alloc_node(*lhs, budget, parent_schema); + let right_node = self.alloc_node(*rhs, budget, parent_schema); let has_sketch = left_node.mode == ExecutionMode::Sketch || right_node.mode == ExecutionMode::Sketch; let stage = if has_sketch { @@ -450,8 +464,8 @@ impl SketchAllocator { // ── Scoping constructs — propagate body's stage ─────────────── QueryExpr::LetBinding { name, expr, body } => { - let expr_node = self.alloc_node(*expr, budget); - let body_node = self.alloc_node(*body, budget); + let expr_node = self.alloc_node(*expr, budget, parent_schema); + let body_node = self.alloc_node(*body, budget, parent_schema); let stage = body_node.stage.clone(); let mode = body_node.mode.clone(); PlanNode { @@ -481,6 +495,10 @@ impl SketchAllocator { col: crate::intent_algebra::legacy_expr::ColumnRef, child: PlanNode, budget: &mut BudgetState, + // Step β: schema in scope at this SketchAgg node. Unused today — + // Step γ wires sketch-placement rules that consult column types + // (e.g. KLL vs DDSketch for `value: Float64`). + _parent_schema: &Schema, ) -> PlanNode { // Exact non-mergeable (Avg) → always Db. if matches!(&op, AggIntent::Avg) { diff --git a/controller/src/physical/planner.rs b/controller/src/physical/planner.rs index d513e9dc..576daa0f 100644 --- a/controller/src/physical/planner.rs +++ b/controller/src/physical/planner.rs @@ -14,6 +14,7 @@ use std::time::Duration; use crate::physical::sketch_catalog; use crate::intent_algebra::legacy_expr::{AggIntent, WindowKind, WindowSpec}; +use crate::intent_algebra::{infer_schema_for_root, Schema}; use crate::types::{SketchParams, SketchType}; // ── PhysicalAggOp (resolved sketch intent) ────────────────────────────────── @@ -230,11 +231,23 @@ pub struct PhysicalPlannerConfig { /// Walks the logical tree bottom-up, assigning each node to a pipeline stage /// (`Placement`), resolving sketch intents to concrete implementations, and /// inserting `Exchange` nodes at stage boundaries. +/// +/// Step β: derives the root-level [`Schema`] from the outermost `Source` +/// leaf and threads it through every recursive [`plan_node`] call. The +/// physical planner is structurally schema-agnostic today (resolution +/// happens against legacy `ColumnRef::Named(_)` strings); the parameter +/// is plumbing for Step γ when sketch-binding decisions start consulting +/// column types. pub fn plan(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { - plan_node(expr, config) + let schema = infer_schema_for_root(expr); + plan_node(expr, config, &schema) } -fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { +fn plan_node( + expr: &QueryExpr, + config: &PhysicalPlannerConfig, + parent_schema: &Schema, +) -> PhysicalNode { match expr { // ── Leaf: scan at Agent ───────────────────────────────────── QueryExpr::Source(s) => PhysicalNode { @@ -249,7 +262,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { // ── Filter: same placement as child ───────────────────────── QueryExpr::Filter { pred, input } => { - let child = plan_node(input, config); + let child = plan_node(input, config, parent_schema); PhysicalNode { placement: child.placement.clone(), op: PhysicalOp::Filter { pred: format!("{pred:?}") }, @@ -260,7 +273,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { // ── SketchAgg: resolve intent → physical, place at Agent or defer ── QueryExpr::SketchAgg { op, col, input } => { - let child = plan_node(input, config); + let child = plan_node(input, config, parent_schema); let resolved = resolve(op); let placement = decide_sketch_placement(&resolved, config); @@ -287,7 +300,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { // ── WindowedAgg: resolve + place with window ──────────────── QueryExpr::WindowedAgg { agg, window, col, input } => { - let child = plan_node(input, config); + let child = plan_node(input, config, parent_schema); let resolved = resolve(agg); let placement = decide_sketch_placement(&resolved, config); let phys_window = resolve_window(window, &placement); @@ -314,7 +327,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { // ── Partition / Merge: Backend stage ──────────────────────── QueryExpr::Partition { keys, input } => { - let child = plan_node(input, config); + let child = plan_node(input, config, parent_schema); let mut node = PhysicalNode { op: PhysicalOp::HashAggregate { keys: keys.keys().to_vec() }, placement: Placement::BackendCollector, @@ -327,7 +340,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { QueryExpr::Merge { inputs } => { let children: Vec = inputs.iter() - .map(|i| plan_node(i, config)) + .map(|i| plan_node(i, config, parent_schema)) .collect(); let sketch_type = children.first() .and_then(|c| match &c.op { @@ -344,7 +357,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { } QueryExpr::Distinct { cols, input } => { - let child = plan_node(input, config); + let child = plan_node(input, config, parent_schema); let pred = format!("distinct({})", display_distinct_cols(cols)); let mut node = PhysicalNode { op: PhysicalOp::Filter { pred }, @@ -358,7 +371,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { // ── TopK / HistogramQuantile / BinaryOp: QueryEngine stage ── QueryExpr::TopK { k, input, .. } => { - let child = plan_node(input, config); + let child = plan_node(input, config, parent_schema); let mut node = PhysicalNode { op: PhysicalOp::TopK { k: *k }, placement: Placement::QueryEngine, @@ -370,7 +383,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { } QueryExpr::HistogramQuantile { phi, input } => { - let child = plan_node(input, config); + let child = plan_node(input, config, parent_schema); let mut node = PhysicalNode { op: PhysicalOp::SketchEval { sketch_type: SketchType::DDSketch, @@ -385,8 +398,8 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { } QueryExpr::BinaryOp { op, lhs, rhs, .. } => { - let left = plan_node(lhs, config); - let right = plan_node(rhs, config); + let left = plan_node(lhs, config, parent_schema); + let right = plan_node(rhs, config, parent_schema); PhysicalNode { op: PhysicalOp::Passthrough, placement: Placement::QueryEngine, @@ -396,7 +409,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { } QueryExpr::PromQLSubquery { input, .. } => { - let child = plan_node(input, config); + let child = plan_node(input, config, parent_schema); let mut node = PhysicalNode { op: PhysicalOp::Passthrough, placement: Placement::QueryEngine, @@ -409,7 +422,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { // ── Aggregate (non-sketch, exact): Database stage ─────────── QueryExpr::Aggregate { keys, input, .. } => { - let child = plan_node(input, config); + let child = plan_node(input, config, parent_schema); let mut node = PhysicalNode { op: PhysicalOp::DbQuery { sql: format!("GROUP BY {:?}", keys) }, placement: Placement::Database, @@ -425,7 +438,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { | QueryExpr::Limit { input, .. } | QueryExpr::Project { input, .. } | QueryExpr::Window { input, .. } => { - let child = plan_node(input, config); + let child = plan_node(input, config, parent_schema); PhysicalNode { op: PhysicalOp::Passthrough, placement: child.placement.clone(), @@ -437,8 +450,8 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { // ── Join: both children, QueryEngine placement ────────────── QueryExpr::Join { left, right, .. } | QueryExpr::SetOp { left, right, .. } => { - let l = plan_node(left, config); - let r = plan_node(right, config); + let l = plan_node(left, config, parent_schema); + let r = plan_node(right, config, parent_schema); PhysicalNode { op: PhysicalOp::Passthrough, placement: Placement::QueryEngine, @@ -448,7 +461,7 @@ fn plan_node(expr: &QueryExpr, config: &PhysicalPlannerConfig) -> PhysicalNode { } // ── LetBinding ────────────────────────────────────────────── - QueryExpr::LetBinding { body, .. } => plan_node(body, config), + QueryExpr::LetBinding { body, .. } => plan_node(body, config, parent_schema), QueryExpr::Ref(_) => PhysicalNode { op: PhysicalOp::Passthrough, placement: Placement::QueryEngine, diff --git a/controller/src/physical/stage_split.rs b/controller/src/physical/stage_split.rs index f33ee6c1..ac3e9f07 100644 --- a/controller/src/physical/stage_split.rs +++ b/controller/src/physical/stage_split.rs @@ -40,6 +40,7 @@ use std::time::Duration; use crate::intent_algebra::legacy_expr::{AggFunc, agg_is_exact, agg_is_mergeable, agg_quantiles, BinaryOpKind, LiteralValue, QueryExpr, ScalarExpr}; use crate::pipeline::format_duration; use crate::intent_algebra::legacy_expr::AggIntent; +use crate::intent_algebra::{infer_schema_for_root, Schema}; use crate::physical::sketch_catalog; use crate::types::{ AgentSubPlan, BackendSubPlan, @@ -56,8 +57,9 @@ use crate::types::{ /// [`crate::types::CollectionPlan::staged_plan`] by the caller /// (`handle_plan` in `main.rs`). pub fn split_expr_by_stage(expr: &QueryExpr, budgets: &StageResourceBudgets) -> StagedPlan { + let schema = infer_schema_for_root(expr); let mut plan = StagedPlan::default(); - walk(expr, &mut plan, budgets); + walk(expr, &mut plan, budgets, &schema); // Build the precompute query_expr from the full tree when the precompute // stage is active (TopK, HistogramQuantile, PromQLSubquery, or deferred ops). @@ -82,14 +84,26 @@ pub fn split_expr_by_stage(expr: &QueryExpr, budgets: &StageResourceBudgets) -> /// Uses recursive descent, so it handles `histogram_quantile`, /// `PromQLSubquery`, and vector `BinaryOp` nodes that the old flat-template /// could not represent. +/// +/// Step β: the root-level [`Schema`] is derived from the outermost +/// `Source` leaf and threaded through the recursive descent. The +/// serialiser is schema-agnostic today (it emits `ColumnRef::Named(s)` +/// verbatim into the PromQL output); Step γ wires schema-aware column +/// resolution at the points that need it. pub fn expr_to_promql(expr: &QueryExpr) -> String { + let schema = infer_schema_for_root(expr); let mut ctx = PromQLCtx::default(); - promql_from_qe(expr, &mut ctx) + promql_from_qe(expr, &mut ctx, &schema) } // ── Tree walker ─────────────────────────────────────────────────────────────── -fn walk(expr: &QueryExpr, plan: &mut StagedPlan, budgets: &StageResourceBudgets) { +fn walk( + expr: &QueryExpr, + plan: &mut StagedPlan, + budgets: &StageResourceBudgets, + parent_schema: &Schema, +) { match expr { // Source — always Agent; populate metric name. QueryExpr::Source(_) => {} @@ -97,19 +111,19 @@ fn walk(expr: &QueryExpr, plan: &mut StagedPlan, budgets: &StageResourceBudgets) // Filter — push label predicates to Agent. QueryExpr::Filter { pred, input } => { collect_label_filters_into(pred, &mut plan.agent.label_filters); - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } // Window — time window lives at Agent. QueryExpr::Window { duration, input, .. } => { plan.agent.window_secs = Some(duration.as_secs()); - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } // SketchAgg — the key sketch assignment decision. QueryExpr::SketchAgg { op, input, .. } => { assign_sketch_agg(op, plan, budgets); - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } // WindowedAgg — bundles window + sketch agg intent. @@ -118,7 +132,7 @@ fn walk(expr: &QueryExpr, plan: &mut StagedPlan, budgets: &StageResourceBudgets) plan.agent.window_secs = Some(size.as_secs()); } assign_sketch_agg(agg, plan, budgets); - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } // Partition — GROUP BY / `by (dims)` always assigned to Backend. @@ -128,7 +142,7 @@ fn walk(expr: &QueryExpr, plan: &mut StagedPlan, budgets: &StageResourceBudgets) plan.backend.group_by.push(k.clone()); } } - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } // Aggregate — SQL GROUP BY + agg functions. @@ -144,20 +158,44 @@ fn walk(expr: &QueryExpr, plan: &mut StagedPlan, budgets: &StageResourceBudgets) for agg in aggs { assign_agg_func(&agg.func, plan, budgets); } - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } // Distinct — absorbed at Backend (HLL dedup elimination is upstream). - QueryExpr::Distinct { input, .. } => { + // + // Step β proof-of-use: opportunistically resolve named-column + // dedup keys against the inherited schema and record their + // positional ids on the staged plan's deferral log. The legacy + // `BackendSubPlan` doesn't carry `ColumnId`s yet (Step γ adds + // them), so failed lookups are logged as TODOs rather than + // surfaced as errors. Wildcard / SampleValue cols are skipped — + // they have no positional id by construction. + QueryExpr::Distinct { cols, input } => { plan.backend.has_dedup = true; - walk(input, plan, budgets); + for c in cols { + if let crate::intent_algebra::legacy_expr::ColumnRef::Named(_) = c { + match crate::intent_algebra::resolve_column_ref(c, parent_schema) { + Ok(_id) => { + // Step γ TODO: thread `_id` through to a + // `BackendSubPlan::distinct_cols: Vec` + // field once the staged plan grows one. + } + Err(e) => { + plan.deferral_log.push(format!( + "Distinct col resolution deferred to Step γ: {e}" + )); + } + } + } + } + walk(input, plan, budgets, parent_schema); } // Merge — Backend merges N agent sketches. QueryExpr::Merge { inputs } => { plan.backend.has_merge = true; for i in inputs { - walk(i, plan, budgets); + walk(i, plan, budgets, parent_schema); } } @@ -165,58 +203,58 @@ fn walk(expr: &QueryExpr, plan: &mut StagedPlan, budgets: &StageResourceBudgets) QueryExpr::TopK { k, input, .. } => { plan.precompute.topk = Some(*k); plan.precompute.active = true; - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } // Sort + Limit — maps to topk semantics at Precompute. QueryExpr::Limit { n, input, .. } => { plan.precompute.topk = Some(*n); plan.precompute.active = true; - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } QueryExpr::Sort { input, .. } => { - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } // HistogramQuantile — precompute engine applies it over ingested histograms. QueryExpr::HistogramQuantile { input, .. } => { plan.precompute.active = true; - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } // PromQLSubquery — precompute engine evaluates the sub-query. QueryExpr::PromQLSubquery { input, .. } => { plan.precompute.active = true; - walk(input, plan, budgets); + walk(input, plan, budgets, parent_schema); } // BinaryOp between two instant vectors — precompute evaluates. QueryExpr::BinaryOp { lhs, rhs, .. } => { plan.precompute.active = true; - walk(lhs, plan, budgets); - walk(rhs, plan, budgets); + walk(lhs, plan, budgets, parent_schema); + walk(rhs, plan, budgets, parent_schema); } // Join — Backend. QueryExpr::Join { left, right, .. } => { plan.backend.has_merge = true; - walk(left, plan, budgets); - walk(right, plan, budgets); + walk(left, plan, budgets, parent_schema); + walk(right, plan, budgets, parent_schema); } // SetOp — treat as Backend merge. QueryExpr::SetOp { left, right, .. } => { plan.backend.has_merge = true; - walk(left, plan, budgets); - walk(right, plan, budgets); + walk(left, plan, budgets, parent_schema); + walk(right, plan, budgets, parent_schema); } // Transparent / passthrough nodes — recurse into child. - QueryExpr::Project { input, .. } => walk(input, plan, budgets), + QueryExpr::Project { input, .. } => walk(input, plan, budgets, parent_schema), QueryExpr::LetBinding { expr, body, .. } => { - walk(expr, plan, budgets); - walk(body, plan, budgets); + walk(expr, plan, budgets, parent_schema); + walk(body, plan, budgets, parent_schema); } // Ref — nothing to assign (resolved externally). @@ -335,7 +373,7 @@ struct PromQLCtx { /// Returns the PromQL string fragment for this node. Inner nodes (Source, /// Filter) return their selector string; outer nodes (SketchAgg, TopK, etc.) /// wrap it. -fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { +fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx, parent_schema: &Schema) -> String { match expr { // ── Leaf ───────────────────────────────────────────────────────────── QueryExpr::Source(s) => s.name.clone(), @@ -343,7 +381,7 @@ fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { // ── Filter — append label matchers to the selector ─────────────────── QueryExpr::Filter { pred, input } => { - let inner = promql_from_qe(input, ctx); + let inner = promql_from_qe(input, ctx, parent_schema); let matchers = scalar_to_label_matchers(pred); if matchers.is_empty() { inner @@ -357,7 +395,7 @@ fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { if ctx.window.is_none() { ctx.window = Some(*duration); } - promql_from_qe(input, ctx) + promql_from_qe(input, ctx, parent_schema) } // ── Partition — store group_by keys for enclosing aggregate ────────── @@ -367,7 +405,7 @@ fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { ctx.group_by.push(k.clone()); } } - promql_from_qe(input, ctx) + promql_from_qe(input, ctx, parent_schema) } // ── WindowedAgg — bundled window + sketch agg ──────────────────────── @@ -377,7 +415,7 @@ fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { ctx.window = Some(*size); } } - let selector = promql_from_qe(input, ctx); + let selector = promql_from_qe(input, ctx, parent_schema); let window_s = window_str(ctx.window); let by = by_clause(&ctx.group_by); sketch_op_to_promql(agg, &selector, &window_s, &by) @@ -385,7 +423,7 @@ fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { // ── SketchAgg — the main aggregation node ──────────────────────────── QueryExpr::SketchAgg { op, input, .. } => { - let selector = promql_from_qe(input, ctx); + let selector = promql_from_qe(input, ctx, parent_schema); let window = window_str(ctx.window); let by = by_clause(&ctx.group_by); sketch_op_to_promql(op, &selector, &window, &by) @@ -399,7 +437,7 @@ fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { ctx.group_by.push(k.clone()); } } - let selector = promql_from_qe(input, ctx); + let selector = promql_from_qe(input, ctx, parent_schema); let window = window_str(ctx.window); let by = by_clause(&ctx.group_by); // Use the first aggregate function to drive the PromQL template. @@ -412,27 +450,27 @@ fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { // ── TopK ───────────────────────────────────────────────────────────── QueryExpr::TopK { k, input, .. } => { - let inner = promql_from_qe(input, ctx); + let inner = promql_from_qe(input, ctx, parent_schema); format!("topk({k}, {inner})") } // ── Sort + Limit — map to topk ─────────────────────────────────────── QueryExpr::Limit { n, input, .. } => { - let inner = promql_from_qe(input, ctx); + let inner = promql_from_qe(input, ctx, parent_schema); format!("topk({n}, {inner})") } - QueryExpr::Sort { input, .. } => promql_from_qe(input, ctx), + QueryExpr::Sort { input, .. } => promql_from_qe(input, ctx, parent_schema), // ── histogram_quantile(φ, rate(selector[w])) ───────────────────────── QueryExpr::HistogramQuantile { phi, input } => { - let selector = promql_from_qe(input, ctx); + let selector = promql_from_qe(input, ctx, parent_schema); let window = window_str(ctx.window); format!("histogram_quantile({phi}, rate({selector}{window}))") } // ── PromQL subquery expr[range:step] ───────────────────────────────── QueryExpr::PromQLSubquery { range, resolution, input } => { - let inner = promql_from_qe(input, ctx); + let inner = promql_from_qe(input, ctx, parent_schema); let step_str = resolution .map(|r| format!(":{}", format_duration(r))) .unwrap_or_default(); @@ -441,8 +479,8 @@ fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { // ── Vector binary op (lhs op rhs) ─────────────────────────────────── QueryExpr::BinaryOp { op, lhs, rhs, vector_match } => { - let lhs_str = promql_from_qe(lhs, ctx); - let rhs_str = promql_from_qe(rhs, &mut PromQLCtx::default()); + let lhs_str = promql_from_qe(lhs, ctx, parent_schema); + let rhs_str = promql_from_qe(rhs, &mut PromQLCtx::default(), parent_schema); let op_str = binop_to_promql(op); let match_str = vector_match .as_ref() @@ -473,19 +511,19 @@ fn promql_from_qe(expr: &QueryExpr, ctx: &mut PromQLCtx) -> String { // ── Merge — serialise first branch (all branches same shape) ───────── QueryExpr::Merge { inputs } => { inputs.first() - .map(|first| promql_from_qe(first, ctx)) + .map(|first| promql_from_qe(first, ctx, parent_schema)) .unwrap_or_default() } // ── Passthrough nodes ───────────────────────────────────────────────── QueryExpr::Distinct { input, .. } - | QueryExpr::Project { input, .. } => promql_from_qe(input, ctx), + | QueryExpr::Project { input, .. } => promql_from_qe(input, ctx, parent_schema), - QueryExpr::LetBinding { body, .. } => promql_from_qe(body, ctx), + QueryExpr::LetBinding { body, .. } => promql_from_qe(body, ctx, parent_schema), // ── Join / SetOp — serialise the outer / left branch ───────────────── - QueryExpr::Join { left, .. } => promql_from_qe(left, ctx), - QueryExpr::SetOp { left, .. } => promql_from_qe(left, ctx), + QueryExpr::Join { left, .. } => promql_from_qe(left, ctx, parent_schema), + QueryExpr::SetOp { left, .. } => promql_from_qe(left, ctx, parent_schema), } } diff --git a/controller/src/query_parser/mod.rs b/controller/src/query_parser/mod.rs index cf950742..9d87971d 100644 --- a/controller/src/query_parser/mod.rs +++ b/controller/src/query_parser/mod.rs @@ -120,9 +120,19 @@ pub fn parse_query(query: &str) -> anyhow::Result { } /// Extract a flat [`ParsedQuery`] by walking a [`QueryExpr`] tree. +/// +/// Step β: the root-level [`crate::intent_algebra::Schema`] is derived +/// from the outermost `Source` leaf (via +/// [`crate::intent_algebra::infer_schema_for_root`]) and threaded through +/// the [`QeCollector::visit`] walk. The collector is schema-agnostic +/// today (it reads metric / aggregate / label names from the legacy +/// `ColumnRef::Named(_)` shape directly); Step γ migrates the column-name +/// reads onto positional `ColumnId` lookups via +/// [`crate::intent_algebra::resolve_column_ref`]. fn qe_to_parsed_query(qe: &QueryExpr) -> ParsedQuery { + let schema = crate::intent_algebra::infer_schema_for_root(qe); let mut c = QeCollector::default(); - c.visit(qe); + c.visit(qe, &schema); c.build() } @@ -141,7 +151,11 @@ struct QeCollector { } impl QeCollector { - fn visit(&mut self, expr: &QueryExpr) { + /// Step β: `parent_schema` is the [`crate::intent_algebra::Schema`] in + /// scope at this node. Collectors that need to convert a + /// `ColumnRef::Named` to a `ColumnId` (Step γ migration) call + /// [`crate::intent_algebra::resolve_column_ref`] with it. + fn visit(&mut self, expr: &QueryExpr, parent_schema: &crate::intent_algebra::Schema) { use crate::intent_algebra::legacy_expr::{FilterOp, FilterVal, LiteralValue, ScalarExpr}; match expr { QueryExpr::Source(s) => { @@ -152,13 +166,13 @@ impl QeCollector { QueryExpr::Filter { pred, input } => { // Extract equality label filters from the predicate tree. collect_filters_from_scalar(pred, &mut self.label_filters); - self.visit(input); + self.visit(input, parent_schema); } QueryExpr::Window { duration, input, .. } => { if self.time_window.is_none() { self.time_window = Some(*duration); } - self.visit(input); + self.visit(input, parent_schema); } QueryExpr::Partition { keys, input } => { for k in keys.keys() { @@ -166,11 +180,11 @@ impl QeCollector { self.group_by_labels.push(k.clone()); } } - self.visit(input); + self.visit(input, parent_schema); } QueryExpr::SketchAgg { op, input, .. } => { self.collect_op(op); - self.visit(input); + self.visit(input, parent_schema); } QueryExpr::WindowedAgg { agg, window, input, .. } => { if self.time_window.is_none() { @@ -179,18 +193,18 @@ impl QeCollector { } } self.collect_op(agg); - self.visit(input); + self.visit(input, parent_schema); } QueryExpr::TopK { k, input, .. } => { self.topk = Some(*k); let prev = self.inside_topk; self.inside_topk = true; - self.visit(input); + self.visit(input, parent_schema); self.inside_topk = prev; } - QueryExpr::Distinct { input, .. } => self.visit(input), + QueryExpr::Distinct { input, .. } => self.visit(input, parent_schema), QueryExpr::Merge { inputs } => { - for i in inputs { self.visit(i); } + for i in inputs { self.visit(i, parent_schema); } } QueryExpr::Aggregate { keys, aggs, input, .. } => { for k in keys { @@ -202,22 +216,22 @@ impl QeCollector { for agg in aggs { self.collect_agg_func_with_group(&agg.func, has_group_by); } - self.visit(input); + self.visit(input, parent_schema); } QueryExpr::Project { input, .. } | QueryExpr::Sort { input, .. } | QueryExpr::Limit { input, .. } | QueryExpr::HistogramQuantile { input, .. } - | QueryExpr::PromQLSubquery { input, .. } => self.visit(input), + | QueryExpr::PromQLSubquery { input, .. } => self.visit(input, parent_schema), QueryExpr::Join { left, right, .. } | QueryExpr::SetOp { left, right, .. } | QueryExpr::BinaryOp { lhs: left, rhs: right, .. } => { - self.visit(left); - self.visit(right); + self.visit(left, parent_schema); + self.visit(right, parent_schema); } QueryExpr::LetBinding { expr, body, .. } => { - self.visit(expr); - self.visit(body); + self.visit(expr, parent_schema); + self.visit(body, parent_schema); } QueryExpr::Ref(_) => {} }