fix: give a multi-set aggregate the grouping set index Substrait defines by namanjain24-sudo · Pull Request #25527 · apache/datafusion

namanjain24-sudo · 2026-09-20T05:43:43Z

Which issue does this PR close?

Closes Substrait: the grouping set column holds __grouping_id, not the set's index #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.

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.

codecov-commenter · 2026-09-20T06:08:43Z

Codecov Report

❌ Patch coverage is 86.44068% with 40 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.82%. Comparing base (2951e12) to head (c8a5a7d).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...afusion/substrait/src/logical_plan/grouping_set.rs	87.81%	13 Missing and 11 partials ⚠️
...ait/src/logical_plan/producer/rel/aggregate_rel.rs	87.34%	2 Missing and 8 partials ⚠️
...ait/src/logical_plan/consumer/rel/aggregate_rel.rs	68.42%	0 Missing and 6 partials ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main   #25527    +/-   ##
========================================
  Coverage   82.82%   82.82%            
========================================
  Files        1149     1150     +1     
  Lines      452179   452399   +220     
  Branches   452179   452399   +220     
========================================
+ Hits       374505   374704   +199     
+ Misses      54873    54871     -2     
- Partials    22801    22824    +23

☔ 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.

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.

namanjain24-sudo · 2026-09-24T15:24:32Z

FYI, cargo test datafusion-cli (amd64) is failing here on a MinIO image pull (quay.io 401 unauthorized on quay.io/minio/minio:RELEASE.2025-02-28T09-55-16Z), not on anything in this PR — same registry issue just hit #25529's and #25664's merge-queue runs today. Nothing else is failing.

namanjain24-sudo · 2026-09-28T08:34:52Z

Bumping this — still open for review whenever someone has bandwidth (the earlier MinIO CI failure was unrelated infra flakiness, not this PR).

kosiew

@namanjain24-sudo,

Thanks for working on this. The grouping-set index mapping looks like the right direction, but I found three edge cases that can currently fail valid plans. I also left one optional test suggestion.

kosiew · 2026-10-07T09:24:27Z

+        }
+        let ordinal = masks.iter().filter(|seen| **seen == mask).count() as u64;
+        masks.push(mask);
+        ids.push((ordinal << group_count) | mask);


This can shift by 64 when group_count == 64, which panics even when the ordinal is zero. Please handle the 64-column zero-ordinal case without shifting, reject nonzero ordinals that need more than 64 bits, and add producer and consumer boundary regressions.

Fixed: at group_count==64 the zero-ordinal case now skips the shift, and a repeated set at 64 columns returns a clean error instead of panicking. Added producer/consumer regressions plus unit tests on grouping_set_ids directly.

Also found while testing: DataFusion's own native aggregate execution (physical-plan/src/aggregates/mod.rs, around line 3216) has the exact same unconditional ordinal << n at n == 64, independent of Substrait. Out of scope here, but flagging in case it's worth its own issue.

kosiew · 2026-10-07T09:24:27Z

+        (0..grouping_id_index)
+            .chain(grouping_id_index + 1..schema.fields().len())
+            .map(|index| Arc::clone(schema.field(index)))
+            .chain(std::iter::once(Arc::new(index_field)))


from_unqualified_fields drops qualifiers, so valid joined columns with the same base name can fail serialization as duplicate fields. Please preserve the qualified identities here, or use unique positional placeholders for this temporary schema, and add a joined-column regression.

Fixed: builds the temporary schema from qualified field specs (new_with_metadata) instead of from_unqualified_fields, so joined columns keep their qualifiers. Added a joined-column regression.

kosiew · 2026-10-07T09:24:27Z

        }
+        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)?


The fixed grouping_set_index alias can collide with a real user column of the same name on both the consumer and producer paths. Please generate a collision-free synthetic name and use that same chosen name for the producer reference, with regressions for both boundaries.

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.

kosiew · 2026-10-07T09:24:27Z

+/// 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<()> {


Suggestion: it would be useful to add a roundtrip with exactly eight distinct grouping expressions and a repeated set, so widening to UInt16 is caused only by the ordinal bit. This would complement the existing duplicate-set and wider-ID tests.

Added: 8 distinct expressions with one set repeated, isolating that the UInt16 widening follows from the duplicate ordinal bit alone, not the column count.

- grouping_set_ids: at exactly 64 grouping columns the mask alone fills the u64, so `ordinal << 64` panicked even when ordinal was 0. Now the zero-ordinal case at 64 columns skips the shift, and a repeated set at 64 columns (which would need a 65th bit) is a clean error instead. - Producer: the grouping-set-index projection's temporary schema used DFSchema::from_unqualified_fields, dropping qualifiers and failing on two joined columns that only collide once reduced to a bare name. Built from qualified field specs instead. - Consumer and producer: both used the fixed literal name "grouping_set_index" for their synthetic column, which could collide with a real column of that name. Added unique_grouping_set_index_name to pick a free name in the schema at hand, used by both sides. Also adds the suggested coverage for 8 distinct grouping expressions with one repeated, isolating that the UInt16 widening follows from the duplicate ordinal bit alone, not the column count.

