Repository navigation
Generate GroupByHash output in multiple RecordBatches rather than one large one #9562
Description
Is your feature request related to a problem or challenge?
- The
AggregateExecgenerates one single (giant)RecordBatchon output (source) - Which is then emitted in parts (via
RecordBatch::slice(), which does not actually allocate any additional memory) (source)
This has at least two potential downsides:
- No memory is freed until the GroupByHash has output every output row
- As we see in Further refine the Top K sort operator #9417, if there are upstream operators like TopK that hold references to any of these sliced
RecordBatchs, those slices are treated as though they were an additional allocation that needs to be tracked (source)
Something like this in pictures:
Output
▲ RecordBatches are
│ slices into a
┌────────────────┐ single large
│ RecordBatch │─ ─ ─ ┐ output batch
└────────────────┴ ─ ─ ┐
▲ │
│ │ ┌────────────────┐
┌────────────────┐ └ ─ ─ ─▶│ │
output stream │ RecordBatch │ │ ├ ─ ─ ─ ─ ─ ─ ─ ─│
└────────────────┘ ─ ─ ─ ─▶│ │
├ ─ ─ ─ ─ ─ ─ ─ ─│
... │ │
│ │
▲ │ ... │
│ │ │
┌────────────────┐ │ │
│ RecordBatch ├ ─ ─ ┐ │ │
└────────────────┘ │ │
▲ │ ├ ─ ─ ─ ─ ─ ─ ─ ─│
│ ─ ─ ─ ─▶│ │
│ └────────────────┘
│
┏━━━━━━━━━━━━━━━━━━━━━━━┓
┃ ┃ Single RecordBatch
┃ ┃
┃ ┃
┃ ┃
┃ ┃
┃ GroupByHashExec ┃
┃ ┃
┃ ┃
┃ ┃
┃ ┃
┃ ┃
┃ ┃
┗━━━━━━━━━━━━━━━━━━━━━━━┛
Describe the solution you'd like
If we had infinite time / engineering hours I think a better approach would actually be to change GroupByHash so it didn't create a single giant contiguous RecordBatch
Instead it would be better if GroupByHash produced a Vec<RecordBatch> and then incrementally fed those batches out
Doing this would allow the GroupByHash to release memory incrementally as it output. This is analogous to how @korowa made join output incremental in #8658
Perhaps something like
▲
│
┌────────────────┐ Output
output stream │ RecordBatch │ RecordBatches are
└────────────────┘ created in smaller
▲ chunks and emitted
│ one by one
│
│
┏━━━━━━━━━━━━━━━━━━━━━━━┓ ┌────────────────┐
┃ ┃ │ RecordBatch │
┃ ┃ └────────────────┘
┃ ┃ ┌────────────────┐
┃ ┃ │ RecordBatch │
┃ ┃ └────────────────┘
┃ GroupByHashExec ┃ ┌────────────────┐
┃ ┃ │ RecordBatch │
┃ ┃ └────────────────┘
┃ ┃ ...
┃ ┃ ┌────────────────┐
┃ ┃ │ RecordBatch │
┃ ┃ └────────────────┘
┗━━━━━━━━━━━━━━━━━━━━━━━┛
Vec<RecordBatch>
Describe alternatives you've considered
No response
Additional context
@yjshen notes:
To improve AggExec's mono output pattern, #7065 might be similar to the idea of incremental output.
Reacted by Marko Grujic, guojidan, Huaijin and Georgi Krastev
Activity
guojidan commentedon Mar 12, 2024
take
guojidan commentedon Mar 12, 2024
This issue is interesting, let me try implement it
alamb commentedon Mar 13, 2024
FYI this issue may be tricky -- as it will be performance critical -- I will be happy to assist
guojidan commentedon Mar 22, 2024
this issue is a bit tricky for me 😢 , I can only think of the following approaches:
change the
GroupedHashAggregateStream::emit()function return a vectorResult<Vec<RecordBatch>>, eachRecordBatchnum_rows equalGroupedHashAggregateStream::batch_size, and changeExecutionState::ProducingOutput(RecordBatch)toExecutionState::ProducingOutput(Vec<RecordBatch>), thenGroupedHashAggregateStream::poll_next()function returns one element of a vector at one loop, likelet output = batchs.pop()sourcealamb commentedon Mar 22, 2024
I think this approach sounds good -- nice proposal
One thing that could help keep the PRs small and manageable would be to switch the APIs as described above but you could avoid having to change all the
GroupAccumulatorsin one PR by returning return aVec<>of size 1 for most of them.Then we can make subsequent PRs to switch over other groups accumulators as needed
Rachelint commentedon Jul 30, 2024
Actually, I found the
slicefunction is not trivial because the eager computation for the null count.One way I could think to solve it is making the computation lazy, see arrow-rs pr(q32 1.10~1.15x faster without the computation for null count):
apache/arrow-rs#6155
The alternative may be the way mentioned in this issue.
Rachelint commentedon Jul 30, 2024
Hi @guojidan, seems there is no update for some months, can I pick this up?
I want to try it and compare the effect with
lazy null count computation.guojidan commentedon Jul 31, 2024
Thanks. You can take this. I am working on another project now.
joroKr21 commentedon May 16, 2025
I reached the same conclusion. Had a hash join after aggregation and it wanted to allocated > 16GB memory for couple hundred MB of aggregation state 😢
@alamb why are you talking about accumulator APIs?
ExecutionState::ProducingOutputis purely local toGroupedHashAggregateStream. I already tried implementing this change and it worked great for my use case - the join no longer ran out of memory 😄alamb commentedon May 16, 2025
I was thinking that the output of the aggregate exec is both the group columns and the aggregate columns. If we want to produce the output without one giant array, I think we would need to update both the group columns and aggregates
Note that @Rachelint is actively working on a PR
joroKr21 commentedon May 16, 2025
Well the hacky solution is just
EmitTo::First(batch_size)until done 😄I guess the concern is that it's not efficient because it's using
Vec::split_offunder the hood in most cases?Although that's not mandated by the API. But currently some plans simply don't work because of the memory overcounting.
alamb commentedon May 16, 2025
I think that will effectively copy all the values potentially total_size/batch_size times.
But it will result in smaller allocations of record batches for sure
alamb commentedon Aug 27, 2026
I filed a ticket / epic to track the idea of blocked state management here