alexandrefimov commented on code in PR #25527:
URL: https://github.com/apache/datafusion/pull/25527#discussion_r4375642909
##########
datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs:
##########
@@ -160,11 +175,19 @@ fn reorder_grouping_set_output(
measure_count
);
}
+ 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)?
Review Comment:
Previously serialized DataFusion plans silently return different results
here. A plan exported on main at `350ca765` for `GROUPING SETS ((a), (b))` with
`GROUPING(a), GROUPING(b)` returns `(1, NULL, 0, 0)` instead of `(1, NULL, 0,
1)` on this branch. Its saved expressions still expect a bitmask.
Could we add an old-plan regression test and decide how to handle this
compatibility break?
##########
datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs:
##########
@@ -39,35 +47,100 @@ pub fn from_aggregate(
.iter()
.map(|e| to_substrait_agg_measure(producer, e, agg.input.schema()))
.collect::<datafusion::common::Result<Vec<_>>>()?;
- let common = (groupings.len() > 1)
- .then(|| grouping_set_output_mapping(grouping_expressions.len(),
measures.len()));
-
- Ok(Box::new(Rel {
+ 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<Rel>,
+ grouping_count: usize,
+ measure_count: usize,
+) -> datafusion::common::Result<Box<Rel>> {
+ 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. 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(
Review Comment:
This adds an unresolved `equal` reference even to `SELECT a, b FROM t GROUP
BY GROUPING SETS ((a), (b))`. Java/Isthmus 0.107.0 accepts main's protobuf plan
but rejects this branch's plan because URN anchor `4294967295` is not
registered.
Could we register the generated comparison correctly and test the export
independently of a DataFusion round trip?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
