perf: cache expected schemas for remote shuffle decoding

[pingzh]

Which issue does this PR close?

Closes #5535.

Rationale for this change

The JVM remote shuffle reader copies and parses the same expected schema on every block. Retain the parsed Spark types for the iterator's lifetime so subsequent blocks avoid JNI schema copies and repeated protobuf/type conversion.

What changes are included in this PR?

• Lazily create a native decoder on the first remote block and reuse it across subsequent blocks. Release it once on iterator close, including early completion and cleanup failures. Empty and local streams allocate no decoder.
• Coordinate decoder creation, decoding, Arrow import, and batch ownership with concurrent task cleanup. Keep transport reads and failure-reporting RPCs outside the monitor so cleanup can proceed while they are blocked.
• Keep Arrow buffer/offset and logical type validation mandatory. Decoder initialization failures remain outside remote corruption reporting.

Remote decode measurements

Linux x86_64, AMD EPYC 9V74, Rust 1.97.1, DataFusion 55.0.0, Arrow 59.3.0; default release profile, 100 samples, 3-second warm-up, 5-second target measurement time. Each block contains 8192 rows and 20 columns (10 Int32, 10 Utf8). Times are Criterion's estimated mean in microseconds per block.

Encoding	Codec	Frame bytes	Local	Remote	Remote/local
plain	None	2,075,868	106.7	290.5	2.72x
plain	LZ4	704,026	410.4	602.9	1.47x
plain	Snappy	785,894	577.1	774.1	1.34x
plain	Zstd(1)	248,249	1577.5	1776.2	1.13x
dictionary	None	709,660	67.0	751.5	11.21x
dictionary	LZ4	342,473	191.3	877.8	4.59x
dictionary	Snappy	370,097	259.6	943.4	3.63x
dictionary	Zstd(1)	205,611	359.3	1053.5	2.93x

The remote dictionary path also expands values for the JVM importer, so that gap includes normalization. Measurements include decompression, IPC decoding, validation/normalization, and batch destruction; they exclude transport, JNI, Arrow FFI, and schema initialization. They do not measure the schema-cache speedup or establish an end-to-end advantage over Celeborn's row decoder. Fixtures differ from the issue's original measurements. Validation remains required despite its cost.

How are these changes tested?

• CometCelebornShuffleReaderSuite on Spark 4.1 and Spark 3.4.3: 69 passed on each against the new release JNI library, including multi-frame reuse, later-frame schema mismatch, initialization failure classification, and cleanup lifecycle coverage.
• Five new latch-based concurrency regressions failed before the lifetime fix and passed afterward. They cover cleanup during decoder creation, active decoding (including interrupted cleanup), Arrow import and batch publication, transport reads, and failure reporting.
• cargo test --release -p datafusion-comet-shuffle --lib: 118 passed.
• Native release build, full JVM reactor test compilation, Spotless, Scalastyle, cargo fmt --check, and git diff --check passed.

[sunchao]

Review result

No actionable defects found in PR #5722 at cedc3bd99, against base 75fdddc92.

Correctness

The cache preserves the existing Spark type conversion and per-block validation. It retains expected types, so later blocks still undergo logical-type and Arrow validation.

I traced normal exhaustion, early task completion, initialization failure, decode failure, and cleanup exceptions. Ownership is cleared before release, remaining cleanup proceeds after failures, and the production callers decode synchronously. I found no reachable release race or introduced leak. Iterator implementation

Performance

For an iterator containing N blocks, schema copying, protobuf parsing, and type construction decrease from N times to once. The cost is two additional lifecycle JNI calls and schema metadata retained until close. Single-block iterators cannot amortize that overhead.

The expected benefit is strongest with many small blocks and wide or nested schemas. The actual speedup remains unmeasured: the PR’s benchmark explicitly excludes JNI and schema initialization. It measures remote decoding cost, not this cache’s benefit or whole-query performance. A before/after benchmark including those costs would establish the net gain.

Design

One immutable native object owned by one iterator is an appropriate lifetime. Lazy creation avoids allocation for empty or unused streams. Keeping initialization failures outside remote-corruption reporting also preserves the right failure classification. Native implementation

Abstraction & complexity

The wrapper is small and serves a clear ownership purpose. I see no unnecessary framework or abstraction, and no evidence that a global cache or combined initialization/decode API would justify additional complexity.

Validation

• 64/64 reader tests passed locally on Spark 4.1, including reuse, later-frame validation, and cleanup cases.
• JVM reactor build, formatting checks, and Rust formatting passed.
• Local native compilation was blocked by the dependency mirror. Tests used a checksum-verified CI library whose source tree matches this commit.
• Latest CI: 59 passed, 6 pending, 7 skipped; no failures. CI run

