feat: add spark.comet.shuffle.maxBufferBytes to cap native shuffle writer memory

[andygrove]

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:

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 uint64 max_buffer_bytes field of the ShuffleWriter proto message (64-bit rather than 32-bit because a valid setting can exceed 2GB), set in CometNativeShuffleWriter.buildUnifiedPlan and threaded through planner.rs and ShuffleWriterExec into the repartitioner.

Zero on the wire means the limit is disabled. planner.rs normalizes that to None at the proto boundary, so the writer holds an Option<usize> and Some(0) is never constructed. That keeps the sentinel in one place and matches the existing memory_limit: Option<usize> knob in the same crate.

The trigger itself is at the existing spill site:

if self.reservation.try_grow(mem_growth).is_err()
    || self
        .max_buffer_bytes
        .is_some_and(|limit| self.reservation.size() >= limit)
{
    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
• an unset limit 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 unset 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.

[comphead]

Suggested change

	spillCount = find(shuffled.queryExecution.executedPlan) {
	case _: CometShuffleExchangeExec => true
	case _ => false
	}.map(_.metrics("spill_count").value).get
	}
	val spillCount =
	shuffled.queryExecution.executedPlan.collectFirst {
	case e: CometShuffleExchangeExec => e.metrics("spill_count").value
	}.get

[andygrove]

Calling collectFirst on the plan directly resolves to TreeNode.collectFirst, which stops at AdaptiveSparkPlanExec since that is a LeafExecNode, so it would miss the exchange whenever AQE is on.

I went with the AdaptiveSparkPlanHelper version instead, which unwraps AQE and query stage nodes in allChildren. Same single partial function you suggested, just the curried form:

spillCount = collectFirst(shuffled.queryExecution.executedPlan) {
  case e: CometShuffleExchangeExec => e.metrics("spill_count").value
}.get

Kept it as an assignment to the outer var so the enclosing helper method still returns the count. Applied in 44f2f0f.

[comphead]

perhaps we can use an Option<usize> ?

[andygrove]

Done

[comphead]

some non blocking minors

[mbutrovich]

Thanks @andygrove! Do we have any guidance on how we might tune this?