From 8d1637807487d70c2b807365686f5291e507b7d9 Mon Sep 17 00:00:00 2001 From: naman Date: Sun, 20 Sep 2026 11:13:05 +0530 Subject: [PATCH 1/4] fix: give a multi-set aggregate the grouping set index Substrait defines Substrait ends an AggregateRel with more than one grouping set with an i32 holding the zero-based index of the set that produced the row. DataFusion ends the same aggregate with __grouping_id, which packs a bitmask of the columns a set leaves out with an ordinal that separates repeated sets. The consumer mapped one onto the other, so a consumed plan returned the bitmask where the spec asks for the index, and the producer wrote the bitmask into the column another engine reads as the index. Both identify the set a row came from, so each is now written as a map of the other: the consumer projects the index, and the producer projects __grouping_id back from the index above an AggregateRel that carries what the spec defines. --- .../consumer/rel/aggregate_rel.rs | 53 +++-- .../src/logical_plan/grouping_set.rs | 204 ++++++++++++++++++ datafusion/substrait/src/logical_plan/mod.rs | 1 + .../producer/rel/aggregate_rel.rs | 193 ++++++++++------- .../tests/cases/aggregation_tests.rs | 64 ++++++ .../tests/cases/roundtrip_logical_plan.rs | 126 ++++++++++- .../duplicate_grouping_sets.json | 97 +++++++++ .../grouping_set_index.json | 108 ++++++++++ 8 files changed, 746 insertions(+), 100 deletions(-) create mode 100644 datafusion/substrait/src/logical_plan/grouping_set.rs create mode 100644 datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/duplicate_grouping_sets.json create mode 100644 datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/grouping_set_index.json 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 982a87d6d5e83..c8b48dd341552 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs @@ -17,10 +17,14 @@ use crate::logical_plan::consumer::{NameTracker, SubstraitConsumer}; use crate::logical_plan::consumer::{from_substrait_agg_func, from_substrait_sorts}; +use crate::logical_plan::grouping_set::{ + GROUPING_SET_INDEX, grouping_id_column, grouping_set_columns, grouping_set_ids, + grouping_sets_of, index_from_grouping_id, +}; 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, + Aggregate, Expr, ExprSchemable, GroupingSet, LogicalPlan, LogicalPlanBuilder, }; use substrait::proto::AggregateRel; use substrait::proto::aggregate_function::AggregationInvocation; @@ -125,33 +129,44 @@ pub async fn from_aggregate_rel( .map(|e| name_tracker.get_uniquely_named_expr(e)) .collect::, _>>()?; + let set_ids = (agg.groupings.len() > 1) + .then(|| { + grouping_set_ids( + &grouping_set_columns(&group_exprs)?, + grouping_sets_of(&group_exprs)?, + ) + }) + .transpose()?; 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) + match set_ids { + Some(set_ids) => grouping_set_output(plan, agg.measures.len(), &set_ids), + None => 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( +/// Shapes DataFusion's `[groups, grouping_id, measures]` aggregate schema into +/// the direct output Substrait gives a multi-set aggregate: +/// `[groups, measures, grouping set index]`. +/// +/// The trailing column is not DataFusion's `__grouping_id`. Substrait defines it +/// as "the zero-based index of the grouping set that yielded the record", while +/// `__grouping_id` packs a bitmask of the columns the set leaves out together +/// with an ordinal that separates repeated sets. Both identify the set that +/// produced a row, so the column is replaced here by an expression mapping one +/// to the other. +/// +/// [Aggregate Operation]: https://substrait.io/relations/logical_relations/#aggregate-operation +fn grouping_set_output( plan: LogicalPlan, measure_count: usize, + set_ids: &[u64], ) -> 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 - ); - }; + let (grouping_id_index, grouping_id) = grouping_id_column(schema)?; if grouping_id_index + measure_count + 1 != schema.fields().len() { return internal_err!( "Grouping set aggregate schema has {} fields after {}, expected {} measures", @@ -160,11 +175,15 @@ fn reorder_grouping_set_output( measure_count ); } + let grouping_id = Expr::Column(grouping_id); + let grouping_id_type = grouping_id.get_type(schema)?; + let set_index = index_from_grouping_id(&grouping_id, &grouping_id_type, set_ids)? + .alias(GROUPING_SET_INDEX); (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)))) + .chain(std::iter::once(set_index)) .collect() }; project(plan, exprs) diff --git a/datafusion/substrait/src/logical_plan/grouping_set.rs b/datafusion/substrait/src/logical_plan/grouping_set.rs new file mode 100644 index 0000000000000..1e3253572058f --- /dev/null +++ b/datafusion/substrait/src/logical_plan/grouping_set.rs @@ -0,0 +1,204 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The column a multi-set aggregate ends with, which Substrait and DataFusion +//! fill differently. +//! +//! Substrait gives an [`AggregateRel`] with more than one grouping set a +//! trailing `i32` holding "the zero-based index of the grouping set that +//! yielded the record" ([Aggregate Operation]). DataFusion ends the same +//! aggregate with `__grouping_id`, which packs a bitmask of the columns the set +//! leaves out together with an ordinal separating repeated sets. Both identify +//! the set a row came from, so each side can be written as a map of the other, +//! which is what the consumer and the producer apply. +//! +//! [`AggregateRel`]: substrait::proto::AggregateRel +//! [Aggregate Operation]: https://substrait.io/relations/logical_relations/#aggregate-operation + +use datafusion::arrow::datatypes::DataType; +use datafusion::common::{ + Column, ScalarValue, internal_datafusion_err, internal_err, not_impl_err, +}; +use datafusion::logical_expr::utils::grouping_set_to_exprlist; +use datafusion::logical_expr::{Aggregate, Case, Expr, GroupingSet, lit}; + +/// The name the grouping set index is given, which Substrait leaves to the +/// plan's root names. +pub(crate) const GROUPING_SET_INDEX: &str = "grouping_set_index"; + +/// The `__grouping_id` value DataFusion gives each grouping set, in the order +/// the sets are listed. +/// +/// The value is `(ordinal << group_count) | mask`: a bit is set in `mask` for +/// every grouping column the set leaves out, counting from the last column, and +/// `ordinal` counts the sets before this one holding the same columns. Both +/// parts follow from the set alone, so no two sets share a value. +pub(crate) fn grouping_set_ids( + columns: &[&Expr], + sets: &[Vec], +) -> datafusion::common::Result> { + let group_count = columns.len(); + if group_count > 64 { + return not_impl_err!( + "Grouping sets with more than 64 columns are not supported" + ); + } + + let mut ids = Vec::with_capacity(sets.len()); + let mut masks = Vec::with_capacity(sets.len()); + for set in sets { + let mut mask = 0u64; + for (position, column) in columns.iter().enumerate() { + if !set.contains(column) { + mask |= 1 << (group_count - 1 - position); + } + } + let ordinal = masks.iter().filter(|seen| **seen == mask).count() as u64; + masks.push(mask); + ids.push((ordinal << group_count) | mask); + } + Ok(ids) +} + +/// The grouping sets of an aggregate DataFusion built from `GROUPING SETS`. +pub(crate) fn grouping_sets_of( + group_exprs: &[Expr], +) -> datafusion::common::Result<&Vec>> { + let [Expr::GroupingSet(GroupingSet::GroupingSets(sets))] = group_exprs else { + return internal_err!( + "Expected a single GROUPING SETS expression, got {group_exprs:?}" + ); + }; + Ok(sets) +} + +/// The grouping columns, in the order DataFusion's aggregate schema holds them. +pub(crate) fn grouping_set_columns( + group_exprs: &[Expr], +) -> datafusion::common::Result> { + grouping_set_to_exprlist(group_exprs) +} + +/// `CASE WHEN = 0 THEN ids[0] ... ELSE ids[last] END`, mapping the +/// grouping set index to DataFusion's `__grouping_id`. +pub(crate) fn grouping_id_from_index( + index: &Expr, + grouping_id_type: &DataType, + ids: &[u64], +) -> datafusion::common::Result { + let ids = ids + .iter() + .map(|id| grouping_id_literal(*id, grouping_id_type)) + .collect::>>()?; + let Some((last, rest)) = ids.split_last() else { + return internal_err!("Grouping set aggregate has no grouping sets"); + }; + case_over( + rest.iter() + .enumerate() + .map(|(position, id)| { + Ok((index.clone().eq(lit(index_literal(position)?)), id.clone())) + }) + .collect::>>()?, + last.clone(), + ) +} + +/// `CASE WHEN = ids[0] THEN 0 ... ELSE last END`, mapping +/// DataFusion's `__grouping_id` to the grouping set index. +pub(crate) fn index_from_grouping_id( + grouping_id: &Expr, + grouping_id_type: &DataType, + ids: &[u64], +) -> datafusion::common::Result { + let Some((_, rest)) = ids.split_last() else { + return internal_err!("Grouping set aggregate has no grouping sets"); + }; + let when_then = rest + .iter() + .enumerate() + .map(|(position, id)| { + let id = grouping_id_literal(*id, grouping_id_type)?; + Ok((grouping_id.clone().eq(id), lit(index_literal(position)?))) + }) + .collect::>>()?; + case_over(when_then, lit(index_literal(rest.len())?)) +} + +/// The `CASE` both maps are written as. The last arm is the `ELSE`: the values +/// are exhaustive, and an `ELSE` keeps the result non-nullable, which is what +/// Substrait requires of the index and DataFusion of `__grouping_id`. +fn case_over( + when_then: Vec<(Expr, Expr)>, + else_expr: Expr, +) -> datafusion::common::Result { + if when_then.is_empty() { + // A single grouping set carries no index column, so both callers stop + // before reaching this. + return Ok(else_expr); + } + Ok(Expr::Case(Case { + expr: None, + when_then_expr: when_then + .into_iter() + .map(|(when, then)| (Box::new(when), Box::new(then))) + .collect(), + else_expr: Some(Box::new(else_expr)), + })) +} + +fn index_literal(position: usize) -> datafusion::common::Result { + i32::try_from(position).map_err(|_| { + internal_datafusion_err!("More grouping sets than an i32 index can hold") + }) +} + +/// A literal of the integer type [`Aggregate::grouping_id_type`] sized to the +/// number of grouping columns. +fn grouping_id_literal( + id: u64, + grouping_id_type: &DataType, +) -> datafusion::common::Result { + let value = match grouping_id_type { + DataType::UInt8 => ScalarValue::UInt8(Some(id as u8)), + DataType::UInt16 => ScalarValue::UInt16(Some(id as u16)), + DataType::UInt32 => ScalarValue::UInt32(Some(id as u32)), + DataType::UInt64 => ScalarValue::UInt64(Some(id)), + other => { + return internal_err!( + "Unexpected {} type: {other}", + Aggregate::INTERNAL_GROUPING_ID + ); + } + }; + Ok(lit(value)) +} + +/// The column DataFusion's aggregate schema holds `__grouping_id` in. +pub(crate) fn grouping_id_column( + schema: &datafusion::common::DFSchema, +) -> datafusion::common::Result<(usize, Column)> { + let Some(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 + ); + }; + Ok((index, Column::from(schema.qualified_field(index)))) +} diff --git a/datafusion/substrait/src/logical_plan/mod.rs b/datafusion/substrait/src/logical_plan/mod.rs index 6f8b8e493f529..5034a024e872a 100644 --- a/datafusion/substrait/src/logical_plan/mod.rs +++ b/datafusion/substrait/src/logical_plan/mod.rs @@ -16,4 +16,5 @@ // under the License. pub mod consumer; +pub(crate) mod grouping_set; pub mod producer; 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 7b6c113ccec0f..a9b436d5ed8d5 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs @@ -15,17 +15,25 @@ // specific language governing permissions and limitations // under the License. +use crate::logical_plan::grouping_set::{ + GROUPING_SET_INDEX, grouping_id_column, grouping_id_from_index, grouping_set_columns, + grouping_set_ids, +}; use crate::logical_plan::producer::{ SubstraitProducer, from_aggregate_function, substrait_field_ref, }; -use datafusion::common::{DFSchemaRef, internal_err, not_impl_err}; +use datafusion::arrow::datatypes::{DataType, Field}; +use datafusion::common::{Column, DFSchema, DFSchemaRef, internal_err, not_impl_err}; use datafusion::logical_expr::expr::Alias; use datafusion::logical_expr::utils::powerset; use datafusion::logical_expr::{Aggregate, Distinct, Expr, GroupingSet}; +use std::sync::Arc; use substrait::proto::aggregate_rel::{Grouping, Measure}; use substrait::proto::rel::RelType; use substrait::proto::rel_common::EmitKind; -use substrait::proto::{AggregateRel, Expression, Rel, RelCommon, rel_common}; +use substrait::proto::{ + AggregateRel, Expression, ProjectRel, Rel, RelCommon, rel_common, +}; pub fn from_aggregate( producer: &mut impl SubstraitProducer, @@ -39,35 +47,93 @@ 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 { + let is_grouping_set = groupings.len() > 1; + let grouping_count = grouping_expressions.len(); + let measure_count = measures.len(); + let aggregate = Box::new(Rel { rel_type: Some(RelType::Aggregate(Box::new(AggregateRel { - common, + common: None, input: Some(input), grouping_expressions, groupings, measures, advanced_extension: None, }))), - })) + }); + + if is_grouping_set { + grouping_set_projection(producer, agg, aggregate, grouping_count, measure_count) + } else { + Ok(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; +/// Puts a multi-set aggregate back into DataFusion's +/// `[groups, grouping_id, measures]` schema. +/// +/// Substrait's own output is `[groups, measures, grouping set index]`, so +/// besides the reordering the trailing column has to be mapped back to +/// `__grouping_id`; see [`crate::logical_plan::grouping_set`]. That map is a +/// projected expression, which leaves the `AggregateRel` itself holding the +/// index the spec defines, for a consumer that reads it. +fn grouping_set_projection( + producer: &mut impl SubstraitProducer, + agg: &Aggregate, + aggregate: Box, + grouping_count: usize, + measure_count: usize, +) -> datafusion::common::Result> { + let schema = agg.schema.as_ref(); + let (grouping_id_index, _) = grouping_id_column(schema)?; + if grouping_id_index != grouping_count { + return internal_err!( + "Aggregate has {grouping_id_index} grouping columns but {grouping_count} grouping expressions were written" + ); + } + let grouping_id_type = schema.field(grouping_id_index).data_type(); + let ids = grouping_set_ids( + &grouping_set_columns(&agg.group_expr)?, + &expand_grouping_sets(&agg.group_expr)?, + )?; + + // The aggregate's output as Substrait orders it, which is what the + // expression below is written against. + let index_field = Field::new(GROUPING_SET_INDEX, DataType::Int32, false); + let substrait_output = DFSchema::from_unqualified_fields( + (0..grouping_id_index) + .chain(grouping_id_index + 1..schema.fields().len()) + .map(|index| Arc::clone(schema.field(index))) + .chain(std::iter::once(Arc::new(index_field))) + .collect(), + schema.metadata().clone(), + )?; + let index = Expr::Column(Column::from_name(GROUPING_SET_INDEX)); + let expression = producer.handle_expr( + &grouping_id_from_index(&index, grouping_id_type, &ids)?, + &Arc::new(substrait_output), + )?; + + // A Substrait project emits its input's fields followed by its + // expressions, so the map sits one past the aggregate's own output. + let index_of_map = grouping_count + measure_count + 1; let output_mapping = (0..grouping_count) - .chain(std::iter::once(grouping_id_index)) - .chain(grouping_count..grouping_id_index) + .chain(std::iter::once(index_of_map)) + .chain(grouping_count..grouping_count + measure_count) .map(|index| index as i32) .collect(); - RelCommon { - emit_kind: Some(EmitKind::Emit(rel_common::Emit { output_mapping })), - hint: None, - advanced_extension: None, - } + + Ok(Box::new(Rel { + rel_type: Some(RelType::Project(Box::new(ProjectRel { + common: Some(RelCommon { + emit_kind: Some(EmitKind::Emit(rel_common::Emit { output_mapping })), + hint: None, + advanced_extension: None, + }), + input: Some(aggregate), + expressions: vec![expression], + advanced_extension: None, + }))), + })) } pub fn from_distinct( @@ -108,68 +174,37 @@ pub fn to_substrait_groupings( schema: &DFSchemaRef, ) -> datafusion::common::Result<(Vec, Vec)> { let mut ref_group_exprs = vec![]; - let groupings = match exprs.len() { - 1 => match &exprs[0] { - Expr::GroupingSet(gs) => match gs { - GroupingSet::Cube(set) => { - // Generate power set of grouping expressions - let cube_sets = powerset(set)?; - cube_sets - .iter() - .map(|set| { - parse_flat_grouping_exprs( - producer, - &set.iter().map(|v| (*v).clone()).collect::>(), - schema, - &mut ref_group_exprs, - ) - }) - .collect::>>() - } - GroupingSet::GroupingSets(sets) => sets - .iter() - .map(|set| { - parse_flat_grouping_exprs( - producer, - set, - schema, - &mut ref_group_exprs, - ) - }) - .collect::>>(), - GroupingSet::Rollup(set) => { - let mut sets: Vec> = vec![vec![]]; - for i in 0..set.len() { - sets.push(set[..=i].to_vec()); - } - sets.iter() - .rev() - .map(|set| { - parse_flat_grouping_exprs( - producer, - set, - schema, - &mut ref_group_exprs, - ) - }) - .collect::>>() + let groupings = expand_grouping_sets(exprs)? + .iter() + .map(|set| parse_flat_grouping_exprs(producer, set, schema, &mut ref_group_exprs)) + .collect::>>()?; + Ok((ref_group_exprs, groupings)) +} + +/// The grouping sets an aggregate is written as, in the order they are emitted. +/// +/// Substrait has no `ROLLUP` or `CUBE`, so both become a list of sets, and the +/// grouping set index of a row follows this order. +fn expand_grouping_sets(exprs: &[Expr]) -> datafusion::common::Result>> { + let sets = match exprs { + [Expr::GroupingSet(gs)] => match gs { + // Generate power set of grouping expressions + GroupingSet::Cube(set) => powerset(set)? + .into_iter() + .map(|set| set.into_iter().cloned().collect()) + .collect(), + GroupingSet::GroupingSets(sets) => sets.clone(), + GroupingSet::Rollup(set) => { + let mut sets: Vec> = vec![vec![]]; + for i in 0..set.len() { + sets.push(set[..=i].to_vec()); } - }, - _ => Ok(vec![parse_flat_grouping_exprs( - producer, - exprs, - schema, - &mut ref_group_exprs, - )?]), + sets.into_iter().rev().collect() + } }, - _ => Ok(vec![parse_flat_grouping_exprs( - producer, - exprs, - schema, - &mut ref_group_exprs, - )?]), - }?; - Ok((ref_group_exprs, groupings)) + exprs => vec![exprs.to_vec()], + }; + Ok(sets) } pub fn parse_flat_grouping_exprs( diff --git a/datafusion/substrait/tests/cases/aggregation_tests.rs b/datafusion/substrait/tests/cases/aggregation_tests.rs index 8af4cb7fa5c39..a426edfc85403 100644 --- a/datafusion/substrait/tests/cases/aggregation_tests.rs +++ b/datafusion/substrait/tests/cases/aggregation_tests.rs @@ -22,6 +22,8 @@ mod tests { use crate::utils::test::{add_plan_schemas_to_ctx, read_json}; use datafusion::arrow::array::record_batch; use datafusion::arrow::datatypes as arrow_schema; + use datafusion::arrow::datatypes::DataType; + use datafusion::assert_batches_sorted_eq; use datafusion::common::{Result, ScalarValue, TableReference}; use datafusion::dataframe::DataFrame; use datafusion::prelude::SessionContext; @@ -260,6 +262,68 @@ mod tests { Ok(()) } + /// Substrait ends a multi-set aggregate with the zero-based index of the + /// grouping set that produced the row, which is not DataFusion's + /// `__grouping_id` bitmask. The sets below are `(a)` then `(b)`, so the two + /// differ: the bitmask is 1 then 2, while the index is 0 then 1. + #[tokio::test] + async fn multiple_grouping_sets_emit_the_set_index() -> Result<()> { + let proto_plan = read_json( + "tests/testdata/test_plans/aggregate_groupings/grouping_set_index.json", + ); + let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?; + let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + + let field = plan.schema().field_with_unqualified_name("grouping_set")?; + assert_eq!(field.data_type(), &DataType::Int32); + assert!(!field.is_nullable(), "the set index is always known"); + + let results = DataFrame::new(ctx.state(), plan).collect().await?; + assert_batches_sorted_eq!( + [ + "+---+----+--------------+", + "| a | b | grouping_set |", + "+---+----+--------------+", + "| | 10 | 1 |", + "| | 20 | 1 |", + "| 1 | | 0 |", + "| 2 | | 0 |", + "+---+----+--------------+", + ], + &results + ); + + Ok(()) + } + + /// The same grouping set listed twice is two sets, so each occurrence gets + /// its own index and its own copy of the rows. + #[tokio::test] + async fn duplicate_grouping_sets_are_separate_indexes() -> Result<()> { + let proto_plan = read_json( + "tests/testdata/test_plans/aggregate_groupings/duplicate_grouping_sets.json", + ); + let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?; + let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + + let results = DataFrame::new(ctx.state(), plan).collect().await?; + assert_batches_sorted_eq!( + [ + "+---+--------------+", + "| a | grouping_set |", + "+---+--------------+", + "| 1 | 0 |", + "| 1 | 1 |", + "| 2 | 0 |", + "| 2 | 1 |", + "+---+--------------+", + ], + &results + ); + + Ok(()) + } + #[tokio::test] async fn multiple_grouping_sets_follow_substrait_output_order() -> Result<()> { let proto_plan = read_json( diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 6d783af16e461..61e7ef6d0e4e2 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -17,6 +17,7 @@ use crate::utils::test::read_json; use datafusion::arrow::array::ArrayRef; +use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::config::Dialect; use datafusion::functions_nested::map::map; use datafusion::logical_expr::{ @@ -416,7 +417,15 @@ async fn aggregate_grouping_sets() -> Result<()> { 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 + // A multi-set aggregate is followed by the projection that maps Substrait's + // grouping set index back to DataFusion's `__grouping_id`. + let Some(RelType::Project(grouping_set_map)) = + &project.input.as_ref().unwrap().rel_type + else { + panic!("expected the grouping set projection"); + }; + let Some(RelType::Aggregate(aggregate)) = + &grouping_set_map.input.as_ref().unwrap().rel_type else { panic!("expected aggregate relation"); }; @@ -426,15 +435,23 @@ async fn aggregate_grouping_sets() -> Result<()> { 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 + // The aggregate keeps Substrait's own output, which ends with the index + assert!( + aggregate.common.is_none(), + "the aggregate should not remap its output" + ); + assert_eq!(grouping_set_map.expressions.len(), 1); + let output_mapping = match grouping_set_map .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"), + _ => panic!("expected the grouping set projection output mapping"), }; - assert_eq!(output_mapping, &[0, 1, 2, 5, 3, 4]); + // 3 grouping columns, the map at 6 (after the aggregate's 6 fields), then + // the 2 measures + assert_eq!(output_mapping, &[0, 1, 2, 6, 3, 4]); let plan = from_substrait_plan(&ctx.state(), &proto).await?; let results = DataFrame::new(ctx.state(), plan).collect().await?; @@ -458,6 +475,107 @@ async fn aggregate_grouping_sets() -> Result<()> { Ok(()) } +/// `GROUPING()` reads the aggregate's `__grouping_id`, so it only survives a +/// round trip if the grouping set index Substrait carries is mapped back to it. +#[tokio::test] +async fn aggregate_grouping_sets_keep_grouping_function() -> Result<()> { + let ctx = create_context().await?; + // `(a)` first, then `(c)`: the bitmask is 1 then 2 while the index is 0 then + // 1, so a test that listed `(a, c)` first would pass either way + let sql = "SELECT a, c, GROUPING(a) AS ga, GROUPING(c) AS gc, avg(b) \ + FROM data GROUP BY GROUPING SETS ((a), (c), (a, c)) ORDER BY a, c, ga, gc"; + let plan = ctx.sql(sql).await?.into_optimized_plan()?; + let expected = DataFrame::new(ctx.state(), plan.clone()).collect().await?; + + let proto = to_substrait_plan(&plan, &ctx.state())?; + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + let actual = DataFrame::new(ctx.state(), plan2).collect().await?; + + assert_eq!( + format!("{}", pretty_format_batches(&expected)?), + format!("{}", pretty_format_batches(&actual)?) + ); + // The values differ per set, so this is not vacuous + assert_snapshot!( + pretty_format_batches(&expected)?, + @r" + +---+------------+----+----+-------------+ + | a | c | ga | gc | avg(data.b) | + +---+------------+----+----+-------------+ + | 1 | 2020-01-01 | 0 | 0 | 2.000000 | + | 1 | | 0 | 1 | 2.000000 | + | 3 | 2020-01-01 | 0 | 0 | 4.500000 | + | 3 | | 0 | 1 | 4.500000 | + | | 2020-01-01 | 1 | 0 | 3.250000 | + +---+------------+----+----+-------------+ + " + ); + Ok(()) +} + +/// The same set listed twice is two sets in Substrait, and DataFusion separates +/// them with an ordinal packed above the `__grouping_id` bitmask. +#[tokio::test] +async fn aggregate_duplicate_grouping_sets() -> Result<()> { + let ctx = create_context().await?; + let sql = "SELECT a, GROUPING(a) AS ga, avg(b) \ + FROM data GROUP BY GROUPING SETS ((a), (a), ()) ORDER BY a, ga"; + let plan = ctx.sql(sql).await?.into_optimized_plan()?; + let expected = DataFrame::new(ctx.state(), plan.clone()).collect().await?; + + let proto = to_substrait_plan(&plan, &ctx.state())?; + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + let actual = DataFrame::new(ctx.state(), plan2).collect().await?; + + assert_eq!( + format!("{}", pretty_format_batches(&expected)?), + format!("{}", pretty_format_batches(&actual)?) + ); + // Each occurrence of `(a)` keeps its own rows + assert_snapshot!( + pretty_format_batches(&expected)?, + @r" + +---+----+-------------+ + | a | ga | avg(data.b) | + +---+----+-------------+ + | 1 | 0 | 2.000000 | + | 1 | 0 | 2.000000 | + | 3 | 0 | 4.500000 | + | 3 | 0 | 4.500000 | + | | 1 | 3.250000 | + +---+----+-------------+ + " + ); + Ok(()) +} + +/// With more than eight grouping columns `__grouping_id` is a `UInt16` rather +/// than a `UInt8`, so the map back to it has to carry the wider literal. +#[tokio::test] +async fn aggregate_grouping_sets_wider_grouping_id() -> Result<()> { + let ctx = create_context().await?; + let columns = (0..9) + .map(|offset| format!("a + {offset}")) + .collect::>() + .join(", "); + let sql = format!( + "SELECT a, GROUPING(a) AS ga, avg(b) FROM data \ + GROUP BY GROUPING SETS (({columns}), (a), ()) ORDER BY a, ga" + ); + let plan = ctx.sql(&sql).await?.into_optimized_plan()?; + let expected = DataFrame::new(ctx.state(), plan.clone()).collect().await?; + + let proto = to_substrait_plan(&plan, &ctx.state())?; + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + let actual = DataFrame::new(ctx.state(), plan2).collect().await?; + + assert_eq!( + format!("{}", pretty_format_batches(&expected)?), + format!("{}", pretty_format_batches(&actual)?) + ); + Ok(()) +} + #[tokio::test] async fn aggregate_grouping_rollup() -> Result<()> { let plan = generate_plan_from_sql( diff --git a/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/duplicate_grouping_sets.json b/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/duplicate_grouping_sets.json new file mode 100644 index 0000000000000..d103e1371d213 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/duplicate_grouping_sets.json @@ -0,0 +1,97 @@ +{ + "relations": [ + { + "root": { + "names": [ + "a", + "grouping_set" + ], + "input": { + "aggregate": { + "input": { + "read": { + "baseSchema": { + "names": [ + "a", + "b" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "virtualTable": { + "expressions": [ + { + "fields": [ + { + "literal": { + "i64": "1" + } + }, + { + "literal": { + "i64": "10" + } + } + ] + }, + { + "fields": [ + { + "literal": { + "i64": "2" + } + }, + { + "literal": { + "i64": "20" + } + } + ] + } + ] + } + } + }, + "groupingExpressions": [ + { + "selection": { + "directReference": { + "structField": { + "field": 0 + } + }, + "rootReference": {} + } + } + ], + "groupings": [ + { + "expressionReferences": [ + 0 + ] + }, + { + "expressionReferences": [ + 0 + ] + } + ] + } + } + } + } + ] +} diff --git a/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/grouping_set_index.json b/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/grouping_set_index.json new file mode 100644 index 0000000000000..ccb93b723a9f0 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/grouping_set_index.json @@ -0,0 +1,108 @@ +{ + "relations": [ + { + "root": { + "names": [ + "a", + "b", + "grouping_set" + ], + "input": { + "aggregate": { + "input": { + "read": { + "baseSchema": { + "names": [ + "a", + "b" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "virtualTable": { + "expressions": [ + { + "fields": [ + { + "literal": { + "i64": "1" + } + }, + { + "literal": { + "i64": "10" + } + } + ] + }, + { + "fields": [ + { + "literal": { + "i64": "2" + } + }, + { + "literal": { + "i64": "20" + } + } + ] + } + ] + } + } + }, + "groupingExpressions": [ + { + "selection": { + "directReference": { + "structField": { + "field": 0 + } + }, + "rootReference": {} + } + }, + { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "rootReference": {} + } + } + ], + "groupings": [ + { + "expressionReferences": [ + 0 + ] + }, + { + "expressionReferences": [ + 1 + ] + } + ] + } + } + } + } + ] +} From 25e78f391d95dc0b4bf2ca55577304b8c19e4e5f Mon Sep 17 00:00:00 2001 From: naman Date: Thu, 24 Sep 2026 21:09:01 +0530 Subject: [PATCH 2/4] chore: retrigger CI after transient MinIO/quay.io registry outage From 95dc41979cf1e010a0347a965644580072550213 Mon Sep 17 00:00:00 2001 From: namanjain24-sudo <180642416+namanjain24-sudo@users.noreply.github.com> Date: Wed, 7 Oct 2026 23:24:10 +0530 Subject: [PATCH 3/4] fix: handle the three grouping-set-index edge cases review flagged - grouping_set_ids: at exactly 64 grouping columns the mask alone fills the u64, so `ordinal << 64` panicked even when ordinal was 0. Now the zero-ordinal case at 64 columns skips the shift, and a repeated set at 64 columns (which would need a 65th bit) is a clean error instead. - Producer: the grouping-set-index projection's temporary schema used DFSchema::from_unqualified_fields, dropping qualifiers and failing on two joined columns that only collide once reduced to a bare name. Built from qualified field specs instead. - Consumer and producer: both used the fixed literal name "grouping_set_index" for their synthetic column, which could collide with a real column of that name. Added unique_grouping_set_index_name to pick a free name in the schema at hand, used by both sides. Also adds the suggested coverage for 8 distinct grouping expressions with one repeated, isolating that the UInt16 widening follows from the duplicate ordinal bit alone, not the column count. --- .../consumer/rel/aggregate_rel.rs | 10 +- .../src/logical_plan/grouping_set.rs | 116 +++++++++++++- .../producer/rel/aggregate_rel.rs | 33 ++-- .../tests/cases/roundtrip_logical_plan.rs | 151 ++++++++++++++++++ 4 files changed, 288 insertions(+), 22 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 c8b48dd341552..9dd13d85f059c 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs @@ -18,8 +18,8 @@ use crate::logical_plan::consumer::{NameTracker, SubstraitConsumer}; use crate::logical_plan::consumer::{from_substrait_agg_func, from_substrait_sorts}; use crate::logical_plan::grouping_set::{ - GROUPING_SET_INDEX, grouping_id_column, grouping_set_columns, grouping_set_ids, - grouping_sets_of, index_from_grouping_id, + grouping_id_column, grouping_set_columns, grouping_set_ids, grouping_sets_of, + index_from_grouping_id, unique_grouping_set_index_name, }; use datafusion::common::{Column, DFSchemaRef, internal_err, not_impl_err}; use datafusion::logical_expr::builder::project; @@ -177,8 +177,12 @@ fn grouping_set_output( } let grouping_id = Expr::Column(grouping_id); let grouping_id_type = grouping_id.get_type(schema)?; + // A real column already named like the preferred alias - most likely + // a user column literally called `grouping_set_index` - would + // otherwise collide with it as a duplicate unqualified field. + let index_name = unique_grouping_set_index_name(schema); let set_index = index_from_grouping_id(&grouping_id, &grouping_id_type, set_ids)? - .alias(GROUPING_SET_INDEX); + .alias(index_name); (0..grouping_id_index) .chain(grouping_id_index + 1..schema.fields().len()) diff --git a/datafusion/substrait/src/logical_plan/grouping_set.rs b/datafusion/substrait/src/logical_plan/grouping_set.rs index 1e3253572058f..402ccd2da707b 100644 --- a/datafusion/substrait/src/logical_plan/grouping_set.rs +++ b/datafusion/substrait/src/logical_plan/grouping_set.rs @@ -31,14 +31,45 @@ use datafusion::arrow::datatypes::DataType; use datafusion::common::{ - Column, ScalarValue, internal_datafusion_err, internal_err, not_impl_err, + Column, DFSchema, ScalarValue, internal_datafusion_err, internal_err, not_impl_err, }; use datafusion::logical_expr::utils::grouping_set_to_exprlist; use datafusion::logical_expr::{Aggregate, Case, Expr, GroupingSet, lit}; -/// The name the grouping set index is given, which Substrait leaves to the -/// plan's root names. -pub(crate) const GROUPING_SET_INDEX: &str = "grouping_set_index"; +/// The preferred name for the grouping set index, which Substrait leaves to +/// the plan's root names. Both the consumer and the producer only use this +/// literal name as the base candidate passed to +/// [`unique_grouping_set_index_name`]; neither hard-codes it as the name a +/// real schema is guaranteed to accept. +const GROUPING_SET_INDEX: &str = "grouping_set_index"; + +/// A name for the grouping set index column that is not already taken in +/// `schema`. +/// +/// `schema` is a real aggregate's output (the consumer) or that aggregate's +/// `DFSchema` sans `__grouping_id` (the producer), so it can legitimately +/// already contain a column named [`GROUPING_SET_INDEX`] - either a user +/// column with that literal name, or (on the producer side specifically) two +/// joined columns that only collide once reduced to this bare, unqualified +/// name. Falling back to the fixed name regardless would make the synthetic +/// column indistinguishable from that real one, or fail schema construction +/// outright with a duplicate-field error. +pub(crate) fn unique_grouping_set_index_name(schema: &DFSchema) -> String { + if schema + .index_of_column_by_name(None, GROUPING_SET_INDEX) + .is_none() + { + return GROUPING_SET_INDEX.to_string(); + } + let mut suffix = 0u32; + loop { + let candidate = format!("{GROUPING_SET_INDEX}_{suffix}"); + if schema.index_of_column_by_name(None, &candidate).is_none() { + return candidate; + } + suffix += 1; + } +} /// The `__grouping_id` value DataFusion gives each grouping set, in the order /// the sets are listed. @@ -47,6 +78,12 @@ pub(crate) const GROUPING_SET_INDEX: &str = "grouping_set_index"; /// every grouping column the set leaves out, counting from the last column, and /// `ordinal` counts the sets before this one holding the same columns. Both /// parts follow from the set alone, so no two sets share a value. +/// +/// At `group_count == 64` the mask alone already fills every bit of the `u64`, +/// leaving no room for an ordinal: `ordinal << 64` panics on the shift amount +/// even when `ordinal` is `0`, and any `ordinal != 0` cannot be represented at +/// all. The first set at that width is still representable - its ordinal is +/// always `0` - so only a repeated set at exactly 64 columns is rejected. pub(crate) fn grouping_set_ids( columns: &[&Expr], sets: &[Vec], @@ -69,7 +106,17 @@ pub(crate) fn grouping_set_ids( } let ordinal = masks.iter().filter(|seen| **seen == mask).count() as u64; masks.push(mask); - ids.push((ordinal << group_count) | mask); + let id = if group_count == 64 { + if ordinal != 0 { + return not_impl_err!( + "A grouping set with 64 columns cannot be repeated: there are no bits left for a duplicate ordinal" + ); + } + mask + } else { + (ordinal << group_count) | mask + }; + ids.push(id); } Ok(ids) } @@ -190,7 +237,7 @@ fn grouping_id_literal( /// The column DataFusion's aggregate schema holds `__grouping_id` in. pub(crate) fn grouping_id_column( - schema: &datafusion::common::DFSchema, + schema: &DFSchema, ) -> datafusion::common::Result<(usize, Column)> { let Some(index) = schema.index_of_column_by_name(None, Aggregate::INTERNAL_GROUPING_ID) @@ -202,3 +249,60 @@ pub(crate) fn grouping_id_column( }; Ok((index, Column::from(schema.qualified_field(index)))) } + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::logical_expr::col; + + fn distinct_columns(count: usize) -> Vec { + (0..count).map(|i| col(format!("c{i}"))).collect() + } + + /// A single grouping set naming every one of 64 columns has mask `0` and + /// is always the first (ordinal `0`) occurrence of that mask, so the + /// `u64` id is representable without shifting by the full 64-bit width. + #[test] + fn grouping_set_ids_at_64_columns_zero_ordinal() -> datafusion::common::Result<()> { + let exprs = distinct_columns(64); + let columns: Vec<&Expr> = exprs.iter().collect(); + let sets = vec![exprs.clone()]; + + let ids = grouping_set_ids(&columns, &sets)?; + + assert_eq!(ids, vec![0u64]); + Ok(()) + } + + /// The same 64-column set repeated needs ordinal `1`, which does not fit + /// in a `u64` alongside a full 64-bit mask. This must be a clean error, + /// not a shift-overflow panic. + #[test] + fn grouping_set_ids_rejects_repeated_set_at_64_columns() { + let exprs = distinct_columns(64); + let columns: Vec<&Expr> = exprs.iter().collect(); + let sets = vec![exprs.clone(), exprs.clone()]; + + let err = grouping_set_ids(&columns, &sets).unwrap_err(); + + assert!( + err.to_string().contains("cannot be repeated"), + "unexpected error: {err}" + ); + } + + /// 65 columns is rejected on its own, independent of any repeated set. + #[test] + fn grouping_set_ids_rejects_more_than_64_columns() { + let exprs = distinct_columns(65); + let columns: Vec<&Expr> = exprs.iter().collect(); + let sets = vec![exprs.clone()]; + + let err = grouping_set_ids(&columns, &sets).unwrap_err(); + + assert!( + err.to_string().contains("more than 64 columns"), + "unexpected error: {err}" + ); + } +} 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 a9b436d5ed8d5..e8a90ece27615 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs @@ -16,8 +16,8 @@ // under the License. use crate::logical_plan::grouping_set::{ - GROUPING_SET_INDEX, grouping_id_column, grouping_id_from_index, grouping_set_columns, - grouping_set_ids, + grouping_id_column, grouping_id_from_index, grouping_set_columns, grouping_set_ids, + unique_grouping_set_index_name, }; use crate::logical_plan::producer::{ SubstraitProducer, from_aggregate_function, substrait_field_ref, @@ -97,17 +97,24 @@ fn grouping_set_projection( )?; // The aggregate's output as Substrait orders it, which is what the - // expression below is written against. - let index_field = Field::new(GROUPING_SET_INDEX, DataType::Int32, false); - let substrait_output = DFSchema::from_unqualified_fields( - (0..grouping_id_index) - .chain(grouping_id_index + 1..schema.fields().len()) - .map(|index| Arc::clone(schema.field(index))) - .chain(std::iter::once(Arc::new(index_field))) - .collect(), - schema.metadata().clone(), - )?; - let index = Expr::Column(Column::from_name(GROUPING_SET_INDEX)); + // expression below is written against. Built from qualified field specs, + // not `DFSchema::from_unqualified_fields`, because dropping the original + // qualifiers here would make two joined columns that only share a bare + // name (e.g. `a.x` and `b.x`) collide as duplicate unqualified fields, + // even though `schema` itself keeps them distinct. + let index_name = unique_grouping_set_index_name(schema); + let index_field = Field::new(index_name.as_str(), DataType::Int32, false); + let qualified_fields = (0..grouping_id_index) + .chain(grouping_id_index + 1..schema.fields().len()) + .map(|index| { + let (qualifier, field) = schema.qualified_field(index); + (qualifier.cloned(), Arc::clone(field)) + }) + .chain(std::iter::once((None, Arc::new(index_field)))) + .collect(); + let substrait_output = + DFSchema::new_with_metadata(qualified_fields, schema.metadata().clone())?; + let index = Expr::Column(Column::from_name(index_name.as_str())); let expression = producer.handle_expr( &grouping_id_from_index(&index, grouping_id_type, &ids)?, &Arc::new(substrait_output), diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 61e7ef6d0e4e2..c4d26a35f6258 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -576,6 +576,157 @@ async fn aggregate_grouping_sets_wider_grouping_id() -> Result<()> { Ok(()) } +/// Eight distinct grouping expressions alone keep `__grouping_id` in `UInt8` +/// (8 bits), unlike [`aggregate_grouping_sets_wider_grouping_id`], which +/// widens it with a ninth *column*. Repeating one of those eight sets instead +/// needs a ninth bit for the duplicate *ordinal*, which widens it the same +/// way - this isolates that the widening follows from the ordinal bit alone. +#[tokio::test] +async fn aggregate_grouping_sets_wider_grouping_id_from_duplicate_ordinal() -> Result<()> +{ + let ctx = create_context().await?; + let columns = (0..8) + .map(|offset| format!("a + {offset}")) + .collect::>() + .join(", "); + let sql = format!( + "SELECT a, GROUPING(a) AS ga, avg(b) FROM data \ + GROUP BY GROUPING SETS (({columns}), ({columns}), (a)) ORDER BY a, ga" + ); + let plan = ctx.sql(&sql).await?.into_optimized_plan()?; + let expected = DataFrame::new(ctx.state(), plan.clone()).collect().await?; + + let proto = to_substrait_plan(&plan, &ctx.state())?; + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + let actual = DataFrame::new(ctx.state(), plan2).collect().await?; + + assert_eq!( + format!("{}", pretty_format_batches(&expected)?), + format!("{}", pretty_format_batches(&actual)?) + ); + Ok(()) +} + +/// Two joined columns that share a bare name (`d1.a` and `d2.a`) are distinct +/// fields in the aggregate's qualified schema, but collide as duplicate +/// unqualified fields if the grouping-set-index projection's temporary schema +/// drops their qualifiers. +#[tokio::test] +async fn aggregate_grouping_sets_joined_columns_same_name() -> Result<()> { + let ctx = create_context().await?; + let sql = "SELECT d1.a, d2.a, GROUPING(d1.a) AS g1, avg(d1.b) \ + FROM data AS d1 JOIN data AS d2 ON d1.c = d2.c \ + GROUP BY GROUPING SETS ((d1.a, d2.a), (d1.a), ()) \ + ORDER BY d1.a, d2.a, g1"; + let plan = ctx.sql(sql).await?.into_optimized_plan()?; + let expected = DataFrame::new(ctx.state(), plan.clone()).collect().await?; + + let proto = to_substrait_plan(&plan, &ctx.state())?; + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + let actual = DataFrame::new(ctx.state(), plan2).collect().await?; + + assert_eq!( + format!("{}", pretty_format_batches(&expected)?), + format!("{}", pretty_format_batches(&actual)?) + ); + Ok(()) +} + +/// A real column literally named `grouping_set_index` must not collide with +/// the synthetic column the consumer and producer both use that name for by +/// default; each must fall back to a different, collision-free name instead. +/// +/// The column has to be a grouping expression itself, not merely aliased in +/// an outer `SELECT`: an alias there lands in a `Projection` on top of the +/// `Aggregate`, so by the time the grouping-set-index map is built the +/// `Aggregate`'s own schema still names the field after the un-aliased +/// expression and never collides. +#[tokio::test] +async fn aggregate_grouping_sets_index_name_collision() -> Result<()> { + let ctx = create_context().await?; + ctx.sql( + "CREATE TABLE idx_collision (grouping_set_index INT, val INT) \ + AS VALUES (1, 10), (2, 20), (1, 30)", + ) + .await? + .collect() + .await?; + let sql = "SELECT grouping_set_index, GROUPING(grouping_set_index) AS ga, avg(val) \ + FROM idx_collision GROUP BY GROUPING SETS ((grouping_set_index), ()) \ + ORDER BY grouping_set_index, ga"; + let plan = ctx.sql(sql).await?.into_optimized_plan()?; + let expected = DataFrame::new(ctx.state(), plan.clone()).collect().await?; + + let proto = to_substrait_plan(&plan, &ctx.state())?; + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + let actual = DataFrame::new(ctx.state(), plan2).collect().await?; + + assert_eq!( + format!("{}", pretty_format_batches(&expected)?), + format!("{}", pretty_format_batches(&actual)?) + ); + Ok(()) +} + +/// A single grouping set naming every one of 64 columns is representable +/// (mask fills the `u64`, ordinal is always `0`), so the logical round trip +/// must succeed rather than panic on the shift that builds its +/// `__grouping_id`. +/// +/// This checks the round-tripped schema rather than executing the plan: +/// DataFusion's own physical aggregate execution +/// (`physical-plan/src/aggregates/mod.rs`) has the same unconditional +/// `ordinal << n` at `n == 64`, independent of Substrait and outside this +/// crate's scope, so `.collect()` on either plan panics regardless of this +/// fix. +#[tokio::test] +async fn aggregate_grouping_sets_64_columns() -> Result<()> { + let ctx = create_context().await?; + let columns = (0..64) + .map(|offset| format!("a + {offset}")) + .collect::>() + .join(", "); + // `()` adds no new column, so this stays at exactly 64 distinct grouping + // columns; selecting a bare `a` would not, since `a` itself never + // appears as one of the `a + ` expressions in either set. + let sql = format!( + "SELECT GROUPING(a + 0) AS g0, avg(b) FROM data \ + GROUP BY GROUPING SETS (({columns}), ())" + ); + let plan = ctx.sql(&sql).await?.into_optimized_plan()?; + + let proto = to_substrait_plan(&plan, &ctx.state())?; + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + + assert_eq!(plan.schema(), plan2.schema()); + Ok(()) +} + +/// A 64-column grouping set that is also repeated has no room left for a +/// duplicate ordinal. The producer must reject it with a clean error, not +/// panic while building the Substrait plan. +#[tokio::test] +async fn aggregate_grouping_sets_64_columns_repeated_is_rejected() -> Result<()> { + let ctx = create_context().await?; + let columns = (0..64) + .map(|offset| format!("a + {offset}")) + .collect::>() + .join(", "); + let sql = format!( + "SELECT GROUPING(a + 0) AS g0, avg(b) FROM data \ + GROUP BY GROUPING SETS (({columns}), ({columns})) ORDER BY g0" + ); + let plan = ctx.sql(&sql).await?.into_optimized_plan()?; + + let err = to_substrait_plan(&plan, &ctx.state()).unwrap_err(); + + assert!( + err.to_string().contains("cannot be repeated"), + "unexpected error: {err}" + ); + Ok(()) +} + #[tokio::test] async fn aggregate_grouping_rollup() -> Result<()> { let plan = generate_plan_from_sql( From 473ecb813eeadaafd032eab1c9838eaf9e2e244d Mon Sep 17 00:00:00 2001 From: namanjain24-sudo <180642416+namanjain24-sudo@users.noreply.github.com> Date: Fri, 9 Oct 2026 12:07:54 +0530 Subject: [PATCH 4/4] fix: check ordinal fits in the remaining bits before shifting in grouping_set_ids The group_count == 64 guard only caught the one case where the shift amount itself is out of range. Below 64 columns the shift amount is always valid, but the *value* can still overflow out of the u64 and silently lose high bits: at 63 columns a third occurrence of the same set needs ordinal 2 (0b10), and 2u64 << 63 == 0, colliding with the first occurrence's id instead of erroring. checked_shl does not catch this either, since the shift amount is in range. Now the representable range is computed from the actual bits left above the mask (64 - group_count) and checked before shifting, for every group_count rather than only 64. Adds boundary unit tests at 62 and 63 columns (last representable ordinal vs. first that overflows) and producer round-trip/rejection tests for the same boundary. --- .../src/logical_plan/grouping_set.rs | 111 ++++++++++++++++-- .../tests/cases/roundtrip_logical_plan.rs | 57 +++++++++ 2 files changed, 156 insertions(+), 12 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/grouping_set.rs b/datafusion/substrait/src/logical_plan/grouping_set.rs index 402ccd2da707b..6cd0d99eb4955 100644 --- a/datafusion/substrait/src/logical_plan/grouping_set.rs +++ b/datafusion/substrait/src/logical_plan/grouping_set.rs @@ -77,13 +77,18 @@ pub(crate) fn unique_grouping_set_index_name(schema: &DFSchema) -> String { /// The value is `(ordinal << group_count) | mask`: a bit is set in `mask` for /// every grouping column the set leaves out, counting from the last column, and /// `ordinal` counts the sets before this one holding the same columns. Both -/// parts follow from the set alone, so no two sets share a value. -/// -/// At `group_count == 64` the mask alone already fills every bit of the `u64`, -/// leaving no room for an ordinal: `ordinal << 64` panics on the shift amount -/// even when `ordinal` is `0`, and any `ordinal != 0` cannot be represented at -/// all. The first set at that width is still representable - its ordinal is -/// always `0` - so only a repeated set at exactly 64 columns is rejected. +/// parts follow from the set alone, so no two sets share a value - *as long as +/// `ordinal` fits in the `64 - group_count` bits above the mask*. `ordinal` +/// grows with how many times a set is repeated, not with `group_count`, so a +/// set repeated enough times can overflow those bits well before +/// `group_count` reaches 64 on its own: with `group_count = 63` only one bit +/// is left for the ordinal, so a third occurrence of the same set (ordinal +/// `2`, i.e. `0b10`) would need to shift that single `1` bit out of the `u64` +/// entirely, silently losing it (`2u64 << 63 == 0`, not an out-of-range +/// shift) and colliding with the first occurrence's id. Checking the shift +/// amount alone (`checked_shl`) does not catch this: `2u64.checked_shl(63)` +/// still returns `Some(0)`, because the amount is in range - only the *value* +/// overflows out of the type, which `checked_shl` does not consider. pub(crate) fn grouping_set_ids( columns: &[&Expr], sets: &[Vec], @@ -106,14 +111,26 @@ pub(crate) fn grouping_set_ids( } let ordinal = masks.iter().filter(|seen| **seen == mask).count() as u64; masks.push(mask); - let id = if group_count == 64 { - if ordinal != 0 { + let id = if ordinal == 0 { + // No shift at all, so this is always representable, including + // the `group_count == 64` case where the mask alone already + // fills every bit and no ordinal bits are available. + mask + } else { + let ordinal_bits = 64 - group_count; + // `1u64 << 64` would itself be an out-of-range shift; at + // `ordinal_bits == 64` (`group_count == 0`) every `u64` ordinal + // is representable anyway, since shifting by 0 loses nothing. + let max_ordinal = if ordinal_bits >= 64 { + u64::MAX + } else { + (1u64 << ordinal_bits) - 1 + }; + if ordinal > max_ordinal { return not_impl_err!( - "A grouping set with 64 columns cannot be repeated: there are no bits left for a duplicate ordinal" + "A grouping set with {group_count} columns cannot be repeated more than {max_ordinal} time(s): the duplicate ordinal does not fit in the {ordinal_bits} bit(s) left above the mask" ); } - mask - } else { (ordinal << group_count) | mask }; ids.push(id); @@ -305,4 +322,74 @@ mod tests { "unexpected error: {err}" ); } + + /// At 63 columns exactly one bit is left for the ordinal, so a second + /// occurrence of the same set (ordinal `1`) is representable: both ids + /// fit in the `u64` without losing any bits. + #[test] + fn grouping_set_ids_ordinal_fits_at_63_columns() -> datafusion::common::Result<()> { + let exprs = distinct_columns(63); + let columns: Vec<&Expr> = exprs.iter().collect(); + let sets = vec![exprs.clone(), exprs.clone()]; + + let ids = grouping_set_ids(&columns, &sets)?; + + assert_eq!(ids, vec![0u64, 1u64 << 63]); + Ok(()) + } + + /// A third occurrence of the same 63-column set needs ordinal `2` + /// (`0b10`), which does not fit in the single bit left above the mask: + /// `2u64 << 63` silently discards the bit instead of producing an + /// out-of-range-shift panic, so this must be a checked, clean error - not + /// a value that happens to collide with another set's id. + #[test] + fn grouping_set_ids_rejects_ordinal_overflow_at_63_columns() { + let exprs = distinct_columns(63); + let columns: Vec<&Expr> = exprs.iter().collect(); + let sets = vec![exprs.clone(), exprs.clone(), exprs.clone()]; + + let err = grouping_set_ids(&columns, &sets).unwrap_err(); + + assert!( + err.to_string().contains("cannot be repeated"), + "unexpected error: {err}" + ); + } + + /// At 62 columns two bits are left for the ordinal, so up to four + /// occurrences (ordinals `0..=3`) of the same set are representable. + #[test] + fn grouping_set_ids_ordinal_fits_at_62_columns() -> datafusion::common::Result<()> { + let exprs = distinct_columns(62); + let columns: Vec<&Expr> = exprs.iter().collect(); + let sets = vec![exprs.clone(), exprs.clone(), exprs.clone(), exprs.clone()]; + + let ids = grouping_set_ids(&columns, &sets)?; + + assert_eq!(ids, vec![0u64, 1u64 << 62, 2u64 << 62, 3u64 << 62]); + Ok(()) + } + + /// A fifth occurrence at 62 columns needs ordinal `4` (`0b100`), which + /// overflows the two bits available. + #[test] + fn grouping_set_ids_rejects_ordinal_overflow_at_62_columns() { + let exprs = distinct_columns(62); + let columns: Vec<&Expr> = exprs.iter().collect(); + let sets = vec![ + exprs.clone(), + exprs.clone(), + exprs.clone(), + exprs.clone(), + exprs.clone(), + ]; + + let err = grouping_set_ids(&columns, &sets).unwrap_err(); + + assert!( + err.to_string().contains("cannot be repeated"), + "unexpected error: {err}" + ); + } } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index c4d26a35f6258..c3c52b721528c 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -727,6 +727,63 @@ async fn aggregate_grouping_sets_64_columns_repeated_is_rejected() -> Result<()> Ok(()) } +/// At 63 columns exactly one bit is left for the duplicate ordinal, so a +/// second occurrence of the same set (ordinal `1`) is still representable. +/// Unlike the 64-column cases above, executing this one does not hit the +/// unrelated shift-amount panic in DataFusion's native aggregate execution +/// (that panic is specifically `shift by 64`; a 63-bit shift is in range), +/// so this checks real query results end to end. +#[tokio::test] +async fn aggregate_grouping_sets_63_columns_repeated_twice() -> Result<()> { + let ctx = create_context().await?; + let columns = (0..63) + .map(|offset| format!("a + {offset}")) + .collect::>() + .join(", "); + let sql = format!( + "SELECT GROUPING(a + 0) AS g0, avg(b) FROM data \ + GROUP BY GROUPING SETS (({columns}), ({columns})) ORDER BY g0" + ); + let plan = ctx.sql(&sql).await?.into_optimized_plan()?; + let expected = DataFrame::new(ctx.state(), plan.clone()).collect().await?; + + let proto = to_substrait_plan(&plan, &ctx.state())?; + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + let actual = DataFrame::new(ctx.state(), plan2).collect().await?; + + assert_eq!( + format!("{}", pretty_format_batches(&expected)?), + format!("{}", pretty_format_batches(&actual)?) + ); + Ok(()) +} + +/// A third occurrence of the same 63-column set needs ordinal `2`, which +/// does not fit in the single bit left above the mask. The producer must +/// reject it with a clean error, not silently collide it with the first +/// occurrence's id. +#[tokio::test] +async fn aggregate_grouping_sets_63_columns_repeated_thrice_is_rejected() -> Result<()> { + let ctx = create_context().await?; + let columns = (0..63) + .map(|offset| format!("a + {offset}")) + .collect::>() + .join(", "); + let sql = format!( + "SELECT GROUPING(a + 0) AS g0, avg(b) FROM data \ + GROUP BY GROUPING SETS (({columns}), ({columns}), ({columns})) ORDER BY g0" + ); + let plan = ctx.sql(&sql).await?.into_optimized_plan()?; + + let err = to_substrait_plan(&plan, &ctx.state()).unwrap_err(); + + assert!( + err.to_string().contains("cannot be repeated"), + "unexpected error: {err}" + ); + Ok(()) +} + #[tokio::test] async fn aggregate_grouping_rollup() -> Result<()> { let plan = generate_plan_from_sql(