andygrove opened a new pull request, #2238:
URL: https://github.com/apache/datafusion-ballista/pull/2238
# Which issue does this PR close?
Backport of #2061 to `branch-54`. The issue it fixes is #2060.
# Rationale for this change
`SortShuffleWriterExec` registers its `MemoryConsumer` with
`.with_can_spill(true)` but then discards the result of the grow:
```rust
// Best-effort: if the pool is bounded and rejects the grow, that's fine —
// the absolute counter below still triggers a spill.
let _ = reservation.try_grow(growth);
```
The comment's reasoning does not hold. `buffered_bytes` is a private
per-task counter compared against `memory_limit_per_task_bytes`; it knows
nothing about the shared `MemoryPool`, so a rejected grow produces no spill at
all. The batch has already been pushed into `BufferedBatches` at that point, so
the writer holds memory the pool has no record of, and other consumers on the
same pool then see bytes as free that are not.
Registering with `can_spill(true)` is a contract: the pool rejects a grow to
ask the consumer to spill. This writer declines to answer, which is a direct
contributor to executor OOMs on shuffle-heavy stages.
# What changes are included in this PR?
A clean cherry-pick of 173e28f6, unmodified.
Treats a rejected grow as a spill trigger alongside the existing per-task
budget:
```rust
let pool_rejected_growth = reservation.try_grow(growth).is_err();
buffered_bytes = buffered_bytes.saturating_add(growth);
if pool_rejected_growth || buffered_bytes >= memory_limit {
```
No retry of the grow is needed after spilling. `push_batch` appends to
`batches` unconditionally, so at the point of rejection the batch is always
spillable; `spill_all_partitions` flushes it and calls `reservation.free()`,
leaving the writer holding nothing.
Adds `spills_when_memory_pool_rejects_growth`, which sets the per-task
budget to 1 GiB (far above the payload, so the private counter can never fire)
against a 64 KiB pool that cannot admit even one ~128 KiB batch, making pool
rejection the only possible spill signal. The commit also renames the
spill-trigger constants in the existing tests so each test reads as one trigger
armed and the other disarmed.
# Are there any user-facing changes?
No API changes. Behaviourally, a sort shuffle writer running against a
contended bounded pool now spills where it previously buffered without limit,
which is the intended behaviour for a consumer registered with
`can_spill(true)`.
---
Verified locally on the `branch-54` base: `cargo fmt --all -- --check` is
clean, and `cargo check --workspace --all-targets --locked` completes with no
warnings on a combined stack of the six backports being proposed together. Test
execution is left to CI.
--
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]
