andygrove opened a new pull request, #7234: URL: https://github.com/apache/datafusion-comet/pull/7234
## Which issue does this PR close? Closes #7071. ## 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`. 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, so the copy buys nothing. With many partitions almost every chunk is a tail. 2. **A flush and an `lseek` per partition.** `finish_partition` called `stream_position()` on the `BufWriter<File>`, which flushes the buffer, so every partition left as its own `write(2)` plus a seek and the 1 MiB output buffer never coalesced small partitions. 3. **A fresh `IpcWriteContext` per writer**, so the first block of every partition regrew arrow-ipc's body scratch and rebuilt the flatbuffer builder. This supersedes the draft #5908. ## What changes are included in this PR? - `BufBatchWriter` gains a passthrough mode and `BufBatchWriter::new_for_partition`, which the multi-partition writers (`LocalPartitionWriter` and `PartitionedSpill`) use. The single-partition writer still coalesces, since its inputs can genuinely be small. - Schemas that contain a view type (`Utf8View`, `BinaryView`, `ListView`, `LargeListView`, at any depth) keep the coalescer. The IPC writer serializes every data buffer a view array references, so writing each small chunk of an `interleave`d batch verbatim would repeat all of the source batches' data buffers in every block, while arrow's coalescer compacts sparse views. `ShuffleBlockWriter::has_view_types` records this once at construction. - `BufBatchWriter::drain` hands buffered bytes to the underlying writer without flushing it; `flush` is `drain` plus a flush. - The task-scoped scratch now carries the IPC context alongside the byte buffer. - `LocalPartitionWriter` tracks the output offset as a running byte total and flushes the output once in `finish_all`, where the total is checked against the real file position and a 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. ## How are these changes tested? - `passthrough_writes_each_batch_as_its_own_block` pins the passthrough block boundaries and row order. - `for_partition_coalesces_only_view_schemas` checks that a view schema still coalesces and a plain schema passes through. - `offsets_match_data_file_with_spilled_and_in_memory_batches` writes partitions with a spilled prefix and in-memory tail, then decodes every block inside each index range, so a wrong arithmetic offset fails to parse. - The existing `datafusion-comet-shuffle` tests (193) and clippy pass. Benchmark, `shuffle_writer_high_partition` (81,920 rows, 4-column schema, codec None, hash partitioning), measured on the draft #5908 before the view-type guard was added; the guard only changes behaviour for view schemas: | 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% | -- 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]