kosiew

@namanjain24-sudo,

Thanks for the follow-up changes. The 64-column guard, qualified schema handling, and collision-free index naming address the earlier concerns.

One correctness issue remains in grouping_set_ids: duplicate ordinals can still overflow when there are fewer than 64 grouping columns. Please add a representability check before shifting and cover the boundary cases with regression tests.

Once that is addressed, I'll be happy to take another look.

kosiew · 2026-10-09T06:23:27Z

+            }
+            mask
+        } else {
+            (ordinal << group_count) | mask


The new guard handles 64 columns, but ordinals can still overflow below that limit. With 63 columns and three identical sets, 2 << 63 wraps to 0, giving the first and third sets the same ID and breaking the reverse mapping. Could you check that the ordinal fits in the remaining 64 - group_count bits before shifting, while still allowing valid duplicates and the zero-ordinal 64-column case? Please add boundary tests and producer/consumer rejection tests for this case. checked_shl alone won't catch discarded high bits.

Fixed: computed the representable range from the actual bits left above the mask (64 - group_count) and check it before shifting, for every group_count rather than only 64. Verified by reverting: got exactly the predicted collision, [0, 1<<63, 0] for 3x at 63 columns. Added boundary unit tests (last-ok vs first-overflow at 62 and 63 columns) plus a producer round-trip at the ok boundary and a rejection test at the overflow boundary.

…ping_set_ids The group_count == 64 guard only caught the one case where the shift amount itself is out of range. Below 64 columns the shift amount is always valid, but the *value* can still overflow out of the u64 and silently lose high bits: at 63 columns a third occurrence of the same set needs ordinal 2 (0b10), and 2u64 << 63 == 0, colliding with the first occurrence's id instead of erroring. checked_shl does not catch this either, since the shift amount is in range. Now the representable range is computed from the actual bits left above the mask (64 - group_count) and checked before shifting, for every group_count rather than only 64. Adds boundary unit tests at 62 and 63 columns (last representable ordinal vs. first that overflows) and producer round-trip/rejection tests for the same boundary.

kosiew

@namanjain24-sudo,

Thanks for addressing the earlier review comments. The grouping ID overflow checks, qualified schema handling, and synthetic index naming issues look resolved. I have two non-blocking suggestions to strengthen the regression tests, but neither needs to hold up approval.

kosiew · 2026-10-12T09:00:42Z

+        .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"


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.

kosiew · 2026-10-12T09:00:42Z

+
+    let err = to_substrait_plan(&plan, &ctx.state()).unwrap_err();
+
+    assert!(


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.

alamb

This PR feels quite large / overly complicated to me for the problem it seems to be solving. Can someone explain why we need 1150 lines of code to support grouping set indexes in substrait?

How about we just unpack the __grouping_id into an index rather than create an expression?

namanjain24-sudo · 2026-10-15T19:11:35Z

Breakdown of the 1256 lines: ~600 is tests/fixtures (330 in roundtrip_logical_plan.rs, 205 in two JSON test plans, 64 in aggregation_tests.rs). Actual logic is grouping_set.rs (395 lines, shared by producer and consumer) plus net +42 producer / +23 consumer, so ~460 lines of real code, not 1150.

On the expression — yeah, it's a lookup: CASE WHEN __grouping_id = ids[0] THEN 0 ... END (and the inverse on the producer side). Can't do this as a plain bit-unpack because Substrait's index is the set's position in the aggregate's own declared list — whatever order the query wrote it in, or ROLLUP/CUBE expanded it into. That position isn't derivable from __grouping_id's (ordinal << group_count) | mask by formula; two aggregates with the same masks can order their sets differently. So recovering the index needs this aggregate's specific set list, which is what the CASE is encoding. Open to a cheaper encoding than Expr::Case if you've got one in mind, just don't think arithmetic alone gets us there.

alamb · 2026-10-15T19:39:57Z

so ~460 lines of real code, not 1150.

Since we have to maintain both code and tests so I think evaluating the entire size of the diff is the right metric, includig tests

That position isn't derivable from __grouping_id's (ordinal << group_count) | mask by formula;

Why not?

If the aggregates order the sets differently, wouldn't there be a projection to reorder the output columns correctly?

Or perhaps we could introduce a Projection to do so

alamb · 2026-10-15T19:40:17Z

Maybe you can provide some example queries

namanjain24-sudo · 2026-10-15T19:45:57Z

Fair point on the metric, happy to count tests in too.

Examples — these are two real ones out of the test suite (aggregate_grouping_sets_keep_grouping_function, aggregate_duplicate_grouping_sets), already passing in CI:

-- (a) first, then (c): __grouping_id's mask comes out 1 then 2, but the index Substrait wants is 0 then 1
SELECT a, c, GROUPING(a) AS ga, GROUPING(c) AS gc, avg(b)
FROM data GROUP BY GROUPING SETS ((a), (c), (a, c));

mask(a)=1 -> index 0, mask(c)=2 -> index 1, mask(a,c)=0 -> index 2. Reading the mask straight as the index is off on every row here. This is the actual bug: the old code wrote the bitmask into the column Substrait defines as the index.

-- (a) repeated: both copies have mask 0, only the ordinal tells them apart
SELECT a, GROUPING(a) AS ga, avg(b)
FROM data GROUP BY GROUPING SETS ((a), (a), ());

ids work out to [0, 2, 1] for indices [0, 1, 2] (1 grouping column, so the second (a)'s ordinal lands in bit 1: (1<<1)|0 = 2). Two different indices, same mask - you need the ordinal and the aggregate's actual set list to tell them apart, not just the bits.

On the projection: there already is one. grouping_set_output does exactly that - project(plan, exprs), swapping __grouping_id for index_from_grouping_id(...), a CASE built over this aggregate's own (mask, ordinal) list. So it's not projection vs. no projection, it's whether that projection's expression has to be this per-aggregate CASE table or if there's a closed form I'm missing - genuinely open to that, but the two examples above are why I don't think mask+ordinal alone gets you there without the list.

…'s overflow path aggregate_grouping_sets_wider_grouping_id_from_duplicate_ordinal claimed to show that repeating a set, rather than adding a column, is what widens __grouping_id. It did not: the eight `a + <offset>` expressions plus the trailing bare `a` are nine distinct grouping expressions, so UInt16 was already forced by the column count and the duplicate ordinal proved nothing. Use bare `a` plus `a + 1` through `a + 7`, so the trailing `(a)` reuses one of the eight, and assert both shapes: eight columns without the repeat keep UInt8, and the same eight with it widen to UInt16. The overflow rejection tests only covered the producer, so the consumer ran the same representability check untested. Add two tests that produce a representable plan, edit the AggregateRel groupings into a shape the producer refuses to emit but another producer could send, and assert the consumer's specific error at 63 and 64 columns.

namanjain24-sudo · 2026-10-22T15:22:04Z

@kosiew both addressed.

Test isolation — you were right, and it was worse than a naming nit: the eight a + <offset> expressions plus the trailing bare a are nine distinct grouping expressions, so UInt16 was already forced by column count and the duplicate ordinal proved nothing. Switched to bare a plus a + 1..a + 7 so the trailing (a) reuses one of the eight. Measured on the branch:

eight columns, no repeat → __grouping_id is UInt8
same eight with the set repeated → UInt16

Both are asserted now, so the widening is pinned to the ordinal rather than the column count.

Consumer coverage — added consume_grouping_sets_63_columns_repeated_thrice_is_rejected and consume_grouping_sets_64_columns_repeated_is_rejected. Each produces a representable plan, edits the AggregateRel groupings in the proto into the shape the producer refuses to emit but another producer could send, then asserts the consumer's specific error ("A grouping set with 63 columns cannot be repeated ... does not fit in the 1 bit(s) left above the mask", and the 64-column / 0-bit variant). Since to_substrait_plan succeeds before the edit, the error provably comes from the consumer path.

Tests only, no production changes.

@alamb — this adds ~145 lines of tests, which cuts against your point that total diff size is the right metric, so happy to drop either of these if you'd rather keep it tight. The open question is still yours though: did the examples I posted above answer why the index can't come out of __grouping_id by formula? I'd rather settle that direction before putting anything more into this PR.

github-actions Bot added the substrait Changes to the substrait crate label Sep 20, 2026

namanjain24-sudo force-pushed the fix-substrait-grouping-set-index branch from 710f66e to 8d16378 Compare September 24, 2026 13:45

chore: retrigger CI after transient MinIO/quay.io registry outage

25e78f3

kosiew requested changes Oct 7, 2026

View reviewed changes

kosiew requested changes Oct 9, 2026

View reviewed changes

namanjain24-sudo and others added 2 commits October 9, 2026 12:07

Merge branch 'main' into fix-substrait-grouping-set-index

c8a5a7d

kosiew approved these changes Oct 12, 2026

View reviewed changes

alamb reviewed Oct 15, 2026

View reviewed changes


		let err = to_substrait_plan(&plan, &ctx.state()).unwrap_err();

		assert!(

Conversation

namanjain24-sudo commented Sep 20, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

What is the testing strategy for this PR?

Are there any user-facing changes?

Uh oh!

codecov-commenter commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

namanjain24-sudo commented Sep 24, 2026

Uh oh!

namanjain24-sudo commented Sep 28, 2026

Uh oh!

kosiew left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kosiew left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kosiew left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

namanjain24-sudo commented Oct 15, 2026

Uh oh!

alamb commented Oct 15, 2026

Uh oh!

alamb commented Oct 15, 2026

Uh oh!

namanjain24-sudo commented Oct 15, 2026

Uh oh!

namanjain24-sudo commented Oct 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

codecov-commenter commented Sep 20, 2026 •

edited

Loading