andygrove opened a new pull request, #2061:
URL: https://github.com/apache/datafusion-ballista/pull/2061
# Which issue does this PR close?
Closes #2060.
Related to #2031 — see the rationale below.
# 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 — 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.
This also matters for #2031, whose proposed remedy is to *"reject
memory-pool growth before real usage exceeds the budget, so DataFusion spills
instead of OOMing"*. That guard cannot work on this operator while it buffers
straight through a rejection. This PR does not close #2031 — it does not add
allocator-level tracking or make `ResourcesExhausted` retriable — but it is a
prerequisite for the rejection-based half of it to have any effect here.
# What changes are included in this PR?
Treat 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.
Tests: 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. It fails on `main` with
`spill_count=0` and passes here. A `pool_bytes` parameter is threaded through
the existing `run_round_trip` helper rather than duplicating it, so the new
test also inherits its assertions that every input key survives the spill path
and that `pool.reserved() == 0` at the end.
# 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)`.
--
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]