I would not request code changes on this revision.

[sunchao]

Approved the reviewed commit cedc3bd. The full review is posted at #5722 (comment).

CI refreshed before publication: 60 passed, 5 pending, 7 skipped, with no failures.

[sunchao]

Correctness

No additional verified P1/P2 finding from this follow-up at cedc3bd99ec4580226ffe77154a18fc7058a1547, against 75fdddc9285ec61c0cd326977c61dd41fca39a8b. The existing review and approval already cover this revision. The maintained Spark 3.5 and 4.0 task-completion and interruption paths support the reviewed iterator cleanup behavior.

Updated validation

The exact-head CI refresh at 2026-09-05 21:46 UTC reports 65 successful / 9 skipped checks, with none pending or failing. I verified that the Spark 3.5 shuffle job and Spark 4.0 shuffle job each passed all 64 reader-suite tests, including the new reuse, later-frame mismatch, and lifecycle cases. Their complete shards passed 426 and 469 tests respectively.

Both jobs checked out merge da40304adc8455b91ce0803f66c27fcd6c592751. Its parents are this base and head, and its entire source tree matches the reviewed head. The native library build succeeded on that merge using the CI profile. These reader tests execute Spark and JNI with a recording Celeborn client and in-memory frames. They do not exercise a deployed Celeborn service. I ran no local build or benchmark. Maintained Spark 3.4 and 4.1 sources remain unavailable for independent source qualification.

Performance

No additional performance finding. The existing review already identifies the missing measurement of the cache's net benefit. The supplied remote/local decode timings exclude JNI and schema initialization, so this follow-up adds no cache-speedup or whole-query performance claim.

Design

No new design concern after checking the production callers, task-completion order, and initialization-failure classification. There has been no code revision since the existing approval that needs a separate design summary.

Abstraction & complexity

No additional abstraction or complexity finding in this unchanged revision. No repeated inline feedback or duplicate approval is proposed.

[jayanth86]

Reviewed cedc3bd99ec4580226ffe77154a18fc7058a1547 against 75fdddc9285ec61c0cd326977c61dd41fca39a8b. I found one additional native-lifetime issue in the supported Spark 3.4.3 Python-consumer cancellation path, detailed inline. The per-frame validation and per-iterator schema scope otherwise remain intact.

Validation included independent source review and three local latch-based probes of the unmodified iterator with mocked JNI. The probes reproduce in-flight release, late handle publication after close, and release following an interrupted join. They do not execute Spark or native Rust.

I also verified that CI built the same source tree as this head: all 64 reader-suite tests passed on Spark 3.5, 4.0, and 4.1, and the Rust job passed 1,163 tests with 4 skipped. Those tests do not cover this cancellation interleaving. A local native build was not run because the required dependencies were not cached.

[jayanth86]

[P1] Keep the decoder alive while a read is in flight

Could creation and decoding be coordinated with this release? On supported Spark 3.4.3, a Python UDF downstream of a Celeborn exchange consumes this iterator on PythonRunner's writer thread. Its completion listener interrupts and joins that writer, but task cancellation with interruptThread=true can interrupt the join. TaskContextImpl catches the listener exception and continues cleanup, so Comet can reach this release while the writer is still in native decompression.

fetchNext() does not take the close() lock. This release therefore frees the boxed decoder while the active JNI call still borrows its expected_types, which are read after decompression in decode_remote_shuffle_batch. That introduces a use-after-free capable of crashing the executor. Previously, the types were owned locally by the JNI invocation and could not be freed by task cleanup.

Please protect the handle across creation, active decoding, and release, and add coverage for interrupted cleanup overlapping a read. A latch-based test using this exact iterator and mocked JNI reproduces release while decoding is still in flight. This is source-verified reachability plus a lifecycle reproduction, not a reproduced native crash.

[andygrove]

The P1 lifetime concern looks right to me. I reproduced the JVM half of it locally with a latch-based probe on this iterator and a mocked Native: releaseRemoteShuffleDecoder runs to completion while decodeShuffleBlockWithValidation is still inside the JNI call, so nothing in the iterator serializes the release against an active decode. remoteDecoderHandle is published from fetchNext() without the monitor that close() takes, and decode_remote_shuffle_batch reads expected_types after read_ipc_compressed_validated has already decompressed the block, so the borrow outlives the free in that ordering. The mirror-image ordering is a leak rather than a crash, since close() reads 0L and the box is never freed.

