perf: reuse zstd compression contexts across shuffle blocks - #5565
perf: reuse zstd compression contexts across shuffle blocks#5565dwsmith1983 wants to merge 9 commits into
Conversation
Every shuffle block previously created and destroyed its own zstd context: a fresh CCtx per encoded block and a fresh DCtx per decoded frame. Context setup is pure overhead that scales with block count, so high-partition shuffles with small blocks pay the most. Encode paths now share one context per task, threaded from the task-level owner so codec memory stays bounded regardless of partition count. The remote shuffle path still frees the zstd workspace with each admitted encode, keeping its memory accounting accurate. Decode reuses a per-thread context behind the existing entry points. Wire format is unchanged. On a 4M-row hash shuffle with 10,000 partitions at zstd level 3, encode time drops ~10% and wall time ~7%; larger-block shapes are within noise. lz4 and snappy keep per-block encoders: no reset API in the pinned crates and much smaller setup cost. Part of apache#5002.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed exact head 8da6e6b445d8801b9e915c59c778cb7034851a70 against base 918bc7d7ac4123a02ef4bba8ddae920f76db9cbc with five independent specialist scopes.
I found two P2 native-memory lifetime regressions that should be addressed before merge:
- The thread-local decode context can retain a 128 MiB zstd workspace per executor worker thread across tasks.
- The local writer can retain up to an 834 MiB zstd compression workspace per active task outside DataFusion memory accounting.
The RSS release path is correct on success and checked error paths. I found no wire-format, row-correctness, or Rust API compatibility defect.
CI snapshot during this review: 36 passed, 27 pending, 7 skipped, and no failures. PR Benchmark Check is skipped.
Local validation:
datafusion-comet-shufflelibrary: 98 passeddatafusion-cometcore library: 201 passed, 4 ignored- Focused IPC tests: 9 passed
- Focused shuffle-scan tests: 7 passed
- Codec-context and multi-partition spill tests passed
- Exact-version zstd workspace reproducer passed
git diff --checkpassed and the worktree remained clean
Validation limits: I did not run a Spark/Celeborn end-to-end workload or reproduce the author's M-series performance numbers. The remaining CI jobs were still running when this review was submitted.
| thread_local! { | ||
| /// Backs the entry points below. They're called from many JVM task threads; a | ||
| /// thread-local gets each thread context reuse without changing any caller. | ||
| static DECODE_CONTEXT: RefCell<ShuffleDecodeContext> = |
There was a problem hiding this comment.
[P2] Could we move this decode context under reader or task ownership, or release it when it exceeds a bounded size? ResetDirective::SessionOnly preserves zstd's allocated window. With the locked zstd 1.5.7 build and the same context sequence used here, a valid 17-byte level-22 frame made DCtx::sizeof() grow from 95,992 to 134,707,000 bytes, and reset left it at that size. Both production entry points use this thread-local context, so executor worker threads retain the native allocation across tasks. At 32 threads that is about 4 GiB. The base path dropped the decoder per frame. It might be worth adding a regression that verifies the workspace is released when a reader closes and after a decode failure.
There was a problem hiding this comment.
Good catch — I had not realized SessionOnly keeps the window allocation. Went with the bounded-size option since the thread-local has no close hook: after every decode, error paths included, the context is dropped if its measured size exceeds 8 MiB (common levels sit at ~1-5 MiB, so they keep reuse; a wide-window frame pays per-frame creation like before). Added the regressions you suggested — one decodes a wide-window level-22 frame and asserts the workspace is released, one does the same through a decode failure.
| data_output: DataOutput, | ||
| /// Compression state shared by every block this task writes; the per-partition | ||
| /// `BufBatchWriter`s borrow it (see [`ShuffleCodecContext`]). | ||
| codec_context: ShuffleCodecContext, |
There was a problem hiding this comment.
[P2] Could we account for this context's retained workspace for the full local-writer lifetime, or release it at a spill or idle boundary? LocalPartitionWriter has no memory reservation for the CCtx, while the repartitioner frees its tracked reservation after spilling. With the locked zstd build, the same Encoder::with_context path retained 72,082,969 bytes at level 15 and 874,070,679 bytes at level 22 after SessionOnly reset. Those levels are accepted by the current configuration. Concurrent tasks can therefore keep large native allocations after the pool reports their buffered memory as released. A size-based regression would also catch this because the current test only checks that the context is present.
There was a problem hiding this comment.
Same treatment here, plus boundary releases: the writer drops its context when each spill event completes and after the final flush, and any block that leaves the context above 8 MiB drops it immediately. So retained memory between phases is zero and the worst case anywhere is the cap, matching the base path profile rather than adding a new reservation surface. The size-based regression drives a real repartitioner through spill and finish and checks the context is gone at both points; level 22 locally measures ~834 MB retained without the cap, which lines up with your numbers.
SessionOnly reset preserves zstd's allocated window, so a retained context grows to the largest workspace it has seen (~128 MiB for the decoder after one wide-window frame, ~834 MiB for the encoder at level 22) and stays there. Cap retained contexts at 8 MiB -- covering the commonly configured levels -- and drop anything larger after each decode (errors included) and each local block encode; higher levels fall back to per-frame creation, the pre-existing cost. The local writer also releases its context when a spill event or the final flush completes, so nothing is retained between write phases.
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
Thanks for taking sunchao's feedback on board. The release-at-boundary handling reads carefully and I like that the error paths are covered.
I checked the pinned zstd 0.13.3 / zstd-safe 7.2.4 / zstd-sys 2.0.16+zstd.1.5.7 sources and Encoder::with_context / Decoder::with_context really are drop-ins for Encoder::new / Decoder::with_buffer. Both land in the same raw + zio path, and the only extra work the owned constructors do is setting the compression level (which zstd_cctx replicates) plus DCtx::init and load_dictionary(&[]), which are no-ops when no dictionary is ever set. So I have no concerns about the wire format claim.
I also measured CCtx::sizeof() and DCtx::sizeof() against that exact zstd build after one streaming frame, since most of my comments turn on those numbers:
| encode level | CCtx::sizeof() |
|---|---|
| 1 / 3 / 6 | 1.31 / 3.49 / 5.24 MiB |
| 7 / 8 | 7.74 MiB |
| 9 / 12 / 15 | 14.74 / 44.74 / 68.74 MiB |
| 22 | 833.58 MiB |
On the decode side a level 19 frame leaves the DCtx at 8.47 MiB and level 22 at 128.47 MiB, while everything up to level 15 stays at or under 4.47 MiB.
One point that is not tied to a line. The benchmark table in the description was measured at 8da6e6b, before abf5807 added the per-block sizeof() check and the release at every spill boundary. Levels 1, 3 and 6 all sit well under the 8 MiB cap so I would expect the gains to survive, but the 10,000 partition level 3 case is both the headline result and the one where the new per-spill release actually lands. Could you re-run against 9afe4eb and update the table? It would be good for whoever merges this to be judging the numbers the code actually produces.
| use std::cell::RefCell; | ||
| use std::io::{Error, ErrorKind, Read}; | ||
|
|
||
| thread_local! { |
There was a problem hiding this comment.
The 8 MiB cap turns the unbounded retention into a bounded one, which is a good fix, but the memory is still held for the life of the thread and still is not visible to any reservation. Both production callers go through the thread-local: the static JNI decodeShuffleBlock and ShuffleScanStream. Those run on JVM task threads and tokio workers, all of which live as long as the executor. A 16 core executor that decodes one zstd shuffle block ends up sitting on roughly 128 MiB of native memory for the rest of its life, including during stages that never shuffle.
ShuffleScanStream looks like a natural owner here. Could decode_shuffle_batch take a &mut ShuffleDecodeContext held by the stream and go through read_ipc_compressed_with? That would bound retention to the operator rather than the thread, and it would leave the thread-local for Java_org_apache_comet_Native_decodeShuffleBlock, which really has no handle to hang a context off. Right now read_ipc_compressed_with and read_ipc_compressed_validated_with are exported from lib.rs but only ever called from tests, so this would also give them a real caller.
There was a problem hiding this comment.
Done — the scan operator owns the decode context now (on the exec rather than the stream since the decode loop runs exec-side, where JNI calls are allowed; retention still dies with the operator). The thread-local is down to one production caller, the static decodeShuffleBlock entry, and the _with variants have a real caller.
| /// Largest zstd workspace worth caching between frames. Covers the commonly configured | ||
| /// levels; higher levels (tens to hundreds of MiB of window) fall back to a fresh context | ||
| /// per frame, which is what per-block encoding paid anyway. | ||
| const MAX_RETAINED_ZSTD_CONTEXT_BYTES: usize = 8 * 1024 * 1024; |
There was a problem hiding this comment.
I measured CCtx::sizeof() against the pinned zstd 1.5.7 build after one streaming frame and the cap is closer to the edge than the comment suggests. Levels 7 and 8 come in at 8,119,825 bytes against a cap of 8,388,608, so about 3% of headroom. Levels 1 through 6 are 1.31 to 5.24 MiB and level 9 jumps to 14.74 MiB, so anything at 9 or above never reuses at all. On the decode side a level 19 frame leaves the DCtx at 8.47 MiB, which also misses the cap.
Two things that would help. Could the measured level to size table go in the comment next to the constant, so the choice of 8 MiB is traceable and someone bumping zstd-sys can see what they are moving? And could a test pin where the boundary actually falls, say level 6 retains and level 9 does not? As it stands a routine dependency bump could push levels 7 and 8 over the line and silently disable the optimization for those users with every test still passing.
There was a problem hiding this comment.
Measured the full table against the pinned build and got the same numbers you did — it is in the comment next to the constant now, with the zstd-sys version. Boundary tests pin levels 6 and 8 retained (8 being the ~3% edge, so a bump that crosses it fails loudly) and level 9 recreated per block; the decode side pins level 1 retained and level 19 released.
| // write header | ||
| output.write_all(&self.header_bytes)?; | ||
|
|
||
| let encode_result = |
There was a problem hiding this comment.
This is right given that rss_codec_workspace is charged per admitted invocation, and I checked that it fires on the error paths as well as the success path. The consequence though is that RSS allocates and frees a CCtx per block exactly as it does on main, so the remote path gets none of the benefit this PR is after. Small frames pushed to Celeborn are arguably the shape where per-block context setup hurts most.
Was reserving the workspace once for the lifetime of the RssPartitionWriter rather than per invocation considered? There is already one writer per task, so the accounting would be a single up-front charge instead of a repeated one. If that turns out to be awkward against the pusher's admission model it would be worth saying so in the description, which currently reads as though both paths benefit.
There was a problem hiding this comment.
I looked at a lifetime reservation: the pusher contract treats reserve, push, and release as one synchronous invocation with a single outstanding reservation, and release takes no amount, so a standing charge would need the JNI contract extended on both sides. Kept RSS per-invocation (same cost as main) and updated the description so it no longer reads as if both paths benefit. Happy to look at extending the pusher contract as a follow-up if there is appetite.
| /// spill event and the final shuffle write each end with the context released. | ||
| #[tokio::test] | ||
| #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` | ||
| async fn local_writer_releases_zstd_context_at_burst_boundaries() { |
There was a problem hiding this comment.
This test only ever asserts !holds_zstd_cctx(), so it would pass just as happily if the context were never created in the first place. That is true of the suite generally. The tests establish that blocks round-trip correctly under a shared context and that release happens at the right boundaries, but nothing observes that N blocks produce fewer than N context creations, which is the actual claim of the PR.
Would a test-only creation counter on ShuffleCodecContext work? Asserting that one spill burst over two partitions creates exactly one context would pin the reuse behaviour directly, and the same counter would let you pin the level boundary from my comment on codec_context.rs.
There was a problem hiding this comment.
Added the counter, test-gated on both context types. The burst test asserts exactly one creation across a two-partition spill burst and two after the finish burst, and the boundary tests from your other comment use the same counter.
The shuffle scan operator now owns its zstd decode context and passes it through the caller-owned decode entry points, so retained memory dies with the operator instead of living as long as the executor thread; the thread-local remains only for the static JNI decode entry, which has nothing to own a context. The retained-size cap gains a measured level-to-workspace table next to the constant and boundary tests that fail loudly if a zstd upgrade moves levels across the cap, and test-only creation counters pin that a multi-partition burst creates one context rather than one per block.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed 7c907a84 against 199a910b. I found no new P1/P2 correctness finding in the current source. The writer benchmark rerun remains unanswered, and the current-head workflows are still action_required, so the performance and execution evidence is not yet sufficient for approval.
Could you also add a focused decoder-reuse microbenchmark against the PR base? Using identical prebuilt frames and fixed total decoded bytes, compare repeated small zstd frames with a large-frame control and a sequence where the measured DCtx exceeds 8 MiB before returning to small frames. Representative numeric and nullable-string data, plus a NONE control, would help distinguish reuse savings from reset/locking overhead. Please report decode time, allocation/context-creation counts, and peak versus retained native memory, including after scan-owner release and static JNI worker reuse. Matching builds/dependencies, repeated warmups, identical decoded results and confirmation of the native reader path would make the comparison useful. This is a request to validate the tradeoff, not a claim that a regression has been measured.
Prebuilt shuffle frames (numeric plus nullable string data) decoded through one reused context and through a fresh context per frame: repeated small zstd frames, a large-frame control, a sequence where a wide-window frame pushes the retained workspace past the cap before small frames resume, and an uncompressed control. Decoded results are asserted identical across variants before anything is measured.
|
Re-ran the writer benchmark at the current head against the merge base, per @andygrove's request, same input and shapes; description table updated:
The retained-size checks and boundary releases cost nothing measurable at 2,000 partitions, and the 10,000-partition result holds at the head the code actually produces (~4% wall, ~9% encode). @sunchao added
Decode reuse is throughput-neutral here, about 56 ns/frame of context setup against ~114 us/frame of decode work, and the over-cap drop/recreate adds nothing measurable, so no regression from the session reset or the operator-owned locking either. The decode-side case for the context API is the bounded retained workspace, not speed. Creation counts are pinned by the test counters and retained-vs-released sizes by the measured table in codec_context.rs; I did not instrument peak RSS beyond sizeof. I can add if you would like. |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed the update at 6c8bae4f; no new findings.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed d8fa1615c9f82555df8e31a85992685a2488dd59 against 8729f6e6adf7091e18a48670e790d4ba8fd41e51. The PR-only patch is unchanged from the previously reviewed 6c8bae4f update, and I found no new P1/P2 issues. The existing approval stands. Current-head workflows report action_required, and no local tests were run in this re-review.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Thanks for syncing the base. I checked the upstream Arrow-export changes against the unchanged codec patch at dc8bb14e and found no new issue. The existing approval stands. No tests were rerun; current-head workflows require action.
3 participants
Which issue does this PR close?
Part of #5002 (the compression-context reuse item; the issue stays open for its remaining items).
Rationale for this change
Every shuffle block currently creates and destroys its own zstd context: a fresh
CCtxper encoded block inShuffleBlockWriter, and a freshDCtxper decoded frame inread_ipc_compressed. Context setup is pure overhead that scales with block count, so high-partition shuffles with small blocks pay the most.What changes are included in this PR?
ShuffleCodecContext/ShuffleDecodeContext(native/shuffle/src/codec_context.rs) wrapping a lazily-createdzstd_safe::CCtx/DCtx, reused viaEncoder::with_context/Decoder::with_context. The session is reset and the level re-applied per frame, so writers with different levels can share a context and a failed encode/decode can't poison the next one.LocalPartitionWriterowns the encode context and the per-partitionBufBatchWriters /SpillWriterborrow it;RssPartitionWriteris already one per task. The shuffle scan operator owns its decode context and goes through the_withentry points, so decode retention is scoped to the operator; the thread-local backs only the static JNI decode entry, which has no handle to own a context.Benchmarks (M-series macOS, 4M-row hash shuffle via
shuffle_bench, 3 iterations after warmup, measured at the current head against the merge base; Linux numbers may differ — happy to see re-runs):The saving is per block, so it grows with partition count / shrinking block size; large-block shapes are compression-bound and within noise, and the retained-size checks and boundary releases cost nothing measurable there. A decode-side microbenchmark (
benches/ipc_decode.rs) shows context reuse is throughput-neutral on decode — the decode-side value is the bounded retained workspace, not speed.How are these changes tested?
Six new tests alongside the existing suites (98 total in the shuffle crate, all passing, plus the core crate's 201):
cargo clippy --all-targets -- -D warningsandcargo fmtclean.