sunchao opened a new pull request, #25172: URL: https://github.com/apache/datafusion/pull/25172
## Which issue does this PR close? Fixes consumer and pool admission when a spilling operator owns multiple reservations. Related to the existing merge-workspace accounting work in #24740. ## Rationale for this change A spilling operator can exceed its fair memory allowance by creating sibling reservations. In a 100-byte pool, one consumer with reservations holding 60 and 40 bytes can currently request another byte from an empty sibling. `new_empty`, `split`, and `take` share a consumer registration but maintain independent reservation sizes, so checking only the requesting reservation gives the operator multiple allowances. Registration also changes fair shares without reclaiming existing reservations. A new consumer can therefore fit within its own share while growing the total beyond the configured pool size. Enforcing these limits exposes another dependency: aggregate spill replay shares one consumer between its merge input and its replay table. The merge must leave working memory for replay instead of taking the entire consumer allowance. ## What changes are included in this PR? - Track aggregate usage by spillable consumer ID and enforce both the consumer allowance and remaining pool capacity during fallible growth. Infallible `grow` retains its existing contract and remains accounted for. - Add an advisory `MemoryPool::memory_limit_for` method, defaulting to `Unknown`, and forward it through the built-in tracking, peak-recording, and merge-workspace pools. - Bound merge fan-in to half the known consumer allowance in all four aggregate replay paths. An indivisible minimum merge may exceed that preference only if the real pool admits it, with one-row output and one-batch read-ahead. - Inspect an unsplittable spill run without rewriting it, preserving progress when the existing files already fill disk quota. - Cover reservation lifetimes, mixed consumers, wrapper budgets, all replay variants, indivisible rows, and permanent memory pressure. Adjust the ordering fixture to force an actual spill while giving merge and replay enough shared capacity. The generic spill-file backend and existing public spill APIs remain in place. This contribution does not add a spilling join or change temporary-file ownership/accounting. ## What is the testing strategy for this PR? Validation runs on upstream base `ac7b18de63515ad15672f4b596b2f7e571e9e352` from September 8, with its lockfile unchanged. The September 10 main lockfile requires versions not yet present in the approved dependency registry (`indexmap` 2.14.2 and `zstd` 0.14.0); this recent upstream base resolves through the same registry. The PR targets main and contains only the accounting/replay changes. The final patch also passed `git apply --3way --check` against September 10 main `1ec9ede5da5c67cbc2d3e0110a7ddbc2a47796a6`. That is a clean-application check, not runtime validation of the merged main source; merged-main CI remains required. In particular, main has since added partial-hash state-batch accounting in #24852, outside the final-aggregate replay code changed here. - `cargo test --locked --profile ci -p datafusion-execution -p datafusion-physical-plan --lib --jobs 16 -- --test-threads=8`: **2,081 passed** (120 execution, 1,961 physical-plan). - `cargo clippy --locked --all-targets --all-features --jobs 16 -- -D warnings`: **passed** on the final source. - `cargo fmt --all -- --check`, `git diff --check`, and `cargo doc --locked -p datafusion-execution --no-deps --all-features --jobs 8`: **passed**. - Required extended workspace command: **passed**, with **11,298 Rust tests passed, 8 ignored**, and **all 512 SQL logic files completed**. It used `--profile ci --workspace --lib --tests --bins --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption`, excluded the examples/benchmarks/CLI packages as required by `AGENTS.md`, and retained the upstream lockfile. - The successful extended run used a 65,536 per-process file-descriptor limit. The initial run hit the host's 1,024-descriptor soft limit in the existing `test_sort_10k_mem` fuzzer (`Too many open files`); that test and the full suite passed after raising the process limit. No source change was needed for that environment issue. No targeted performance benchmark was run. The local sort-preserving-merge benchmark uses in-memory streams and an unbounded pool, and the vectorized-aggregate benchmark measures group-value builders; neither exercises the changed fair-pool admission or spill-replay paths. This PR makes no performance claim. ## Are there any user-facing changes? Fallible allocations that exceed a spill consumer's combined allowance or the pool's remaining capacity are now rejected. Under permanent memory starvation, replay reports a resource error and releases its reservations and spill files rather than exceeding the configured memory capacity. `MemoryPool::memory_limit_for` is an additive advisory API. Its answer may change as other consumers register or allocate, and every allocation still goes through `try_grow`. External transparent wrappers should forward the method to make their inner consumer allowance available to aggregate replay planning. -- 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]
