perf: reuse zstd compression contexts across shuffle blocks - #5565
perf: reuse zstd compression contexts across shuffle blocks#5565dwsmith1983 wants to merge 28 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.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Rechecked 7a120d1e after the base update. The codec patch is byte-for-byte unchanged. I checked the imported Celeborn reservation/transport and shared codec-enum interfaces and found no new P1/P2. The existing approval stands.
No tests were rerun. This head has no check runs and three workflows awaiting approval.
…or-reuse # Conflicts: # native/core/src/execution/operators/shuffle_scan.rs # native/shuffle/src/lib.rs
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed 1372d786575c2512b08ce0cc1ea0eea9fa0bfac8 against ef62b463, including the merge adaptation that passes the operator-owned decode context through the shared remote dictionary-normalization path. I found no remaining P1/P2.
Four focused decoder-component tests passed, covering reused/fresh/static entry points, codec transitions, dictionary values and nulls, error recovery, and oversized-context release. This is not full Comet/JNI/Spark validation or a benchmark. The three current-head workflows still require authorization, with no executed check results reported.
|
Reviewed
The remaining abstractions mostly earn their place:
One optional performance follow-up: zstd’s Rust wrapper still allocates a 32 KiB output buffer per frame. This predates the PR. Profile that allocation before adding lower-level streaming machinery to reuse it. Allocation (line 49) This was a fresh source/design review. The measured codec paths and lockfile are unchanged from our earlier benchmarked revision; I did not rerun benchmarks or tests. Current CI has 64 successful checks, with the performance check skipped. |
…put state ShuffleScanExec now keeps the input batch and its decoder in one ScanInputState behind a single mutex, still shared by plan and stream clones, and releases the decoder's retained zstd workspace when the input reaches EOF instead of holding it until the plan drops. The ipc_decode bench adds a 400-row case, zstd levels 1 and 3, an LZ4 control, a per-reader-lifetime variant that counts context startup and cleanup, and a mutex-locked variant mirroring the operator's lock. The decode-context reuse test now decodes its reference with a new ShuffleDecodeContext per frame; read_ipc_compressed uses the thread-local cached context, so the old reference was itself a reused context.
|
Thanks for the careful pass. All four are in d21b24e:
The 32 KiB per-frame output buffer in the zstd writer is noted as a separate follow-up; I did not touch it here. |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Re-reviewed 89aed721 against 719cba11. The complete feature patch and all 19 feature-file blobs are identical to the previously reviewed f2df122c pair. The entire native tree is unchanged. The only head update is CometExplodeBenchmark.scala, copied exactly from the new base. No new or remaining P1/P2 was found, and the existing approval stands.
The unchanged framing and resource-lifetime paths were rechecked against the maintained Spark 3.5/4.0 sources. This update changes no expression, null, ANSI, error or fallback behavior. Maintained Spark 3.4/4.1 source coverage remains unavailable.
Source/dependency identity and three diff checks passed. No product tests, JVM/native execution, formatting or benchmarks were rerun. At 05:24 UTC, all three current-head workflows were action_required, with zero check runs. The current synthetic merge matches this head's tree, which establishes source identity but not executed validation.
Performance
There is no shuffle hot-path, allocation, locking or dependency change in this update. Existing decoder benchmark coverage is unchanged, and the inherited explode benchmark does not measure codec reuse. No new performance measurement is claimed.
Design
The merge preserves the already-reviewed ownership and reservation boundaries. The imported benchmark adds no production registration or configuration, so I found no new design interaction.
Abstraction & complexity
No production helper or abstraction changes here. The synchronization introduces no additional complexity into the codec feature.
|
This seems lot a lot of additional complexity for little performance gain? |
|
Here's my AI review:
Digging into the numbers, I think that instinct is right, and I'd narrow the scope rather than land this as is. The measured gain is about 3 µs per shuffle block, and the defaults never reach it. Two things from Working backwards from the benchmark table gives a consistent mechanism:
(4M rows / 10k partitions is ~400 rows each, under the 8192 batch size, so blocks ≈ partitions.) Worth noting that level 6 has a larger workspace than level 3 and still showed nothing, which confirms the variable is block count, not level. So the honest claim is ~3 µs saved per block: 29 ms at 10,000 blocks, 0.6 ms at Spark's default 200 partitions. Given how sensitive these shuffle benchmarks are to machine load, I'd also want the 10k result reproduced on Linux with more than 3 iterations before treating the 4% as real. The decode half has no measured benefit, and it introduces a retention hazard main doesn't have. The comment above puts it at 56 ns/frame of setup against ~114 µs/frame of work. The stated justification — "the decode-side value is the bounded retained workspace, not speed" — is circular: on main, The thread-local in Where the complexity actually lands. It's ~290 net production lines and ~1,080 test/bench lines, but the part that costs us long term is four things:
What I'd suggest instead: keep the encode-side reuse for local shuffle, and drop the decode half entirely — Then, could the One hypothesis I checked and discarded, in case anyone else wonders the same: hoisting |
Fair question. the gain is on the encode side, 7 percent wall and 11 percent encode time at 10k partitions with zstd level 3 on the writer benchmark, and flat at large blocks. The decode side measured throughput neutral; its value is bounding retained memory. About half of the added lines are tests, and a good part of the surface came from review asks (the exec-side decode context, the retained-size cap, the spill boundary hook, RSS releasing per invocation). If the trade off looks wrong to you, the cleanest is to drop the decode side entirely and keep only writer-side context reuse. That removes the scan operator changes, the decode context type, and the thread-local, roughly half the diff, while keeping the measured win. |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Re-reviewed a23d1ed6278dd753f32dc386a000bb0f9ece65a0 against 7190df631afe3795914839203c7afe57ea23903c, since accepted 89aed721. This is not a byte-identical synchronization. Seventeen of the 19 feature files are identical, while two RSS files integrate main's typed size-limit errors and tests. The PR's added and removed lines remain identical after excluding hunk positions and unchanged context.
I traced the affected admission, encode, error, growth-retry and split paths. write_rss_batch still releases its zstd context before returning, so output/compaction cleanup and reservation release occur afterward. Early admission failures allocate no context. Scratch draining, local spill release before reservation.free(), configured-level initialization, independent-frame completion, decoder reset and EOF release remain intact. No new P1/P2 was found in this integration.
The maintained Spark 3.5/4.0 compression sources preserve configured levels and distinguish ordinary completed frames from continuous streams. Comet still completes each independent IPC frame. This update changes no expression, null or ANSI semantics and does not imply interchangeability with Spark serializer bytes. Maintained Spark 3.4/4.1 sources remain unavailable. Locked Arrow 58.4 and zstd dependencies are unchanged. The newer Arrow 59.2 work described in #5446 is outside this exact pair.
Three diff checks and fresh source/dependency/archive verification passed. No product tests, native/JNI/Spark execution, formatting or benchmarks ran. At 2026-09-05T21:35:07.916Z, all three workflows were action_required, with zero head or merge checks. GitHub reported dirty and no synthetic merge SHA, although the assigned base is a parent of this head. There is no executed CI checkout to qualify as a pass. An existing approval is associated with this head, so I am leaving a follow-up comment rather than another approval.
Performance
The published writer table implies about 4.1% wall-time and 8.9% encode-time improvement for 10,000 partitions at zstd level 3, rather than the later reply's 7%/11%. These are historical author measurements, not a fresh result. The defaults are LZ4 and zstd level 1. This does not establish that a non-default zstd level is necessary for a gain, or that the added context bookkeeping has zero default-path cost.
The decoder reports range from roughly 2% setup/drop cost at 400 rows to neutral at 8,192 rows. The benchmark includes finite-reader and mutex-proxy controls but not the production JNI path or measured peak memory. The existing request for a repeated Linux writer comparison is appropriate. Any retained decoder optimization also needs evidence that its benefit justifies the retained memory and lifecycle overhead. Source equality does not supply that measurement.
Design
The encoder-only option offered by the author is a concrete simplification. I agree that decoder retention is not a memory improvement over this base: the base drops an owned decoder after every frame. The PR instead adds operator-owned and thread-local retention, then bounds it. EOF release fixes the operator lifetime, but the JNI thread-local can still retain up to the cap across tasks. Please resolve the existing decoder-scope discussion before merge. The merge update has not removed that part of the implementation.
Abstraction & complexity
Task-owned encoder state and mutable borrowing avoid multiplying workspaces by output partition. The spill hook and RSS per-invocation release express real accounting boundaries. Dropping decoder reuse would remove a separate context type, wrapper entry points, thread-local state and owner lifecycle work.
I would not remove the measured-size guard or level initialization solely from the discussion's uniform-level argument. with_context does not initialize the configured level, and every replacement context needs that initialization. A level-based retention policy would also need to preserve the actual memory bound across supported configurations and dependency updates. The current 8 MiB check limits retention between frames, not peak memory.
Drop the decode-side context reuse: it measured throughput neutral and retained a decompression workspace that main never held. Shuffle reads go back to a fresh decoder per frame. The write path keeps the task-owned ShuffleCodecContext, its 8 MiB retained-size cap, and the per-encode level initialization.
|
Trimmed to the write path in 27c2033: the decoder context, the thread-local, the scan operator changes and the decode bench are gone, and ipc.rs is back to main. The retained-size cap and the level initialization stay, for the reasons sunchao gave. Remaining diff is about 260 production lines. One correction to my earlier reply: the posted table implies 4.1 percent wall and 8.9 percent encode at 10,000 partitions, not 7 and 11. Linux run, since you asked: rust 1.94 on linux/arm64 in Docker, 10,000 partitions, zstd level 3, ten timed iterations, base and head alternated. Round one gives 3.0 percent wall and 6 percent encode with the min/max ranges not overlapping. A second round showed more, but the base run drifted, so I would not read anything into it beyond the direction. That is in line with your per-block estimate. On macOS the same shape over five iterations gives 5.6 percent wall and 6.6 percent encode. |
1e60c6a to
81828df
Compare
andygrove
added
enhancement
#5038 moved Arrow's |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Re-reviewed 5a1ead6e8914ceaf18d95cb53b597ebe19013032 against 7f1e00189b1ed86f1cb5acd872d97fce694482b1, following review 5123119243. The decoder-scope concern is addressed: ipc.rs and remote decoding match the base, the production scan code matches the base, and the decoder context, thread-local cache and decode benchmark are removed. The remaining scan changes adapt tests to the writer API. I found no new or remaining P1/P2.
I rechecked the encoder paths through the merged source. Each zstd frame resets its session, applies the configured level and finishes independently. The size check runs after successful and failed compression; spill completion releases zstd before the input reservation is freed. RSS releases it before returning to the caller's reservation cleanup. Successful local finish releases it explicitly; production finish errors unwind the task-owned repartitioner. Mutable borrowing keeps the context exclusive while a frame is encoded. Dictionary-bearing schemas still use a fresh Arrow stream writer per frame.
The maintained Spark 3.5/4.0 compression sources confirm configured-level handling and the distinction between completed frames and continuous streams. Comet preserves its existing independent IPC framing; this does not imply interchangeability with Spark serializer bytes. There are no new expression, null, ANSI or fallback semantics. Maintained Spark 3.4/4.1 source coverage remains unavailable.
This pair locks Arrow 59.3.0 and DataFusion 55.0.0, with identical dependencies in head/base; zstd remains 0.13.3 / zstd-safe 7.2.4 / zstd 1.5.7. I verified the relevant library source against the locked crate checksums, refreshed the maintained Spark source checks, and passed both diff checks and formatting checks for all 15 changed Rust files. No native build, product tests, JNI/Spark execution or benchmarks ran in this follow-up. At 2026-09-08T02:57:34.298Z, all three current-head workflows were action_required, with zero head or merge check results. The test counts in the description are author reports, not current-head CI validation.
Performance
The September 6 report addresses the request for a repeated Linux writer comparison: Linux/arm64 Docker, ten timed iterations with base/head alternation, reporting 3.0% wall and 6% encode improvement in the first round. The author appropriately discounts the drifting second round and corrects the earlier 7%/11% claim to the table's 4.1%/8.9%. These remain author measurements of the trimmed revision, without raw iterations or complete commit/build provenance in the discussion; they do not establish performance at this merged head. The description should also be reconciled with that report because it still says no Linux run was available.
Arrow's shared context now contains a reusable FlatBuffer builder, so hoisting it can avoid repeated builder allocations even when outer compression is LZ4 or NONE. However, on this base the multi-partition spill/finish loops create temporary BufBatchWriters sequentially. The source does not establish the latest comment's implication of 10,000 simultaneously retained builders. A paired allocation measurement through that local multi-partition path would substantiate the additional benefit; the existing fresh/reused context benchmark alone does not measure the ownership change. RSS still creates a zstd context per admitted block, as documented.
Design
The update takes the concrete simplification discussed previously: retain writer-side reuse and restore ordinary per-frame decoder ownership. This removes the decoder retention introduced by the earlier revision. The remaining task-owned encoder context fits the sequential local write path, and the spill hook expresses a real lifetime boundary. RSS's shorter lifetime remains justified by its existing admission contract; extending reuse there would require a separate reservation change.
Abstraction & complexity
Removing the decoder type, wrapper APIs, thread-local cache and scan-state lifecycle work substantially narrows the maintenance surface. The remaining context groups state borrowed by the same encode operation without adding a pool or synchronization layer. Arrow copies the finished metadata into an owned buffer before resetting its builder, so sequential reuse does not leave prior frames borrowing mutable builder storage.
The configured-level initialization and measured-size guard remain justified. The 8 MiB limit bounds retained zstd workspace between blocks, not peak encode memory or the Arrow builder. It should not be described as a total task-memory cap. The level-8 boundary test remains useful protection against a dependency update silently disabling reuse.
|
The multi-partition path builds its writers one at a time, so "ten thousand retained builders" was wrong; the effect is per flush, not simultaneous retention. Measured with the allocation counter #5727 added, driving
The 8 allocations and 552 bytes are the |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Rechecked d4fec191 against bb9e7402, following review 5136951863. All 15 files in the PR contribution are byte-identical to the previously reviewed 5a1ead6e head. The entire incremental source diff equals the base update, and the lockfile is unchanged across the prior head, current head and current base. The decoder files and production scan code still match the base. I found no new or remaining P1/P2, and the existing approval stands.
The unchanged callers preserve the reviewed lifetime boundaries. Spill completion releases zstd before freeing the input reservation, including after a write error. RSS releases it before returning to the caller's reservation cleanup. Successful local finish releases it explicitly, while production finish errors drop the task-owned repartitioner. Session reset, configured-level initialization, independent frame completion and fresh dictionary tracking remain intact. The maintained Spark 3.5/4.0 compression sources still agree on configured-level and completed-frame behavior. This update adds no expression, null, ANSI or fallback behavior, and does not imply that Comet IPC bytes are interchangeable with Spark serializer bytes. Maintained Spark 3.4/4.1 source coverage remains unavailable.
Both diff checks and fresh source/dependency verification passed. No native build, product tests, JNI/Spark execution or benchmarks ran in this follow-up. At 2026-09-08T09:25:23.897Z, all three workflows required action, with zero head or merge checks. The synthetic merge has the assigned parents and the head's tree, which establishes source identity but supplies no executed CI result. The test counts in the description remain author reports.
Performance
The description now reconciles the Linux benchmark report. It reports ten timed iterations on Linux/arm64, a first-round change from 1.349s to 1.309s wall time and 0.713s to 0.672s encode time, and explicitly discounts the drifting second round. These are historical author measurements of the trimmed revision, with incomplete exact commit/build and raw-iteration provenance. They do not establish performance at this rebased head.
The new local-writer allocation report also corrects the earlier simultaneous-retention claim. The multi-partition writers are sequential. The reported probe exercises finish and spill passes over 10,000 partitions after warm-up, and reports eight fewer allocations and 552 fewer allocated bytes per flushed partition for NONE, LZ4 and zstd. That targets the production ownership change rather than only the existing fresh/reused-context microbenchmark. These remain author-reported results: the probe sources and raw output are offered but not included in the discussion, so I have not independently reproduced the counts. The counter excludes zstd's C allocations. The Arrow allocation result and zstd workspace retention should therefore remain separate from each other and from an end-to-end throughput claim.
Design
The base update preserves the writer-only scope and the existing reservation interfaces. It introduces no new interaction with decoder ownership. Task-scoped local state and per-invocation RSS release remain appropriate for their respective callers, so I found no new design change required by the rebase.
Abstraction & complexity
No helper, pool or synchronization layer was added in this update. The measured-size guard still bounds retained zstd workspace between blocks, not peak encode memory or Arrow builder capacity. The level-8 boundary test and explicit level initialization retain their existing purpose without another abstraction.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Rechecked c314c320 against 8e684685, following published review 5139980456 at d4fec191. All 15 authored files and the complete contribution diff are unchanged. The 43-file increment from that published head exactly equals the base update. The further 17-file increment since 3298b72c also equals its base update. All 27 previously audited source units, including writer, decoder, JNI, iterator and lockfile files, are unchanged. I found no new or remaining P1/P2. The existing approval stands.
I inspected the newly inherited Variant normalization and concat_ws integration boundaries. They do not modify the shuffle interfaces, frame completion, writer ownership or release paths. The previous lifetime and error-cleanup conclusions remain applicable by exact source identity. The maintained Spark 3.5/4.0 compression-source refs and hashes were rechecked. Configured-level and completed-frame behavior remains unchanged, with no new expression, null, ANSI or fallback semantics in the compression contribution. Maintained Spark 3.4/4.1 coverage remains unavailable. This is not a separate qualification of the inherited Variant or expression PRs.
Three diff checks passed. No build, product tests, JNI/Spark execution or benchmarks ran. The lockfile remains identical across the published head, intervening head, current head and base. The complete public snapshot ends at 2026-09-08T16:31:17.975Z. All three workflows require action and the fresh 16:45:21 UTC check found zero jobs. The synthetic merge has the assigned parents and head tree, without an executed CI result. Prior test results are not current-head validation.
Performance
The description and every discussion body are unchanged since the preceding audit. There is no new writer measurement. The published review's Linux and allocation-result qualifications still apply. Main's added hash benchmark and Variant/string work do not measure this PR's unchanged compression reuse. No current-head speedup is inferred from source equality.
Design
The base merges preserve the existing writer-only contribution and reservation interfaces. Local task ownership and RSS per-invocation release are unchanged. I found no new design interaction requiring a change to this PR.
Abstraction & complexity
The compression contribution adds no helper, pool or synchronization layer in this update. The existing level initialization and measured-size guard retain their reviewed purposes. The 8 MiB guard bounds retained zstd workspace between blocks, not peak encode memory or total task memory.
4 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 written through the local shuffle path creates and destroys its own zstd context: a fresh
CCtxper encoded block inShuffleBlockWriter. Context setup is pure overhead that scales with block count, so high-partition shuffles with small blocks pay the most. The saving is a few microseconds per block, which is noise at Spark's default 200 partitions and a measurable few percent of encode time at 10,000.What changes are included in this PR?
ShuffleCodecContext(native/shuffle/src/codec_context.rs) wrapping a lazily createdzstd_safe::CCtx, reused viaEncoder::with_context. The session is reset and the configured level applied on every frame, sincewith_contextdoes not initialize the level and a failed encode must not poison the next block.LocalPartitionWriterowns the context and the per-partitionBufBatchWriters andSpillWriterborrow it;RssPartitionWriteris already one per task.write_burst_completereleases the workspace at spill and finish boundaries.CCtxsizes per level, with a test that fails if a zstd bump moves level 8 across it) bounds what a task keeps between bursts: zstd's session reset preserves the allocated window, so a context that once saw a wide-window frame would otherwise stay that large.Benchmarks with
shuffle_bench(4M-row hash shuffle, single task, M-series macOS, release builds). The first table is the earlier three-iteration run at the head of that revision against its merge base; the second is a five-iteration rerun of the trimmed head against current main on a different generated input, so only base-versus-head within a table is meaningful.The wall delta is about the size of the run-to-run spread, so encode time is the more direct signal. Large-block shapes are compression-bound and within noise.
On Linux (arm64, Docker, rust 1.94, four CPUs pinned), same shape and input, ten timed iterations per run with base and head alternated over two rounds: round one gives 1.349s to 1.309s wall (3.0 percent) and 0.713s to 0.672s encode (6 percent), with the min/max ranges not overlapping. The second round's base run drifted slower than every other run, so only the direction is taken from it.
How are these changes tested?
Tests in the shuffle crate (130 passing) and the core crate (266 passing, 4 ignored):
Scala on Spark 3.5: CometNativeShuffleSuite, CometShuffleSuite, and CometCelebornShuffleReaderSuite, 150 passing.
cargo clippy --all-targets -- -D warningsandcargo fmtclean.