feat(functions-aggregate): support nested types (List, Struct, Map) in first_value / last_value GroupsAccumulator - #23628
Conversation
…lue GroupsAccumulator `first_value(x ORDER BY o)` and `last_value(x ORDER BY o)` currently fall back to the per-group `Accumulator` path whenever the value column is a nested type (List, LargeList, Struct, Map, ...), because `groups_accumulator_supported()` whitelists only scalar primitives and byte types. The per-group `Accumulator` path stores state as `ScalarValue` and calls `ScalarValue::try_from_array` on every candidate row. For nested types this allocates on each row and, for `List` / `Struct` / `Map`, the resulting `ScalarValue` holds an Arc slice into the source batch buffer — pinning every batch that produced a current winner. On wide payloads (e.g. `List<Struct>` with kilobyte-per-row values) this compounds into much larger memory usage than the data itself and can trigger OOM in high-cardinality dedup queries. This PR extends the `GroupsAccumulator` fast path to nested types via a new `GenericValueState` — a `Vec<Option<ScalarValue>>`-backed `ValueState` that calls `ScalarValue::compact()` on each stored scalar so the winner value becomes an owned copy instead of an Arc slice into the source batch. The `compact()` step is essential: without it the fast path would still pin source batches even though the accumulator's own reported size looks small. # Types now supported - `List`, `LargeList`, `ListView`, `LargeListView` - `FixedSizeList` - `Struct` - `Map` Directly addresses the primitive needed by the `DISTINCT ON` performance report (apache#16620) and unblocks a `ROW_NUMBER() = 1 -> aggregate` logical rewrite. # Tests - 9 unit tests in `first_last/state.rs` covering `List<Utf8>`, `Struct`, `LargeList`, `FixedSizeList`, `Map`, `EmitTo::First(n)` partial emit, null handling, and — importantly — a regression test that verifies `ScalarValue::compact()` releases the Arc reference to the source batch (`test_generic_value_state_compact_releases_parent_batch`). - 2 integration tests in `first_last.rs` running the full `FirstLastGroupsAccumulator<GenericValueState>` for `List<Int32>`, one functional and one asserting the accumulator's reported size stays O(#groups) rather than O(#rows) on a 10-group / 10k-row workload. Follow-up: add a benchmark comparing the two paths on wide-payload aggregates so the memory improvement is measurable in-tree. Closes apache#23601 Part of apache#23600
There was a problem hiding this comment.
Pull request overview
This PR extends the first_value / last_value ordered GroupsAccumulator fast path to support nested/composite Arrow types (e.g., List, Struct, Map) by introducing a ScalarValue-backed GenericValueState that compacts stored winners to avoid pinning input batches in memory.
Changes:
- Added
GenericValueState(Vec<Option<ScalarValue>>+ cached heap size) and used it for nested/composite output types. - Expanded
groups_accumulator_supported()andcreate_groups_accumulator()to route nested types through theGroupsAccumulatorpath. - Added unit + integration tests covering nested types, partial emits, null handling, and bounded size behavior.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| datafusion/functions-aggregate/src/first_last/state.rs | Adds GenericValueState and extensive unit tests for nested-type state behavior and size accounting. |
| datafusion/functions-aggregate/src/first_last.rs | Wires GenericValueState into ordered groups accumulators for nested types and adds end-to-end integration tests. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| let array: ArrayRef = make_list_utf8_array(); | ||
| let weak_before = Arc::downgrade(&array); | ||
| assert_eq!(weak_before.strong_count(), 1); | ||
|
|
||
| state.update(0, &array, 0)?; | ||
| drop(array); | ||
|
|
||
| // If compact() did its job, dropping the source array must release | ||
| // the only strong reference. If the stored ScalarValue kept a slice | ||
| // of the parent, strong_count would stay >= 1 and the batch would | ||
| // remain pinned. | ||
| assert_eq!( | ||
| weak_before.strong_count(), | ||
| 0, | ||
| "compact() failed to break the Arc reference to the source batch" | ||
| ); |
There was a problem hiding this comment.
Good catch — addressed in 2d8e8bd: the pinning tests now record each source batch's underlying value-data buffer pointer (to_data().buffers()[1].as_ptr(), exactly as suggested) and assert the emitted output aliases none of them, instead of checking the outer Arc's strong_count. The stale comments describing the old Weak-handle mechanism were cleaned up in 582ae48.
Streams 50 independent batches of `List<Int32>` payload through
`FirstLastGroupsAccumulator<GenericValueState>`, dropping each batch's local
`Arc` immediately after `update_batch`. Asserts:
1. Every source batch's `Weak` handle reports `strong_count == 0` after
its local `Arc` is dropped — proves the accumulator does not pin any
source batch via a stored `Arc` slice.
2. `group_acc.size()` stays under a small bound (well under
`#batches * #rows * per-list-cost`).
3. The accumulator can still emit correct winners after every source
batch has been dropped — proves stored values are owned copies.
This is the direct regression test for the wide-payload pinning behaviour
that motivated the `compact()` step in `GenericValueState::update`.
Per Copilot's review: checking `Arc::strong_count` on the outer `ArrayRef` does not actually prove that `compact()` broke buffer sharing. For nested types, `ScalarValue::try_from_array` goes through `list_array.value(idx)`, which returns a sliced *child* array whose Arc chain is independent of the outer `ArrayRef`. The outer array can be dropped (strong_count == 0) while the child value-data buffer remains shared with the source batch. The correct check is to compare the raw pointer of the value-data buffer before and after storing. - `test_generic_value_state_compact_releases_parent_batch`: capture the source `Utf8` value-data buffer pointer, then compare against the stored `ScalarValue::List`'s value-data buffer pointer. `assert_ne!` catches the case where `compact()` is removed. - `test_first_group_acc_list_no_source_batch_pinning`: capture each source batch's `Int32` value-data buffer pointer, then check that the emitted result's value-data buffer pointer matches none of them.
saadtajwar
left a comment
saadtajwar
left a comment
There was a problem hiding this comment.
Not a maintainer, but this change set makes sense and LGTM!
… struct aggregate Add the CoalesceFirstLast logical optimizer rule, gated behind optimizer.enable_coalesce_first_last (default off). Peer first_value / last_value aggregates that share an ORDER BY key are rewritten into one struct-valued aggregate plus a get_field projection, so the input is scanned once instead of once per expression and one per-group state slot is kept instead of N. Skips DISTINCT / FILTER / IGNORE NULLS, expressions without an ORDER BY, and grouping-set aggregates, and no-ops when named_struct / get_field are not registered. Closes apache#23602. Part of apache#23600; the struct fast path depends on apache#23628.
… struct aggregate Add the CoalesceFirstLast logical optimizer rule, gated behind optimizer.enable_coalesce_first_last (default off). Peer first_value / last_value aggregates that share an ORDER BY key are rewritten into one struct-valued aggregate plus a get_field projection, so the input is scanned once instead of once per expression and one per-group state slot is kept instead of N. Skips DISTINCT / FILTER / IGNORE NULLS, expressions without an ORDER BY, and grouping-set aggregates, and no-ops when named_struct / get_field are not registered. Registers the config and rule in the generated artifacts that enumerate every setting and optimizer rule: configs.md, information_schema.slt, optimizer_rule_reference.md, and the EXPLAIN VERBOSE rule list in explain.slt. Closes apache#23602. Part of apache#23600; the struct fast path depends on apache#23628.
There was a problem hiding this comment.
@zhuqi-lucas
Looks good overall. I found no blocking issues and left two follow-up suggestions for test coverage and comment clarity.
| // avoids the per-row ScalarValue churn of the per-group `Accumulator` | ||
| // path: winner extraction happens once per group per batch, not once | ||
| // per candidate row. | ||
| DataType::List(_) |
There was a problem hiding this comment.
The new generic path has good accumulator-level coverage for first_value(List<Int32>), and the existing SLT already covers SQL-level first_value and last_value with list payloads. As a follow-up, it would be useful to add SQL coverage for a distributed partial/final merge shape or another nested payload such as Struct or Map. That would help confirm the planner selects this path beyond the current list case. This is not blocking since the direct FirstLastGroupsAccumulator tests cover the core state and update behavior, and list SQL coverage already exists.
There was a problem hiding this comment.
Agreed — and this is already in motion as part of the #23600 epic work: I've prepared a nested-type benchmark suite (struct/list update/evaluate/merge plus a coalesce-peers comparison) that I'll open as a follow-up PR once this merges, and the review on #23682 asks for exactly the SQL-level matrix you describe (struct/map payloads + partial/final merge shapes, baseline-vs-enabled result equality). Early numbers from the bench: a struct(i64,utf8,f64) accumulator costs ~1.07x a single primitive one — the ordering comparison dominates and is shared.
| const BATCHES: usize = 50; | ||
| const ROWS_PER_BATCH: usize = 256; | ||
|
|
||
| // Keep a `Weak` handle to each batch's payload array. After we drop |
There was a problem hiding this comment.
The test comment still mentions keeping Weak handles and checking strong_count, but the implementation now records raw child-buffer pointers. Could you update the comment to describe the current assertion? That should make the test easier for future readers to follow.
There was a problem hiding this comment.
Good catch — fixed in 582ae48. Both stale spots (the doc-comment's item 2 and the inline note before source_value_ptrs) now describe the raw-pointer assertion the test actually performs.
The pinning regression tests were switched from Weak-handle strong_count checks to raw child-buffer pointer comparison, but two comments still described the old mechanism. Align them with what the assertions actually do. Addresses kosiew's review comment.
Thanks @kosiew for review and good suggestions, addressed now. |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #23628 +/- ##
=======================================
Coverage ? 80.87%
=======================================
Files ? 1101
Lines ? 375920
Branches ? 375920
=======================================
Hits ? 304010
Misses ? 53778
Partials ? 18132 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Which issue does this PR close?
Closes #23601.
Part of the epic #23600.
Rationale for this change
first_value(x ORDER BY o)/last_value(x ORDER BY o)currently fall back to the per-groupAccumulatorpath wheneverxis a nested type (List,LargeList,Struct,Map, ...), becausegroups_accumulator_supported()whitelists only scalar primitives and byte types.Two consequences we hit:
Accumulatorpath callsScalarValue::try_from_arrayon every candidate row, and storesScalarValue::List(orScalarValue::Struct, ...) as anArcslice into the source batch. For wide payloads this both allocates heavily per row and pins the source batch in memory for every group whose current winner came from it. On high-cardinality dedup queries (SELECT DISTINCT ON,ROW_NUMBER() = 1, single-passFIRST_VALUE(...ORDER BY)over a wideList<Struct>) that combination can OOM even when the final output would fit comfortably.distinct on (columns)#16620 (Performance of DISTINCT ON (columns)) reports at the SQL layer.What changes are included in this PR?
Adds a new
GenericValueStateinfirst_last/state.rs— aVec<Option<ScalarValue>>-backedValueState— and wires it intocreate_groups_accumulatorfor nested types. The state usesScalarValue::compact()aftertry_from_arrayso the stored winner is an owned copy instead of anArcslice into the source batch (otherwise the fast path would still pin source batches even though the accumulator's own reported size looks small).Types now supported by the
GroupsAccumulatorfast path:List,LargeList,ListView,LargeListViewFixedSizeListStructMapThe existing
PrimitiveValueState(scalar primitives) andBytesValueState(Utf8 / Binary variants) paths are unchanged.Are these changes tested?
Yes.
first_last/state.rscoveringList<Utf8>,Struct,LargeList,FixedSizeList,Map,EmitTo::First(n)partial emit, null handling, size accounting on overwrite and shrink, and a dedicated regression test (test_generic_value_state_compact_releases_parent_batch) that verifiescompact()releases theArcreference to the source batch.first_last.rsexercising the fullFirstLastGroupsAccumulator<GenericValueState>forList<Int32>: one functional multi-batch test, and one that asserts the accumulator's reportedsize()stays proportional to#groupsrather than#rowson a 10-group × 10 000-row workload.cargo fmtandcargo clippy -p datafusion-functions-aggregate --lib --tests -- -D warningsare clean.Are there any user-facing changes?
Yes, but only in the sense that a previously-slow / OOM-prone path becomes fast for the affected data types. No SQL or API surface changes; the same queries run without modification and previously-passing tests continue to pass.
Follow-ups
FIRST_VALUE(... ORDER BY o)expressions into a single struct accumulator.Filter(row_number() = 1) → Aggregate(FIRST_VALUE(... ORDER BY o)). Depends on this PR to emit the fast path.