apache / apache/datafusion

Hash aggregation produces batches reporting huge memory size

Open
#22,526 21 comments 0 reactions 0 assignees View on GitHub
bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.