perf: reuse per-partition scratch in the shuffle write path - #5568
perf: reuse per-partition scratch in the shuffle write path#5568dwsmith1983 wants to merge 6 commits into
Conversation
The multi-partition write path pays two allocation costs per partition per event: BufBatchWriter starts from an empty byte buffer and regrows it toward the 1MB write-buffer size, and the interleave iterator materializes the whole partition index list as a fresh usize-pair vector (16 bytes per row) plus a batch-ref vector over all buffered batches. The task-level writer now owns one byte buffer and recycles it through the sequential partition loops, drained between partitions with its capacity kept, so an event allocates the buffer once instead of once per partition and holds exactly one regardless of partition count. The interleave path builds the batch-ref list once per write cycle and converts indices per output chunk into a small reusable scratch instead of widening the full list up front. Wire format and produced batches are unchanged, pinned by byte-level tests. Write time drops ~27% on a 2000-partition zstd shuffle bench; the larger effect is allocator pressure under concurrent tasks, which a single-task bench understates. Part of apache#5002.
|
@sunchao the Iceberg 1.11 / Spark 4.1 runtime leg failed before any tests ran — Gradle could not fetch the shadow plugin from plugins.gradle.org (TLS handshake dropped during project configuration). Same leg passed on #5565 shortly before. Could you rerun that job when you get a chance? I don't have permissions to. |
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
I checked this out and ran the output through the real ShuffleWriterExec path on both main and this branch across 15 configurations, hashing the data and index files each time: 2000-partition hash, single partition, forced spilling via max_buffer_bytes, a 64-byte write buffer to force a flush per block, across None, Lz4Frame and Zstd(1). Every digest matched, which covers the spill path that the new unit tests don't reach. All 94 shuffle tests pass and clippy is clean.
The win reproduces and it's bigger than you claimed. Running PR, main, main, PR alternating to control for drift, this branch was fastest in every configuration: roughly 10 to 15% end to end at 2000 partitions and 4 to 8% at 8000, flat at 200. Those are whole-plan numbers including hashing, so your write-phase measurement looks right. Nice result.
A few things I'd like to sort out before this goes in, mostly around the buffer's lifetime and settling the convention with #5565 while both are still open. Those are inline.
One more that doesn't anchor anywhere. Could you add a high-partition case to native/shuffle/benches/shuffle_writer.rs? It tops out at 16 partitions today, so "the criterion suite shows no regressions" is accurate but the suite can't actually see this change in either direction, and the whole argument for the PR is high partition counts. I added a throwaway bench at 200/2000/8000 and the win is clear at 2000 and up, so a case there would both demonstrate it and guard it, and it would give PR Benchmark Check something to watch.
| /// Partitions are written strictly one at a time, so a single buffer keeps its | ||
| /// grown capacity across the whole task instead of every partition regrowing a | ||
| /// fresh allocation toward the write buffer size. | ||
| recycled_buffer: Vec<u8>, |
There was a problem hiding this comment.
One thing I want to make sure we've thought about. write_batch_to_buffer appends a whole block and only then checks pos >= buffer_max_size, so the buffer's high-water mark is the write buffer size plus one full block, not the write buffer size. Previously the writer was dropped at the end of each partition and that peak went back to the allocator. Now it lives in recycled_buffer until the LocalPartitionWriter drops, and nothing charges it to the shuffle reservation. The spill path is where I'd worry, since spills happen under memory pressure and we'd now be holding an extra untracked allocation between them. Would a shrink_to(write_buffer_size) when you store it back be enough to bound that? #5565 bounds retention for its codec context, so it would be good for these two to agree.
Related, and worth fixing since we squash-merge: the description says this is "a strict improvement for memory at high partition counts". Partitions are already written one at a time, so at most one buffer was ever live and peak is unchanged. The real change is that the peak is retained rather than released. Could you reword that so the tradeoff is on the record?
Separately, into_buffer() will happily hand back a buffer with bytes still in it, and new() clears whatever it's given. Both callers flush first so this is fine today, but if someone later reuses a buffer from an unflushed writer the encoded bytes just disappear and it surfaces as missing rows two stages downstream. Could into_buffer() carry a debug_assert!(self.buffer.is_empty())? The doc comment already states the invariant.
There was a problem hiding this comment.
[P2] Could you extend the requested benchmark with a wide-then-narrow forced-spill case to quantify this retained-buffer tradeoff? At 3b4425cb, a successful spill returns the byte buffer to LocalPartitionWriter before freeing the reservation, and the write threshold is checked only after a whole block is encoded. This establishes a lifetime/accounting change, not a measured RSS regression.
One bounded case would use an early hot-key batch of 8,192 rows with 8-KiB Binary values, followed by small values with the same schema. Compare that with reversed-order and narrow-only controls, using None compression and a forced-spill threshold. Compare base 199a910bd8623cac3db82d70f9bfb24865613c2c against head 3b4425cb4b0c7e0688248f6c2571c08307c3fdb1 with the same release build, allocator, hardware, input seed, warmup and repeated runs.
Could you report total time, allocation/growth bytes, retained buffer capacity and tracked reservation after each spill, peak/live memory, spill bytes, and matching decoded rows/order? This would complement the existing high-partition throughput request and help choose whether to cap, release or account for the retained buffer.
There was a problem hiding this comment.
All three taken: flush now ends with shrink_to(write_buffer_size), so retained capacity is bounded at the configured size (mid-partition the high-water can still reach the cap plus the largest block, same as before, only what is retained changed); the description is reworded to say the peak is retained rather than released, and now capped, instead of claiming an improvement; and the drained-scratch invariant is debug-asserted with a test. The into_buffer question dissolved with the borrow change from your other comment.
There was a problem hiding this comment.
Ran it with your shape: an 8,192-row hot batch of 8 KiB Binary payloads followed by ~500k rows of 16-byte payloads, plus reversed-order and narrow-only controls; codec none, 2,000 partitions, max-buffer-bytes 4 MiB (about 10 spills per run), 3 iterations after warmup, base vs head on the same build, machine, and input seed:
| input | base avg | head avg |
|---|---|---|
| wide-then-narrow | 0.401s | 0.384s |
| narrow-then-wide | 0.421s | 0.413s |
| narrow-only | 0.398s | 0.397s |
Head is slightly faster on all three, most on the wide-then-narrow case, so no time regression from retention. On the memory side the answer changed while you were writing the comment: flush now shrinks the scratch back to the configured write-buffer size, so retained capacity after the wide batch is bounded by construction and pinned by a test rather than instrumented at runtime, which I think lands on the "cap" option among cap/release/account. Spilled bytes and outputs match within the run-to-run variation of the spill trigger (Arrow's memory estimates shift block boundaries by a spill or so even between runs of the same binary, so byte-exact cross-tree diffs are not meaningful under forced spilling; row-level equality is pinned by the unit tests and the digest matrix in the review above). I did not instrument peak RSS or live reservation traces.
| writer: W, | ||
| buffer_max_size: usize, | ||
| batch_size: usize, | ||
| mut buffer: Vec<u8>, |
There was a problem hiding this comment.
How do you feel about matching the shape #5565 uses? It threads its task-scoped state as a &mut ShuffleCodecContext parameter on write and flush, where this moves a Vec<u8> in through the constructor and back out through into_buffer(). If both land we end up with two conventions for the same idea in the same struct, and whichever rebases second is unlikely to revisit it.
Borrowing would also clean up a small wart. The mem::take empties the caller's buffer up front and the write-back only runs on success, so any error in between quietly ends recycling for the rest of the task. Not a correctness issue, but it goes away on its own if the buffer is borrowed.
There was a problem hiding this comment.
Done, write and flush borrow the scratch per call now, same shape as the codec context, and the constructor param and into_buffer are gone. You are right that it also fixes the take/put-back wart; an error no longer ends recycling. The one thing ownership gave for free was buffer identity, so the writer now records the scratch address on first use and debug-asserts every later call passes the same one, a swapped buffer would silently drop un-flushed bytes otherwise.
| pub(super) fn produce<'a>( | ||
| &'a mut self, | ||
| &'a self, | ||
| refs: &'a [&'a RecordBatch], |
There was a problem hiding this comment.
The indices are positions into self.buffered_batches, but refs here is unconstrained, so a slice built from the wrong producer or a filtered one compiles fine and shuffles the wrong rows instead of failing. Would you mind adding a debug_assert_eq!(refs.len(), self.buffered_batches.len()) at the top? Cheap, and it catches the realistic version of that mistake.
There was a problem hiding this comment.
Added, with a debug-build test pinning that a truncated refs slice panics.
…city The write scratch is now threaded through write and flush as a borrow instead of moving through the constructor, so one convention covers all task-scoped shuffle scratch, and an error mid-partition can no longer strand the buffer inside a dropped writer. The writer asserts a fresh scratch is drained and that every call passes the same buffer, since a swapped buffer would silently lose unflushed bytes. Flush shrinks the scratch back to the configured write-buffer size, bounding retained capacity that could previously reach the buffer size plus the largest block. The producer asserts the batch-ref slice covers every buffered batch, and the criterion suite gains a high-partition group so partition scaling stays visible to benchmark checks.
|
Thanks for running the digest matri, good to have the spill path independently covered, and glad the win held up end to end. The criterion suite now has a high-partition group (200/2,000/8,000, hash partitioning, no compression) so PR Benchmark Check can actually watch this; on my machine it lands at 2.75/9.44/21.5 ms respectively. The retention and convention comments are addressed inline, and the description is reworded per your squash-merge point. |
|
@sunchao |
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
Took another look at 08ec7be48, since the reshaping commit moved enough that my earlier pass no longer covered it. I re-ran an end-to-end comparison against the merge base, hashing the data and index files across 144 configurations (three input shapes including the 8 KiB Binary hot batch, partition counts 1 through 2000, None/Lz4Frame/Zstd(1), write buffers of 1 MiB and 64 bytes, and max_buffer_bytes unset and at 256 KiB). Every digest matched, and the criterion high-partition group reproduces a 4 to 5% win against the base once I port it to the baseline tree. Two things inline.
| for batch in iter.by_ref() { | ||
| let batch = batch?; | ||
| buf_batch_writer.write(&batch, &metrics.encode_time, &metrics.write_time)?; | ||
| buf_batch_writer.write( |
There was a problem hiding this comment.
The borrow change fixed the mem::take wart I raised, but I think it opened a different one on the same path. If the iterator or the writer errors part way through a partition, finish_partition returns before flush runs, so recycled_buffer keeps whatever was already encoded. I confirmed it with a probe: inject an error after one good batch and the buffer comes back holding 668 bytes. The next partition's BufBatchWriter then seeks to the end and appends after those bytes, and its flush writes them into that partition's byte range, so a reader sees a corrupt block rather than an error.
Nothing reaches that state today, because every error here aborts the task and check_scratch catches it in debug builds. But it is the one invariant in this design that the code does not enforce, and the release-build failure mode is silent wrong data rather than a crash. Both earlier shapes were immune for free: the owned buffer died with the writer, and the move-in/move-out version in 3b4425cb lost recycling rather than corrupting anything. The borrow shape I pushed you toward is what introduced it, so this one is on me.
Would you mind draining the scratch when the write loop or flush errors, in both SpillWriter::write and the Multi arm of finish_partition? Binding the loop and flush to a result and calling recycled_buffer.clear() before propagating would make the invariant true by construction instead of by assertion.
There was a problem hiding this comment.
Good catch, this one is on me too. Both call sites now bind the loop and flush to a result and clear the scratch before propagating, so an errored partition always hands back an empty buffer. Added a failing test at each site first; they showed 1244 retained bytes and now come back empty.
| let mut writer = BufBatchWriter::new(block_writer, &mut output, large_cap, 8192); | ||
| writer.write(&batch, &mut scratch, &time, &time).unwrap(); | ||
| writer.flush(&mut scratch, &time, &time).unwrap(); | ||
| assert!(scratch.capacity() > 0 && scratch.capacity() <= large_cap); |
There was a problem hiding this comment.
The second half of flush_caps_retained_scratch_capacity cannot fail as written. With large_cap at 1 MiB and a block around a kilobyte, scratch.capacity() > 0 && scratch.capacity() <= large_cap holds for any implementation, including one where flush shrank to a single byte or did not shrink at all. The comment says it is checking that there is no over-shrinking, which is the property worth having, but with batch_size at 8192 the coalescer defers serialization to flush and shrink_to ends up a no-op, so nothing is being observed.
Could you use a batch_size below the row count the way the first half does, so write actually serializes into the scratch, then capture the capacity and assert flush leaves it alone? assert_eq!(scratch.capacity(), cap_after_write) would fail if this ever became shrink_to_fit, which is the regression the case exists to catch.
There was a problem hiding this comment.
You are right, the coalescer was buffering everything so the second half was not proving anything. The test now uses a batch size below the row count so the writes actually serialize, captures the capacity right after them, and asserts flush leaves it exactly unchanged. I checked it fails if the flush path shrinks the buffer.
An error mid-partition left encoded bytes in the recycled scratch buffer, and the next partition to reuse it would flush those bytes into its own range. Both call sites now clear the scratch before propagating the error. Also rewrites the capacity test to serialize for real and assert the exact capacity is preserved across a flush.
3 participants
Which issue does this PR close?
Part of #5002 (the two per-partition allocation items in the medium tier; the issue stays open for its remaining items). Independent of #5565 — whichever merges second gets a small mechanical rebase over the shared
BufBatchWriterconstructor.Rationale for this change
The multi-partition write path pays two allocation costs per partition per event:
BufBatchWriterstarts from an empty byte buffer and regrows it toward the write-buffer size (1MB default).finish_partitionand the spill path construct one per partition per event, so a 2,000-partition pass runs 2,000 growth cycles — roughly an extra copy of the payload through doubling-growth memcpy, plus the transient allocations.PartitionedBatchIteratormaterializes the whole partition's(u32,u32)index list as a fresh(usize,usize)vector (16 bytes per row, re-materialized every write cycle) and rebuilds a batch-ref vector over all buffered batches for every non-empty partition.What changes are included in this PR?
BufBatchWriter::newtakes a caller-owned buffer andinto_buffer()hands it back drained (capacity kept). The task-levelLocalPartitionWriterowns one buffer and recycles it through both sequential partition loops (spill and finish), so an event allocates it once and holds exactly one regardless of partition count. Peak memory is unchanged (partitions were already written one at a time, so at most one buffer was ever live); the change is that the peak is retained across partitions rather than released, andflushcaps the retained capacity at the configured write-buffer size. The scratch is borrowed per call (write/flushtake&mut Vec<u8>), matching the convention perf: reuse zstd compression contexts across shuffle blocks #5565 uses for its codec context; the long-lived single-partition writer carries its own scratch.PartitionedBatchesProducer::producetakes&selfplus a caller-built batch-ref slice (batch_refs(), built once per write cycle); the iterator borrows the raw(u32,u32)indices and converts per output chunk into a small reusable scratch (capacity ≤ batch size) instead of widening the full list up front. The empty-partition path stays allocation-free.Wire format and produced batches are unchanged, pinned by byte-level tests (recycled vs fresh buffers produce identical bytes; chunked conversion interleaves identically to full up-front conversion, tail chunk included).
Benchmarks (M-series macOS, 4M-row hash shuffle via
shuffle_bench, 3 iterations after warmup): write time drops ~27% at 2,000 partitions (0.015s → 0.011s) with the other phases flat, and the criterion end-to-end suite shows no regressions (one config improved). Honest framing: a single-task bench understates this change — the removed churn is per-partition allocator traffic, which matters most under concurrent tasks sharing the allocator, and the bounded buffer count is the memory-profile win.How are these changes tested?
Two new tests alongside the existing suites (94 in the shuffle crate, all passing, plus the core crate's 201):
cargo clippy --all-targets -- -D warningsandcargo fmtclean.