perf: return native shuffle partition offsets over JNI instead of an index file

[peterxcli]

Which issue does this PR close?

Closes #5790.

Rationale for this change

The native writer knows every partition offset when a map task finishes, but handed them to the JVM through a temp file: finish_all wrote numPartitions + 1 offsets to an index file, and CometNativeShuffleWriter read it back with Files.readAllBytes, converted offsets to lengths, deleted it, then passed the lengths to writeMetadataFileAndCommit — which writes Spark's real index file. The temp file existed only to move an array of longs across JNI.

Measuring the removed sequence directly (create, write offsets, close, readAllBytes, parse, delete), and the parse separately with bytes already in memory, 2000 iterations after 300 warmup, each count measured twice in both orders:

partitions	round trip (two runs)	parse alone	file syscalls
200	111.8 us / 214.1 us	4.1 us	~110-210 us
2000	209.7 us / 181.5 us	38.7 us	~145-170 us
16000	588.7 us / 694.8 us	292.0 us	~300-400 us

A ~150-200 us syscall floor per map task plus ~18 ns per partition of parsing. What replaces it is one JNI call returning a long[], which the JVM allocated before anyway.

This is fixed per-task overhead, not throughput, so the job-level effect is the saving times the map-task count. Measured on macOS/APFS on a loaded laptop; the syscall half is filesystem-sensitive and would grow on EBS. The 200-partition row is the noisiest, so treat it as a range. The parse allocations are item 5 of #5198, which this removes along with the file.

What changes are included in this PR?

• Offsets are published through a PartitionOffsets slot shared by LocalPartitionWriter and its ShuffleWriterDestination, set once in finish_all.
• Native.getShufflePartitionOffsets returns them as a long[], following the existing writeSortedFileNative precedent.
• LocalPartitionWriter.output_index_file and the legacy ShuffleWriter.output_index_file are removed, their field numbers reserved, and the planner validation for them dropped.
• native_shuffle.md and the in-tree comments describe the single data file and the offsets returned over JNI.

Two things worth a reviewer's eye:

Offsets are captured at end of stream, not after draining. CometExecIterator.hasNext closes itself when the stream ends, and close calls releasePlan, freeing the context that owns the writer. Reading after drainAndClose returned freed memory and garbage lengths. The iterator captures them before close, when built with capturePartitionOffsets — set only for local output, since RSS reports lengths through its pusher.

Lengths are sized from the returned offsets. A range partitioning whose sampled bounds are empty is serialized as a single partition, so native can write fewer partitions than the declared output partitioning reports. The old index file was sized by what the writer produced, and the lengths still are.

How are these changes tested?

Existing coverage, updated to assert on published offsets rather than index file bytes — a more direct check of the same property. The empty-schema tests assert offset count, leading zero and trailing data length; the round-robin determinism test compares offsets across runs; rss_execution_tests asserts [0, dataFileLength] for a single partition; planner tests assert a fresh destination has published nothing. Tests covering the removed validation are deleted, and the proto round-trip tests decode a plan that still carries the reserved tag 4.

datafusion-comet-proto 4, datafusion-comet-shuffle 125, and the datafusion-comet shuffle and planner tests (21) pass; workspace clippy clean. CometNativeShuffleSuite 53, CometCelebornNativeShuffleWriterSuite 16, CometShuffleSuite 44.

The use-after-free above was caught by CometNativeShuffleSuite, which failed every execution test until the capture moved inside the iterator.

[sunchao]

Correctness

Prior behavior and implementation

The old native writer emitted a temporary index solely to transfer partition offsets to the JVM. The JVM read and parsed it, deleted it, and passed lengths to Spark, which wrote the committed index. This PR removes that intermediate file. finish_all flushes the data output, publishes checked signed offsets through PartitionOffsets, and the JVM obtains a long[] through JNI.

The ordering is correct in the reviewed path. Native execution completes the shuffle write before returning its empty output stream. CometExecIterator.hasNext captures offsets before close releases the execution context. The JNI result is a JVM-owned copy, so it remains valid after native teardown. Later hasNext calls see the closed iterator. An exception during capture still reaches drainAndClose's safe-finally cleanup before any shuffle commit.

The local writer retains one offset per output partition plus the trailing file length. Empty partitions retain equal adjacent offsets, and zero-row output publishes all zeros. The existing checked u64 to i64 conversion still rejects overflow. Lengths are derived from the published array rather than the number of input partitions. Hash, range and round-robin routing, compression and row serialization are unchanged. RSS leaves capture disabled and continues obtaining lengths through its task-owned pusher.

