From 5563a08421943287d9a414c0d80af9613397f40f Mon Sep 17 00:00:00 2001 From: jonasHanhan Date: Wed, 4 Nov 2026 10:25:40 +0800 Subject: [PATCH 1/2] fix(substrait): dedupe repeated GROUP BY keys in the producer DataFusion's Aggregate schema keeps each grouping key once, but the Substrait producer was emitting a repeated expression_reference for GROUP BY k, k. The consumer then expanded those refs into extra columns, so projections built from the DataFusion plan indexed the wrong fields and a round trip swapped the grouping and measure columns. Skip duplicate grouping expressions within each grouping so the AggregateRel output layout matches the DataFusion schema. Closes #27058. --- .../producer/rel/aggregate_rel.rs | 11 +++ .../tests/cases/roundtrip_logical_plan.rs | 82 ++++++++++++++++++- 2 files changed, 92 insertions(+), 1 deletion(-) 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 f3064b9477383..60d729e1e94e3 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs @@ -167,8 +167,19 @@ pub fn parse_flat_grouping_exprs( ) -> datafusion::common::Result { let mut expression_references = Vec::with_capacity(exprs.len()); let mut grouping_expressions = Vec::with_capacity(exprs.len()); + // DataFusion's Aggregate schema keeps each grouping key once + // (`grouping_set_to_exprlist`). Emit the same uniqueness here so the + // Substrait AggregateRel's output column count matches that schema; + // otherwise a projection built from the DataFusion plan indexes the + // measure columns as if the repeated keys were absent, and a round trip + // swaps columns (see https://github.com/apache/datafusion/issues/27058). + let mut seen: Vec<&Expr> = Vec::with_capacity(exprs.len()); for e in exprs { + if seen.contains(&e) { + continue; + } + seen.push(e); let rex = producer.handle_expr(e, schema)?; grouping_expressions.push(rex.clone()); let reference = ref_group_exprs.iter().position(|existing| existing == &rex); diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 4285b7a708fe7..35b1e3c583422 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -40,7 +40,7 @@ use std::hash::Hash; use std::sync::Arc; use substrait::proto::extensions::simple_extension_declaration::MappingType; use substrait::proto::rel::RelType; -use substrait::proto::{Plan, Rel, plan_rel}; +use substrait::proto::{AggregateRel, Plan, Rel, plan_rel}; #[derive(Debug)] struct MockSerializerRegistry; @@ -432,6 +432,86 @@ async fn simple_aggregate() -> Result<()> { roundtrip("SELECT a, sum(b) FROM data GROUP BY a").await } +#[tokio::test] +async fn aggregate_duplicate_group_by_key_roundtrip() -> Result<()> { + // https://github.com/apache/datafusion/issues/27058 + // Repeated GROUP BY keys are legal SQL and do not change the result. The + // eliminate_duplicated_expr optimizer removes them, so use the unoptimized + // plan to exercise the producer path that used to emit duplicate grouping + // expression references and shift projection column indices on round trip. + let ctx = create_context().await?; + let sql = "SELECT a, sum(b) AS s FROM data GROUP BY a, a ORDER BY a"; + let plan = ctx.sql(sql).await?.into_unoptimized_plan(); + + let LogicalPlan::Sort(sort) = &plan else { + panic!("expected Sort, got {plan}"); + }; + let LogicalPlan::Projection(projection) = sort.input.as_ref() else { + panic!("expected Projection under Sort, got {}", sort.input); + }; + let LogicalPlan::Aggregate(agg) = projection.input.as_ref() else { + panic!( + "expected Aggregate under Projection, got {}", + projection.input + ); + }; + assert_eq!( + agg.group_expr.len(), + 2, + "fixture must keep duplicate group keys" + ); + assert_eq!( + agg.schema.fields().len(), + 2, + "Aggregate schema still has one column per unique group key plus measures" + ); + + let proto = to_substrait_plan(&plan, &ctx.state())?; + let aggregate = find_aggregate_rel(&proto).expect("plan should contain AggregateRel"); + assert_eq!( + aggregate.grouping_expressions.len(), + 1, + "producer should declare the repeated key once" + ); + assert_eq!( + aggregate.groupings[0].expression_references, + vec![0], + "producer should not repeat the grouping expression reference" + ); + + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + let expected = DataFrame::new(ctx.state(), plan.clone()).collect().await?; + let actual = DataFrame::new(ctx.state(), plan2).collect().await?; + assert_eq!( + format!("{expected:?}"), + format!("{actual:?}"), + "round trip should preserve column order and values" + ); + + // Optimized plans already drop the duplicate key via eliminate_duplicated_expr. + roundtrip("SELECT a, sum(b) AS s FROM data GROUP BY a, a").await?; + Ok(()) +} + +fn find_aggregate_rel(plan: &Plan) -> Option<&AggregateRel> { + fn walk(rel: &Rel) -> Option<&AggregateRel> { + match rel.rel_type.as_ref()? { + RelType::Aggregate(a) => Some(a.as_ref()), + RelType::Project(p) => walk(p.input.as_ref()?), + RelType::Sort(s) => walk(s.input.as_ref()?), + RelType::Filter(f) => walk(f.input.as_ref()?), + RelType::Fetch(f) => walk(f.input.as_ref()?), + _ => None, + } + } + plan.relations + .iter() + .find_map(|pr| match pr.rel_type.as_ref()? { + plan_rel::RelType::Root(r) => walk(r.input.as_ref()?), + _ => None, + }) +} + #[tokio::test] async fn aggregate_distinct_with_having() -> Result<()> { let ctx = create_context_without_single_distinct_to_group_by().await?; From 8193c563303ea9e7f4a2dc3fcb0b3dc84e9b4455 Mon Sep 17 00:00:00 2001 From: jonasHanhan Date: Fri, 6 Nov 2026 10:15:32 +0800 Subject: [PATCH 2/2] ci: retrigger Rust fork workflow