Skip to content

perf: write every partition of a JVM shuffle spill in one native call #7065

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andygrove wants to merge 4 commits into
base: main
Choose a base branch
from
Open

perf: write every partition of a JVM shuffle spill in one native call #7065

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5c53def
perf: write every partition of a JVM shuffle spill in one native call
andygrove Oct 16, 2026
20a9dc3
test: pin dictionary changes and decoded rows in the sorted spill tests
andygrove Oct 16, 2026
a0f90de
test: run the dictionary change test with zstd offsets as lengths
andygrove Oct 16, 2026
c4ae0a8
refactor: simplify the sorted row encoder and the spill metrics
andygrove Oct 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 17 additions & 9 deletions docs/source/contributor-guide/jvm_shuffle.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,13 +152,14 @@ Selection logic in `CometShuffleManager.shouldBypassMergeSort()`:

### Writers

| Class | Location | Description |
| ----------------------------------- | ---------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------- |
| `CometBypassMergeSortShuffleWriter` | `.../shuffle/CometBypassMergeSortShuffleWriter.java` | Hash-based writer. Creates one `CometDiskBlockWriter` per partition. Supports async writes. |
| `CometUnsafeShuffleWriter` | `.../shuffle/CometUnsafeShuffleWriter.java` | Sort-based writer. Uses `CometShuffleExternalSorter` to buffer and sort records, then merges spill files. |
| `CometDiskBlockWriter` | `.../shuffle/CometDiskBlockWriter.java` | Buffers rows in memory pages for a single partition. Spills to disk via native encoding. Used by bypass writer. |
| `CometShuffleExternalSorter` | `.../shuffle/sort/CometShuffleExternalSorter.java` | Buffers records across all partitions, sorts by partition ID, spills sorted data. Used by unsafe writer. |
| `SpillWriter` | `.../shuffle/SpillWriter.java` | Base class for spill logic. Manages memory pages and calls `Native.writeSortedFileNative()` for Arrow IPC encoding. |
| Class | Location | Description |
| ----------------------------------- | ---------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------- |
| `CometBypassMergeSortShuffleWriter` | `.../shuffle/CometBypassMergeSortShuffleWriter.java` | Hash-based writer. Creates one `CometDiskBlockWriter` per partition. Supports async writes. |
| `CometUnsafeShuffleWriter` | `.../shuffle/CometUnsafeShuffleWriter.java` | Sort-based writer. Uses `CometShuffleExternalSorter` to buffer and sort records, then merges spill files. |
| `CometDiskBlockWriter` | `.../shuffle/CometDiskBlockWriter.java` | Buffers rows in memory pages for a single partition. Spills to disk via native encoding. Used by bypass writer. |
| `CometShuffleExternalSorter` | `.../shuffle/sort/CometShuffleExternalSorter.java` | Buffers records across all partitions, sorts by partition ID, spills sorted data. Used by unsafe writer. |
| `SpillSorter` | `.../shuffle/sort/SpillSorter.java` | Holds one spill's records for `CometShuffleExternalSorter`. Writes every partition of the spill with one `Native.writeSortedPartitionsNative()` call. |
| `SpillWriter` | `.../shuffle/SpillWriter.java` | Base class for spill logic. Manages memory pages. `doSpilling()` encodes one partition's rows with `Native.writeSortedFileNative()`. |

### Reader

Expand All @@ -174,8 +175,15 @@ Selection logic in `CometShuffleManager.shouldBypassMergeSort()`:
1. `ShuffleWriteProcessor` calls `CometShuffleManager.getWriter()`
2. Writer receives `Iterator[Product2[K, V]]` where V is `UnsafeRow`
3. Rows are serialized and buffered in off-heap memory pages
4. When memory threshold or batch size is reached, `SpillWriter.doSpilling()` is called
5. Native code (`Native.writeSortedFileNative()`) converts rows to Arrow arrays and writes IPC format
4. When memory threshold or batch size is reached, the buffered rows are written. `CometDiskBlockWriter`
calls `SpillWriter.doSpilling()` for its partition. `CometShuffleExternalSorter` calls
`SpillSorter.writeSortedFileNative()` once for each spill and once more for the final file, and each
call writes every partition
5. Native code converts the rows to Arrow arrays and writes them as Arrow IPC blocks. `doSpilling()`
passes one partition's row addresses and sizes to `Native.writeSortedFileNative()`. `SpillSorter`
sorts its record pointers by partition with `Native.sortRowPartitionsNative()`, then passes them,
with the address of each data page, to `Native.writeSortedPartitionsNative()`. That call reads the
rows in place, writes every partition of the spill, and returns each partition's length and checksum
6. For bypass writer: partition files are concatenated into final output
7. For sort writer: spill files are merged

Expand Down
Loading
Loading