Repository navigation
fix: give a multi-set aggregate the grouping set index Substrait defines - #25527
namanjain24-sudo wants to merge 4 commits into
Conversation
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25527 +/- ##
==========================================
+ Coverage 82.50% 82.77% +0.26%
==========================================
Files 1140 1148 +8
Lines 438050 450800 +12750
Branches 438050 450800 +12750
==========================================
+ Hits 361410 373135 +11725
- Misses 54833 54932 +99
- Partials 21807 22733 +926 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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.
710f66e to
8d16378
Compare
|
Bumping this — still open for review whenever someone has bandwidth (the earlier MinIO CI failure was unrelated infra flakiness, not this PR). |
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
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.
| } | ||
| let ordinal = masks.iter().filter(|seen| **seen == mask).count() as u64; | ||
| masks.push(mask); | ||
| ids.push((ordinal << group_count) | mask); |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
| (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))) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
| } | ||
| 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)? |
There was a problem hiding this comment.
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.
There was a problem hiding this 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.
| /// 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<()> { |
There was a problem hiding this comment.
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.
There was a problem hiding this 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.
- 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
left a comment
kosiew
left a comment
There was a problem hiding this comment.
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.
| } | ||
| mask | ||
| } else { | ||
| (ordinal << group_count) | mask |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
3 participants
Which issue does this PR close?
__grouping_id, not the set's index #25208.Rationale for this change
Substrait ends an
AggregateRelwith more than one grouping set with an extrai32whose 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:UInt8rather thanInt32;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 newlogical_plan::grouping_setmodule shared by the two directions:CASE WHEN __grouping_id = <id of set 0> THEN 0 ... ELSE <last index> END, so the plan ends with the index, as a requiredInt32.AggregateRelholding what the spec defines and projects__grouping_idback 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
ELSEarm rather than aWHEN: the values are exhaustive, and anELSEkeeps the column non-nullable, which Substrait requires of the index and DataFusion of__grouping_id.Two details the maps have to respect:
ROLLUPorCUBE, so the producer writes both as a list of sets, reversed forROLLUPand as a powerset forCUBE. The index follows that list, so the expansion is now one function used both to write the groupings and to compute the ids.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_indexis the issue's plan: sets(a)then(b), where the index differs from the bitmask. It checks the column is a requiredInt32holding 0 and 1, wheremaingives aUInt8holding 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_idis aUInt16and the map has to carry the wider literal.aggregate_grouping_setsnow asserts that theAggregateRelno 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:
multiple_grouping_sets_emit_the_set_index,aggregate_duplicate_grouping_setsandaggregate_grouping_sets_keep_grouping_functionfail;aggregate_grouping_sets;UInt16arm of the literal narrowed toUInt8, onlyaggregate_grouping_sets_wider_grouping_idfails.cargo test -p datafusion-substraitpasses (58 unit, 215 integration with the 6 already ignored, 3 doc tests), as docargo fmt --all -- --check,cargo clippy -p datafusion-substrait --all-targets --features physical -- -D warningsandcargo xtask ci step test substrait. A Substrait round trip ofaggregate.slt, which is not in that job, reports the same 10 pre-existing failures asmain.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 theAggregateRelthat maps that index back to__grouping_id, so theAggregateRelitself holds the column the spec describes. No Rust API changes.