apache / apache/datafusion-comet
Return native shuffle partition offsets over JNI instead of through a temp index file
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 198
Description
### What is the problem the feature request solves?
The native shuffle writer already knows every partition offset by the time it finishes a map task, but it hands them to the JVM through a temporary file rather than returning them directly.
`LocalPartitionWriter::finish_all` creates a temp index file and writes `num_output_partitions + 1` little-endian i64 offsets into it. Back in the JVM, `CometNativeShuffleWriter` reads the whole file with `Files.readAllBytes`, converts the offsets to partition lengths, deletes the file, and passes the lengths to `IndexShuffleBlockResolver.writeMetadataFileAndCommit`, which writes Spark's real index file. The temp file exists only to move an array of longs across the JNI boundary, and it costs every map task a file create, write, read, and unlink on top of the index file Spark writes anyway.
The parse is allocation-heavy too: `grouped(OFFSET_LENGTH)` allocates an intermediate array per partition and the `map` allocates a `ByteBuffer` per partition, so `2 * numPartitions` short-lived objects per map task. That part is already recorded as item 5 of #5198, but it goes away entirely if the file does.
### Describe the potential solution
Return the offsets across the existing JNI boundary instead, as a `long[]` or a direct buffer filled by native code, and drop the temp index file. The JVM side then walks a primitive array once to turn offsets into lengths and commits as it does today.
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.