You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/03 03:10:35 UTC
(datafusion) branch gh-readonly-queue/main/pr-23468-7a1f468643f5fcc33166b05a6cf2210675cab449 created (now 6ffc45d081)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-23468-7a1f468643f5fcc33166b05a6cf2210675cab449
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 6ffc45d081 fix(substrait): preserve grouping set output order (#23468)
This branch includes the following new commits:
new 6ffc45d081 fix(substrait): preserve grouping set output order (#23468)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix(substrait): preserve grouping set output order (#23468)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23468-7a1f468643f5fcc33166b05a6cf2210675cab449
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 6ffc45d0815f5f1e0b912cf5f42b34a5ab236dce
Author: Bruno Volpato <br...@gmail.com>
AuthorDate: Thu Sep 3 03:10:08 2026 +0000
fix(substrait): preserve grouping set output order (#23468)
## Which issue does this PR close?
- Closes #17910.
## Rationale for this change
For multiple grouping sets, Substrait orders aggregate output as
grouping expressions, measures, then grouping-set ID. DataFusion puts
its internal `__grouping_id` before measures. The consumer applied
`RelCommon.emit` to DataFusion's order, so a mapping such as `[0, 1, 2]`
returned `__grouping_id` where the plan requested its first measure.
See the [Substrait AggregateRel output mapping
specification](https://substrait.io/relations/logical_relations/#aggregate-operation).
## What changes are included in this PR?
- Reorder consumer output to Substrait's direct order before applying
emit mappings.
- Emit a producer mapping back to DataFusion's aggregate order.
- Deduplicate producer grouping expressions used by grouping references.
- Test both issue reproducer and two-measure roundtrip.
## Are these changes tested?
- `cargo test -p datafusion-substrait --features protoc`
- `cargo clippy --all-targets --all-features -- -D warnings`
## Are there any user-facing changes?
Substrait plans with multiple grouping sets now return emitted columns
in specified order. No public API change.
---
.../src/logical_plan/consumer/rel/aggregate_rel.rs | 48 +++++++-
.../src/logical_plan/producer/rel/aggregate_rel.rs | 31 ++++-
.../substrait/tests/cases/aggregation_tests.rs | 38 ++++++
.../tests/cases/roundtrip_logical_plan.rs | 55 ++++++++-
.../aggregate_groupings/multiple_groupings.json | 127 +++++++++++++++++++++
5 files changed, 289 insertions(+), 10 deletions(-)
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 8c0114b90e..982a87d6d5 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::<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,
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 dec94b0422..7b6c113cce 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::<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 {
diff --git a/datafusion/substrait/tests/cases/aggregation_tests.rs b/datafusion/substrait/tests/cases/aggregation_tests.rs
index 92a41850b2..e572023f17 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 c9f874dd9b..3716b0feba 100644
--- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs
+++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs
@@ -344,10 +344,59 @@ 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 ctx = create_context().await?;
+ let proto = roundtrip_with_ctx(
+ "SELECT a, c, d, avg(b), sum(e) FROM data GROUP BY GROUPING SETS ((a, c), (a), (d), ())",
+ ctx.clone(),
)
- .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]);
+
+ let plan = from_substrait_plan(&ctx.state(), &proto).await?;
+ let results = DataFrame::new(ctx.state(), plan).collect().await?;
+ datafusion::assert_batches_sorted_eq!(
+ [
+ "+---+------------+-------+-------------+-------------+",
+ "| a | c | d | avg(data.b) | sum(data.e) |",
+ "+---+------------+-------+-------------+-------------+",
+ "| | | | 3.250000 | 6442450943 |",
+ "| | | false | 2.000000 | 4294967295 |",
+ "| | | true | 4.500000 | 2147483648 |",
+ "| 1 | | | 2.000000 | 4294967295 |",
+ "| 1 | 2020-01-01 | | 2.000000 | 4294967295 |",
+ "| 3 | | | 4.500000 | 2147483648 |",
+ "| 3 | 2020-01-01 | | 4.500000 | 2147483648 |",
+ "+---+------------+-------+-------------+-------------+",
+ ],
+ &results
+ );
+
+ 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 0000000000..3779ee32c6
--- /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"
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org