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, plus one try_grow per spill while the spill file buffer is below full size.
• 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.
• A spill under pool pressure opens the spill file while its batches are still charged, so its buffer usually starts at 8 KiB. spill() now calls write_burst_complete after reservation.free(), and the local writer then asks for the full size again (PartitionedSpill::restore_buffer): try_grow of the difference, then a new BufWriter at full size that takes the old buffer's bytes. The bytes always fit the larger empty buffer, so the swap writes nothing to the file and cannot fail, which matters because write_burst_complete returns nothing. A refusal leaves the buffer as it was, and the next spill asks again. Once grown, the spill buffer stays charged until the final write, so under steady pressure the batches have that much less room, as they would with a full-size buffer from the start; a pool smaller than about one write buffer keeps the 8 KiB buffer. The data file buffer is sized when the writer attaches, before this consumer holds anything, so pressure from the writer's own batches does not shrink it. The copy buffer is sized during the final write while the last batches are still charged and is not grown back; a range larger than it goes through io::copy instead of one pread.
• 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

shuffle_bench over TPC-H lineitem (the first 6,001,215 rows of the SF10 file, the SF1 row count), 200 partitions, lz4, no buffer limit, ci profile, 7 rounds per pool with the three builds' order rotated each round, in a Linux container. Medians of the shuffle write time, with spill counts:

pool	main	this PR before the restore	this PR
16 MiB	0.128 s (67 spills)	0.130 s (73 to 77)	0.130 s (76 to 77)
32 MiB	0.127 s (34 to 35)	0.146 s (35 to 37)	0.129 s (36 to 37)
64 MiB	0.126 s (17)	0.130 s (18)	0.120 s (18)

Charging the buffers takes about 2 MiB of the pool, which is the few extra spills. Before the restore, the spill buffer stayed at 8 KiB after a pressure spill for the rest of the task; with it, the buffer is back to full size after each spill and write time is back at main's.

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; with no buffer limit and pools of 1, 4 and 16 MiB, the spill buffer is back to full size after every pressure spill, the reservation equals the bytes held, and every row is written once; at the spill file level, a buffer the full pool cut to 8 KiB stays there while the pool is full, then grows by exactly the difference with its buffered bytes moved and nothing written until the flush, and a spill file that failed a write is not grown back.

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.

[andygrove]

With maxBufferBytes at its default of 0, every spill is a pressure spill, so ensure_spill_file_created asks for the spill file buffer while the batch reservation that just failed to grow is still full. That try_grow is usually refused, and it always is when a batch is smaller than the write buffer, so the buffer drops to 8 KiB. It then stays at 8 KiB for the rest of the task, even though the spill frees the whole batch reservation a moment later. I checked this with a GreedyMemoryPool probe at 4, 16 and 64 MiB with no buffer limit. The spill buffer finished the task at 8 KiB every time, including the 16 MiB run that had 15 MiB free before the final write.

On TPC-H SF1 lineitem through shuffle_bench (200 partitions, lz4, 16 MiB pool, ci profile, 7 interleaved runs), shuffle write time went from 0.140s on main to 0.185s with this PR, and the total went up about 1%. At 32 and 64 MiB I couldn't see a difference.

Could spill() call write_burst_complete after self.reservation.free(), and have the local writer retry the full-size spill buffer there? PartitionedSpill can try_grow the difference, flush the small buffer, take the writer back with BufWriter::into_parts, and wrap it in a full-size one. I tried that locally. Write time went back to 0.142s and all 179 shuffle tests still pass. A test could spill under pool pressure with no buffer limit and check that the spill buffer is back to full size once the spill has freed its batches.

[dwsmith1983]

Could spill() call write_burst_complete after self.reservation.free(), and have the local writer retry the full-size spill buffer there?

Done in adb9daa. spill() now calls write_burst_complete after reservation.free(), and PartitionedSpill::restore_buffer grows the reservation by the difference with try_grow, takes the writer back with BufWriter::into_parts, and wraps it in a full-size BufWriter that takes over the buffered bytes. It moves those bytes instead of flushing them: they always fit the larger empty buffer, so the swap writes nothing and cannot fail, which matters because write_burst_complete returns nothing. A refusal leaves the 8 KiB buffer in place and the next spill asks again. A spill file that already failed a write is not grown.

The data file buffer is sized when the writer attaches, before the consumer holds anything, so the writer's own batches don't shrink it. The copy buffer is sized during the final write while the last batches are still charged and is not grown back; a range larger than it goes through io::copy.

pressure_spill_restores_the_spill_buffer spills with no buffer limit in 1, 4 and 16 MiB pools and checks that the spill buffer is back to full size after every spill, that the reservation equals the bytes held, and that every row is written once.

shuffle_bench on lineitem at the SF1 row count, 200 partitions, lz4, no buffer limit, ci profile, 7 rounds with the build order rotated, in a Linux container. Median shuffle write time:

Pool	main	before	after
16 MiB	0.128 s	0.130 s	0.130 s
32 MiB	0.127 s	0.146 s	0.129 s
64 MiB	0.126 s	0.130 s	0.120 s

On this data the 8 KiB buffer showed up at 32 MiB rather than 16.

