diff --git a/crates/frontend-sql/src/sql/mod.rs b/crates/frontend-sql/src/sql/mod.rs index eaee3fe5..9dd87bec 100644 --- a/crates/frontend-sql/src/sql/mod.rs +++ b/crates/frontend-sql/src/sql/mod.rs @@ -17,7 +17,7 @@ use datafusion::logical_expr::{ }; use datafusion::prelude::SessionContext; -use asap_ir::intent_algebra::schema::Schema; +use asap_ir::intent_algebra::schema::{DataType, Schema}; use asap_ir::intent_algebra::{ ColumnRef, CompareOp, JoinKind, L2Expr, L3Scalar, SetOpKind, WindowFuncKind, }; @@ -33,7 +33,7 @@ mod types; pub use types::SqlCatalog; use self::expr::df_expr_to_l2; -use self::types::schema_to_arrow; +use self::types::{arrow_to_l3, schema_to_arrow}; /// Lowers SQL strings to the Layer-2 [`relational::QueryExpr`] over a table /// [`SqlCatalog`]. Call [`convert_root`](asap_l2::convert_root) @@ -341,18 +341,11 @@ impl<'a> SqlLowerer<'a> { fn lower_aggregate(&self, agg: &logical_expr::Aggregate) -> Result { 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). + // `GROUPING SETS`/`ROLLUP`/`CUBE` emit several grouping levels from one + // scan. `Aggregate.by` is a single key set, so each level becomes its own + // `Aggregate` and they are merged (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}" - ))); + return self.lower_grouping_sets(agg, gs, input); } // `Aggregate.by` and the reducers index *columns*, so a grouping or @@ -432,6 +425,134 @@ impl<'a> SqlLowerer<'a> { }) } + /// `GROUP BY ROLLUP/CUBE/GROUPING SETS` — multi-level grouping (issue #118). + /// + /// One scan produces several grouping levels; `Aggregate.by` holds a single + /// key set. So each level becomes its own `Aggregate`, and the levels are + /// `Merge`d. A level that omits a key still has to *emit* it — as `NULL`, per + /// SQL — so each branch is wrapped in a `Project` that reinstates the missing + /// keys as typed nulls and restores the canonical column order. That keeps + /// the branches union-compatible, which `Merge` requires (it derives its + /// schema from the first child). + /// + /// `Aggregate.child` is duplicated per level. `plan::cse` hoists it back into + /// a single producer — the same trade `histogram_quantiles` makes (#109). + /// + /// DataFusion's `__grouping_id` discriminator is dropped: it only exists to + /// tell a subtotal's `NULL` apart from a data `NULL`, which is observable + /// solely through `GROUPING(col)` — an aggregate this front end rejects. + fn lower_grouping_sets( + &self, + agg: &logical_expr::Aggregate, + gs: &logical_expr::GroupingSet, + input: L2, + ) -> Result { + // DataFusion normalizes every mixed form (`GROUP BY g, ROLLUP(d)`) into a + // single `GroupingSets`, so one grouping expression is the only shape. + if agg.group_expr.len() != 1 { + return Err(LoweringError::UnsupportedFeature( + "a grouping set alongside plain GROUP BY keys".into(), + )); + } + + // `distinct_expr()` is ordered exactly like the aggregate's leading + // schema fields, which is the column order the enclosing Projection + // expects. The field after them is `__grouping_id`. + let distinct = gs.distinct_expr(); + for e in &distinct { + if !matches!(unalias(e), Expr::Column(_)) { + return Err(LoweringError::UnsupportedFeature(format!( + "non-column key inside a multi-level grouping: {e}" + ))); + } + } + let keys: Vec<(String, DataType)> = agg + .schema + .fields() + .iter() + .take(distinct.len()) + .map(|f| Ok((f.name().to_string(), arrow_to_l3(f.data_type())?))) + .collect::>()?; + + let out_names: Vec = agg + .schema + .fields() + .iter() + .skip(distinct.len() + 1) // + `__grouping_id` + .map(|f| f.name().to_string()) + .collect(); + + // Reducer arguments still materialize as derived columns (#110); the + // grouping keys are plain columns, so they only need carrying through. + let mut derived = DerivedCols::default(); + for e in &distinct { + derived.passthrough(e)?; + } + let aggr_expr = agg + .aggr_expr + .iter() + .map(|e| derived.rewrite_agg(e)) + .collect::, LoweringError>>()?; + let aggs = aggr_expr + .iter() + .enumerate() + .map(|(i, e)| { + let mut item = lower_agg_item(e)?; + if let Some(name) = out_names.get(i) { + item.alias = Some(name.clone()); + } + Ok(item) + }) + .collect::, LoweringError>>()?; + let input = derived.wrap(input)?; + + let branches = expand_grouping_set(gs) + .iter() + .map(|level| { + let level_keys = distinct + .iter() + .filter(|e| level.contains(e)) + .map(|e| expr_to_group_ref(e)) + .collect::, LoweringError>>()?; + let aggregate = L2::Aggregate { + keys: level_keys, + without: false, + aggs: aggs.clone(), + having: None, + input: Box::new(input.clone()), + }; + // Reinstate omitted keys as typed nulls, in canonical order. + let cols = keys + .iter() + .zip(&distinct) + .map(|((name, dtype), e)| L2ProjectItem { + alias: Some(name.clone()), + expr: if level.contains(e) { + L2Expr::Column(ColumnRef::Named(name.clone())) + } else { + L2Expr::Cast { + expr: Box::new(L2Expr::Literal(L3Scalar::Null)), + to: dtype.clone(), + try_cast: false, + } + }, + }) + .chain(out_names.iter().map(|n| L2ProjectItem { + alias: Some(n.clone()), + expr: L2Expr::Column(ColumnRef::Named(n.clone())), + })) + .collect(); + Ok(L2::Project { + cols, + qualifier: None, + input: Box::new(aggregate), + }) + }) + .collect::, LoweringError>>()?; + + Ok(L2::Merge { inputs: branches }) + } + fn lower_sort(&self, sort: &logical_expr::Sort) -> Result { // A count-ranked `ORDER BY … LIMIT k` is the frequency heavy-hitter the // `TopK` intent represents, but that promotion now happens in the shared @@ -559,6 +680,35 @@ fn as_grouping_set(expr: &Expr) -> Option<&logical_expr::GroupingSet> { } } +/// The grouping levels a `GroupingSet` stands for, widest first (issue #118). +/// +/// `ROLLUP(a, b)` → `(a,b), (a), ()` — the prefixes. +/// `CUBE(a, b)` → `(a,b), (a), (b), ()` — the power set. +/// `GROUPING SETS` is already the explicit list. +fn expand_grouping_set(gs: &logical_expr::GroupingSet) -> Vec> { + match gs { + logical_expr::GroupingSet::Rollup(exprs) => (0..=exprs.len()) + .rev() + .map(|n| exprs[..n].to_vec()) + .collect(), + logical_expr::GroupingSet::Cube(exprs) => { + // Bitmask descending, so the full set leads and `()` trails. + (0..(1u32 << exprs.len())) + .rev() + .map(|mask| { + exprs + .iter() + .enumerate() + .filter(|(i, _)| mask & (1 << i) != 0) + .map(|(_, e)| e.clone()) + .collect() + }) + .collect() + } + logical_expr::GroupingSet::GroupingSets(sets) => sets.clone(), + } +} + /// Derived columns materialized in a `Project` beneath an `Aggregate` (#110). /// /// `Aggregate.by` holds positional `ColumnId`s and each reducer holds one input diff --git a/crates/frontend-sql/tests/sql_lowering.rs b/crates/frontend-sql/tests/sql_lowering.rs index 19146bd6..f07ccb5d 100644 --- a/crates/frontend-sql/tests/sql_lowering.rs +++ b/crates/frontend-sql/tests/sql_lowering.rs @@ -990,21 +990,180 @@ async fn a_shared_expression_is_materialized_once() { assert_eq!(aggs[0].input_col(), aggs[1].input_col()); } +// ── Issue #118: multi-level grouping expands into one Aggregate per level ─── + +/// The branches of the first `Merge` along the single-child spine. +fn merge_branches(qe: &QueryExpr) -> &Vec { + fn find(qe: &QueryExpr) -> Option<&Vec> { + match qe { + QueryExpr::Merge { children } => Some(children), + QueryExpr::Project { child, .. } + | QueryExpr::Filter { child, .. } + | QueryExpr::Sort { child, .. } + | QueryExpr::Limit { child, .. } => find(child), + _ => None, + } + } + find(qe).expect("expected a Merge") +} + +/// `(group keys, column names)` of each merged grouping level. +fn grouping_levels(qe: &QueryExpr) -> Vec<(GroupKeys, Vec)> { + merge_branches(qe) + .iter() + .map(|b| { + let QueryExpr::Project { child, .. } = b else { + panic!("expected a Project per level, got {b:?}"); + }; + let QueryExpr::Aggregate { by, .. } = child.as_ref() else { + panic!("expected an Aggregate under the Project, got {child:?}"); + }; + let names = b + .output_schema() + .expect("level schema") + .columns + .iter() + .map(|c| c.name.clone()) + .collect(); + (by.clone(), names) + }) + .collect() +} + #[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"); +async fn rollup_expands_to_one_aggregate_per_prefix() { + // ROLLUP(a, b) → (a,b), (a), () — three levels, widest first. + let qe = + lower("SELECT service, bytes, SUM(latency) FROM metrics GROUP BY ROLLUP(service, bytes)") + .await; + let levels = grouping_levels(&qe); + let keys: Vec<_> = levels.iter().map(|(by, _)| by.clone()).collect(); + assert_eq!( + keys, + vec![ + GroupKeys::by(vec![1, 3]), + GroupKeys::by(vec![1]), + GroupKeys::none(), + ] + ); +} + +#[tokio::test] +async fn cube_expands_to_the_power_set() { + // CUBE(a, b) → (a,b), (a), (b), () — four levels. + let qe = + lower("SELECT service, bytes, SUM(latency) FROM metrics GROUP BY CUBE(service, bytes)") + .await; + assert_eq!(grouping_levels(&qe).len(), 4); +} + +#[tokio::test] +async fn a_mixed_grouping_set_is_normalized_by_datafusion() { + // `GROUP BY g, ROLLUP(d)` arrives as one GroupingSets, not a plain key + // alongside a grouping set — so there is only one shape to handle. + let qe = + lower("SELECT service, bytes, SUM(latency) FROM metrics GROUP BY service, ROLLUP(bytes)") + .await; + assert_eq!(grouping_levels(&qe).len(), 2); +} + +#[tokio::test] +async fn omitted_grouping_keys_become_typed_nulls() { + // Every level must emit every key — as NULL where the level omits it — or + // `Merge` (which takes the first child's schema) would misdescribe the rest. + // The null is *cast*: a bare Null literal infers as Float64. + let qe = lower("SELECT service, SUM(bytes) FROM metrics GROUP BY ROLLUP(service)").await; + let levels = grouping_levels(&qe); + assert_eq!(levels.len(), 2); + for (_, names) in &levels { + assert_eq!( + names, + &["service".to_string(), "sum(metrics.bytes)".to_string()] + ); + } + + // The `()` level projects `service` as a Utf8 null, not a Float64 one. + let schema = merge_branches(&qe)[1] + .output_schema() + .expect("level schema"); + assert_eq!(schema.columns[0].name, "service"); + assert_eq!( + schema.columns[0].dtype, + DataType::Utf8, + "the omitted key must keep its declared type" + ); +} + +#[tokio::test] +async fn grouping_levels_are_union_compatible() { + let qe = lower( + "SELECT service, bytes, SUM(latency) FROM metrics GROUP BY GROUPING SETS ((service),(bytes),())", + ) + .await; + let shapes: Vec<_> = merge_branches(&qe) + .iter() + .map(|b| { + b.output_schema() + .expect("level schema") + .columns + .iter() + .map(|c| (c.name.clone(), c.dtype.clone())) + .collect::>() + }) + .collect(); + assert!( + shapes.windows(2).all(|w| w[0] == w[1]), + "levels disagree: {shapes:?}" + ); +} + +#[tokio::test] +async fn grouping_function_is_rejected() { + // `__grouping_id` is dropped when the levels are expanded. It is observable + // only through `GROUPING(col)`, so dropping it loses nothing representable — + // this test is what makes that true. + let err = lower_sql( + "SELECT service, SUM(bytes), GROUPING(service) FROM metrics GROUP BY ROLLUP(service)", + &catalog(), + AccuracyTarget::Exact, + ) + .await + .expect_err("GROUPING() must be rejected while __grouping_id is dropped"); + assert!(format!("{err}").contains("grouping"), "got {err}"); +} + +#[tokio::test] +async fn a_non_column_key_inside_a_grouping_set_is_rejected() { + // The #110 derived-column machinery covers plain `GROUP BY `; inside a + // grouping set the key also has to be reinstatable as a typed null. + let err = lower_sql( + "SELECT date_trunc('minute', ts) AS m, SUM(bytes) FROM metrics GROUP BY ROLLUP(m)", + &catalog(), + AccuracyTarget::Exact, + ) + .await + .expect_err("expression key inside ROLLUP must be rejected"); + assert!( + format!("{err}").contains("non-column key inside a multi-level grouping"), + "got {err}" + ); +} + +#[tokio::test] +async fn multi_level_grouping_composes_with_a_derived_reducer_argument() { + // #110's materializing Project sits beneath every level's Aggregate. + let qe = lower("SELECT service, SUM(bytes * 8) FROM metrics GROUP BY ROLLUP(service)").await; + for b in merge_branches(&qe) { + let QueryExpr::Project { child, .. } = b else { + panic!("expected a Project per level"); + }; + let QueryExpr::Aggregate { aggs, child, .. } = child.as_ref() else { + panic!("expected an Aggregate"); + }; + assert!(matches!(aggs.as_slice(), [AggIntent::Sum { col: Some(_) }])); assert!( - format!("{err}").contains("multi-level grouping"), - "{q} gave {err}" + matches!(**child, QueryExpr::Project { .. }), + "the derived-column projection should sit under each level" ); } }