From 710f66e9e9ab652b9dd1a7c00d57ed041f9dfead Mon Sep 17 00:00:00 2001 From: naman Date: Sun, 20 Sep 2026 11:13:05 +0530 Subject: [PATCH] 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 e572023f17a92..1cedb91a1ac7c 100644 --- a/datafusion/substrait/tests/cases/aggregation_tests.rs +++ b/datafusion/substrait/tests/cases/aggregation_tests.rs @@ -20,6 +20,8 @@ #[cfg(test)] mod tests { use crate::utils::test::{add_plan_schemas_to_ctx, read_json}; + use datafusion::arrow::datatypes::DataType; + use datafusion::assert_batches_sorted_eq; use datafusion::common::Result; use datafusion::dataframe::DataFrame; use datafusion::prelude::SessionContext; @@ -69,6 +71,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 813d0ed6c3489..65aea2c646fae 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 + ] + } + ] + } + } + } + } + ] +}