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