andygrove opened a new pull request, #4989:
URL: https://github.com/apache/datafusion-comet/pull/4989
## Which issue does this PR close?
Closes #.
## Rationale for this change
`MultiPartitionShuffleRepartitioner` buffers input batches and per-partition
row indices in memory, and spills to disk only when the DataFusion memory pool
denies an allocation:
```rust
if self.reservation.try_grow(mem_growth).is_err() {
self.spill()?;
}
```
That is the only spill trigger. When the pool is large or unbounded, the
writer can hold an arbitrary amount of data resident before any spill happens,
and there is no way to cap the writer's buffered footprint independently of the
pool.
This PR adds an optional config that caps how many bytes the writer buffers.
Spilling is then triggered by memory pressure or by hitting the configured
limit, whichever comes first.
## What changes are included in this PR?
New config `spark.comet.shuffle.maxBufferBytes`, defaulting to 0, which
disables the limit and preserves today's behavior exactly. `ByteUnit.BYTE`, so
a bare number means bytes and suffixed values such as `512m` still parse.
The value is carried to the native side on a new `int64 max_buffer_bytes`
field of the `ShuffleWriter` proto message (`int64` rather than `int32` because
a valid setting can exceed 2GB), set in
`CometNativeShuffleWriter.buildUnifiedPlan` and threaded through `planner.rs`
and `ShuffleWriterExec` into the repartitioner.
The trigger itself is at the existing spill site:
```rust
if self.reservation.try_grow(mem_growth).is_err()
|| (self.max_buffer_bytes > 0 && self.reservation.size() >=
self.max_buffer_bytes)
{
self.spill()?;
}
```
The comparison value is `reservation.size()`, the same running estimate
already used when requesting from the pool: deduped capacity of the buffers
pinned by buffered batches, plus the allocated size of `partition_indices`.
`spill()` calls `reservation.free()`, so it resets after each spill and always
reflects bytes currently buffered. `try_grow` is evaluated first so the
reservation accounts for the batch on both paths. The check runs after the
batch is buffered, so the writer can overshoot by at most one batch, which is
how the memory-pressure trigger already behaves.
Scope is limited to the multi-partition writer.
`SinglePartitionShufflePartitioner` writes straight through and never spills,
so it takes no threshold. The JVM columnar shuffle path is unaffected.
`shuffle_bench.rs` gains a matching `--max-buffer-bytes` arg.
## How are these changes tested?
Three Rust unit tests in `native/shuffle/src/shuffle_writer.rs`, all against
a memory pool large enough that `try_grow` never fails, so any spill observed
must come from the new limit:
- a small limit spills
- a limit of 0 does not spill
- end-to-end through `ShuffleWriterExec`, decoding the written shuffle
blocks and asserting a spilling run produces the same rows in the same order as
a non-spilling run
The first two are a deliberate pair: the zero case pins that the spills in
the first come from the new limit rather than the pre-existing memory-pressure
trigger. I verified the tests are not vacuous by temporarily disabling the new
condition, which fails the first test with `got 0` while the others still pass.
One JVM test in `CometNativeShuffleSuite` covers the plumbing the Rust tests
cannot reach (config to proto to native), asserting that a 64k limit spills
more than the default of 0. It compares against the default rather than
asserting an absolute zero, since the default path is still free to spill under
memory pressure and how much of that happens depends on the pool size.
The design for this change was worked out using the brainstorming skill.
--
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]
