kosiew commented on code in PR #25527:
URL: https://github.com/apache/datafusion/pull/25527#discussion_r4246267966
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -458,6 +475,315 @@ 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::<Vec<_>>()
+ .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(())
+}
+
+/// 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::<Vec<_>>()
+ .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::<Vec<_>>()
+ .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 + <offset>` 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::<Vec<_>>()
+ .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(())
+}
+
+/// 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::<Vec<_>>()
+ .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::<Vec<_>>()
+ .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!(
Review Comment:
The overflow rejection tests currently exercise only the producer, so the
consumer's error path is not covered. Could you add consumer tests that
deserialize grouping sets with overflowing duplicate ordinals at 63 and 64
columns and verify the specific representability error? The shared validation
looks correct, so this is just a coverage suggestion.
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -458,6 +475,315 @@ 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::<Vec<_>>()
+ .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(())
+}
+
+/// 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::<Vec<_>>()
+ .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"
Review Comment:
This test has nine distinct grouping expressions because the final `(a)`
adds bare `a` to `a + 0` through `a + 7`, so UInt16 is already needed without
the duplicate ordinal. Could you use bare `a` plus `a + 1` through `a + 7` and
assert that there are exactly eight distinct grouping expressions and a UInt16
grouping ID? That would isolate widening caused by the duplicate ordinal.
--
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]
