You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/01/31 12:29:21 UTC
[I] Use interleave_record_batch to avoid tiny batches in sort-based shuffle [datafusion-ballista]
andygrove opened a new issue, #1432:
URL: https://github.com/apache/datafusion-ballista/issues/1432
## Problem
The current sort-based shuffle writer uses DataFusion's `BatchPartitioner::partition()` which calls `take_arrays()` to split each input batch into per-partition sub-batches. With `N` output partitions, a single input batch of 8192 rows produces up to `N` sub-batches averaging `8192/N` rows each. With 200 output partitions, that's ~41 rows per batch.
These tiny batches are:
- Appended individually to `PartitionBuffer`
- Spilled individually to IPC files
- Written individually to the final output file
Each batch carries per-batch Arrow IPC overhead (metadata, alignment padding), and the large number of small batches hurts compression ratios and read performance.
## Proposed Solution
Adopt the approach used by [Apache DataFusion Comet](https://github.com/apache/datafusion-comet/blob/main/native/core/src/execution/shuffle/shuffle_writer.rs) which avoids this problem using `arrow::compute::interleave_record_batch`:
1. **Buffer input batches whole** — store unmodified input batches in a `Vec<RecordBatch>` instead of immediately splitting them per partition
2. **Track indices, not data** — for each row, record `(batch_index, row_index)` into per-partition index vectors (`Vec<Vec<(u32, u32)>>`)
3. **Deferred materialization** — when it's time to write (spill or final output), call `interleave_record_batch(&buffered_batches, &indices)` to gather rows into well-sized output batches (up to the configured `batch_size`)
This produces properly-sized output batches regardless of the number of output partitions, with a single efficient gather operation instead of many small `take_arrays` calls.
### Comparison
| Aspect | Current | Proposed |
|--------|---------|----------|
| When rows are partitioned | Immediately per input batch | Deferred — only indices tracked |
| Output batch sizes | `input_rows / N` (tiny) | Up to `batch_size` (well-sized) |
| Data copies | `take_arrays` per partition per batch | One `interleave` call at write time |
### Additional Consideration
Comet also integrates with DataFusion's `MemoryReservation` / `MemoryConsumer` for memory-aware spilling instead of a fixed byte threshold. This could be adopted as a follow-up improvement.
--
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: github-unsubscribe@datafusion.apache.org.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [I] Use interleave_record_batch to avoid tiny batches in sort-based shuffle [datafusion-ballista]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on issue #1432:
URL: https://github.com/apache/datafusion-ballista/issues/1432#issuecomment-3828441839
https://docs.rs/arrow/latest/arrow/compute/struct.BatchCoalescer.html#method.push_batch_with_indices
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [I] Use interleave_record_batch to avoid tiny batches in sort-based shuffle [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on issue #1432:
URL: https://github.com/apache/datafusion-ballista/issues/1432#issuecomment-3855079153
> Hey [@andygrove](https://github.com/andygrove), I'd love to give this a go. Happy to compare interleave against coalesce as well as suggested.
That would be great, thank you. I am hoping that whatever we learn here can also be applied to Comet, where the shuffle design is very similar.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [I] Use interleave_record_batch to avoid tiny batches in sort-based shuffle [datafusion-ballista]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on issue #1432:
URL: https://github.com/apache/datafusion-ballista/issues/1432#issuecomment-3828438963
`interleave` is relatively slow compared to `coalesce` for in order results.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [I] Use interleave_record_batch to avoid tiny batches in sort-based shuffle [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove closed issue #1432: Use interleave_record_batch to avoid tiny batches in sort-based shuffle
URL: https://github.com/apache/datafusion-ballista/issues/1432
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [I] Use interleave_record_batch to avoid tiny batches in sort-based shuffle [datafusion-ballista]
Posted by "mattcuento (via GitHub)" <gi...@apache.org>.
mattcuento commented on issue #1432:
URL: https://github.com/apache/datafusion-ballista/issues/1432#issuecomment-3854915802
Hey @andygrove, I'd love to give this a go. Happy to compare interleave against coalesce as well as suggested.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [I] Use interleave_record_batch to avoid tiny batches in sort-based shuffle [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on issue #1432:
URL: https://github.com/apache/datafusion-ballista/issues/1432#issuecomment-4329723722
Hi @mattcuento. Just wanted to check in - are you still working on this? I have some time now and could take this on.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [I] Use interleave_record_batch to avoid tiny batches in sort-based shuffle [datafusion-ballista]
Posted by "mattcuento (via GitHub)" <gi...@apache.org>.
mattcuento commented on issue #1432:
URL: https://github.com/apache/datafusion-ballista/issues/1432#issuecomment-4331816969
> Hi @mattcuento. Just wanted to check in - are you still working on this? I have some time now and could take this on.
Hey @andygrove, apologies I've been busier than anticipated and haven't been able to make progress past my staged draft. I should have more time next week but please feel free to go ahead and take it. I see you have a referenced PR, happy to help review next week when I free up
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org