The maintained Spark 3.5 and 4.0 IndexShuffleBlockResolver implementations expect partition lengths, reconstruct the final index with a leading zero, and may overwrite the supplied lengths when reusing an existing successful attempt. This PR still passes the mutable lengths array to that resolver and builds MapStatus afterward, preserving that behavior. It also preserves the existing empty checksum array. I found no verified P1/P2 correctness issue in the handoff or its interaction with the base's aggregate-planning changes.

Validation and limits

Reviewed head 7e4c5e4139561f013480a39d263a936f541f5284 against base 424c31aa79d13fddf743ffa29bae3c6f146e6c5e. Local work comprised source, diff and provenance checks. I did not run a new local Rust or JVM build.

The Rust CI job passed 1,260 tests with five skipped, including local offset publication, empty-schema output, zero rows and round-robin determinism. The Linux Spark 3.5 shuffle job passed 438 tests with six canceled. The Linux Spark 4.0 shuffle job and macOS Spark 4.0 shuffle job each passed 487 with none canceled.

Those jobs checked out synthetic merge 4a2f1ba0b0ad7987991c458ea1aba76391b8d229, whose first parent is 80c8ec37af422db1952a1128cb23043c9bbfbb72. All 15 authored files and the complete native core, shuffle and protocol trees match the reviewed head. Native artifact producer and consumer IDs and SHA-256 digests match on Linux and macOS. The whole trees and workflow files differ, so these are qualified merge-context results, not an execution of the exact assigned head/base pair. At 06:04 UTC on September 11, the fresh check snapshot showed 70 successes and nine skips. Maintained Spark 3.4 and 4.1 sources were unavailable.

Performance

Removing the temporary file eliminates its create/write/read/delete work and the JVM's grouped-byte parsing. The replacement still converts the native offset vector, allocates and fills a JVM long[], and allocates the final lengths array. PartitionOffsets.get borrows the published slice without another Rust clone. These operations are linear in the output partition count and happen once at completion.

The reported 200/2,000/16,000-partition measurements time the removed path and its parser, rather than both complete implementations. They establish an avoided cost, but do not measure net saving through the replacement JNI path. Could you add a paired microbenchmark of the old and new complete handoffs at one partition and those three counts, with the same offsets, warmup and filesystem, reporting latency and allocation? Include the native conversion, JNI array allocation/copy and JVM differencing in the new measurement. That would substantiate the per-task improvement before extrapolating it to aggregate task cost or job time. I found no measured regression to assign P1/P2 severity.

Design

The data file remains owned by the existing Spark commit flow, while completion metadata moves through the execution context. Capturing at end of stream is the key lifecycle decision because the generic iterator releases that context immediately afterward. The shared one-shot slot lets the completed writer expose metadata without extending native plan lifetime or adding another temporary-file protocol.

The protobuf field renumbering and new JNI entry point require matching JVM and native builds. Retaining the legacy data-path field does not establish compatibility with an older native binary. The reviewed CI uses matching builds, and no mixed-version execution is claimed here.

Abstraction & complexity

PartitionOffsets is a small OnceLock<Vec<i64>> wrapper with an explicit duplicate-publication error. Its shared ownership survives destination cloning and child replacement, while each planned task receives a fresh slot. This is an appropriate scope for the result.

The iterator's capture flag adds shuffle-specific completion handling, but defaults to false and is enabled only for local native shuffle output. The implementation stays limited to that lifecycle boundary. It removes index-file fields, planner validation and parser code without introducing a general callback or metadata framework.

[andygrove]

Thanks for this. The core change looks correct to me. The offsets are captured in hasNext before releasePlan runs, the long[] is a JVM-owned copy, and deriving the partition lengths from offsets.length - 1 is a small improvement over sizing them from numParts. The error paths still close the iterator and clean up the temp data file.

I do have one thing I would like changed before this merges. I do not think we should renumber protobuf fields. The last commit moves codec through partition_writer down by one to close the gap left by output_index_file, and drops the reserved declarations. Every other retired field in operator.proto uses reserved, and renumbering is what the comment on output_data_file two lines up says we are trying to avoid. Details inline.

Two smaller things in the same pass:

• docs/source/contributor-guide/native_shuffle.md (steps 6 and 7, around lines 154 to 159) still describes the index file as one of two output files and says CometNativeShuffleWriter reads it back for partition lengths. Could you update that section to describe the data file plus the offsets returned through Native.getShufflePartitionOffsets?
• The "What changes are included" section of the PR description still says the field numbers are reserved, which no longer matches the diff. Once the proto is back to reserved that will be accurate again.

[andygrove]

