jayzhan211 opened a new issue, #26705:
URL: https://github.com/apache/datafusion/issues/26705
### Describe the bug
For aggregates whose state holds every input value, such as `median`,
`percentile_cont` and `array_agg`, the state moves from the partial aggregate
to the final aggregate without a memory reservation. The partial aggregate
emits its whole state as one batch, then its stream ends and releases its
reservation. The final aggregate reserves memory only after it has merged a
batch. In between, states as large as the input are held unreserved, so a plan
with more than one partition uses far more memory than the pool allows.
### To Reproduce
Run each query in its own `datafusion-cli` process (`/usr/bin/time -l` is
macOS; use `-v` on Linux):
```sh
# Query 1
/usr/bin/time -l datafusion-cli -m 300m \
-c "SET datafusion.execution.target_partitions = 1;" \
-c "SELECT median(v) FROM generate_series(1, 30000000) AS t(v);"
# Query 2
/usr/bin/time -l datafusion-cli -m 300m \
-c "SET datafusion.execution.target_partitions = 4;" \
-c "SELECT median(v) FROM generate_series(1, 30000000) AS t(v);"
```
### Expected behavior
| Query | Plan | Peak RSS, 6 runs each (300 MB limit) | Expected |
|---|---|---|---|
| 1 | `AggregateExec: mode=Single` | 519 MB every run | (baseline) |
| 2 | `Partial` -> `Final` | 656 MB to 939 MB | About the same as Query 1:
the extra memory is the partial states, which should be reserved while held |
Query 1 is above the limit for a separate reason: its value buffer
reallocates into one twice its size before freeing the old one, the same shape
as #25606.
### Additional context
- The partial `state()` moves all values into a single list value:
[percentile_cont.rs#L469-L495](https://github.com/apache/datafusion/blob/039c0b99d5b05e73ff5def029c0131b76b481c6f/datafusion/functions-aggregate/src/percentile_cont.rs#L469-L495).
- `AggregateStream` grows its reservation only after a batch has been
aggregated:
[aggregate_stream.rs#L407-L411](https://github.com/apache/datafusion/blob/039c0b99d5b05e73ff5def029c0131b76b481c6f/datafusion/physical-plan/src/aggregates/aggregate_stream.rs#L407-L411).
- Counting allocations against reservations (the approach of #25651) at 10M
rows, the final aggregate had 96 bytes reserved while 147 MB was allocated.
`array_agg(s) ... GROUP BY v % 10` shows the same gap.
- One direction: the final aggregate reserves an incoming state batch before
merging it, since the contract in `memory_pool/mod.rs` makes the operator that
holds batches responsible for them.
- Related: #25851, #26701.
### Willingness to contribute
I'm not planning to fix this myself, so anyone is welcome to pick it up.
Feel free to ping me with questions about the cause or when you'd like a review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
