namanjain24-sudo opened a new pull request, #25527: URL: https://github.com/apache/datafusion/pull/25527
## Which issue does this PR close? - Closes #25208. ## Rationale for this change Substrait ends an `AggregateRel` with more than one grouping set with an extra `i32` whose value is "the zero-based index of the grouping set that yielded the record" ([Aggregate Operation]). DataFusion ends the same aggregate with `__grouping_id`, which packs two things: a bitmask with a bit set for every grouping column the set leaves out, counting from the last column, and an ordinal that separates repeated sets. The consumer and the producer both treated that column as Substrait's index, so: - a consumed plan returned the bitmask where the spec asks for the index, typed `UInt8` rather than `Int32`; - a produced plan wrote the bitmask into the column another engine, such as substrait-java, reads as the index. The two coincide for some lists of sets, which is why this went unnoticed: for `(a, b)` then `(a)` both are 0 then 1. For `(a)` then `(b)` the bitmask is 1 then 2 while the index is 0 then 1. [Aggregate Operation]: https://substrait.io/relations/logical_relations/#aggregate-operation ## What changes are included in this PR? Both values identify the set a row came from, and every set has its own `__grouping_id`, so each side can be written as a map of the other. The maps live in a new `logical_plan::grouping_set` module shared by the two directions: - The **consumer** projects `CASE WHEN __grouping_id = <id of set 0> THEN 0 ... ELSE <last index> END`, so the plan ends with the index, as a required `Int32`. - The **producer** leaves the `AggregateRel` holding what the spec defines and projects `__grouping_id` back from the index above it, which is also where the reordering to DataFusion's `[groups, grouping_id, measures]` now happens. A plan DataFusion writes therefore reads correctly in another engine, and still round trips here. The last set is the `ELSE` arm rather than a `WHEN`: the values are exhaustive, and an `ELSE` keeps the column non-nullable, which Substrait requires of the index and DataFusion of `__grouping_id`. Two details the maps have to respect: - **Set order.** Substrait has no `ROLLUP` or `CUBE`, so the producer writes both as a list of sets, reversed for `ROLLUP` and as a powerset for `CUBE`. The index follows that list, so the expansion is now one function used both to write the groupings and to compute the ids. - **Repeated sets.** `GROUPING SETS ((a), (a))` is two sets with two indexes, which DataFusion separates with the ordinal packed above the bitmask, so the map stays one-to-one. A single grouping set has no index column and is untouched, as is `SELECT DISTINCT`. ## What is the testing strategy for this PR? Consumer, in `aggregation_tests.rs`: - `multiple_grouping_sets_emit_the_set_index` is the issue's plan: sets `(a)` then `(b)`, where the index differs from the bitmask. It checks the column is a required `Int32` holding 0 and 1, where `main` gives a `UInt8` holding 1 and 2. - `duplicate_grouping_sets_are_separate_indexes`: the same set twice gets index 0 and 1, with its rows once per index. Round trip, in `roundtrip_logical_plan.rs`: - `aggregate_grouping_sets_keep_grouping_function`: `GROUPING(a)` reads `__grouping_id`, so this only holds if the index is mapped back to it. The sets are `(a)`, `(c)`, `(a, c)`, chosen so the index and the bitmask differ; with `(a, c)` first the test would pass either way. - `aggregate_duplicate_grouping_sets`: `GROUPING SETS ((a), (a), ())` keeps each occurrence's rows. - `aggregate_grouping_sets_wider_grouping_id`: nine grouping columns, so `__grouping_id` is a `UInt16` and the map has to carry the wider literal. - `aggregate_grouping_sets` now asserts that the `AggregateRel` no longer remaps its output and that the projection above it carries the map and the mapping. Each half was reverted on its own to check the tests pin it: - without the consumer change, `multiple_grouping_sets_emit_the_set_index`, `aggregate_duplicate_grouping_sets` and `aggregate_grouping_sets_keep_grouping_function` fail; - without the producer change, the last two fail, together with `aggregate_grouping_sets`; - with the `UInt16` arm of the literal narrowed to `UInt8`, only `aggregate_grouping_sets_wider_grouping_id` fails. `cargo test -p datafusion-substrait` passes (58 unit, 215 integration with the 6 already ignored, 3 doc tests), as do `cargo fmt --all -- --check`, `cargo clippy -p datafusion-substrait --all-targets --features physical -- -D warnings` and `cargo xtask ci step test substrait`. A Substrait round trip of `aggregate.slt`, which is not in that job, reports the same 10 pre-existing failures as `main`. ## Are there any user-facing changes? Plans consumed from Substrait end a multi-set aggregate with the grouping set index, a required `Int32`, instead of DataFusion's `__grouping_id`. Plans produced for a multi-set aggregate now carry a projection above the `AggregateRel` that maps that index back to `__grouping_id`, so the `AggregateRel` itself holds the column the spec describes. No Rust API changes. -- 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]