Please keep the original tag numbers and mark the retired one as reserved 4; reserved "output_index_file"; instead of shifting codec through partition_writer down by one. That is what the rest of this file does for retired fields, it costs nothing, and it keeps the comment above about output_data_file being retained for older native binaries true. With the renumbering, an older binary would decode a wire-type mismatch on tag 4.

[peterxcli]

Thanks for reviewing this. My understanding is that users normally deploy the released Comet JAR, which bundles the JVM code and its matching native library. Both protobuf producers and consumers therefore change together. Do we support any deployment where those components intentionally use different Comet versions?
With matching builds, removing the fields and renumbering should not affect runtime correctness. I understand the value of keeping reserved declarations for consistency and preventing accidental reuse, but I think that is a separate consideration from supporting older native binaries. Could we clarify that compatibility requirement and update the existing comment accordingly?

[andygrove]

Same for LocalPartitionWriter. Tag 2 is now silently free, so a reserved 2; reserved "output_index_file"; here would keep someone from reusing it for a different field later.

[peterxcli]

same reply in #5807 (comment)

[sunchao]

Rechecked unchanged head 7e4c5e41 against 424c31aa after Andy’s review. I agree with retaining ShuffleWriter tags 5–11 and reserving the removed tag 4 and name, plus reserving LocalPartitionWriter tag 2 and name. These changes preserve field identities without increasing encoded tag size. My earlier matched-JVM/native-build qualification described the execution we validated, not a reason to reuse those tags. Protobuf’s field-number guidance supports keeping them stable.

Restoring the tags alone does not establish compatibility with older native binaries. This JVM no longer supplies the old index-file path and calls the new getShufflePartitionOffsets entry point. Please make the legacy-compatibility comments reflect that boundary.

The documentation and description requests also remain: native_shuffle.md should describe offsets captured through JNI before native-plan release, followed by Spark creating the committed index. The description’s claim that the removed tags are reserved will become accurate when the schema is corrected. I have not duplicated the existing inline comments.

The earlier paired old/new handoff benchmark request remains unanswered. CI is still green, with 70 successful checks and nine skips. I reverified that all 15 authored files match its executed merge 4a2f1ba0, while the complete trees differ. No new local runtime tests or benchmarks were run. This follow-up preserves my existing approval without submitting another one.

[sunchao]

Responding to your compatibility question, also referenced in the other thread: the policy at this head already requires the JVM JAR and native library to come from the same Comet release and explicitly excludes cross-release mixing. It takes effect from 1.0.0, without retroactive guarantees for 0.x. Your point about matching builds is correct: renumbering alone does not break a producer and consumer rebuilt with the same schema.

The packaging configuration bundles the native libraries. NativeBase also tries java.library.path before the bundled resource. That is a loading mechanism, not an additional compatibility guarantee. This PR adds getShufflePartitionOffsets and stops supplying the old index-file path, so retaining protobuf tags alone cannot make the old native implementation usable. The older-binary comments should be aligned with the documented coupling policy.

My field-number/reservation request is therefore a schema-maintenance convention, following protobuf guidance, not a promise of a stable mixed-version JNI ABI. The same Comet policy explicitly excludes internal plan protobuf from its compatibility guarantees.

No source changes or new verified P1/P2 findings since my last review. The existing documentation and paired-handoff benchmark requests remain open. CI remains at 70 successes and nine skips, with the same qualified merge-context coverage described previously. No new runtime tests or benchmarks were run.

[peterxcli]

My field-number/reservation request is therefore a schema-maintenance convention, following protobuf guidance, not a promise of a stable mixed-version JNI ABI. The same Comet policy explicitly excludes internal plan protobuf from its compatibility guarantees.

@sunchao @andygrove my point is: if there’s no compatibility concern, this is a great opportunity for a code cleanup. If we keep the old code around indefinitely, the codebase will only grow harder to maintain. And we don’t lose anything by removing it -- the schema evolution is still fully traceable through git blame and the PR history.

That said, I’m happy to keep it if there’s a compatibility case I’m missing. Would love to hear your thoughts.

[andygrove]

Thanks for the reply. You are right that we do not support a jar and a native library from different Comet releases. docs/source/about/versioning_policy.md requires them to come from the same release and explicitly excludes the internal plan protobuf from the compatibility guarantees. So renumbering will not break a matched build and the wire type mismatch I described cannot actually be hit. Please treat that part of my earlier comment as withdrawn.

The request still stands for a different reason. reserved is what operator.proto already does for every other retired field. HashAggregate carries reserved 3, 8; with the names, DeltaScanCommon carries three separate declarations, and IcebergDeleteFile has reserved 5, 6, 7;. It is one line, it does not change the encoding, and it keeps the file self describing for anyone reading a captured plan or an older log. Renumbering does the opposite. It reassigns six live tags in a message that is under active change, and nothing in the file now records that tag 4 ever meant output_index_file. Could we leave codec through partition_writer on 5 through 11 and add reserved 4; reserved "output_index_file"; here, plus reserved 2; reserved "output_index_file"; on LocalPartitionWriter?

