andygrove opened a new pull request, #5908:
URL: https://github.com/apache/datafusion-comet/pull/5908
## Which issue does this PR close?
Part of #5905 (findings W2 and W3). Does not close it.
## Rationale for this change
In a multi-partition native shuffle, every output partition gets its own
short-lived `BufBatchWriter` per spill or finish cycle, so anything that writer
does per instance is multiplied by `partitions x cycles`. Three such costs were
found in the review:
1. **Every sub-`batch_size` chunk was materialized twice.**
`PartitionedBatchIterator` already emits maximal `batch_size` chunks plus one
tail per partition. Full chunks bypass the `BatchCoalescer`, but the tail is
`copy_rows`'d into in-progress builders and re-emitted as the same block. There
is never a second batch for the tail to coalesce with in this path, so the copy
buys nothing. With many partitions almost every chunk is a tail, so this was a
full second copy of the shuffle payload plus one coalescer (with a boxed
in-progress builder per column) per partition per cycle.
2. **A fresh `IpcWriteContext` per writer.** arrow-ipc only retains capacity
inside the context, so the first block of every partition regrew the body
scratch from empty and rebuilt the flatbuffer builder.
3. **A flush and an `lseek` per partition.** `finish_partition` called
`stream_position()` on the `BufWriter<File>` to record the partition offset,
which flushes the buffer, and then `BufBatchWriter::flush` flushed it again.
Every partition therefore left as its own `write(2)` plus a seek, so the 1 MiB
output buffer never coalesced small partitions.
## What changes are included in this PR?
- `BufBatchWriter` gains a passthrough mode (`new_passthrough`) that
serializes every batch as its own block, and a `drain` method that hands
buffered bytes to the underlying writer without flushing it. `flush` is now
`drain` plus a flush. The coalescing mode is unchanged and still used by the
single-partition writer, whose inputs can genuinely be small.
- The per-call scratch `Vec<u8>` becomes a `ShuffleScratch { buffer,
ipc_context }`, so the task-scoped recycling that already existed for the byte
buffer now also covers the IPC context.
- `LocalPartitionWriter` multi-partition mode and `SpillWriter` use
passthrough writers with the shared scratch. `LocalPartitionWriter` tracks the
output offset as a running byte total (spill copy bytes plus
`BufBatchWriter::bytes_written`) instead of asking the `BufWriter` for its
position, and flushes the output once in `finish_all`, where the running total
is checked against the actual file position and any mismatch fails the task.
Block boundaries and the file format are unchanged: the multi-partition path
wrote one block per chunk before too, just after an extra copy.
## Benchmark
`shuffle_writer_high_partition` (81,920 rows of a 4-column schema, codec
None, hash partitioning), Apple Silicon, criterion `--baseline`:
| Partitions | Before | After | Change |
|---:|---:|---:|---:|
| 200 | 3.68 ms | 2.91 ms | -20% |
| 2000 | 11.95 ms | 6.78 ms | -43% |
| 8000 | 26.75 ms | 13.25 ms | -51% |
The main `shuffle_writer` group (16 partitions, all codecs, range
partitioning) improved by 3 to 6 percent and the single-partition cases are
unchanged.
## How are these changes tested?
- New unit test `passthrough_writes_each_batch_as_its_own_block` in
`buf_batch_writer.rs` pins the passthrough mode's block boundaries and row
order.
- New unit test `offsets_match_data_file_with_spilled_and_in_memory_batches`
in `local_partition_writer.rs` writes two partitions, one with a spilled prefix
and in-memory tail, then decodes every block inside each index range and checks
the row counts, so a wrong arithmetic offset would fail to parse.
- Existing `datafusion-comet-shuffle` tests (127) and clippy pass.
- `CometNativeShuffleSuite` (end-to-end native shuffle writes and reads,
including spilling cases) passes against the rebuilt native library.
--
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]
