jayzhan211 opened a new pull request, #26986: URL: https://github.com/apache/datafusion/pull/26986
## Which issue does this PR close? - Closes #26707. ## Rationale for this change `WindowAggExec` and `BoundedWindowAggExec` hold input rows, but neither reserves memory, so a memory limit does not bound them. A window query neither spills nor fails with `ResourcesExhausted`; it uses as much memory as its input needs. #26707 measures a `WindowAggExec` query over 50M rows peaking at about 1.2 GiB with `-m 200m`, and a `BoundedWindowAggExec` query with a 20M-row frame at 2.9 to 3.7 GiB, the same as without a limit. ## What changes are included in this PR? Three commits: a `bounded_window` benchmark case with a `Utf8View` column, then one commit per operator. **`WindowAggExec`.** `WindowAggStream` registers a `MemoryConsumer` and grows its reservation for each input batch it buffers, for the concatenation of those batches, and for the window result columns. Once the output is computed it drops the buffered batches and frees the reservation; like other operators' output, the emitted batch is not reserved. One `RecordBatchMemoryCounter` counts each buffer once, because these batches often share buffers: the scan splits a large batch into `batch_size` slices, aggregates emit their output as slices of one batch, and `concat_batches` of a single batch returns that batch's own arrays. **`BoundedWindowAggExec`.** `BoundedWindowAggStream` counts the buffers held by `input_buffer` and the partition buffers, and resizes its reservation to that count after each input batch. The count changes only where a batch is replaced or dropped: when a partition is created or extended, when `input_buffer` is replaced, and when an ended partition is removed. Pruning keeps slices of the same buffers, so it needs no update. Batches are counted in one of two ways: - Batches that can share buffers go in a `RecordBatchMemoryTracker`, which counts each buffer once. These are `input_buffer` and each new partition's batch, which is a slice of an input batch. In `Linear` mode all of a batch's new partitions are slices of one batch, and none of them ends before the input does. - Once a partition is extended, `concat_batches` has copied its rows into new buffers that no other batch holds, so their size goes in a plain byte count. Counting them in the tracker as well cost 3% to 18% on the dense `bounded_window` cases below, where nearly every partition is extended by about one row per batch. Schemas with view, dictionary, run-end encoded or nested types stay on the tracker, because concatenating them can keep the input's buffers. When a partition's batch is released, the new `RecordBatchMemoryTracker::is_counted` tells which count it is in, by looking up the batch's first buffer. Not reserved: - Memory inside the window functions, such as accumulator and frame state, and in `BoundedWindowAggExec` the results computed but not yet emitted (`out_col`). Both live inside the `WindowExpr` implementations, out of the operator's sight, so reserving them needs either an API change or, in `BoundedWindowAggExec`, a scan of every partition per batch. - The copies alive while `BoundedWindowAggExec` re-concatenates a partition or `input_buffer`: the reservation is resized after each input batch, when only the new batches are left. - In `Linear` mode, the hash of each buffered row and `most_recent_row`, a one-row slice that keeps its input batch alive. - The result array `WindowAggExec` builds for each partition before concatenating them. With many small partitions these arrays are mostly array objects, which a buffer count barely sees, so they need their own fix: not building one array per partition. - Array objects: the reservations count Arrow buffers, not the array structs that hold them, which matters only when the data is split into many tiny arrays. ## What is the testing strategy for this PR? New tests in `datafusion/core/tests/memory_limit/mod.rs`: - `oom_window_agg_exec` and `oom_bounded_window_agg_exec`: a query that holds its whole input returns `ResourcesExhausted` naming the window stream. On `main` both queries succeed. - `window_agg_exec_reserves_input_concatenation`: two 400 KB input batches, their 800 KB concatenation and the 800 KB result need more than 2 MB. - `window_agg_exec_reserves_shared_buffers_once`: 2,000 one-row slices of one ~280 KB batch run under 16 MB (reserving each slice separately needs ~560 MB), and a single 100,000-row batch runs under 2 MB, which holds its input and result only if its concatenation, the same batch, is not counted again. - `bounded_window_agg_exec_reserves_held_buffers`: three queries that hold part of their input. The first needs ended partitions released (`Sorted` mode, 10-row partitions, under 2 MB). The second needs the 1,000 partitions that start as slices of one batch counted once (`Linear` mode under 2 MB, where counting each slice separately reserves ~130 MB). The third adds a `Utf8View` column, whose data buffers every extended partition keeps, so they must stay in the tracker (under 16 MB; counting them per partition needs ~6 GB). And unit tests: `test_window_agg_releases_reservation_after_output` in `window_agg_exec.rs` checks that nothing is reserved once the output is returned, and `test_is_counted` in `datafusion-common` covers the new method. Benchmarks, this PR against its base on an M4 Pro laptop, medians of paired runs (above 1 is slower): | Benchmark | Cases | Ratio | |---|---|---| | `window_query_sql` | the 6 `OVER ()` and `PARTITION BY` cases (`WindowAggExec`) | 0.97 to 1.03 | | `bounded_window` | linear dense, 100 partitions | 1.00 | | `bounded_window` | linear dense, 10,000 partitions (5 cases) | 1.01 to 1.07 | | `bounded_window` | linear sparse, 32,768 partitions (3 cases) | 0.97 to 1.05 | | `bounded_window` | sorted, 10,000 partitions | 1.00 | The dense 10,000-partition cases are the worst case for this change: each 8,192-row batch extends nearly every partition by about one row, and each extend now sizes the partition's old and new batch. In a probe build that skips accounting on extend these cases run at 0.99 to 1.01, and in one that only sizes the new batch at 1.02 to 1.03, so most of what remains is the cost of sizing each extend exactly. The sparse cases never extend a partition, so their spread is noise between builds. ## Are there any user-facing changes? Yes. With a memory limit, a window query that holds more than the pool allows now fails with `ResourcesExhausted` instead of going over the limit. Neither operator can spill yet (#22946), so there is no fallback. In a shared pool, window operators now take memory that other operators could use before. The 56.0.0 upgrade guide gets an entry for this. Public API: `RecordBatchMemoryTracker::is_counted` is added in `datafusion-common`. Nothing is changed or removed. -- 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]
