apache / apache/datafusion-comet

Return native shuffle partition offsets over JNI instead of through a temp index file

Open
#5,790 1 comment 0 reactions 1 assignee Claimed by @peterxcli View on GitHub
requires-triage
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.