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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.