Skip to content

fix: give a multi-set aggregate the grouping set index Substrait defines - #25527

Open
namanjain24-sudo wants to merge 1 commit into
apache:mainfrom
namanjain24-sudo:fix-substrait-grouping-set-index
Open

namanjain24-sudo wants to merge 1 commit into
apache:mainfrom
namanjain24-sudo:fix-substrait-grouping-set-index

Conversation

@namanjain24-sudo

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

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.

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.

Substrait ends an AggregateRel with more than one grouping set with an i32
holding the zero-based index of the set that produced the row. DataFusion
ends the same aggregate with __grouping_id, which packs a bitmask of the
columns a set leaves out with an ordinal that separates repeated sets. The
consumer mapped one onto the other, so a consumed plan returned the bitmask
where the spec asks for the index, and the producer wrote the bitmask into
the column another engine reads as the index.

Both identify the set a row came from, so each is now written as a map of
the other: the consumer projects the index, and the producer projects
__grouping_id back from the index above an AggregateRel that carries what
the spec defines.
@github-actions github-actions Bot added the substrait Changes to the substrait crate label Sep 20, 2026
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 81.86275% with 37 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.38%. Comparing base (b4a8c82) to head (710f66e).

Files with missing lines Patch % Lines
...afusion/substrait/src/logical_plan/grouping_set.rs 81.81% 12 Missing and 8 partials ⚠️
...ait/src/logical_plan/producer/rel/aggregate_rel.rs 85.52% 4 Missing and 7 partials ⚠️
...ait/src/logical_plan/consumer/rel/aggregate_rel.rs 66.66% 0 Missing and 6 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #25527      +/-   ##
==========================================
- Coverage   82.38%   82.38%   -0.01%     
==========================================
  Files        1138     1139       +1     
  Lines      434328   434457     +129     
  Branches   434328   434457     +129     
==========================================
+ Hits       357824   357918      +94     
- Misses      54872    54885      +13     
- Partials    21632    21654      +22     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

substrait Changes to the substrait crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Substrait: the grouping set column holds __grouping_id, not the set's index

2 participants