fix(substrait): preserve grouping set output order

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::<Result<Vec<Expr>, _>>()?;
 
-        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<LogicalPlan> {
+    let exprs: Vec<Expr> = {
+        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,

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::<datafusion::common::Result<Vec<_>>>()?;
+    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 {

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(())
+    }
 }

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]

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"
+  }
+}