Repository navigation
Preallocate RowsGroupColumn buffers during take_n rebuilds #23994
Description
Related PRs
Problem
Currently, RowsGroupColumn::take_n rebuilds the retained rows with zero-capacity Rows:
let mut remaining = self.row_converter.empty_rows(0, 0);At that point the retained row count is known (self.group_values.num_rows() - n), and Arrow Rows also exposes row lengths. Rebuilding from zero capacity can force avoidable growth and copying while pushing the remaining rows.
RowsGroupColumn::try_new also initializes with empty_rows(0, 0), but construction currently has no reliable row-count hint. Treat constructor preallocation as optional follow-up, not required scope.
Why it matters
take_n runs on emit paths for row-backed group columns. In high-cardinality or repeated partial-emit workloads, rebuilding retained row buffers from zero capacity can add avoidable CPU and allocator pressure without changing semantics.
Invariant / desired behavior
RowsGroupColumn should preserve exact grouping semantics while avoiding known reallocations in rebuild paths:
- No change to equality or group-id behavior.
- No change to emitted values, output type reconstruction, or emit ordering.
take_nshould allocate enough capacity up front when retained row count and byte size are known.
Proposed direction
Refactor only the take_n rebuild path first.
-
Compute exact retained capacity before rebuilding:
remaining_rows = self.group_values.num_rows() - nremaining_bytes = self.group_values.lengths().skip(n).sum()
-
Create the replacement buffer with those capacities:
self.row_converter.empty_rows(remaining_rows, remaining_bytes)
-
Push the same retained rows in the same order.
-
Leave constructor capacity hinting out unless a trustworthy local hint already exists.
- Do not add speculative planner/batch-size plumbing just to pass an estimate.
Scope
In
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rstake_npreallocation.- Focused test coverage proving unchanged outputs across
take_n. - Targeted perf/allocation comparison if practical.
Out
- Rewriting
Rowsinternals or Arrow row implementation. - Changing grouping semantics, output schema behavior, dictionary re-encoding, or fallback selection logic.
- Adding broad constructor/planner capacity plumbing without an existing reliable hint.
- Broader aggregation architecture changes unrelated to
RowsGroupColumn::take_n.
Acceptance criteria
-
RowsGroupColumn::take_nno longer rebuilds retained rows with zero-capacityRows. - Retained row and byte capacities are derived from existing
Rowsdata, not guesses. - Existing row-backed group-value tests pass unchanged.
- Added or updated tests exercise
take_nand validate unchanged emitted arrays / remaining rows. - If benchmarked, before/after results show no regression and expected allocation or throughput improvement on representative row-backed workloads.
Tests / verification
- Run existing row-backed group-column tests:
cargo test -p datafusion-physical-plan row_backed --lib
- Add or update a focused unit test that:
- Appends multiple row-backed group values.
- Calls
take_n. - Verifies emitted output and later
buildoutput are unchanged.
- Optional perf check:
- Compare before/after on a row-backed aggregation case with high and medium cardinality inputs.
Activity
saadtajwar commentedon Jul 30, 2026
take
saadtajwar commentedon Jul 30, 2026
Happy to take a stab at this!
perf: preallocate RowsGroupColumn buffers in take_n (apache#24070)
perf: preallocate RowsGroupColumn buffers in take_n (apache#24070)