Conversation
…index file The native shuffle writer knew every partition offset by the time it finished a map task, but handed them to the JVM through a temporary file: LocalPartitionWriter::finish_all created an index file and wrote num_output_partitions + 1 little-endian i64 offsets into it, and CometNativeShuffleWriter read the whole file back with Files.readAllBytes, converted the offsets to lengths, deleted it, and passed the lengths to IndexShuffleBlockResolver.writeMetadataFileAndCommit, which writes Spark's real index file. The temp file existed only to move an array of longs across the JNI boundary, and cost every map task a create, write, read and unlink on top of the index file Spark commits anyway. The parse also allocated an intermediate array and a ByteBuffer per partition (item 5 of apache#5198), which goes away with the file. The offsets are now published in memory through a PartitionOffsets slot shared by the writer and its ShuffleWriterDestination, and read back over JNI by Native.getShufflePartitionOffsets. The index path no longer travels in the plan, so LocalPartitionWriter.output_index_file and the legacy ShuffleWriter.output_index_file are removed and their field numbers reserved. The offsets have to be read while the native plan is still alive. CometExecIterator closes itself when its stream reaches the end, and close releases the execution context that owns the writer, so reading after drainAndClose returned freed memory and produced garbage lengths. The iterator instead captures the offsets at end of stream, before close, when built with capturePartitionOffsets, which only the local destination sets: RSS reports its partition lengths through its pusher. Partition lengths are derived from effectivePartitionCount, the output partition count, not the numParts constructor argument, which is the input partition count. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
- The shuffle crate does not know about JNI, so PartitionOffsets and the local writer describe the handover in terms of the caller driving the plan rather than the JVM reading over JNI. - ShuffleWriterExec::try_new says what it writes where: partition data to a local file, offsets in memory. - CometExecIterator had three lookalike names for one thing. The read is now a named method, readPartitionOffsetsBeforeClose, whose name and doc carry the constraint that made the placement surprising: the offsets live in the native execution context, close releases it, and hasNext closes as soon as the plan runs out of output, so the final hasNext is the last point they can be read. The field is partitionOffsets and the constructor flag is documented. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Two CI failures, both from assuming more than the writer guarantees. The proto crate's own tests still referenced ShuffleWriter.output_index_file and LocalPartitionWriter.output_index_file, so datafusion-comet-proto failed to compile its test target. I had only checked the shuffle and core crates locally rather than the whole workspace. The round-trip tests now assert that a new plan carries no index path, and that a plan still carrying the retired tag 4 decodes cleanly because the tag is reserved rather than reused. partitionLengths was sized by effectivePartitionCount, which is not what the writer produces. isSinglePartitioning serializes a range partitioning whose sampled bounds came out empty as SinglePartition, so native writes one partition while the declared output partitioning still reports several, and the require failed with "returned 2 partition offsets for 10 output partitions". The index file was always sized by what the writer produced, so deriving the length count from the returned offsets restores the previous behaviour exactly. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Requested in review. The plan proto is built by the JVM and consumed by native in the same process from the same artifact, so no old plan ever meets a new reader and there is nothing for a reserved tag to protect against. Decoding is unaffected either way: an undeclared tag is skipped as an unknown field, and reserved only stops protoc from later reusing the number. The proto round-trip test covering a plan that still carries the retired tag 4 keeps passing, and its comment no longer credits reserved for that. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Comments cut back to what they need to say, per review: the JNI entry point, the PartitionOffsets type and its set/get, the destination field, try_new, the partition_offsets accessor, the zero-offset test, the finish_all note, and the two comments in CometNativeShuffleWriter. ShuffleWriter field numbers 5 through 11 shift down to 4 through 10, closing the gap the retired output_index_file left. Both sides of the plan are generated from this file and ship together, so no encoded plan outlives the change. That does mean tag 4 now belongs to codec, and the LegacyShuffleWriter test struct claimed it for a string. Decoding a plan carrying it would be a wire type mismatch rather than a skipped unknown field, so the struct drops that field. The test still covers a legacy plan decoding without a partition writer. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The native shuffle writer spilled each output partition to its own temporary file, created on that partition's first spill and held open until the task finished. A task with P partitions that spilled held up to P spill files open at once, each with a create and an unlink, and every spill scattered its bytes across P files. A task now spills to a single file. Each write appends the partition's blocks and records the range they occupy, and finish_partition copies a partition's ranges into the output in write order. Correctness does not depend on the order partitions are written in, which PartitionWriter leaves unspecified. A single spill round gives contiguous ascending ranges, so the merge reads sequentially. A write that fails partway can leave uncounted bytes in the shared file, which would shift every later range, so the spill refuses further writes and range reads after a failure. The merge also checks each copy's length, so a spill file shorter than its ranges fails the partition instead of writing it short. Closes apache#3859. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Copying a spill range with io::copy costs an lseek, two statx calls and a copy_file_range. Ranges that fit in the write buffer are now read with one read_exact_at into a scratch buffer; longer ranges still use io::copy. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Each partition's spill write ended with a flush, so a spill round issued one write syscall per partition. The spill file's writer is now buffered across partitions and flushed before the merge reads it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
enhancement
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Reviewed a524afb2 against 481aefea, separating the three spill commits from the stacked offset/JNI changes in #5807. I found no P1/P2 issue in this review.
Previously, each spilled output partition retained its own file and writer. The new task-owned PartitionedSpill appends complete encoded blocks to one file and records each partition's ranges in append order. Finalization copies those ranges before that partition's remaining in-memory batches. Empty partitions keep equal adjacent offsets, and the final offset is published only after the output is flushed. Deferring the underlying flush does not defer completion of an IPC or codec frame.
The maintained Spark 3.5 and 4.0 branches concatenate each partition's spill segments before publishing lengths, and their IndexShuffleBlockResolver creates the final index and commits the data. This PR retains that JVM commit boundary. The byte-copy change does not reinterpret values, nulls, dictionary data, hash/range assignments, or fallback decisions. Failed spill writes or flushes poison the shared spill. Both copy paths reject truncation, and errors propagate before offsets or map status are published. Each writer owns its spill and read handle. It adds no cross-task shared writer or lock. Spark checksum-file behavior remains the existing empty-checksum-array behavior.
Validation
The Linux Rust job executed merge 4a437438, whose parents are the reviewed base and head: 1,448 tests passed, five skipped. The three files changed by the spill commits are byte-identical between that merge and the head. This includes the new interleaving, buffering, one-file, poison, and truncation tests.
The Spark 3.5 shuffle job passed 442 Scala tests with six canceled. Spark 4.0 and Spark 4.1 each passed 491. All ran the same merge, including the forced-spill test. Six additional local probes passed for the verbatim range-copy and flush-deferral helpers: mixed copy sizes, pending output, truncation, read/write errors, and deferred flush failure. Those are macOS component tests with an error-type shim, not an Arrow/JNI or Linux zero-copy run. Maintained Spark 3.4/4.1 source branches were unavailable, so their source-level compatibility was not independently verified.
Performance
The design removes per-partition file creation/open/unlink work and repeated block-writer clones. Small ranges trade kernel-copy setup for a positional read and buffered write. Large ranges keep bounded io::copy. The shared spill buffer combines small partition writes. These changes match the bottleneck described in the PR.
The author reports 6.6–19.7% lower times for 4,000–16,000 partitions, but 2.0–4.3% higher times for the lighter 200-partition spill cases. These are rotating-order microbenchmarks on one rotational-disk host, not independently reproduced results. The first-commit TPC-H/ClickBench runs did not spill and do not validate the final spill path's performance.
The principal memory tradeoff is one range per nonempty partition per spill, retained until writer teardown, plus a spill-write buffer and merge-read buffer. The range metadata is outside the partitioner's batch reservation. Could you add a current-head microbenchmark with many spill rounds and concurrent tasks on SSD/NVMe, reporting peak memory as well as elapsed time? The unreserved range index grows with spill count, so this would check memory scaling and whether the reported light-spill slowdown also appears on non-rotational storage.
Design
The ownership boundary is clear: one append-only spill belongs to one local writer, while final output remains contiguous by partition. Keeping the existing encoder and Spark commit path limits the change to temporary storage and merge mechanics. A range records only a completed write, and poisoning the entire shared file is the appropriate response once a partial write makes later offsets unreliable.
A file per spill round is a plausible alternative, but would reintroduce file and merge-handle growth with spill count. The current range index is a reasonable tradeoff for the demonstrated many-partition workload. There is no evidence here requiring a more elaborate merge strategy.
Abstraction & complexity
PartitionedSpill earns its scope by keeping file ownership, append offsets, partition ranges, and the failure state together. The small DeferFlush adapter preserves per-partition batch finalization while batching physical writes. The copy helper isolates the two I/O paths and their exact-length checks. Both remain private to local shuffle. I found no unnecessary general-purpose abstraction to remove.
2 participants
Which issue does this PR close?
Closes #3859.
Stacked on #5807; only the last three commits belong to this PR.
Rationale for this change
The native shuffle writer spilled each output partition to its own file and kept every file open until the merge. With many partitions a task creates, opens, closes and unlinks thousands of files, and under the default 1024 soft
nofilelimit a task spilling 1000 partitions fails withToo many open files.What changes are included in this PR?
PartitionedSpillreplaces the per-partitionSpillWriters: one spill file per task and the byte ranges each partition's blocks occupy. A failed write makes the spill unusable.finish_partitioncopies a partition's ranges into the output in write order. A range that fits in the write buffer is read with oneread_exact_at; longer ranges keepio::copy. A spill file shorter than its ranges fails the task.BufWriteracross partitions, flushed before the merge reads the file.How are these changes tested?
New unit tests: partitions interleaved across spill rounds read back in write order through both copy paths, one spill file for any partition count, truncated spill file, spill unusable after a failed write, writes buffered until flush.
datafusion-comet-shuffle132 passed, clippy clean.shuffle_benchoutput is byte-identical to #5807's at 64, 1000 and 4000 partitions with 8 spills and at 200 partitions with 49 spills.Benchmarks
shuffle_bench, 8M rows (5 numeric columns), lz4, 16-core Linux host with a rotational disk, 3 rounds in rotating order. Mean time change against #5807:--max-buffer-bytes(spills)At 200 partitions with 49 spills, syscalls drop from 12,152 to 10,375 (
write9,846 to 263, no per-partitionopenat/unlink).TPC-H SF1/SF10/SF100 and ClickBench (Spark 4.1, 2 executors x 8 cores) never spill at default settings and show no difference. These end-to-end runs used the first commit only.