Hash aggregation produces batches reporting huge memory size
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### Describe the bug
For aggregations (both partial and final) with `GroupOrdering::None`, a huge batch is produced after consuming all the input RecordBatches, which is further sliced in order to produce batches of `batch_size` length.
In row_hash.rs, `ExecutionState::ProducingOutput(batch)` slices the large batch:
```
let remaining = batch.slice(size, num_remaining);
let output = batch.slice(0, size);
```
Unfortunately `get_array_memory_size` for each of these small RecordBatches returns the physical memory of the initial huge batch, causing unnecessary spills in the downstream operator.
Operators such as RepartitionExec use `batch.get_array_memory_size()` for deciding whether or not to spill.
Related to https://github.com/apache/datafusion/issues/19481
### To Reproduce
See the linked PR: https://github.com/apache/datafusion/pull/22527
### Expected behavior
### Option 1. Avoid producing the initial huge RecordBatch
Avoid producing a large RecordBatch and slicing it afterwards, instead emit only batch_size RecordBatches at a time, this would also fix: https://github.com/apache/datafusion/issues/18907
The issue with this approach is described here: https://github.com/apache/datafusion/issues/19906, i.e. the bottleneck of Vec::drain of shifting all the existing elements in the vector (though I'm not sure this is more significant then the performance impact of spilling)
Merging https://github.com/apache/datafusion/pull/15591 would probably make this solution more feasible, alternatively VecDeque could be considered: https://github.com/apache/datafusion/issues/19906#issuecomment-3771879828
### Option 2. Fix memory accounting
Alternatively, fix the memory accounting for these small RecordBatches, so they only report the memory occupied by their slice instead of the entire underlying array.
Another idea:
extend this function to work across multiple batches and deduplicate the arrow buffers: https://github.com/apache/datafusion/blob/7c05b208e115c29d1f014d2e318b974984cea1a6/datafusion/common/src/utils/memory.rs#L133
### Additional context
_No response_
Contributor guide
Research direction
Start in row_hash.rs at ExecutionState::ProducingOutput(batch), then inspect RepartitionExec's use of get_array_memory_size and datafusion/common/src/utils/memory.rs around the linked function. Compare the two proposed approaches with the linked reproduction PR and related issues. Done means sliced output batches report appropriate memory for downstream spill decisions without regressing aggregation output behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- databases
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100