From 2bb35219146660a79b7a62904b5989ef21d51567 Mon Sep 17 00:00:00 2001 From: zz_y Date: Thu, 9 Jul 2026 12:46:13 -0600 Subject: [PATCH] feat(sql): materialize expression GROUP BY and reducer arguments (#110) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Aggregate.by` holds positional ColumnIds and each reducer holds one input column, so neither could hold an expression. `GROUP BY date_trunc('minute', ts)` — the canonical time-series bucketing idiom, and the SQL counterpart of a PromQL range selector — was rejected, which meant the SQL front end could not express windowed aggregation at all. `SUM(bytes * 8)` was rejected for the same reason. Materialize each grouping / reducer expression as a derived column in a Project beneath the Aggregate, then group and reduce over that column. The projection also carries through the plain columns the aggregate still references, since a Project replaces its child's schema rather than extending it. The projection is inserted only when something actually needs deriving, so every query that lowers today keeps its exact tree shape. Two expressions with the same value are materialized once, so `SUM(v*2), MIN(v*2)` share a column. A Project carries one relation qualifier for all its columns, so `a.k` and `b.k` cannot both survive it. That is reported as an ambiguity — but only when the projection is actually inserted, otherwise a join whose keys need no materialization would regress. GROUPING SETS / ROLLUP / CUBE are a different problem: several grouping levels plus a `__grouping_id` discriminator from one scan, which `Aggregate.by` cannot express. They previously failed as "non-column GROUP BY expression", implying they were the same issue; they now fail as "multi-level grouping" and are tracked in #118. Supersedes the rejection this change replaces: `SUM()` and `median()` used to error rather than lower with `col: None` (#115). They now lower with `col: Some()`. The #115 invariant — a SQL reducer never reaches L3 with `col: None` — still holds and is still asserted. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/frontend-sql/src/sql/mod.rs | 197 ++++++++++++++++++- crates/frontend-sql/tests/sql_lowering.rs | 223 ++++++++++++++++++++-- 2 files changed, 391 insertions(+), 29 deletions(-) diff --git a/crates/frontend-sql/src/sql/mod.rs b/crates/frontend-sql/src/sql/mod.rs index 12bf6dcf..8ef6e7e0 100644 --- a/crates/frontend-sql/src/sql/mod.rs +++ b/crates/frontend-sql/src/sql/mod.rs @@ -10,7 +10,7 @@ use std::sync::Arc; -use datafusion::common::ScalarValue; +use datafusion::common::{Column as DfColumn, ScalarValue}; use datafusion::datasource::MemTable; use datafusion::logical_expr::{ self, Distinct, Expr, JoinType, LogicalPlan, WindowFunctionDefinition, @@ -339,12 +339,66 @@ impl<'a> SqlLowerer<'a> { } fn lower_aggregate(&self, agg: &logical_expr::Aggregate) -> Result { - let input = Box::new(self.lower_plan(&agg.input)?); - let keys = agg - .group_expr + let input = self.lower_plan(&agg.input)?; + + // `GROUPING SETS`/`ROLLUP`/`CUBE` emit several grouping levels plus a + // `__grouping_id` discriminator from one scan. `Aggregate.by` is a + // single key set, so there is nothing to lower them onto (issue #118). + if let Some(gs) = agg.group_expr.iter().find_map(as_grouping_set) { + let kind = match gs { + logical_expr::GroupingSet::Rollup(_) => "ROLLUP", + logical_expr::GroupingSet::Cube(_) => "CUBE", + logical_expr::GroupingSet::GroupingSets(_) => "GROUPING SETS", + }; + return Err(LoweringError::UnsupportedFeature(format!( + "multi-level grouping: {kind}" + ))); + } + + // `Aggregate.by` and the reducers index *columns*, so a grouping or + // reducer expression (`GROUP BY date_trunc(…)`, `SUM(a * 8)`) has no + // slot. Materialize each one as a derived column in a `Project` beneath + // the aggregate, then group/reduce over that column (issue #110). + let mut derived = DerivedCols::default(); + + // DataFusion strips `AS m` from a grouping expression, so the aggregate + // schema's field name is what the enclosing Projection references — + // the derived column has to carry exactly that name. + let group_names: Vec = agg + .schema + .fields() .iter() - .map(expr_to_group_ref) - .collect::, _>>()?; + .take(agg.group_expr.len()) + .map(|f| f.name().to_string()) + .collect(); + + let mut keys = Vec::with_capacity(agg.group_expr.len()); + for (i, e) in agg.group_expr.iter().enumerate() { + match unalias(e) { + Expr::Column(_) => { + derived.passthrough(e)?; + keys.push(expr_to_group_ref(e)?); + } + other => { + let name = group_names + .get(i) + .cloned() + .unwrap_or_else(|| other.to_string()); + derived.materialize(name.clone(), df_expr_to_l2(other)?)?; + keys.push(ColumnRef::Named(name)); + } + } + } + + // Reducer arguments get the same treatment; `rewrite_agg` returns the + // aggregate with its argument repointed at the derived column. + let aggr_expr = agg + .aggr_expr + .iter() + .map(|e| derived.rewrite_agg(e)) + .collect::, LoweringError>>()?; + + let input = Box::new(derived.wrap(input)?); // DataFusion names the aggregate outputs in its own schema (e.g. // "sum(metrics.bytes)") — the same names the enclosing Projection // references. The schema is [group fields …, aggregate fields …], so @@ -357,8 +411,7 @@ impl<'a> SqlLowerer<'a> { .skip(agg.group_expr.len()) .map(|f| f.name().to_string()) .collect(); - let aggs = agg - .aggr_expr + let aggs = aggr_expr .iter() .enumerate() .map(|(i, e)| { @@ -491,6 +544,134 @@ fn lower_agg_item(expr: &Expr) -> Result { } } +/// Strip `AS alias` wrappers. +fn unalias(expr: &Expr) -> &Expr { + match expr { + Expr::Alias(a) => unalias(&a.expr), + other => other, + } +} + +/// The `GroupingSet` inside a grouping expression, if any. +fn as_grouping_set(expr: &Expr) -> Option<&logical_expr::GroupingSet> { + match unalias(expr) { + Expr::GroupingSet(gs) => Some(gs), + _ => None, + } +} + +/// Derived columns materialized in a `Project` beneath an `Aggregate` (#110). +/// +/// `Aggregate.by` holds positional `ColumnId`s and each reducer holds one input +/// column, so neither can hold an expression. `GROUP BY date_trunc('minute', t)` +/// and `SUM(bytes * 8)` are therefore rewritten to group/reduce over a projected +/// column that carries the expression's value. +/// +/// The projection also has to carry through the plain columns the aggregate +/// still references, since a `Project` replaces its child's schema rather than +/// extending it. +#[derive(Default)] +struct DerivedCols { + cols: Vec, + /// Whether any column is genuinely derived. Without one the aggregate keeps + /// its original child, so trees that lower today keep their exact shape. + any: bool, + /// First same-name-different-value collision, reported only if the + /// projection is actually inserted (see [`Self::wrap`]). + collision: Option, +} + +impl DerivedCols { + /// Add `alias := expr`, or note a collision if `alias` already means + /// something else. `Project` carries one relation qualifier for all its + /// columns, so `a.k` and `b.k` cannot both survive it — but that only + /// matters when a projection gets inserted at all. + fn push(&mut self, alias: String, expr: L2Expr) { + let existing = self + .cols + .iter() + .find(|c| c.alias.as_deref() == Some(&alias)); + match existing { + // Same name, same value — one projected column serves both uses. + Some(e) if e.expr == expr => {} + Some(_) => { + self.collision.get_or_insert(alias); + } + None => self.cols.push(L2ProjectItem { + alias: Some(alias), + expr, + }), + } + } + + /// A plain column the aggregate references — carried through unchanged. + fn passthrough(&mut self, expr: &Expr) -> Result<(), LoweringError> { + let Expr::Column(c) = unalias(expr) else { + return Ok(()); + }; + self.push(c.name.clone(), df_expr_to_l2(expr)?); + Ok(()) + } + + /// A genuinely derived column: `alias` now names `expr`'s value. + fn materialize(&mut self, alias: String, expr: L2Expr) -> Result<(), LoweringError> { + self.any = true; + self.push(alias, expr); + Ok(()) + } + + /// Repoint a reducer's argument at a derived column when it is an + /// expression; otherwise carry its plain input column through. + fn rewrite_agg(&mut self, expr: &Expr) -> Result { + let Expr::AggregateFunction(agg_fn) = unalias(expr) else { + return Ok(expr.clone()); + }; + // `COUNT(*)` reduces no column; `agg_col_name` covers bare/aliased/cast + // columns, so `None` here means the argument really is an expression. + let counts_rows = agg_fn.func.name().eq_ignore_ascii_case("count") && !agg_fn.distinct; + let Some(arg) = agg_fn.args.first() else { + return Ok(expr.clone()); + }; + if counts_rows { + return Ok(expr.clone()); + } + match agg_col_name(&agg_fn.args) { + Some(name) => { + self.push(name, df_expr_to_l2(arg)?); + Ok(expr.clone()) + } + None => { + let alias = unalias(arg).to_string(); + self.materialize(alias.clone(), df_expr_to_l2(arg)?)?; + let mut agg_fn = agg_fn.clone(); + agg_fn.args[0] = Expr::Column(DfColumn::new_unqualified(alias)); + Ok(Expr::AggregateFunction(agg_fn)) + } + } + } + + /// Wrap `input` in the materializing `Project`, or return it untouched when + /// nothing needed deriving — so a query that lowers today keeps its exact + /// tree, and a name collision that the projection would have flattened only + /// matters once the projection exists. + fn wrap(self, input: L2) -> Result { + if !self.any { + return Ok(input); + } + if let Some(alias) = self.collision { + return Err(LoweringError::UnsupportedFeature(format!( + "ambiguous column `{alias}` beneath an expression GROUP BY / \ + aggregate — alias the relations apart" + ))); + } + Ok(L2::Project { + cols: self.cols, + qualifier: None, + input: Box::new(input), + }) + } +} + /// The first aggregate argument's column name (bare / aliased / cast column), /// or `None` for `*` / a non-column expression. fn agg_col_name(args: &[Expr]) -> Option { diff --git a/crates/frontend-sql/tests/sql_lowering.rs b/crates/frontend-sql/tests/sql_lowering.rs index 5388008a..038693e1 100644 --- a/crates/frontend-sql/tests/sql_lowering.rs +++ b/crates/frontend-sql/tests/sql_lowering.rs @@ -61,6 +61,36 @@ fn find_aggregate(qe: &QueryExpr) -> Option<(&GroupKeys, &Vec)> { } } +/// The first `Aggregate` node itself, for tests that need its child. +fn find_aggregate_node(qe: &QueryExpr) -> Option<&QueryExpr> { + match qe { + QueryExpr::Aggregate { .. } => Some(qe), + QueryExpr::Project { child, .. } + | QueryExpr::Filter { child, .. } + | QueryExpr::Sort { child, .. } + | QueryExpr::Limit { child, .. } => find_aggregate_node(child), + _ => None, + } +} + +/// The names of the columns the first `Aggregate`'s reducers read, resolved +/// against its child's schema, plus whether that child is a materializing +/// `Project` (issue #110). +fn reducer_input_names(qe: &QueryExpr) -> (Vec, bool) { + let QueryExpr::Aggregate { aggs, child, .. } = + find_aggregate_node(qe).expect("expected an Aggregate") + else { + unreachable!() + }; + let schema = child.output_schema().expect("child schema"); + let names = aggs + .iter() + .filter_map(|a| a.input_col()) + .map(|id| schema.columns[id].name.clone()) + .collect(); + (names, matches!(**child, QueryExpr::Project { .. })) +} + /// Find the first `Join` node along the single-child spine. fn find_join(qe: &QueryExpr) -> Option<&QueryExpr> { match qe { @@ -258,16 +288,22 @@ async fn distinct_value_reducer_is_rejected_not_dropped() { } #[tokio::test] -async fn aggregate_over_non_column_expression_is_rejected() { - // L3 reduces a column, not an arbitrary expression — SUM(bytes + 1) must be - // rejected rather than silently reducing a probe column. - let res = lower_sql( - "SELECT SUM(bytes + 1) FROM metrics", - &catalog(), - AccuracyTarget::Exact, - ) - .await; - assert!(res.is_err(), "SUM() should be rejected"); +async fn aggregate_over_an_expression_reduces_a_derived_column() { + // L3 reduces a column, not an arbitrary expression. `SUM(bytes + 1)` used to + // be rejected for that reason; since #110 the expression is materialized as + // a derived column in a `Project` beneath the aggregate, and reduced there. + let qe = lower("SELECT SUM(bytes + 1) FROM metrics").await; + let (_, aggs) = find_aggregate(&qe).expect("expected an Aggregate"); + assert!( + matches!(aggs.as_slice(), [AggIntent::Sum { col: Some(_) }]), + "expected Sum bound to the derived column, got {aggs:?}" + ); + let (names, materialized) = reducer_input_names(&qe); + assert!(materialized, "expected a materializing Project"); + assert!( + names[0].contains("bytes") && names[0].contains('1'), + "the reduced column should be the projected `bytes + 1`, got {names:?}" + ); } #[tokio::test] @@ -769,19 +805,26 @@ async fn count_distinct_carries_its_input_column() { } #[tokio::test] -async fn quantile_and_count_distinct_over_an_expression_are_rejected() { +async fn quantile_and_count_distinct_over_an_expression_bind_the_derived_column() { // A SQL aggregate has no "sample value" to fall back on, so an expression - // argument would lower to `col: None` and silently drop the expression. - // Reject it, exactly as `SUM(a*b)` is rejected. + // argument must never reach L3 as `col: None` (#115). Since #110 it reaches + // L3 as `col: Some(derived)` instead of being rejected. for q in [ "SELECT approx_percentile_cont(bytes * 8, 0.95) FROM metrics", "SELECT COUNT(DISTINCT bytes * 8) FROM metrics", "SELECT approx_distinct(bytes * 8) FROM metrics", ] { - let res = lower_sql(q, &catalog(), AccuracyTarget::Exact).await; + let qe = lower(q).await; + let (_, aggs) = find_aggregate(&qe).expect("expected an Aggregate"); assert!( - res.is_err(), - "aggregate over an expression must be rejected: {q}" + aggs[0].input_col().is_some(), + "{q} must bind a column, never `col: None`, got {aggs:?}" + ); + let (names, materialized) = reducer_input_names(&qe); + assert!(materialized, "{q} expected a materializing Project"); + assert!( + names[0].contains("bytes"), + "{q} should reduce the projected `bytes * 8`, got {names:?}" ); } } @@ -839,13 +882,151 @@ async fn median_threads_the_accuracy_target() { } #[tokio::test] -async fn median_over_an_expression_is_rejected() { - // Inherits the #115 column-binding rule: no column, no intent. - let res = lower_sql( - "SELECT median(bytes * 8) FROM metrics", +async fn median_over_an_expression_binds_the_derived_column() { + // Was rejected when filed (#111); supported since #110 materialized the + // expression. What must still hold is the #115 rule: never `col: None`. + let qe = lower("SELECT median(bytes * 8) FROM metrics").await; + let (_, aggs) = find_aggregate(&qe).expect("expected an Aggregate"); + assert!( + matches!(aggs.as_slice(), [AggIntent::Quantile { col: Some(_), q, .. }] if (*q - 0.5).abs() < 1e-9), + "expected Quantile(0.5) bound to the derived column, got {aggs:?}" + ); +} + +// ── Issue #110: expression GROUP BY (time bucketing) ──────────────────────── + +#[tokio::test] +async fn time_bucketing_group_by_lowers_to_a_derived_key() { + // The canonical time-series shape: `GROUP BY date_trunc(...)`. The bucket + // expression is materialized beneath the aggregate and grouped on. + let qe = + lower("SELECT date_trunc('minute', ts) AS m, SUM(bytes) FROM metrics GROUP BY m").await; + let node = find_aggregate_node(&qe).expect("expected an Aggregate"); + let QueryExpr::Aggregate { + by, aggs, child, .. + } = node + else { + unreachable!() + }; + assert!( + matches!(**child, QueryExpr::Project { .. }), + "expected a materializing Project beneath the Aggregate" + ); + let schema = child.output_schema().expect("child schema"); + assert_eq!(by, &GroupKeys::by(vec![0])); + assert!( + schema.columns[0].name.contains("date_trunc"), + "group key should be the projected bucket, got {:?}", + schema.columns[0].name + ); + // The reducer still binds its own column, not the bucket. + assert!(matches!(aggs.as_slice(), [AggIntent::Sum { col: Some(1) }])); +} + +#[tokio::test] +async fn time_bucketing_keeps_the_scan_predicate() { + // The projection is inserted above the scan, so a WHERE clause still folds + // onto the Scan rather than being stranded. + let qe = lower( + "SELECT date_trunc('minute', ts) AS m, SUM(bytes) FROM metrics \ + WHERE bytes > 10 GROUP BY m", + ) + .await; + fn scan_has_predicate(qe: &QueryExpr) -> bool { + match qe { + QueryExpr::Scan { predicates, .. } => !predicates.is_empty(), + QueryExpr::Project { child, .. } + | QueryExpr::Filter { child, .. } + | QueryExpr::Aggregate { child, .. } + | QueryExpr::Sort { child, .. } + | QueryExpr::Limit { child, .. } => scan_has_predicate(child), + _ => false, + } + } + assert!(scan_has_predicate(&qe), "WHERE should stay on the Scan"); +} + +#[tokio::test] +async fn a_plain_group_by_inserts_no_projection() { + // Queries that lowered before #110 must keep their exact tree shape — the + // projection appears only when something actually needs materializing. + for q in [ + "SELECT service, SUM(bytes) FROM metrics GROUP BY service", + "SELECT SUM(bytes) FROM metrics", + "SELECT COUNT(*) FROM metrics", + ] { + let qe = lower(q).await; + let QueryExpr::Aggregate { child, .. } = + find_aggregate_node(&qe).expect("expected an Aggregate") + else { + unreachable!() + }; + assert!( + !matches!(**child, QueryExpr::Project { .. }), + "{q} should not gain a projection" + ); + } +} + +#[tokio::test] +async fn a_shared_expression_is_materialized_once() { + let qe = lower("SELECT SUM(bytes * 2), MIN(bytes * 2) FROM metrics").await; + let QueryExpr::Aggregate { aggs, child, .. } = + find_aggregate_node(&qe).expect("expected an Aggregate") + else { + unreachable!() + }; + assert_eq!( + child.output_schema().expect("child schema").columns.len(), + 1, + "the two reducers should share one derived column" + ); + assert_eq!(aggs[0].input_col(), aggs[1].input_col()); +} + +#[tokio::test] +async fn multi_level_grouping_is_rejected() { + // ROLLUP/CUBE/GROUPING SETS emit several grouping levels plus a + // `__grouping_id` discriminator; `Aggregate.by` is a single key set. + for q in [ + "SELECT service, SUM(bytes) FROM metrics GROUP BY ROLLUP(service)", + "SELECT service, SUM(bytes) FROM metrics GROUP BY CUBE(service)", + "SELECT service, SUM(bytes) FROM metrics GROUP BY GROUPING SETS ((service), ())", + ] { + let err = lower_sql(q, &catalog(), AccuracyTarget::Exact) + .await + .expect_err("multi-level grouping must be rejected"); + assert!( + format!("{err}").contains("multi-level grouping"), + "{q} gave {err}" + ); + } +} + +#[tokio::test] +async fn an_ambiguous_passthrough_column_is_rejected_only_when_projecting() { + // A `Project` carries one relation qualifier for all its columns, so `a.k` + // and `b.k` cannot both survive it. That only matters once a projection is + // inserted: without a derived column the join keys resolve as before. + let ok = lower_sql( + "SELECT m.service, h.service, SUM(m.bytes) FROM metrics m \ + JOIN hosts h ON m.service = h.service GROUP BY m.service, h.service", &catalog(), AccuracyTarget::Exact, ) .await; - assert!(res.is_err(), "median over an expression must be rejected"); + assert!( + ok.is_ok(), + "no derived column ⇒ no projection ⇒ no ambiguity" + ); + + let err = lower_sql( + "SELECT m.service, h.service, SUM(m.bytes * 2) FROM metrics m \ + JOIN hosts h ON m.service = h.service GROUP BY m.service, h.service", + &catalog(), + AccuracyTarget::Exact, + ) + .await + .expect_err("ambiguous passthrough must be rejected, not silently resolved"); + assert!(format!("{err}").contains("ambiguous column"), "got {err}"); }