perf: charge the local shuffle writer's buffers to the memory pool

[dwsmith1983]

Which issue does this PR close?

Closes #6196.

Rationale for this change

The local shuffle writer allocates buffers of spark.comet.shuffle.native.writeBufferSize bytes that no reservation covers: the data file BufWriter, the spill file BufWriter, the buffer finish_partition reads spilled ranges through, and the recycled scratch that blocks are encoded into, plus the zstd context when zstd is configured. With the 1 MiB default, a multi-partition task that has spilled holds about 4 MiB that neither Comet's pool nor Spark's TaskMemoryManager sees, so it has to fit in spark.executor.memoryOverhead.

What changes are included in this PR?

• MultiPartitionShuffleRepartitioner::try_new hands the writer a second reservation on the consumer it already registers, reservation.new_empty(), through a new PartitionWriter::attach_buffer_reservation with a no-op default. No new consumer, so fair_unified's per-consumer share is unchanged, and the existing reservation, maxBufferBytes and memory_spilled_bytes are untouched: spill() still frees only the main reservation.
• LocalPartitionWriter keeps that reservation equal to the bytes its buffers hold: the data file buffer capacity, the spill file buffer capacity, the copy buffer length, the scratch capacity and the zstd context size. It syncs at the end of write, finish_partition, write_burst_complete and finish_all, on every exit. resize makes no pool call when the size is unchanged, so steady state adds no pool traffic per batch; in practice it is one grow and one shrink per spill when zstd is on, plus the one-time buffer grows.
• Buffers allocated up front ask for the full write buffer size with try_grow. When the pool refuses, the buffer is 8 KiB and that is what gets charged. The data file buffer is built before the consumer registers, so on a refusal it is rebuilt at the fallback size while still empty. The spill file and copy buffers are reserved after their file is open, so a failed open leaves no charge. Reads of a range larger than the copy buffer already go through io::copy.
• Memory that already exists when it is charged (the scratch, the zstd context) uses grow, which Comet's pools record as overcommit. The scratch's transient peak between syncs, the write buffer size plus one block, stays uncharged.
• The single-partition, empty-schema and remote shuffle writers are unchanged; they have no reservation to charge, and registering one would change fair_unified's shares. The tuning guide and the native shuffle contributor doc list them as untracked.

Benchmark

A pressure-triggered spill in the shuffle crate: 200 hash partitions, a 16 MiB GreedyMemoryPool, 1 MiB write buffer, Lz4, 256 MiB of Int64 input in 8192-row batches, release build, medians of three runs.

	spill count	memory_spilled_bytes	spilled bytes	insert (spills)	final write	total
main	36	607,136,768	153,994,145	692 ms	25 ms	717 ms
this PR	39	613,866,496	153,571,301	709 ms	36 ms	745 ms

The buffers take about 2 MiB of the 16 MiB pool, which is the three extra spills. The spill file buffer's grow runs during a pressure spill while the main reservation is still full, so it gets the 8 KiB fallback for the rest of the task, which is the slower final merge. Reserved bytes after the write: data buffer 1 MiB, spill buffer 8 KiB, copy buffer 1 MiB, scratch 38 KiB. Upgrading that buffer to full size right after a spill frees the main reservation, when the pool has room again, is a possible follow-up.

How are these changes tested?

Shuffle crate tests with a GreedyMemoryPool: the buffers are charged after attach, after a spill and after the write, the main reservation is empty after a spill, and the pool reads zero after drop; with pools of 1.5 MiB and 512 KiB the refused buffers fall back to 8 KiB, the reservation equals the bytes held after every step, and the output is byte-identical to a 1 GiB pool run; a refused copy buffer copies a range larger than it through io::copy and a smaller one through pread; with maxBufferBytes set, spill_count and memory_spilled_bytes match a run without buffer charging; with zstd, the reservation grows by the context size after a write and shrinks after write_burst_complete and finish_all; growing only the encode scratch moves the reservation by exactly the scratch's growth.

The shuffle crate, the memory pool tests, clippy, fmt and cargo bench --no-run pass; CometNativeShuffleSuite and CometShuffleSuite pass on Spark 3.5 and CometNativeShuffleSuite on 4.1.

[sunchao]

Summary

• Prior state and problem: Multi-partition local shuffle retained file buffers, encoding scratch, and zstd workspace without charging them to the memory pool.
• Design approach: Give the writer a reservation.new_empty() sibling on the existing consumer, separate from buffered input batches.
• Correctness / compatibility analysis: No introduced P1/P2 issues found within this review. Reservation ownership, error cleanup, spill copying, and partition offsets remain consistent. Compared Spark memory acquisition and shuffle commit contracts across 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Partition assignment and IPC format are unchanged.
• Key design decisions: Reusing the consumer preserves fair-pool shares. Separate reservations preserve maxBufferBytes and spilled-memory metrics. Refused file/copy buffers fall back to at most 8 KiB.
• Implementation sketch: One reservation hook and a shared allocation helper keep accounting with the buffer owner. Capacity-based synchronization includes retained zstd memory and handles failed writes.
• Behavioral changes worth calling out: Accounting for these buffers can trigger earlier spills. Small fallback buffers trade I/O throughput for lower memory usage. Unchanged reservation sizes avoid pool calls. Single-partition, empty-schema, and remote writers retain their existing behavior.
• Suggested improvements: None at the P1/P2 bar.

