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}"); }