What I do not think is true is that this PR introduces the hazard. CometExecIterator has had the same shape for a long time, with a wider window on a much hotter path. It registers close() as a task completion listener at construction (CometExecIterator.scala:189), close() is synchronized and calls nativeLib.releasePlan(plan) (:275 and :305), and getNextBatch calls nativeLib.executePlan(..., plan, ...) outside that monitor (:205). On the native side releasePlan is a Box::from_raw drop of the context (jni_api.rs:1051) while executePlan holds get_execution_context(exec_context), an unbounded &'a mut ExecutionContext derived from the raw handle, for the duration of an entire plan execution (jni_api.rs:848). So if the threaded Python consumer plus an interrupted cleanup wait reaches the shuffle decoder, it already reaches every Comet native plan today.

Given that, I would rather not hold this PR for it. Could we file one issue covering JVM-owned native handle lifetime for both NativeBatchDecoderIterator and CometExecIterator, and fix them the same way? Fixing only the new handle here would leave the wider window untouched and leave us with two different ownership conventions for the same problem.

For what it is worth, the minimal fix for this iterator looks like taking the monitor around just the nativeUtil.getNextBatch call and leaving readNextBlock() outside it, so in.close() can still unblock a transport read. Decoding one block is bounded work, so cleanup cannot hang waiting on the monitor. Whatever we settle on should apply to releasePlan as well.

[pingzh]

Codex acting for the user: Fixed the NativeBatchDecoderIterator lifetime race in 92fa68c5b.

Decoder creation, native decoding, Arrow import, and batch publication now share the monitor used by close(). next() also transfers batch ownership under that monitor. A transport read that finishes after cleanup rechecks isClosed before creating or using a decoder.

readNextBlock() stays outside the monitor so in.close() can unblock it. Native decode failures unwind Arrow cleanup before reporting outside the monitor, since reporting may perform an RPC.

Added five latch-based regressions covering creation, active decoding (including interrupted cleanup), import/publication, blocked transport reads, and failure reporting. All five failed before the fix; all 69 reader-suite tests now pass on both Spark 4.1 and Spark 3.4.3. Spotless and Scalastyle also pass.

The broader CometExecIterator lifetime concern remains a separate follow-up.

[sunchao]

Correctness

[P1] Protect the cached decoder's lifetime during cancellation.

This follow-up at cedc3bd99ec4580226ffe77154a18fc7058a1547, against 75fdddc9285ec61c0cd326977c61dd41fca39a8b, supports the ownership concern in the existing P1 thread. My earlier cleanup assessment was too broad: synchronous decoding within the consumer does not establish exclusive access throughout task cleanup.

The iterator does not coordinate decoder creation and use with closure. The new native ownership contract therefore depends on caller guarantees that are not sufficient under cancellation. Clearing the stored handle before release makes repeated closure idempotent; it does not establish that active use has finished or that creation cannot publish state after closure. Before this change, the expected types belonged to the individual JNI invocation.

Independent inspection of the maintained Spark 3.5 branch also found a threaded Python input consumer and an interruptible cleanup wait. That requires the lifetime protection discussed in the existing thread. The maintained 4.0 branch instead consumes Python input on the task thread; I am not extending that specific threaded-consumer conclusion to 4.0. Maintained 3.4 and 4.1 sources were unavailable, so I have not independently qualified those versions or reproduced the external reviewer's probes.

The final GitHub CI check at 2026-09-06 00:01:47 UTC has 65 successful / 9 skipped checks, with no pending or failing check. Source inspection shows sequential decoder lifecycle coverage and separate concurrent raw-stream-opening tests; neither establishes the new decoder's lifetime during cancellation. This follow-up ran no local Spark, JNI, native build, or concurrency reproduction. No duplicate inline is proposed.

Performance

No additional performance finding from this lifecycle review. A lifetime fix should preserve cancellation responsiveness and avoid unnecessary contention during normal reading. Those effects need validation once the fix exists; this follow-up adds no timing claim.

Design

The iterator needs one explicit ownership rule covering initialization, publication, active native use, and final release. That rule must remain valid when completion callbacks fail or cleanup is interrupted. The native safety preconditions should be enforced by the owning layer, rather than relying solely on normal callback order. Input shutdown must still be able to unblock transport reads.

Abstraction & complexity

Keep the lifetime mechanism local to the iterator and make its relationship with batch and NativeUtil cleanup clear. Merely adding visibility to the handle or another close-only guard would leave the ownership gap. Focused regression coverage should verify the lifetime invariants across cancellation and cleanup failures, while retaining the existing allocation and idempotence checks.