Reviewed the full eight-file diff from base 65a0cda1cf62877a37ec1a1f5f9ebc9dd1405ebd to head f0aa284d8bcdc7c75826c8c409f1bc1ee8329be7. Routed skills: review-comet-pr, review-comet-memory-pr, and review-comet-shuffle-pr. The PR is not a draft. The discussion snapshot contained no existing reviews, issue comments, inline comments, or threads.

Validation: All 177 shuffle unit tests passed. Two disposable fault-injection tests also passed, covering accounting after encoding failures and failed spill-file reads under ample and zero-capacity pools. The checkout was restored clean.

Exact-head CI: Comet CI and CodeQL report action_required. The label workflow remains queued. There is no completed CI test verdict. Local validation did not rerun JVM end-to-end suites, the Spark SQL matrix, production JNI-backed pool tests, or benchmarks.

[andygrove]

This is a light fully automated review since there are so many PRs open.

buffer_bytes_held at native/shuffle/src/writers/local/local_partition_writer.rs:195 charges codec_context.retained_bytes(), and that workspace is allocated by libzstd through libc malloc (zstd-sys only swaps the allocator on wasm). So with this PR the pool tracks C library memory that the AccountingAllocator behind native_allocated never sees, and two docs now say the opposite. The "Non-Rust allocations" bullet at docs/source/contributor-guide/memory_management.md:406 says neither the memory pool nor the allocation counters see such memory. Step 3 of "Sizing the Overhead from the Memory Usage Log" at docs/source/user-guide/latest/tuning/memory.md:175 says neither allocated nor reserved includes memory from native C libraries such as zstd.

When spark.comet.shuffle.compression.codec is zstd, reserved now holds one context per multi-partition shuffle task that is mid-spill or in its final write (about 1.3 MiB at the default level and close to 8 MiB at levels 7 and 8), and that also comes straight off the allocated - reserved difference the tuning page sizes the overhead from. Charging it still seems right to me, since it is real memory Spark should budget for. Could those two passages call out the shuffle writer's zstd context as the exception, alongside the list this PR already updates higher up in the same tuning page?

[dwsmith1983]

Could those two passages call out the shuffle writer's zstd context as the exception, alongside the list this PR already updates higher up in the same tuning page?

Done in 0d23cb4. The "Non-Rust allocations" bullet in memory_management.md now says the pool tracks such memory only where an operator charges it by hand, and names the multi-partition shuffle writer's zstd context as the one case today, in reserved and never in allocated. Step 3 of the sizing section in tuning/memory.md says the same, and that the allocated - reserved difference understates what has to fit in the overhead by one context per such task, with the two sizes you gave.

[dwsmith1983]

@sunchao CI passed on a888fbeb6 and it has your approval. Could it go into the merge queue?

[sunchao]

Summary

• Prior state and problem: Multi-partition local shuffle retained file buffers, encoding scratch, and zstd workspace without charging them to the memory pool.
• Design approach: Give the writer a separate reservation.new_empty() reservation on the existing consumer and synchronize it with retained buffer sizes.
• Correctness / compatibility analysis: No introduced P1/P2 issues found within this review. Spill replay, partition offsets, reservation cleanup, and fallback-buffer output remain consistent. Checked Spark memory-acquisition and shuffle-commit contracts against upstream sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0.
• Key design decisions: Sharing the consumer preserves fair-pool shares. Keeping buffer accounting separate preserves maxBufferBytes and spilled-memory metrics. Refused file and copy buffers fall back to at most 8 KiB.
• Implementation sketch: A reservation hook and shared allocation helper keep accounting with the buffer owner. Capacity synchronization handles retained scratch and zstd memory without pool calls when the size is unchanged.
• Behavioral changes worth calling out: Compared with branch-1.1, the intended changes reduce available pool headroom and can trigger earlier spills. Smaller fallback buffers trade throughput for lower memory use. The author’s pressure benchmark reports approximately 4% higher runtime with three additional spills. Single-partition, empty-schema, and remote writers retain their previous behavior.
• Suggested improvements: None at the P1/P2 bar. The existing zstd accounting documentation concern is addressed.

Reviewed the full nine-file diff from base fef94f6cd78b18151dff57b7a936798385356de5 to head 8274c7a55b6bf44f7f7126941f26f34377685ef2, including existing reviews and discussions. The PR remains open and non-draft. Routed skills: review-comet-pr, review-comet-memory-pr, and review-comet-shuffle-pr.

Validation: All 177 shuffle unit tests passed. Two disposable fault-injection tests also passed, checking accounting after encoding failures and failed spill-file opens with ample and zero-capacity pools. The checkout was restored clean.

Exact-head CI: Comet CI remains in progress. At final inspection, 17 checks succeeded, 14 were skipped, and seven remained running, with no reported failures. Running checks include Rust tests and Spark 4.1 shuffle tests. CodeQL passed. Spark SQL and macOS jobs were skipped. Local validation did not run JVM end-to-end suites, the Spark SQL matrix, production JNI-backed pool tests, or benchmarks.

[dwsmith1983]

@andygrove when you get a chance, could you review this and approve the CI run? The branch is synced with main.