peterxcli opened a new pull request, #5807: URL: https://github.com/apache/datafusion-comet/pull/5807
## Which issue does this PR close? Closes #5790. ## Rationale for this change The native shuffle writer already knows every partition offset by the time it finishes a map task, but it 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; `CometNativeShuffleWriter` then 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 was allocation-heavy too: `grouped(OFFSET_LENGTH)` allocated an intermediate array per partition and the `map` a `ByteBuffer` per partition. That is item 5 of #5198, and it goes away with the file. ## What changes are included in this PR? - The writer publishes its offsets through a new `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 for returning a `jlongArray`. - 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 planner validation for those paths goes with them. - `CometNativeShuffleWriter` derives partition lengths from the returned offsets and no longer writes, reads or deletes a temp index file. Two things worth a reviewer's attention: **The offsets have to be read while the native plan is still alive.** `CometExecIterator.hasNext` closes itself when the stream ends, and `close` calls `releasePlan`, which frees the execution context that owns the writer. Reading after `drainAndClose` returned freed memory and produced garbage lengths. The iterator now captures the offsets at end of stream, before `close`, when constructed with `capturePartitionOffsets` — set only for the local destination, since RSS reports partition lengths through its pusher instead. **Partition lengths come from `effectivePartitionCount`, not `numParts`.** `numParts` is the input partition count (`CometExecIterator` asserts `partitionIndex < numParts`); the offsets are sized by the output partition count. ## How are these changes tested? Existing coverage, updated to assert on the published offsets rather than on index file bytes, which is a more direct check of the same property: - `shuffle_writer.rs` empty-schema tests now assert offset count, the leading zero, and that trailing offsets equal the data length; the round-robin determinism test compares published offsets across runs instead of index files. - `rss_execution_tests.rs` asserts the single-partition writer publishes `[0, dataFileLength]`. - Planner tests assert a freshly built local destination has nothing published yet. The tests covering the removed index-path validation are deleted. - `CometNativeShuffleSuite` and `CometCelebornNativeShuffleWriterSuite` proto round-trip tests updated for the removed fields. Test runs on this branch: - `cargo test -p datafusion-comet-shuffle` — 125 passed - `cargo test -p datafusion-comet --lib` — 330 passed - `CometNativeShuffleSuite` — 53 passed - `CometCelebornNativeShuffleWriterSuite` — 16 passed - `CometShuffleSuite` — 44 passed - `cargo clippy --all-targets -D warnings` clean on both crates The use-after-free above was caught by `CometNativeShuffleSuite`, which failed every execution test until the capture moved inside the iterator. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
