NoahKusaba opened a new issue, #2858:
URL: https://github.com/apache/datafusion-ballista/issues/2858
**Is your feature request related to a problem or challenge?**
Two shuffle writers `take` rows into an intermediate batch before pushing
them into a per-partition `BatchCoalescer`, so every row is copied twice:
- **Pipelined hash writer** (`pipelined/writer.rs`).
`BatchPartitioner::partition_iter` `take`s each partition's rows into a small
batch, which is then pushed into that partition's coalescer.
- **Sort-shuffle split on arrival** (#2774, open). One grouped
`take_record_batch` per input batch, then a slice of it is pushed into each
partition's coalescer.
apache/arrow-rs#11174 (approved) makes
`BatchCoalescer::push_batch_with_indices` copy straight from the source batch
when the indices are non-null `UInt32` and the selection is sparse: at most
half the rows, or at most 1/16 if the schema has a view column. Otherwise it
falls back to `take`. There is no API change.
**Describe the solution you'd like**
Once Ballista is on an arrow release that includes #11174:
1. **Pipelined hash writer.** Route rows with the sort-shuffle writer's
`RowRouter`, which matches `BatchPartitioner`'s hashing, and share it between
the two writers. Then call `push_batch_with_indices` for each partition.
2. **Sort-shuffle split path, after #2774.** Slice the grouped indices per
partition and call `push_batch_with_indices` instead of `take`.
- With a view column and fewer than about 16 partitions, arrow falls back
to a `take` per partition, which costs more than today's single grouped `take`.
Keep the grouped `take` in that case unless measurements show otherwise.
- Also re-check the split budget's 1 MiB-per-view-column estimate: the
direct path sizes view buffers differently.
3. **Measure** `write_time`, `buffer_time` and peak RSS at 12 and 200+
partitions, with and without view columns.
That release is likely 60.x, which comes with the DataFusion 56 upgrade
(#2718), or 61, which needs a later DataFusion release.
**Describe alternatives you've considered**
- Keep `take` + `push_batch`. It works, but keeps the double copy.
- Copy rows by index in Ballista. That duplicates arrow-rs internals.
**Additional context**
Comet dropped a similar immediate-mode shuffle
(apache/datafusion-comet#3855) for two reasons:
- The intermediate small arrays doubled peak RSS. #11174 addresses this.
- Per-partition coalescers don't scale to many partitions. This doesn't
apply to the pipelined writer, which already keeps a coalescer per partition.
#2774 deals with it by splitting only when the coalescers fit a memory budget.
Peak RSS is the number to watch.
--
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]