I do agree with the cleanup instinct, and I think it points at a bigger target than the tag numbers. If there is no supported mixed version deployment then the top level ShuffleWriter.output_data_file is dead weight too. buildUnifiedPlan always sets partition_writer, so the fallback branch in shuffle_writer_destination is unreachable from any JVM we ship, and after this PR an older native library cannot serve a new JVM anyway because it has no getShufflePartitionOffsets symbol. Would you be willing to open an issue for removing that field and keep this PR to the index file removal?

Two things from my earlier pass are still open. docs/source/contributor-guide/native_shuffle.md steps 6 and 7 still describe two output files and say CometNativeShuffleWriter reads the index file back for partition lengths. And the PR description still says the retired field numbers are reserved.

Three in-tree comments also still describe the index file. The doc on LocalPartitionWriter says it writes "a single data file plus an index file recording the byte offset where each partition begins", which is the struct this PR changed. ShufflePartitioner::shuffle_write in partitioners/traits.rs says "Write shuffle data and shuffle index file to disk". And the output_dir doc in bin/shuffle_bench.rs still says "data/index files" though the bench no longer creates one. Could those go in the same pass?

For what it is worth I went through the risky parts and they look right. The external shuffle service contract is intact since the deleted file was Comet's own .index.tmp and Spark's committed .index is still written by writeMetadataFileAndCommit. Offset bounds, empty partitions and the zero-partition case all reproduce the old index-derived behavior. Task retry and speculation are actually improved, since two attempts previously shared one .index.tmp path.

[sunchao]

Rechecked 1f2d90f355c1 against 481aefea9c60, including the changes since 7e4c5e41. The new commit addresses the field-number requests: ShuffleWriter keeps tags 5–11 and reserves tag 4 plus its name, and LocalPartitionWriter reserves tag 2 plus its name. The restored test decodes a plan carrying the retired index-path field. The write-path guide and the three stale index-file comments now describe the in-memory offsets and Spark's committed index correctly.

The iterator, JVM shuffle writer, native offset getter, plan-release function and destination selection are unchanged from the previously reviewed implementation. The local writer's incremental change is documentation only. I found no new or remaining verified P1/P2 in the offset handoff or failure handling.

The current Rust CI passes all four protobuf tests, including the restored tag-4 case, plus the local-offset and empty-schema cases. Linux Spark 3.5 shuffle CI reports 442 passed and six canceled. Spark 4.0 shuffle CI reports 491 passed with none canceled. These jobs checked out aa02f30e5f7d, whose parents are the assigned base/head and whose tree equals this head. The native artifact ID and digest match between its producer and all three inspected Spark shuffle consumers. The current snapshot has 53 successful and ten skipped checks.

The earlier paired old/new handoff benchmark request remains unanswered, so the net latency/allocation saving is still unverified. The existing comments about supporting older native binaries also remain inconsistent with the documented same-release coupling, as discussed in the previous follow-up. I have not duplicated those requests inline. No new local product build or benchmark was run. Maintained Spark 3.5/4.0 resolver sources were reverified, while maintained 3.4/4.1 sources remain unavailable.

[sunchao]

Re-reviewed 2c98b34a against 1f2d90f3 and base 58ab5f61. This revision merges main. The tag 5–11 retention and retired tag 4/tag 2 reservations remain intact. The offset getter, iterator capture-before-release and destination selection are unchanged. I checked the new codec-context and plan-injection integrations: final data flush still precedes offset publication, task-specific output paths are preserved, and Spark commits the index before MapStatus is created. No new or remaining verified P1/P2 findings.

Current Rust CI passed 1,478 tests with five skipped, including the protobuf, offset, empty-output, determinism and new codec-release cases. Linux Spark 4.1 shuffle CI passed 494 tests with none canceled. Both checked out 82896764, whose parents are this base/head and whose tree equals the head. The native artifact ID and SHA-256 digest match between the producer and shuffle consumer. At September 17, 08:46 UTC, the overall snapshot has 20 successful checks, 13 skipped and two still running.

The previously recorded paired-handoff benchmark and same-release wording requests remain unanswered. The supplied timings still measure the removed path, so net latency/allocation savings remain unverified. No new local product build or benchmark was run. Maintained Spark 3.5/4.0 resolver sources were reverified. Maintained 3.4/4.1 sources remain unavailable.