apache / apache/datafusion-comet

Read local shuffle blocks directly in native code instead of copying them through a JVM stream

Open
#5,791 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?

Shuffle blocks that live on the same host as the reduce task are read through the same path as blocks fetched over the network. `CometShuffleBlockIterator` (and `NativeBatchDecoderIterator` on the non-directRead path) wraps the block's `InputStream` in a channel and copies the compressed bytes into a direct `ByteBuffer` before handing the address to native code. For a local block those bytes are already a byte range of a shuffle data file on local disk, so the copy and the per-block stream machinery buy nothing.

Spark's resolver already models a local block as exactly that range: `IndexShuffleBlockResolver.getBlockData` returns a `FileSegmentManagedBuffer` carrying the file, offset and length. Native code could `pread` or map that range straight into the decoder and skip the JVM copy.

The share of reads this affects is roughly one over the number of hosts, so on a wide cluster it is small. On single-node and few-node runs, including the TPC-H and TPC-DS benchmarks Comet is usually measured with, it is most or all of the shuffle read.

### Describe the potential solution

`CometShuffleManager.getReader` already has `blocksByAddress` in hand, so local blocks could be split out and served from their `FileSegmentManagedBuffer` while remote blocks continue through `ShuffleBlockFetcherIterator`. That is the awkward part: the fetcher hands back `Iterator[(BlockId, InputStream)]`, so the `ManagedBuffer` is not visible to the reader today and local blocks would have to be diverted before it rather than unwrapped after.

Worth measuring on a single-node run first, since decompression and IPC decode may dominate the copy this removes.

### 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.