feat(physical-plan): add GroupColumn support for FixedSizeList / List / LargeList / Struct in multi-column GROUP BY#22706
Conversation
There was a problem hiding this comment.
Pull request overview
This PR extends the GroupValuesColumn (column-wise group key storage) fast path in DataFusion’s physical plan to support nested group-by key types, avoiding fallback to the row-encoded GroupValuesRows path when a GROUP BY includes FixedSizeList, List/LargeList, or Struct columns.
Changes:
- Added
GroupColumnbuilders forFixedSizeList<primitive>,List<T>/LargeList<T>, andStruct<...>with recursive child support. - Refactored
GroupValuesColumn::internto use a recursivemake_group_column(Field)factory and updatedsupported_typeto recursively validate nested types. - Added extensive unit/integration tests covering nested builders, slicing behavior, take/drain behavior, and dispatcher/allow-list consistency.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs |
Adds recursive make_group_column factory, extends supported_type for nested types, and adds tests for support/dispatcher behavior. |
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs |
Introduces FixedSizeListGroupValueBuilder<T> for fixed-size list group keys. |
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs |
Introduces ListGroupValueBuilder<O> for List / LargeList group keys with offset-based storage. |
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/struct_.rs |
Introduces StructGroupValueBuilder for struct group keys with per-field child builders. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
… / LargeList / Struct in multi-column GROUP BY Addresses apache#22682 by adding column-native GroupColumn implementations for the nested types, so a multi-column GROUP BY containing any of these no longer falls back from `GroupValuesColumn` to the byte-encoded `GroupValuesRows` path. ## Summary Today `multi_group_by::supported_schema` is all-or-nothing: if any single GROUP BY column is not in the `supported_type` allow-list, the whole grouping falls back to `GroupValuesRows`, even when every other column would have qualified for the fast column-wise + cross-column short-circuit path. The most common offender in real workloads is a single nested column (e.g. `LargeList<Struct<Utf8, LargeUtf8>>`) that drags an otherwise primitive-heavy 30+ column GROUP BY onto the slow path. This PR adds: - `FixedSizeListGroupValueBuilder<T>` for `FixedSizeList<primitive>`. Uses an inner `PrimitiveGroupValueBuilder` for the flat child storage. - `StructGroupValueBuilder` for `Struct<...>` whose fields are themselves `GroupColumn`-supported. Recursive composition: each child field is itself constructed via the factory. - `ListGroupValueBuilder<O>` (`O: OffsetSizeTrait`) for `List<T>` and `LargeList<T>`. Owns an offsets `Vec<O>` and a child `Box<dyn GroupColumn>` built via the factory. Null outer rows occupy a zero-length range so the offset buffer stays monotonic without consuming child slots. - A recursive `make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>>` factory that replaces the giant inline match in `GroupValuesColumn::intern`. The factory walks nested fields and recursively builds children for `Struct` / `List` / `LargeList` cases. - `supported_type` recursively checks children for `List` / `LargeList` / `Struct`. `FixedSizeList` remains restricted to primitive children in this PR; composing FSL with nested children can follow in a separate PR. ## Tests (38 new, 68/68 group_values pass) Per-builder: - append / build round trip - `equal_to` for outer nulls, inner nulls, identical, different (and for List: different lengths, empty) - **sliced input array** (offset != 0) — critical real-world case where the input array carries a non-zero slice offset - `take_n` boundary cases: `n=0`, `n=len`, prefix-contains-nulls - `vectorized_equal_to` / `vectorized_append` match per-row reference - `size()` grows with appends - `build` on an empty builder returns an empty array Composition: - `Struct<Struct<...>>` (two-level nested struct) - `List<List<Int32>>` (two-level nested list) - `LargeList<Struct<Utf8, Int32>>` (mirrors the production-shape nested type that motivated this PR) Dispatcher and invariants: - Each builder is exercised end-to-end via `new_group_values` to prove the dispatcher routes the schema to `GroupValuesColumn` (not `GroupValuesRows`) and dedup is correct across batches. - `supported_type` accepts every primitive + nested combination listed above and rejects `Float16`, `Decimal256`, `Time64(Second)`, `List<Float16>`, `Struct{Float16}`, `FixedSizeList<Utf8>`, etc. - A `supported_type` <-> `make_group_column` consistency fuzz over 21 supported + 9 unsupported types asserts the biconditional. Locks the invariant against future drift between the allow-list and the dispatcher. ## Known limitations - `FixedSizeList<T>` is restricted to primitive `T`. Composing FSL with `Struct` or `List` children can follow in a separate PR. - `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. - `Map` support is not in this PR; ordering semantics need a separate discussion. ## No regression on existing types The factory refactor preserves the existing specialization order: primitive / byte / byte-view / boolean / decimal128 builders are still matched first. Only types that were previously rejected fall through to the new builders. All existing aggregate tests (115/115) pass unchanged.
…n row-encoded For three representative shapes, assert that GroupValuesColumn::size() is strictly smaller than GroupValuesRows::size() after interning the same data. Pins the memory-savings claim of this PR so a future regression in column-native storage trips a unit test before reaching users. Measured savings (printed via eprintln on each run): - List<Int32>, 500 groups: 3.9x smaller - LargeList<Struct<Utf8, Utf8>>, 200 groups: 3.4x smaller - 4 cheap cols + 1 LargeList<Struct>, 300 groups: 2.2x smaller The third shape mirrors the production-shape lock-in described in the issue: a single nested column drags otherwise primitive-heavy keys onto the row-encoded path. The 2.2x figure shows the composite saving when the nested column escapes the Rows fallback alongside its cheap neighbors.
1. fixed_size_list: use `FixedSizeListArray::value_offset(row)` instead of manual `(offset + row) * list_len` arithmetic. Correctness-equivalent today (FSL's `offset()` returns 0 and `values()` is already sliced), but the new form avoids double-applying any future offset behavior and matches the arrow-array convention. 2. list: guard `push_offset` with `checked_add` before the `O::from_usize` range check. Defense in depth against usize wrap for List<i32> at pathological cumulative lengths. 3. supported_type vs dispatcher alignment for Time32 / Time64. The dispatcher only handles Time32(Second/Millisecond) and Time64(Microsecond/Nanosecond), but supported_type previously did not match Time64 at all and matched Time32(_) (including invalid units). Tighten supported_type to the exact set the dispatcher supports, and extend the supported_type <-> make_group_column consistency fuzz to cover all four valid Time variants and the four invalid ones. Locks the alignment going forward.
- mod.rs: remove unnecessary `#[expect(non_snake_case)]` on `mod struct_` (snake_case-compliant module name, attribute was unfulfilled) - mod.rs: `schema.clone()` -> `Arc::clone(&schema)` per project lint - mod.rs: drop redundant `.to_string()` after `format!` - mod.rs: inline format args in eprintln! - fixed_size_list.rs: change `new(data_type: DataType)` -> `new(data_type: &DataType)` (clippy::needless_pass_by_value); the function only inspects the data_type by reference. Update test callers to pass `&data_type(true)`. Update the dispatcher call site in mod.rs to drop `.clone()` (it already had a reference). - list.rs: elide explicit lifetimes on `list_array` helper (clippy::needless_lifetimes) - list.rs: introduce `NotesRow` and `MatrixRow` type aliases in tests to satisfy clippy::type_complexity - struct_.rs: introduce `OuterRow` type alias in nested_struct_of_struct test, same reason.
zhuqi-lucas
force-pushed
the
qizhu/df-nested-group-column
branch
from
854a86e to
af12254
Compare
|
Before I review the 2k LOC a couple of questions:
Is "most common offender" from your production data? Are there any other unsupported cases (e.g. are all dates, etc. supported)? A table / confirmation of what is and what is not supported would be great.
Since you're adding support for multiple things my immediate ask would be: can we split this into 1 PR prepping whatever refactoring is needed + 1 PR per type we are adding support for? You can prototype this by refactoring this PR into stacked commits then we can split it once that looks good. |
Thanks @adriangb taking a look. Q1: source of "most common offender" + supported-types tableYes, from production observation. The shape is a wide multi-column GROUP BY (~40 columns) on a SEC filings materialized view, where a single Below is the multi-column GROUP BY supported set after this PR. (Single-column GROUP BY has separate specializations in Supported
Still falls back to
The PR pins the Q2: split into smaller PRsHappy to. Proposed stacked sequence:
I'll prototype the split as stacked commits on this branch first so you can sanity-check the boundaries before I open the actual sequence of PRs. |
|
Would you be open to creating an EPIC for the other types as well? Interval, Duration, Float16, Decimal256, Dictionary and FixedSizeBinary in particular seem like pretty common primitives that we should be able to handle. You bring up a good point with memory regression tests: would it make sense to also merge some dhat tests in one of the setup PRs so we can measure before / after memory usage? E.g. you write a test and set the limit to 8GB, then when you add support for a type lower that to 2GB. |
|
Both make sense, agree on both. EPIC for remaining type coverage Opened #22715 as the broader tracking EPIC. It includes the existing #22682 (nested types, in flight via this PR sequence) plus the 6 additional primitives you called out ( dhat memory regression tests Good call. The Plan: in the refactor / setup PR (PR 1 of the split), add a small dhat-based harness gated behind a |
| self.outer_len | ||
| } | ||
|
|
||
| fn size(&self) -> usize { |
There was a problem hiding this comment.
size_of_val from std::mem does len * size_of::<O>() not giving current vec allocation (capacity * size_of::<O>())
instead could use:
fn size(&self) -> usize {
self.offsets.allocated_size()
+ self.outer_nulls.allocated_size()
+ self.child.size()
}helper from datafusion: use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt};
There was a problem hiding this comment.
Thanks @nathanb9 for review, addressed in this PR, but we want to split to more PRs, i will address there too.
…counting Per nathanb9 on apache#22706: `size_of_val(slice)` returns `len * size_of::<T>()`, not the underlying Vec's allocation, so size() under-reports the heap footprint that the memory pool uses for spill decisions. Switch to the existing `VecAllocExt::allocated_size()` helper used throughout multi_group_by (primitive.rs, bytes.rs, mod.rs): - list.rs: offsets: Vec<O> -> offsets.allocated_size() - struct_.rs: also account for children: Vec<Box<dyn GroupColumn>>'s own capacity (previously only child contents were counted) fixed_size_list.rs has no owned Vec fields beyond outer_nulls (already correct) so it does not need a change.
Which issue does this PR close?
Closes #22682.
Rationale for this change
Today
multi_group_by::supported_schemais all-or-nothing: if any single GROUP BY column is not in thesupported_typeallow-list, the whole grouping falls back to the byte-encodedGroupValuesRowspath, even when every other column would have qualified for the fast column-wise + cross-column short-circuit path. The most common offender in real workloads is a single nested column (e.g.LargeList<Struct<Utf8, LargeUtf8>>) that drags an otherwise primitive-heavy 30+ column GROUP BY onto the slow path.This PR adds column-native
GroupColumnimplementations forFixedSizeList<primitive>,Struct<...>,List<T>, andLargeList<T>, so a GROUP BY containing any of these no longer falls back fromGroupValuesColumntoGroupValuesRows.What changes are included in this PR?
FixedSizeListGroupValueBuilder<T>forFixedSizeList<primitive>. Uses an innerPrimitiveGroupValueBuilderfor the flat child storage.StructGroupValueBuilderforStruct<...>with recursive child support.ListGroupValueBuilder<O: OffsetSizeTrait>forList<T>(O=i32) andLargeList<T>(O=i64) with recursive child support. Null outer rows occupy a zero-length range so the offset buffer stays monotonic without consuming child slots.make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>>factory inmulti_group_by/mod.rsthat replaces the inline match inGroupValuesColumn::intern. Walks nested fields and recursively builds children forStruct/List/LargeList.supported_typerecursively checks children forList/LargeList/Struct.FixedSizeListis restricted to primitive children in this PR; composing FSL with nested children can follow in a separate PR.Are these changes tested?
Yes. 38 new tests, 68/68 group_values tests pass, 142/142 aggregates regression pass.
Per-builder:
equal_tofor outer nulls, inner nulls, identical, different (and for List: different lengths, empty)take_nboundary cases:n=0,n=len, prefix containing null rowsvectorized_equal_to/vectorized_appendmatch per-row referencesize()grows with appendsbuildon empty builder returns empty arrayComposition:
Struct<Struct<Utf8, Int32>>(two-level nested struct)List<List<Int32>>(two-level nested list)LargeList<Struct<Utf8, Int32>>end-to-end throughnew_group_valuesInvariants:
supported_typeaccepts every primitive + nested combo and rejectsFloat16,Decimal256,Time64(Second),List<Float16>,Struct{Float16},FixedSizeList<Utf8>, etc.supported_type<->make_group_columnconsistency fuzz over 21 supported + 9 unsupported types asserts the biconditional. Pins the invariant against future drift between the allow-list and the dispatcher.Are there any user-facing changes?
No public API changes. The new builders are internal to
physical-plan.supported_typebecomes more permissive (returnstruefor the new types), so existing GROUP BYs that were previously falling back toGroupValuesRowswill now useGroupValuesColumnautomatically; this is a performance improvement, not a behavior change.Known limitations
FixedSizeListis restricted to primitive child types. Composing FSL withStruct/Listchildren is straightforward via the factory but not in this PR.vectorized_equal_to/vectorized_appendfor 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.Map<K, V>not included; ordering semantics need a separate discussion.Happy to split this into 4-5 smaller PRs (factory refactor, then one PR per type) if maintainers prefer reviewing the staging incrementally. The single commit is structured so the factory and the per-type builders are mutually dependent and easier to read together.
Measured memory savings (regression tests)
Three new unit tests assert
GroupValuesColumn::size() < GroupValuesRows::size()after interning the same data and print the ratio. Numbers from the local test run:List<Int32>, 500 groupsLargeList<Struct<Utf8, Utf8>>, 200 groupsLargeList<Struct>, 300 groupsThe third shape mirrors the lock-in pattern described in #22682: a single nested column drags otherwise primitive-heavy keys onto the row-encoded path. The 2.2x composite saving comes from the nested column escaping the Rows fallback alongside its cheap neighbors.
These tests pin the memory claim so future regressions in column-native storage trip a unit-level assertion before reaching users.
Adjacent fix from review
Copilot review caught a pre-existing inconsistency on main: the dispatcher in
make_group_columnhandlesTime32(Second / Millisecond)andTime64(Microsecond / Nanosecond), butsupported_typedid not matchTime64(_)at all and matchedTime32(_)permissively (including the invalid Microsecond / Nanosecond variants). This meantsupported_typeand the dispatcher disagreed on Time types: a schema withTime64(Microsecond)was rejected bysupported_typeeven though the dispatcher would have handled it.Tightened
supported_typeto the exact set the dispatcher supports and extended the supported_type ↔ make_group_column consistency fuzz to cover all four valid + four invalid Time combinations. The alignment is now pinned by tests.