perf: reuse zstd compression contexts across shuffle blocks by dwsmith1983 · Pull Request #5565 · apache/datafusion-comet

dwsmith1983 · 2026-08-31T02:59:56Z

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 CCtx per encoded block in ShuffleBlockWriter. 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?

New ShuffleCodecContext (native/shuffle/src/codec_context.rs) wrapping a lazily created zstd_safe::CCtx, reused via Encoder::with_context. The session is reset and the configured level applied on every frame, since with_context does not initialize the level and a failed encode must not poison the next block.
Contexts are task-scoped, never per output partition: LocalPartitionWriter owns the context and the per-partition BufBatchWriters and SpillWriter borrow it; RssPartitionWriter is already one per task. write_burst_complete releases the workspace at spill and finish boundaries.
A retained-size cap (8 MiB, checked against the measured CCtx sizes 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.
The RSS path frees the zstd workspace at the end of each admitted encode, success or error: its memory accounting charges the workspace per admitted invocation and releases it afterward, so the context must not outlive that window. The remote path therefore keeps per-block context creation as on main; extending reuse there needs the pusher's reserve/push/release contract changed on both sides.
The decode side is unchanged from main. An earlier revision reused the decoder context too; it measured throughput neutral and introduced retention main does not have, so it was dropped after review.
Wire format is unchanged. lz4 and snappy keep per-block encoders; the pinned crates expose no reset for them and their setup cost is far below zstd's workspace.

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.

Shape	base	this PR	encode base -> PR
2,000 partitions, zstd level 1	0.323s	0.328s	0.188s -> 0.187s
2,000 partitions, zstd level 6	0.743s	0.746s	0.606s -> 0.610s
10,000 partitions, zstd level 3	0.581s	0.557s	0.326s -> 0.297s

10,000 partitions, zstd level 3, 5 iterations	base	this PR	delta
wall, average	1.202s	1.135s	-5.6%
wall, min / max	1.179s / 1.228s	1.113s / 1.171s	ranges do not overlap
encode, last iteration	0.653s	0.610s	-6.6%

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. A Linux rerun with more iterations would be welcome; none was available here.

How are these changes tested?

Tests in the shuffle crate (130 passing) and the core crate (266 passing, 4 ignored):

reuse across blocks, writers, and codecs round-trips and every block decodes independently
two writers with different zstd levels sharing one context each keep their own level
a mid-frame write failure does not poison the context for the next block
the retained-size cap releases an oversized context after an encode, with the level-8 boundary sentinel
the RSS release-versus-local-retain contract is pinned via a test accessor
spill and finish boundaries release the workspace, including on the error path

Scala on Spark 3.5: CometNativeShuffleSuite, CometShuffleSuite, and CometCelebornShuffleReaderSuite, 150 passing. cargo clippy --all-targets -- -D warnings and cargo fmt clean.

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

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-shuffle library: 98 passed
datafusion-comet core 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 --check passed 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.

sunchao · 2026-08-31T04:09:11Z

+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> =


[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.

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.

sunchao · 2026-08-31T04:09:11Z

    data_output: DataOutput,
+    /// Compression state shared by every block this task writes; the per-partition
+    /// `BufBatchWriter`s borrow it (see [`ShuffleCodecContext`]).
+    codec_context: ShuffleCodecContext,


[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.

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

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.

andygrove · 2026-08-31T14:28:27Z

+use std::cell::RefCell;
 use std::io::{Error, ErrorKind, Read};

+thread_local! {


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.

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.

andygrove · 2026-08-31T14:28:27Z

+/// 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;


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.

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.

andygrove · 2026-08-31T14:28:27Z

        // write header
        output.write_all(&self.header_bytes)?;

+        let encode_result =


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.

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.

andygrove · 2026-08-31T14:28:27Z

+    /// 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() {


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.

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

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.

dwsmith1983 · 2026-09-01T02:19:57Z

Re-ran the writer benchmark at the current head against the merge base, per @andygrove's request, same input and shapes; description table updated:

shape	base avg	head avg	encode base -> head
2,000 parts, zstd-1	0.323s	0.328s	0.188 -> 0.187s
2,000 parts, zstd-6	0.743s	0.746s	0.606 -> 0.610s
10,000 parts, zstd-3	0.581s	0.557s	0.326 -> 0.297s

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 benches/ipc_decode.rs: frames prebuilt once via ShuffleBlockWriter (Int64 + nullable Utf8), decoded results asserted identical across variants before measuring, four scenarios, repeated small zstd frames, a large-frame control, small frames with a wide-window level-19 frame every 16 (retained workspace passes the cap, context drops and recreates), and a NONE control, each with one reused context vs a fresh context per frame (the per-frame-creation profile of the base path; head-fresh matches a trimmed copy of the bench run on the base tree within noise: 7.289 vs 7.320 ms on the small-frames scenario).

scenario	reused	fresh per frame
64 small zstd-3 frames	7.285 ms	7.289 ms
1 large zstd-3 frame	8.945 ms	8.918 ms
over-cap recovery	7.772 ms	7.779 ms
64 small NONE frames	337 us	342 us

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

Reviewed the update at 6c8bae4f; no new findings.

sunchao

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

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

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

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.

sunchao · 2026-09-04T01:18:39Z

Reviewed 1372d786 against ef62b463. The overall design is sound, and I found no new blocking defect. I would keep context reuse, with these refinements:

Release decoder memory at EOF.
The EOF branch leaves the decoder context populated. Plan and stream clones share it, so a finished input can retain its workspace while other parts of the task continue. Clearing the context at definitive EOF preserves all useful reuse and shortens memory retention. The 8 MiB cap bounds this cost but does not eliminate it. EOF handling (line 185)
Combine the batch and decoder under one mutex.
The new decoder mutex is only acquired while the existing batch mutex is already held. It provides no additional concurrency in the current call graph. A single Arc<Mutex<ScanInputState>>, containing the batch and decoder, would preserve sharing across clones while removing an allocation per operator, an extra lock per fetched block, and a second synchronization invariant. This is a concrete simplification; its performance magnitude remains unmeasured. State and locking (line 141)
Improve the benchmark’s coverage of the actual design.
Its “small” frames contain 8,192 rows, and it calls the decoder directly, bypassing the operator mutex. The reusable context also survives across benchmark iterations. That measures steady-state decoding, but cannot establish reader startup/cleanup costs or operator overhead. Add the previously measured 400-row case, zstd levels 1/3, an LZ4 control, and a finite reader lifetime. An operator-path benchmark would address the remaining overhead question. Benchmark (line 143)
Correct the supposedly fresh decoder in the reuse test.
Its fresh result comes from read_ipc_compressed, which now uses the thread-local cached context. The test therefore compares two reused contexts. Create a new ShuffleDecodeContext for each reference decode so the test exercises its stated comparison. Test (line 331)

The remaining abstractions mostly earn their place:

Task-owned encoder state with mutable borrowing is appropriate. Putting contexts inside per-partition writers would multiply retained memory and lose reuse across small partitions.
The spill-completion hook is reasonable. It expresses a lifecycle boundary without teaching the partitioner about zstd, and current callers invoke it on both success and error.
Resetting the session and applying the level are justified. The zstd constructor does neither. Caching the last level would add bookkeeping for an unmeasured, small saving.
RSS’s per-invocation release should remain. Longer reuse requires changing its memory-reservation contract. The existing context types and thin decoder wrappers are preferable to introducing a general codec pool or framework.

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.

dwsmith1983 · 2026-09-04T04:03:57Z

Thanks for the careful pass. All four are in d21b24e:

EOF now releases the decoder. Every store of an input batch goes through one receive point on the scan input state, and InputBatch::EOF resets the decode context there, so a finished input no longer keeps its zstd workspace while the task continues. Pinned by eof_releases_retained_decoder_workspace.
Batch and decoder live under a single Arc<Mutex<ScanInputState>>. The separate decoder mutex is gone, and the batch field is no longer pub (nothing outside the module used it). Sharing across execute() clones is unchanged.
The decode benchmark now covers {400, 8192} rows x {zstd 1, zstd 3, lz4, none}, each with a context that outlives iterations, one locked per frame behind the same Arc<Mutex<_>> shape the operator uses, one created per reader lifetime of 64 frames and dropped (so startup and cleanup are measured), and a fresh context per frame. On this machine the per-frame context create and drop is about 2 percent of decode time at 400 rows and within noise at 8192 rows, and the lock is not measurable. That matches the earlier framing: the decode side is about retention bounds, not throughput. A benchmark through ShuffleScanExec itself is not possible without a JVM, since its decode path runs inside with_env and fetches blocks through JNI; the locked variant is the closest stand-in and is labeled as such.
The reuse test now builds a fresh ShuffleDecodeContext for every reference decode and asserts the reused side created exactly one context across all frames, so it compares what it says it compares.

The 32 KiB per-frame output buffer in the zstd writer is noted as a separate follow-up; I did not touch it here.

dwsmith1983 · 2026-09-04T04:24:26Z

@sunchao the four refinements from your design pass are in d21b24e, details in the comment above. Ready for another look.

sunchao

Correctness

Re-reviewed f2df122c against 2da32915, including the merge with main's scratch-buffer recycling. Fifteen of the 19 feature patches are unchanged after removing hunk positions. I traced the four adaptations that now pass both scratch and codec context through the writer APIs.

Single-partition writes keep the same scratch/context through final flush. Multi-partition spill and finish paths drain or clear scratch before reuse. The spill-completion hook still releases zstd state before freeing the input reservation, including on write failure. Codec reset, compression-level assignment, independent-frame completion, RSS release and decoder EOF handling are unchanged. I found no new or remaining P1/P2.

The maintained Spark 3.5/4.0 compression sources were checked for frame completion and level handling. This preserves Comet's existing IPC framing, not interchangeability with Spark's serializer. No expression, null or ANSI behavior changes here. Maintained Spark 3.4/4.1 source coverage remains unavailable.

Formatting checks passed for all 18 changed Rust files. No native compilation, product tests, JVM execution or benchmarks were run in this follow-up. Earlier remote results were not treated as current validation. At 03:42 UTC, all three exact-head workflows were action_required, with zero check runs. The synthetic merge also uses a newer first parent than the API base, so it is not evidence of an executed test pass.

Performance

The scratch adaptation adds borrowing, not another payload copy or a codec allocation per partition. Task-level local context reuse and scratch capacity reduction remain intact. RSS still releases the workspace per admitted invocation. The unchanged zstd wrapper retains a 32 KiB output allocation per frame, and the 8 MiB context threshold bounds between-frame retention rather than peak allocation.

The existing decode benchmark covers fresh and reused contexts plus codec controls, but not JNI fetching. Its earlier figures were not reproduced against this merged base. I make no new throughput claim.

Design

Caller-owned scratch and task-owned codec state remain separate resources with explicit lifetimes. Passing both through the existing APIs preserves main's recycling and the established reservation boundaries. Extending RSS reuse would still require changing its reservation contract.

Abstraction & complexity

The update adds no pool, cache hierarchy or synchronization layer. It keeps the small context types and existing writer interfaces. The shared scan-state mutex is unchanged, and I found no additional abstraction needed for this adaptation.

sunchao

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.

andygrove · 2026-09-05T14:31:07Z

This seems lot a lot of additional complexity for little performance gain?

andygrove · 2026-09-05T14:37:32Z

Here's my AI review:

This seems lot a lot of additional complexity for little performance gain?

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 CometConf.scala matter here: the default shuffle codec is lz4 (line 510) and the default zstd level is 1 (line 518). This PR does nothing for lz4 or snappy, so out of the box it's inert — you need a non-default codec and a non-default level and a high partition count before it does anything.

Working backwards from the benchmark table gives a consistent mechanism:

shape	encode base → PR	blocks	implied per-block
10,000 parts, zstd-3	0.326 → 0.297s	~10,000	~2.9 µs
2,000 parts, zstd-1	0.188 → 0.187s	~2,000	predicted ~6 ms, below noise
2,000 parts, zstd-6	0.606 → 0.610s	~2,000	predicted ~6 ms, below noise

(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, zstd::Decoder::with_buffer builds a fresh DCtx per frame, so retention between frames is already zero. The PR introduces retention and then spends the 8 MiB cap, the sizeof() checks, the error-path releases, and the EOF hook bounding retention it created.

The thread-local in ipc.rs also backs the live production path for the JVM shuffle reader (NativeBatchDecoderIterator → Native.decodeShuffleBlock), so every Spark task thread now holds roughly 1-2.6 MB of native zstd workspace that Comet's memory accounting doesn't track. That's exactly the property used to justify releasing per-invocation on the RSS encode path — same hazard, opposite decision, for a 0.05% saving.

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:

MAX_RETAINED_ZSTD_CONTEXT_BYTES is a magic constant backed by a hand-measured table of zstd workspace sizes, with a test whose only job is to fail when a zstd bump silently disables the optimization (level 8 sits 3% under the cap). That's a tax on every dependency bump.
Two divergent lifetime contracts for one object — RSS releases, local retains — enforced by a bool threaded through write_batch_with_codec_limits.
A new write_burst_complete() on the PartitionWriter trait that exists only to drop one codec's allocator, i.e. a lifecycle hook in the partitioner API serving zstd specifically.
The thread-local above.

What I'd suggest instead: keep the encode-side reuse for local shuffle, and drop the decode half entirely — ShuffleDecodeContext, the four _with entry points, decode_remote_shuffle_batch_with, the thread-local, the ScanInputState restructure, benches/ipc_decode.rs, and the decode tests. That's roughly half the production diff, removed for something already measured as neutral.

Then, could the sizeof()-based cap be replaced with a retain/don't-retain decision made once in ShuffleBlockWriter::try_new from the configured level? Codec and level are uniform per query, so the alternating-level test is defending a case that can't occur in production, and reset(SessionOnly) preserves parameters anyway, which makes the per-block set_parameter redundant. That would drop the constant, the measured table, and both release_zstd_if_oversized sites, leaving the mechanical &mut threading as the only real cost.

One hypothesis I checked and discarded, in case anyone else wonders the same: hoisting CompressionContext out of BufBatchWriter is not a hidden memory win at high partition counts. IpcWriteOptions::try_new(64, false, V5) leaves Arrow's inner buffer compression off, so those per-partition contexts were an empty Vec plus a None, about 40 bytes each. Good hygiene, not a saving.

dwsmith1983 · 2026-09-05T14:37:54Z

This seems lot a lot of additional complexity for little performance gain?

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.

peterxcli

Haven't looked through the code yet, but this might relate to #5446 and #5008

sunchao

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.

dwsmith1983 · 2026-09-06T12:23:05Z

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.

…or-reuse

sunchao reviewed Aug 31, 2026

View reviewed changes

dwsmith1983 requested a review from sunchao August 31, 2026 04:45

Merge branch 'main' into perf/shuffle-compressor-reuse

9afe4eb

dwsmith1983 mentioned this pull request Aug 31, 2026

perf: reuse per-partition scratch in the shuffle write path #5568

Merged

andygrove reviewed Aug 31, 2026

View reviewed changes

sunchao reviewed Aug 31, 2026

View reviewed changes

dwsmith1983 and others added 2 commits September 1, 2026 08:42

Merge branch 'main' into perf/shuffle-compressor-reuse

16198d9

Merge branch 'main' into perf/shuffle-compressor-reuse

6c8bae4

dwsmith1983 requested review from andygrove and sunchao September 1, 2026 02:49

sunchao approved these changes Sep 1, 2026

View reviewed changes

Merge branch main into perf/shuffle-compressor-reuse

d8fa161

sunchao reviewed Sep 2, 2026

View reviewed changes

Merge branch 'main' into perf/shuffle-compressor-reuse

dc8bb14

sunchao reviewed Sep 2, 2026

View reviewed changes

Merge branch 'main' into perf/shuffle-compressor-reuse

7a120d1

sunchao reviewed Sep 3, 2026

View reviewed changes

dwsmith1983 requested a review from sunchao September 3, 2026 07:02

Merge remote-tracking branch 'origin/main' into perf/shuffle-compress…

1372d78

…or-reuse # Conflicts: # native/core/src/execution/operators/shuffle_scan.rs # native/shuffle/src/lib.rs

sunchao approved these changes Sep 3, 2026

View reviewed changes

dwsmith1983 added 2 commits September 4, 2026 10:23

Merge branch 'main' into perf/shuffle-compressor-reuse

3fb42e1

sunchao approved these changes Sep 4, 2026

View reviewed changes

Merge branch 'main' into perf/shuffle-compressor-reuse

f2df122

sunchao approved these changes Sep 5, 2026

View reviewed changes

Merge branch 'main' into perf/shuffle-compressor-reuse

89aed72

dwsmith1983 requested a review from sunchao September 5, 2026 04:38

sunchao reviewed Sep 5, 2026

View reviewed changes

Merge branch 'main' into perf/shuffle-compressor-reuse

37b697f

Merge branch 'main' into perf/shuffle-compressor-reuse

495cead

peterxcli reviewed Sep 5, 2026

View reviewed changes

Merge branch 'main' into perf/shuffle-compressor-reuse

a23d1ed

sunchao reviewed Sep 5, 2026

View reviewed changes

dwsmith1983 added 2 commits September 6, 2026 18:46

Merge branch 'main' into perf/shuffle-compressor-reuse

74b57b5

dwsmith1983 requested a review from sunchao September 6, 2026 12:34

Merge remote-tracking branch 'origin/main' into perf/shuffle-compress…

81828df

…or-reuse

dwsmith1983 force-pushed the perf/shuffle-compressor-reuse branch from 1e60c6a to 81828df Compare September 6, 2026 12:38

andygrove added enhancement New feature or request performance area:shuffle Shuffle (JVM and native) labels Sep 6, 2026

Uh oh!

Conversation

dwsmith1983 commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andygrove left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

dwsmith1983 commented Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

sunchao commented Sep 4, 2026

Uh oh!

dwsmith1983 commented Sep 4, 2026

Uh oh!

dwsmith1983 commented Sep 4, 2026

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Correctness

Performance

Design

Abstraction & complexity

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Correctness

Performance

Design

Abstraction & complexity

Uh oh!

dwsmith1983 commented Aug 31, 2026 •

edited

Loading

dwsmith1983 commented Sep 1, 2026 •

edited

Loading