[andygrove]

Thanks for the quick turnaround. restore_buffer is what I was hoping for. spill() frees the batches first, write_burst_complete then asks for the full-size spill buffer again, and the buffered bytes move across without touching the file. Keeping the reservation equal to the bytes held after every spill makes it easy to trust.

I reran this on 09977e3c44 with shuffle_bench on SF1 lineitem (200 partitions, lz4, ci profile, M3 Max, quiet machine, 9 to 11 interleaved rounds). Median write time:

pool	main	before the restore	now
16 MiB	0.166 s (67 spills)	0.206 s (73)	0.169 s (75 to 77)
8 MiB	0.185 s (123 to 125)	0.270 s (147)	0.202 s (179 to 180)

The 8 KiB buffer regression is gone. What is left is the pool space the buffers now occupy. Main at a 14 MiB pool spills 75 to 76 times and has the same 0.169 s median as this PR at 16 MiB, so the remaining cost is the accounting and not the buffer size.

I also checked these.

• All 181 shuffle crate tests pass, and clippy and fmt are clean.
• I broke the new code eight ways. I moved write_burst_complete back before free(), made restore_buffer a no-op, skipped the try_grow, charged the full size instead of the difference, dropped the buffered bytes, ignored failed, kept the small capacity, and never called it from the writer. Each one fails a test.
• With 512 and 2048 partitions a restore moves 4.7 to 7.6 KB of buffered bytes, and every row still lands once in the output.
• CometNativeShuffleSuite and CometNativeShuffleWriterSuite pass against this head on the default Spark profile.

Approving.

[sunchao]

Summary

• Prior state and problem: Multi-partition local shuffle retained file buffers, copy storage, encoding scratch, and cached zstd workspace without charging them to the memory pool, understating its memory use.
• Design approach: Give the local writer a separate reservation.new_empty() reservation on the repartitioner's existing consumer and synchronize it with retained buffer sizes.
• Correctness: Checked reservation ownership, error cleanup, spill replay, partition offsets, and buffer restoration. Restoration preserves pending bytes without writing to the file, and failed spill files remain unusable. No introduced P1/P2 issues found within this review.
• Compatibility analysis: Compared Spark's memory-grant and shuffle-commit contracts against sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Partition assignment, IPC format, and offset-to-length conversion remain unchanged. The six affected Rust files on branch-1.1 match the supplied base, so the release comparison shows the same intended accounting changes.
• Key design decisions: Reusing the consumer preserves fair-pool shares. Separate reservations preserve the meaning of maxBufferBytes and memory_spilled_bytes. Refused file and copy buffers fall back to at most 8 KiB so spilling can continue.
• Implementation sketch: reserve_buffer handles allocation requests, while buffer_bytes_held and sync_buffer_reservation account for retained capacities and zstd workspace. spill() releases batch memory before write_burst_complete retries the full-size spill buffer.
• Performance: Unchanged reservation sizes avoid pool calls, and successful restoration requires only a bounded buffer copy. The existing exact-head benchmark reports write medians of 0.169 versus 0.166 seconds at 16 MiB and 0.202 versus 0.185 seconds at 8 MiB. That discussion demonstrates the persistent 8 KiB regression is fixed and attributes the remaining cost to reduced batch headroom and additional spills. I did not rerun benchmarks.
• Design: Keeping accounting with the buffer owner makes allocation, retention, and release straightforward to follow. The change improves pool accounting without claiming to impose a hard limit on physical memory.
• Abstraction & complexity: The default no-op PartitionWriter hook accommodates backend-specific accounting without registering another consumer. The error-path synchronization wrappers serve a concrete purpose. No unnecessary abstraction meeting the P1/P2 reporting bar was identified.
• Behavioral changes worth calling out: Charged buffers can trigger earlier pressure spills. Scratch growth between synchronization points remains uncharged, and infallible accounting can record overcommit. Single-partition, empty-schema, and remote writer implementations retain their existing behavior. The documentation now distinguishes reserved zstd memory from Rust allocation counters.
• Suggested improvements: No additional improvements meeting the P1/P2 reporting bar were identified. The existing spill-buffer performance concern and zstd documentation concern are addressed at this head.

Reviewed the complete nine-file diff from base 00a4b422f0ed6560ef76dce004b94c3697613108 to head 09977e3c44f4752d98ec941ca50c721a0ce7b213. The PR remains non-draft. Read the supplied reviews and conversation before forming conclusions. Routed skills: review-comet-pr, review-comet-shuffle-pr, and review-comet-memory-pr.

Validation: All 181 shuffle unit tests passed locally. Two additional disposable failure probes passed, covering accounting after encoding failures and failed spill-file opens with ample and zero-capacity pools. Temporary source changes were restored, and the checkout is clean.

Exact-head CI: Comet CI passed. GitHub reports 23 successful checks, 15 skipped checks, and no failures or unfinished checks. Successful jobs include Rust tests, the Spark 4.1 Comet suites, and TPC-H/TPC-DS verification. Spark SQL, macOS, Iceberg, and benchmark checks were skipped. Local validation did not rerun JVM suites, the Spark SQL version matrix, production JNI-backed memory contention, or benchmarks.