feat(functions-aggregate): support nested types (List, Struct, Map) in first_value / last_value GroupsAccumulator#23628
Open
zhuqi-lucas wants to merge 3 commits into
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
Contributor
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.
Comment on lines
+844
to
+859
| 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" | ||
| ); |
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
approved these changes
Jul 16, 2026
saadtajwar
left a comment
saadtajwar
left a comment
Contributor
There was a problem hiding this comment.
Not a maintainer, but this change set makes sense and LGTM!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.