diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs index 413ee4b537c29..106d8c8dfb324 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs @@ -17,8 +17,11 @@ use crate::logical_plan::consumer::{NameTracker, SubstraitConsumer}; use crate::logical_plan::consumer::{from_substrait_agg_func, from_substrait_sorts}; -use datafusion::common::{DFSchemaRef, not_impl_err}; -use datafusion::logical_expr::{Expr, GroupingSet, LogicalPlan, LogicalPlanBuilder}; +use datafusion::common::{Column, DFSchemaRef, internal_err, not_impl_err}; +use datafusion::logical_expr::builder::project; +use datafusion::logical_expr::{ + Aggregate, Expr, GroupingSet, LogicalPlan, LogicalPlanBuilder, +}; use substrait::proto::AggregateRel; use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::aggregate_rel::Grouping; @@ -122,12 +125,51 @@ pub async fn from_aggregate_rel( .map(|e| name_tracker.get_uniquely_named_expr(e)) .collect::, _>>()?; - input.aggregate(group_exprs, aggr_exprs)?.build() + let plan = input.aggregate(group_exprs, aggr_exprs)?.build()?; + if agg.groupings.len() > 1 { + reorder_grouping_set_output(plan, agg.measures.len()) + } else { + Ok(plan) + } } else { not_impl_err!("Aggregate without an input is not valid") } } +/// Reorders DataFusion's `[groups, grouping_id, measures]` aggregate schema to +/// Substrait's direct output order of `[groups, measures, grouping_id]`. +fn reorder_grouping_set_output( + plan: LogicalPlan, + measure_count: usize, +) -> datafusion::common::Result { + let exprs: Vec = { + let schema = plan.schema(); + let Some(grouping_id_index) = + schema.index_of_column_by_name(None, Aggregate::INTERNAL_GROUPING_ID) + else { + return internal_err!( + "Grouping set aggregate schema is missing {}", + Aggregate::INTERNAL_GROUPING_ID + ); + }; + if grouping_id_index + measure_count + 1 != schema.fields().len() { + return internal_err!( + "Grouping set aggregate schema has {} fields after {}, expected {} measures", + schema.fields().len() - grouping_id_index - 1, + Aggregate::INTERNAL_GROUPING_ID, + measure_count + ); + } + + (0..grouping_id_index) + .chain(grouping_id_index + 1..schema.fields().len()) + .chain(std::iter::once(grouping_id_index)) + .map(|index| Expr::Column(Column::from(schema.qualified_field(index)))) + .collect() + }; + project(plan, exprs) +} + #[expect(deprecated)] async fn from_substrait_grouping( consumer: &impl SubstraitConsumer, diff --git a/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs b/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs index dec94b0422257..7b6c113ccec0f 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs @@ -24,7 +24,8 @@ use datafusion::logical_expr::utils::powerset; use datafusion::logical_expr::{Aggregate, Distinct, Expr, GroupingSet}; use substrait::proto::aggregate_rel::{Grouping, Measure}; use substrait::proto::rel::RelType; -use substrait::proto::{AggregateRel, Expression, Rel}; +use substrait::proto::rel_common::EmitKind; +use substrait::proto::{AggregateRel, Expression, Rel, RelCommon, rel_common}; pub fn from_aggregate( producer: &mut impl SubstraitProducer, @@ -38,10 +39,12 @@ pub fn from_aggregate( .iter() .map(|e| to_substrait_agg_measure(producer, e, agg.input.schema())) .collect::>>()?; + let common = (groupings.len() > 1) + .then(|| grouping_set_output_mapping(grouping_expressions.len(), measures.len())); Ok(Box::new(Rel { rel_type: Some(RelType::Aggregate(Box::new(AggregateRel { - common: None, + common, input: Some(input), grouping_expressions, groupings, @@ -51,6 +54,22 @@ pub fn from_aggregate( })) } +/// Maps Substrait's `[groups, measures, grouping_id]` direct output to +/// DataFusion's `[groups, grouping_id, measures]` aggregate schema. +fn grouping_set_output_mapping(grouping_count: usize, measure_count: usize) -> RelCommon { + let grouping_id_index = grouping_count + measure_count; + let output_mapping = (0..grouping_count) + .chain(std::iter::once(grouping_id_index)) + .chain(grouping_count..grouping_id_index) + .map(|index| index as i32) + .collect(); + RelCommon { + emit_kind: Some(EmitKind::Emit(rel_common::Emit { output_mapping })), + hint: None, + advanced_extension: None, + } +} + pub fn from_distinct( producer: &mut impl SubstraitProducer, distinct: &Distinct, @@ -165,8 +184,12 @@ pub fn parse_flat_grouping_exprs( for e in exprs { let rex = producer.handle_expr(e, schema)?; grouping_expressions.push(rex.clone()); - ref_group_exprs.push(rex); - expression_references.push((ref_group_exprs.len() - 1) as u32); + let reference = ref_group_exprs.iter().position(|existing| existing == &rex); + let reference = reference.unwrap_or_else(|| { + ref_group_exprs.push(rex); + ref_group_exprs.len() - 1 + }); + expression_references.push(reference as u32); } #[expect(deprecated)] Ok(Grouping { diff --git a/datafusion/substrait/tests/cases/aggregation_tests.rs b/datafusion/substrait/tests/cases/aggregation_tests.rs index 92a41850b208d..e572023f17a92 100644 --- a/datafusion/substrait/tests/cases/aggregation_tests.rs +++ b/datafusion/substrait/tests/cases/aggregation_tests.rs @@ -68,4 +68,42 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn multiple_grouping_sets_follow_substrait_output_order() -> Result<()> { + let proto_plan = read_json( + "tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json", + ); + let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?; + let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + + assert_snapshot!( + plan, + @r" + Projection: c0, c1, sum(c0) AS summation + Aggregate: groupBy=[[GROUPING SETS ((c0), (c1), (c0, c1))]], aggr=[[sum(c0)]] + Values: (Int64(1), Int64(10)), (Int64(1), Int64(20)), (Int64(2), Int64(10)) + " + ); + + let results = DataFrame::new(ctx.state(), plan).collect().await?; + datafusion::assert_batches_sorted_eq!( + [ + "+----+----+-----------+", + "| c0 | c1 | summation |", + "+----+----+-----------+", + "| | 10 | 3 |", + "| | 20 | 1 |", + "| 1 | | 2 |", + "| 1 | 10 | 1 |", + "| 1 | 20 | 1 |", + "| 2 | | 2 |", + "| 2 | 10 | 2 |", + "+----+----+-----------+", + ], + &results + ); + + Ok(()) + } } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 018e1aef80ea1..c646f8c59a585 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -344,10 +344,39 @@ async fn aggregate_multiple_keys() -> Result<()> { #[tokio::test] async fn aggregate_grouping_sets() -> Result<()> { - roundtrip( - "SELECT a, c, d, avg(b) FROM data GROUP BY GROUPING SETS ((a, c), (a), (d), ())", + let proto = roundtrip_with_ctx( + "SELECT a, c, d, avg(b), sum(e) FROM data GROUP BY GROUPING SETS ((a, c), (a), (d), ())", + create_context().await?, ) - .await + .await?; + + let Some(plan_rel::RelType::Root(root)) = &proto.relations[0].rel_type else { + panic!("expected root relation"); + }; + let Some(RelType::Project(project)) = &root.input.as_ref().unwrap().rel_type else { + panic!("expected project relation"); + }; + let Some(RelType::Aggregate(aggregate)) = &project.input.as_ref().unwrap().rel_type + else { + panic!("expected aggregate relation"); + }; + + assert_eq!(aggregate.grouping_expressions.len(), 3); + assert_eq!(aggregate.groupings[0].expression_references, [0, 1]); + assert_eq!(aggregate.groupings[1].expression_references, [0]); + assert_eq!(aggregate.groupings[2].expression_references, [2]); + assert!(aggregate.groupings[3].expression_references.is_empty()); + let output_mapping = match aggregate + .common + .as_ref() + .and_then(|common| common.emit_kind.as_ref()) + { + Some(substrait::proto::rel_common::EmitKind::Emit(emit)) => &emit.output_mapping, + _ => panic!("expected aggregate output mapping"), + }; + assert_eq!(output_mapping, &[0, 1, 2, 5, 3, 4]); + + Ok(()) } #[tokio::test] diff --git a/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json b/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json new file mode 100644 index 0000000000000..3779ee32c6ab8 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json @@ -0,0 +1,127 @@ +{ + "extensionUris": [ + { + "extensionUriAnchor": 1, + "uri": "https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml" + } + ], + "extensions": [ + { + "extensionFunction": { + "extensionUriReference": 1, + "functionAnchor": 1, + "name": "sum:i64" + } + } + ], + "relations": [ + { + "root": { + "input": { + "aggregate": { + "common": { + "emit": { + "outputMapping": [0, 1, 2] + } + }, + "input": { + "read": { + "baseSchema": { + "names": ["c0", "c1"], + "struct": { + "nullability": "NULLABILITY_REQUIRED", + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + } + ] + } + }, + "common": { + "direct": {} + }, + "virtualTable": { + "expressions": [ + { + "fields": [ + { "literal": { "i64": "1", "nullable": false } }, + { "literal": { "i64": "10", "nullable": false } } + ] + }, + { + "fields": [ + { "literal": { "i64": "1", "nullable": false } }, + { "literal": { "i64": "20", "nullable": false } } + ] + }, + { + "fields": [ + { "literal": { "i64": "2", "nullable": false } }, + { "literal": { "i64": "10", "nullable": false } } + ] + } + ] + } + } + }, + "groupingExpressions": [ + { + "selection": { + "directReference": { "structField": {} }, + "rootReference": {} + } + }, + { + "selection": { + "directReference": { "structField": { "field": 1 } }, + "rootReference": {} + } + } + ], + "groupings": [ + { "expressionReferences": [0] }, + { "expressionReferences": [1] }, + { "expressionReferences": [0, 1] } + ], + "measures": [ + { + "measure": { + "arguments": [ + { + "value": { + "selection": { + "directReference": { "structField": {} }, + "rootReference": {} + } + } + } + ], + "functionReference": 1, + "invocation": "AGGREGATION_INVOCATION_ALL", + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "phase": "AGGREGATION_PHASE_INITIAL_TO_RESULT" + } + } + ] + } + }, + "names": ["c0", "c1", "summation"] + } + } + ], + "version": { + "minorNumber": 29, + "producer": "substrait-go v4.2.0" + } +}