jayzhan211 opened a new issue, #26707:
URL: https://github.com/apache/datafusion/issues/26707
### Describe the bug
`WindowAggExec` and `BoundedWindowAggExec` never reserve memory: there is no
`MemoryConsumer` or `MemoryReservation` anywhere under
`datafusion/physical-plan/src/windows/`. Both hold input rows for as long as
their window needs them. `WindowAggExec` buffers its whole input before it
computes anything, and `BoundedWindowAggExec` buffers the rows its frames still
need.
So a memory limit has no effect on them. A window query neither spills nor
fails with `ResourcesExhausted`; it uses as much memory as its input needs, and
in a container with a memory limit the process can be killed instead. The
contract in `memory_pool/mod.rs` asks an operator whose memory grows with its
input to reserve that memory before allocating it, and to spill or return an
error when the pool refuses.
### To Reproduce
Run each command in its own process (`/usr/bin/time -l` is macOS; use `-v`
on Linux). Drop `-m 200m` for the "no limit" column.
```sh
# Query 1: WindowAggExec, the whole input is one window partition
/usr/bin/time -l datafusion-cli -m 200m \
-c "SELECT max(s) FROM (SELECT sum(v) OVER () AS s FROM generate_series(1,
50000000) AS t(v));"
# Query 2: BoundedWindowAggExec, a frame that reaches 20M rows ahead
/usr/bin/time -l datafusion-cli -m 200m \
-c "SET datafusion.execution.batch_size = 1000000;" \
-c "SELECT max(s) FROM (SELECT sum(v) OVER (ORDER BY v ROWS BETWEEN
CURRENT ROW AND 20000000 FOLLOWING) AS s FROM generate_series(1, 40000000) AS
t(v));"
```
Query 2 uses a large batch size only to keep the run short:
`BoundedWindowAggExec` re-concatenates its buffered rows for every input batch.
### Expected behavior
| Query | Plan | Peak RSS, no limit | Peak RSS, `-m 200m` | Expected with
`-m 200m` |
|---|---|---|---|---|
| 1 | `WindowAggExec` | 1186 MiB (3 runs) | 1193 to 1226 MiB, succeeds (6
runs) | Fails with `ResourcesExhausted` once the buffered rows exceed the limit
(or spills, after #22946) |
| 2 | `BoundedWindowAggExec` | 2.8 to 3.2 GiB (5 runs) | 2.9 to 3.7 GiB,
succeeds (10 runs) | Same as query 1 |
Measured with `datafusion-cli` built from `main` at b93ef03be4 (release
profile without LTO) on macOS arm64.
### Additional context
- `WindowAggStream` keeps every input batch and concatenates them when its
input ends:
[window_agg_exec.rs#L502-L511](https://github.com/apache/datafusion/blob/b93ef03be4d1a5c7d95330dbb971e49804e542be/datafusion/physical-plan/src/windows/window_agg_exec.rs#L502-L511),
[#L610](https://github.com/apache/datafusion/blob/b93ef03be4d1a5c7d95330dbb971e49804e542be/datafusion/physical-plan/src/windows/window_agg_exec.rs#L610),
[#L545](https://github.com/apache/datafusion/blob/b93ef03be4d1a5c7d95330dbb971e49804e542be/datafusion/physical-plan/src/windows/window_agg_exec.rs#L545).
Query 1's peak is about three copies of its 381 MiB input (50M `Int64`
values): the buffered batches, their concatenation, and the result column.
- `BoundedWindowAggExec` keeps `input_buffer` and a buffer per window
partition, and re-concatenates both for every input batch:
[bounded_window_agg_exec.rs#L666-L707](https://github.com/apache/datafusion/blob/b93ef03be4d1a5c7d95330dbb971e49804e542be/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs#L666-L707),
[window_state.rs#L336-L341](https://github.com/apache/datafusion/blob/b93ef03be4d1a5c7d95330dbb971e49804e542be/datafusion/expr/src/window_state.rs#L336-L341).
Query 2's peak grows faster than its frame: 222 to 295 MiB with a 1-row frame,
602 to 691 MiB at 5M rows, 1.1 GiB at 10M, and 2.9 to 3.7 GiB at 20M, while the
frame itself is 8 bytes a row (38 to 153 MiB). Only the window's buffering
depends on the frame, so at its peak the operator holds many times the frame.
The two buffers and their re-concatenation explain about four copies; I haven't
found the rest. A reservation has to count all of it.
- The contract:
[memory_pool/mod.rs#L68-L85](https://github.com/apache/datafusion/blob/b93ef03be4d1a5c7d95330dbb971e49804e542be/datafusion/execution/src/memory_pool/mod.rs#L68-L85).
- One direction: each window stream registers a `MemoryConsumer` that cannot
spill. `WindowAggExec` grows its reservation for every buffered batch and for
the concatenated copy; `BoundedWindowAggExec` resizes its reservation to its
buffers after each batch and after pruning. A test in
`datafusion/core/tests/memory_limit` that expects query 1 to return
`ResourcesExhausted` at a small limit would fail today, because the query
succeeds.
- Trade-off to decide: queries that pass today with a limit that is too
small would start failing with `ResourcesExhausted`, and in a shared pool the
window operators would take memory that spilling operators use now. Comet
closed a change that charged more memory to Spark's pool, because without a way
to reclaim it, native spills turned into task failures
(apache/datafusion-comet#5998).
- Separately, `WindowAggExec` keeps one result array per window partition
until a final `concat`
([window_agg_exec.rs#L558-L570](https://github.com/apache/datafusion/blob/b93ef03be4d1a5c7d95330dbb971e49804e542be/datafusion/physical-plan/src/windows/window_agg_exec.rs#L558-L570)),
so many small partitions cost far more than the input: `sum(v) OVER (PARTITION
BY v)` over the same 50M rows peaks at 7.8 to 9.7 GiB. That needs its own fix.
- Related: #22946 (spilling for `WindowAggExec`, which needs this accounting
first), #23197 (window EPIC, "WindowAggExec Memory / Spilling"), #25637
(spilling coverage gaps), #25650 (processes OOM killed due to inaccurate memory
accounting), #23207 (an earlier accounting PR for `WindowAggExec`, closed
without review).
### 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 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]
