namanjain24-sudo commented on code in PR #25527:
URL: https://github.com/apache/datafusion/pull/25527#discussion_r4210181193
##########
datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs:
##########
@@ -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)?
Review Comment:
Fixed: added unique_grouping_set_index_name (shared helper in
grouping_set.rs) that falls back to a free name when the schema already has one
named grouping_set_index; used by both the producer and the consumer. Added a
regression with a real column of that name.
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -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<()> {
Review Comment:
Added: 8 distinct expressions with one set repeated, isolating that the
UInt16 widening follows from the duplicate ordinal bit alone, not the column
count.
--
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]
