feat(physical-plan): add GroupColumn support for FixedSizeList<primitive> in multi-column GROUP BY#23128
Conversation
… size, drop per-compare allocation in equal_to Four Copilot inline comments handled together: * `group_column_supported_type` and `make_group_column` both now reject `FixedSizeList` with a negative `list_size`. Arrow defines that size as non-negative, but the enum lets callers construct a negative value programmatically; without this guard the `i32 as usize` cast inside the builder would wrap and risk panics / OOM. Guards are duplicated on both sides so direct callers of `make_group_column` (and the `supported_schema` planner gate) both fail safely. Consistency fuzz already covers the `make_group_column` side via the existing unsupported_cases path; an explicit dedicated test below covers both surfaces. * `FixedSizeListGroupValueBuilder::new` asserts `list_len >= 0` so any direct construction that bypasses the factory fails fast with a clear message instead of silently producing a broken builder. `list_len_usize` now uses `usize::try_from(...).expect(...)` instead of `as usize`; with the constructor's assertion the conversion is provably infallible, but going through `try_from` means a future invariant break shows up as an explicit panic rather than a silent wrap. * `equal_to` no longer calls `list_array.value(rhs_row)`, which constructs a sliced child `ArrayRef` on every comparison. The hot path now borrows `list_array.values()` and uses `value_offset(rhs_row)` to compute the child base index, mirroring the existing approach in `append_val`. Zero allocations per equality check. New test: `negative_list_size_is_rejected_by_allow_list_and_dispatcher` locks in the negative-size guard at both surfaces. 13/13 FSL tests pass, 128/128 aggregates pass, clippy clean.
|
cc @adriangb @alamb @Rich-T-kid |
|
taking a look 👀 |
Rich-T-kid
left a comment
There was a problem hiding this comment.
This PR looks mostly good to me. Left a couple nits and comments. 🚀
| // wrap when cast to `usize` and trigger panics / OOM in the | ||
| // builder. Reject explicitly here as well so direct callers | ||
| // of `make_group_column` fail safely. | ||
| if list_size < 0 { |
There was a problem hiding this comment.
since group_column_supported_type doesn't allow for a negative list_size, would it make sense to have this be an assert!()?
There was a problem hiding this comment.
Done in bd4532a5d — assert!(list_size >= 0). Test split into allow-list (bool) + dispatcher (#[should_panic]).
| // FixedSizeList<non-primitive>: only primitive children are | ||
| // covered by this PR; nested children depend on the | ||
| // List / Struct builders that follow in the EPIC sequence. |
There was a problem hiding this comment.
nit: are these comments needed? Maybe we could link to the broader epic instead?
There was a problem hiding this comment.
Trimmed to a one-liner pointing at EPIC #22715.
| fn size(&self) -> usize { | ||
| self.outer_nulls.allocated_size() + self.child.size() | ||
| } |
There was a problem hiding this comment.
nit: this doesn't include the bytes held by list_len or the usize pointer held by the field: FieldRef & outer_bytes
There was a problem hiding this comment.
Kept as-is to match peer builders (Primitive/Bytes/ByteView all report only heap buffers). Happy to do a follow-up sweep if we want all builders exact.
| )) | ||
| } | ||
|
|
||
| fn take_n(&mut self, n: usize) -> ArrayRef { |
There was a problem hiding this comment.
Do you think it'd be worth adding some kind of check for n here? maybe in the case where n >= self.size, you could call self.build() directly. This can avoid the split_vec_min_alloc() call that occurs in PrimitiveGroupValueBuilder::take_n()/
There was a problem hiding this comment.
This may be worth placing in GroupValuesColumn::Emit() 🤔
There was a problem hiding this comment.
nvm on this. Since build takes a boxed reference to the trait this wouldn't work since take_n takes a mutable reference. This would need to be done in GroupValuesColumn::Emit()
There was a problem hiding this comment.
Ack — needs Emit() plumbing, not take_n. Follow-up issue under EPIC #22715 to do this uniformly across GroupColumn impls.
| // `group_column_supported_type` and `make_group_column` both reject | ||
| // negative `list_len` upstream so this branch should be unreachable | ||
| // for any schema that survives the allow-list. Assert defensively | ||
| // to fail fast (and with a clear message) if a direct caller ever | ||
| // bypasses the factory with an invalid Arrow type. |
There was a problem hiding this comment.
nit: I think this comment can also be removed. maybe including a piece of why this is un-reachable in the assert!() would be helpful and concise.
There was a problem hiding this comment.
Done — folded into the assert! message so the rationale shows up in the backtrace.
| for j in 0..list_len { | ||
| if !self.child.equal_to(lhs_base + j, child_array, rhs_base + j) { | ||
| return false; | ||
| } | ||
| } | ||
| true |
There was a problem hiding this comment.
Nice! Avoiding an allocation in exchange for O(list_len) integer comparisons should be a solid performance win
There was a problem hiding this comment.
I know you mentioned the memory saving from transition from GroupValuesRows to GroupValuesColumn, I think it'd also be insightful to have benchmarks as well.
There was a problem hiding this comment.
Fair — follow-up under EPIC #22715 to add a criterion bench (GroupValuesRows vs GroupValuesColumn), shared harness across the six upcoming PRs in the series.
…ive> in multi-column GROUP BY PR 2 of the EPIC apache#22715 sequence. Builds on the dispatcher refactor + factory framework that landed in PR 1 (apache#22751) to bring the first nested type into `GroupValuesColumn`, unblocking the rest of the sequence (PR 3 Struct, PR 4 List, PR 5 LargeList, PR 6 composite FSL, PR 7 Map). ## Which issue does this PR close? Part of apache#22682 (nested type coverage) and apache#22715 (EPIC). ## Rationale for this change Today a single `FixedSizeList<primitive>` column in a GROUP BY drags the whole grouping onto the byte-encoded `GroupValuesRows` fallback, even when every other column would have qualified for the column-wise + cross-column short-circuit fast path in `GroupValuesColumn`. With the recursive `make_group_column` factory + `group_column_supported_type` allow-list in place from PR 1, this PR is now a self-contained addition of one builder and two dispatcher entries. ## What changes are included in this PR? - `FixedSizeListGroupValueBuilder<T: ArrowPrimitiveType>` in a new `fixed_size_list` submodule. Storage: outer null bitmap + a child `PrimitiveGroupValueBuilder<T, true>` that holds every element flat (length = `outer_len * list_len`). Element `j` of outer row `i` lives at child index `i * list_len + j`. - `group_column_supported_type` accepts `FixedSizeList<primitive>` for the primitive subset wired through the dispatcher (Int8..Int64, UInt8..UInt64, Float32, Float64, Date32, Date64). Composite children (Struct / List inside FSL) are deferred to a follow-up after their respective builders land. - `make_group_column` dispatches `DataType::FixedSizeList(...)` to the appropriate `FixedSizeListGroupValueBuilder<T>` via an `instantiate_fsl!` macro mirroring the existing `instantiate_primitive!` pattern. ## Are these changes tested? Yes. 12 new unit tests on `FixedSizeListGroupValueBuilder` plus extensions to the existing `group_column_supported_type` ⇔ `make_group_column` consistency fuzz. Per-builder: - append / build round trip with mixed outer-null and inner-null rows - `equal_to` for identical / different / outer-null / inner-null rows - `take_n` boundary cases (`n=0`, `n=len`, prefix containing null rows) - sliced input array (offset != 0) - `vectorized_equal_to` / `vectorized_append` match per-row reference - `size()` grows with appends - `build` on empty builder returns empty array - end-to-end dispatcher → `GroupValuesColumn` routing Consistency fuzz now covers four primitive FSL children (signed, unsigned, float, date) on the supported side and three non-primitive FSL children (Utf8, Decimal128, Boolean) on the rejected side, locking in the scope boundary for this PR. 127/127 aggregates tests pass, clippy clean, fmt clean. ## What follows in the EPIC - PR 3: `Struct<...>` builder + dispatcher. - PR 4: `List<T>` builder + dispatcher (recursive child via factory). - PR 5: `LargeList<T>` builder + dispatcher. - PR 6: Relax FSL child restriction to allow `FSL<Struct>` / `FSL<List>` once the prerequisite child builders are in. - PR 7: `Map<K,V>` (`List<Struct<keys, values>>` Arrow representation).
… size, drop per-compare allocation in equal_to Four Copilot inline comments handled together: * `group_column_supported_type` and `make_group_column` both now reject `FixedSizeList` with a negative `list_size`. Arrow defines that size as non-negative, but the enum lets callers construct a negative value programmatically; without this guard the `i32 as usize` cast inside the builder would wrap and risk panics / OOM. Guards are duplicated on both sides so direct callers of `make_group_column` (and the `supported_schema` planner gate) both fail safely. Consistency fuzz already covers the `make_group_column` side via the existing unsupported_cases path; an explicit dedicated test below covers both surfaces. * `FixedSizeListGroupValueBuilder::new` asserts `list_len >= 0` so any direct construction that bypasses the factory fails fast with a clear message instead of silently producing a broken builder. `list_len_usize` now uses `usize::try_from(...).expect(...)` instead of `as usize`; with the constructor's assertion the conversion is provably infallible, but going through `try_from` means a future invariant break shows up as an explicit panic rather than a silent wrap. * `equal_to` no longer calls `list_array.value(rhs_row)`, which constructs a sliced child `ArrayRef` on every comparison. The hot path now borrows `list_array.values()` and uses `value_offset(rhs_row)` to compute the child base index, mirroring the existing approach in `append_val`. Zero allocations per equality check. New test: `negative_list_size_is_rejected_by_allow_list_and_dispatcher` locks in the negative-size guard at both surfaces. 13/13 FSL tests pass, 128/128 aggregates pass, clippy clean.
* Dispatcher (`mod.rs`): change the `list_size < 0` guard from a `return not_impl_err!` to a bare `assert!`. `group_column_supported_type` already rejects negative sizes at the schema-check layer, so any reachable path with `list_size < 0` is a programmer bug rather than an unsupported case worth returning gracefully. Split the existing test into an allow-list case (checked via `bool`) and a dispatcher case (verified with `#[should_panic]`) so both surfaces stay covered. * Constructor (`fixed_size_list.rs`): fold the multi-line invariant comment above `assert!` into the assert message itself, which now reads "requires non-negative list size (allow-list / dispatcher should have rejected earlier)". Simpler to read at the panic site, and the rationale surfaces exactly where a future reader would want it (the panic backtrace). Same trim on `list_len_usize`. * Dispatcher consistency-fuzz corpus (`mod.rs`): shorten the FixedSizeList<non-primitive> comment to a one-liner pointing at EPIC apache#22715, instead of the two-line prose explanation. Tests: `cargo test -p datafusion-physical-plan --lib multi_group_by` now runs 42/42 (up from 41), including the split negative-size cases and the new `#[should_panic]` guarantee. Clippy clean.
29114c3 to
bd4532a
Compare
Thanks @Rich-T-kid for review, addressed comments in latest PR. |
|
nice! I'll try and take another look tomorrow |
Rich-T-kid
left a comment
There was a problem hiding this comment.
This PR looks good to me. left a comment.
nice work @zhuqi-lucas 🚀
| for &row in rows { | ||
| self.append_val(array, row)?; | ||
| } |
There was a problem hiding this comment.
in #22706 you mentioned
vectorized_equal_to / vectorized_append for the new builders fall back to per-row loops. They are correctness-equivalent to the specialized vectorized paths used by the primitive builders, but a follow-up can add type-specialized batched comparators where beneficial.
which is fine! did you have a specialized vectorized_append in mind? just curious.
|
Thanks @Rich-T-kid for review, also cc @alamb @adriangb |
|
Thanks for the ping -- I'll try to look at this some time, but I have a bunch of other stuff to review first. Why is this feature important? Do we have users grouping by FixedSizeList ? Specifically, I am worried that adding specialized code for all possible types even when they aren't really used will lead to more code to maintain and large binary sizes. For example, this code // `group_column_supported_type` already restricts the child to
// the primitive subset supported here. Any unsupported child
// type returned `false` upstream and was routed to the
// `GroupValuesRows` fallback, so the wildcard arm below is
// only reachable from the consistency-fuzz test.
macro_rules! instantiate_fsl {
($t:ty) => {{
let b = fixed_size_list::FixedSizeListGroupValueBuilder::<$t>::new(
data_type,
);
v.push(Box::new(b) as _);
}};
}
match child_field.data_type() {
DataType::Int8 => instantiate_fsl!(Int8Type),
DataType::Int16 => instantiate_fsl!(Int16Type),
DataType::Int32 => instantiate_fsl!(Int32Type),
DataType::Int64 => instantiate_fsl!(Int64Type),
DataType::UInt8 => instantiate_fsl!(UInt8Type),
DataType::UInt16 => instantiate_fsl!(UInt16Type),
DataType::UInt32 => instantiate_fsl!(UInt32Type),
DataType::UInt64 => instantiate_fsl!(UInt64Type),
DataType::Float32 => instantiate_fsl!(Float32Type),
DataType::Float64 => instantiate_fsl!(Float64Type),
DataType::Date32 => instantiate_fsl!(Date32Type),
DataType::Date64 => instantiate_fsl!(Date64Type),
other => {
return not_impl_err!(
"FixedSizeList<{other}> not supported in GroupValuesColumn"
);
}
}Is likely going to create 12 copies of the new code |
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { |
There was a problem hiding this comment.
I would recommend (only) slt tests, unless there is some particular case we can't covere with those end to end tests
slt tests are less resource intensive, don't add to compile time, and ensure everything is hooked up end to end.
|
Thanks @alamb, no hurry for this PR. Context: this is PR 2 of #22715, split out from the original mega-PR #22706 at @adriangb's request (that PR added FSL / List / LargeList / Struct all in one, 4 files / +2468). #22751 landed as PR 1 (dispatcher + factory framework, primitives only). #23128 is the smallest additive next step — one new builder + two dispatcher entries. Motivation: production workloads that GROUP BY a mix of primitives plus a nested column (Struct<List> in the case I benchmarked) fall entirely into GroupValuesRows because of the one nested column. In my measurements the fallback ran at ~25 GiB peak vs ~15 GiB with the nested column removed — about 40% of peak was fallback tax, not the nested data itself. The EPIC brings nested types onto the column-wise path so mixed shapes stop falling back. On binary size: the 12 arms mirror the existing |
This makes sense to me (as in as single nested GrOUP BY columns forces a fall back for all columns) My concern is about trying to add per-type support for all remaining types -- it seems like a lot of new code, and as proposed still isn't general purpose (this PR only supports FixedSlizeList -- and not for example FixedSizedList , etc If the goal is to avoid falling back to Rows for multi column grouping, what do you think about implementing a generic This would also avoid having to make specialized versons of We maybe could also remove the |
|
Thanks @alamb — I agree, this is a better direction than adding a native Opened #23404 to track it. |
|
Sorry for not doing this before |
@alamb in this context do nested types encompass dictionary arrays? |
Yes probably -- though I think a case can be made that explicitly dictionary support might be valuable The thing with Dictionaries is that the dictionary indices for hte same logical value can vary from batch to batch, so we can't just do the grouping based on the dictionary indices We could potentially have an optimized dictionary group values that somehow wraps another group values and does the dictionary translation 🤔 |
This is pretty much what #23187 does, benchmarks show good results
I think theres room to optimize for the low cardinality case similar to what #21765 benchmarks showcased. Due to |
PR 2 of the EPIC #22715 sequence. Builds on the dispatcher refactor + factory framework that landed in PR 1 (#22751) to bring the first nested type into `GroupValuesColumn`, unblocking the rest of the sequence (PR 3 Struct, PR 4 List, PR 5 LargeList, PR 6 composite FSL, PR 7 Map).
Which issue does this PR close?
Part of #22682 (nested type coverage) and #22715 (EPIC).
Rationale for this change
Today a single `FixedSizeList` column in a GROUP BY drags the whole grouping onto the byte-encoded `GroupValuesRows` fallback, even when every other column would have qualified for the column-wise + cross-column short-circuit fast path in `GroupValuesColumn`. With the recursive `make_group_column` factory + `group_column_supported_type` allow-list in place from PR 1, this PR is now a self-contained addition of one builder and two dispatcher entries.
What changes are included in this PR?
Are these changes tested?
Yes. 12 new unit tests on `FixedSizeListGroupValueBuilder` plus extensions to the existing `group_column_supported_type` ⇔ `make_group_column` consistency fuzz.
Per-builder:
Consistency fuzz now covers four primitive FSL children (signed, unsigned, float, date) on the supported side and three non-primitive FSL children (Utf8, Decimal128, Boolean) on the rejected side, locking in the scope boundary for this PR.
127/127 aggregates tests pass, clippy clean, fmt clean.
What follows in the EPIC
Are there any user-facing changes?
No behavior change for users whose GROUP BY did not previously contain a `FixedSizeList` column. For users who did, the grouping now uses the column-native fast path instead of falling back to `GroupValuesRows` — same results, less memory and CPU.