Repository navigation
Integrate BatchCoalescer into RepartitionExec #18782
Description
Activity
take
Reacted by Andrew LambHi @alamb just double checking, I noticed this comment https://github.com/apache/datafusion/pull/18630/files#r2526799488, but regardless we still want to integrate to remove dependency on
CoalescerBatchesExec? Thanks!Hi @alamb just double checking, I noticed this comment https://github.com/apache/datafusion/pull/18630/files#r2526799488, but regardless we still want to integrate to remove dependency on
CoalescerBatchesExec? Thanks!If you mean
Yes - I tried to apply this optimization to RepartitionExec but unfortunately got some regressions (my feeling is due to some decreased parallelism / work distribution).
Yes, I guess I am hoping that someone could figure out what was going on / optimize the code so that we could combine them
Reacted by jizezhangHi @alamb wanted to check my understanding and discuss the approach a little bit. Please let me know your thoughts and correct me if anything sounds off.
I thought about two places where we could potentially coalesce batches:
- Coalesce hash-partitioned input batches for the same output partition in an input stream before sending to output channel, e.g. coalesce in
RepartitionExec::pull_from_inputbeforechannel.sender.send.- For order preserving case, this would result in larger input batches for merge-sort, and as merge-sort needs to wait for batches from all input partition streams to be ready, I wonder if this would impact performance.
- Coalesce batches in output stream when
poll_nextis called, e.g.- For non order-preserving case, output stream is a
PerPartitionStream; for order-preserving case, output stream is aSortPreservingMergeStreamwhose input is vector ofPerPartitionStreams. - If coalescing in e.g.
PerPartitionStream::poll_next_inner, it also means larger input batches for merge-sort. - Alternatively, could create a wrapper stream around
PerPartitionStreamthat coalesces while callingPerPartitionStream::poll_next. This would look very much like aCoalesceBatchesbut not as a separate node. - For order preserving case, maybe coalesce inside
SortPreservingMergeStream::poll_next_inner.
- For non order-preserving case, output stream is a
Thanks a lot!
- Coalesce hash-partitioned input batches for the same output partition in an input stream before sending to output channel, e.g. coalesce in
Coalesce batches in output stream when poll_next is called, e.g.
This would be my personal suggestion as a place to start because:
- It (should) be the same effect as running
CoalesceBatchesinRepartitionExec. (if you coalesce on input it may affect the execution order and affect performance in unexpected ways) - It sounds relatively simple
If coalescing in e.g. PerPartitionStream::poll_next_inner, it also means larger input batches for merge-sort.
This is a good thing I think
I don't have any great suggestion about if a new stream would be better than integrating the coalescer into the existing streams -- i think we would have to try them both out and see which looked best
Thank you for working on this @jizezhang 🙏
- It (should) be the same effect as running
Hi @alamb, I wonder if I may confirm the behavior of one test
with you wrt integrating withBatchCoalescer.I noticed that this test runs two versions of the same query, one with Dataframe API and one with
SessionContext::sql. The logical plans resulting from the two are slightly different, in particular theprojectionpart:- With DataFrame API
[2025-11-22T21:29:47Z DEBUG datafusion_optimizer::utils] Final optimized plan: Aggregate: groupBy=[[]], aggr=[[count(?table?.flag)]] TableScan: ?table? projection=[flag], full_filters=[?table?.flag = Int32(0)] - With
sql,[2025-11-22T21:31:15Z DEBUG datafusion_optimizer::utils] Final optimized plan: Projection: count(Int64(1)) AS count(*) Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] TableScan: data projection=[], full_filters=[data.flag = Int32(0)]
The
SessionContext::sqlversion had an issue when using arrowBatchCoalescerkernel viaLimitedBatchCoalescer. The reason (I think) is that the custom table providerCustomProviderin the test has a branching logic on schema depending onprojection
thus the resulting physical plans differ inschema:- With DataFrame API,
AggregateExec: mode=Final, gby=[], aggr=[count(?table?.flag)], schema=[count(?table?.flag):Int64] CoalescePartitionsExec, schema=[count(?table?.flag)[count]:Int64] AggregateExec: mode=Partial, gby=[], aggr=[count(?table?.flag)], schema=[count(?table?.flag)[count]:Int64] RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1, schema=[flag:Int32] CooperativeExec, schema=[flag:Int32] CustomPlan: batch_size=1, schema=[flag:Int32] - With
sql,ProjectionExec: expr=[count(Int64(1))@0 as count(*)], schema=[count(*):Int64] AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))], schema=[count(Int64(1)):Int64] CoalescePartitionsExec, schema=[count(Int64(1))[count]:Int64] AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))], schema=[count(Int64(1))[count]:Int64] RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1, schema=[] CooperativeExec, schema=[] CustomPlan: batch_size=1, schema=[]
However, the batches returned by the associated custom execution plan
CustomPlanis always full (not projected)Arrow
BatchCoalescerchecks that the provided schema to the coalescer matches actual batch schema https://github.com/apache/arrow-rs/blob/a67d49758b1faee7d42fe3b215e226d6d560f237/arrow-select/src/coalesce.rs#L428, thus theSessionContext::sqlversion would panic.The issue went away when I modified
CustomPlan::executeto return projected batches, but I wanted to check whether the behavior is expected and whether the test should be updated, or otherwise what approaches to take. Thanks a lot!- With DataFrame API
Actually the behavior of one other test
test_preserve_order_with_spillingmay probably be also affected by this change. When reserving memory for an array, e.g. herethe buffer size is computed usingdatafusion/datafusion/physical-plan/src/sorts/stream.rs
Lines 247 to 250 in d24eb4a
capacityfor primitive arrays https://github.com/apache/arrow-rs/blob/a8a63c28d14b99d8f50b32f3184ab986bad15e50/arrow-array/src/array/primitive_array.rs#L1242.When arrow
BatchCoalescercoalesces batches, it copies rows from batches to the internalin_progress_arrays, and forInProgressPrimitiveArray, this involves anensure_capacitycall that takesbatch_sizehttps://github.com/apache/arrow-rs/blob/a67d49758b1faee7d42fe3b215e226d6d560f237/arrow-select/src/coalesce/primitive.rs#L58. If using default batch size of 8192, 64B is not enough for this test. I plan to adjust the batch size but wanted to mention here just in case this does not make sense. Thanks.The issue went away when I modified CustomPlan::execute to return projected batches, but I wanted to check whether the behavior is expected and whether the test should be updated, or otherwise what approaches to take. Thanks a lot!
I don't fully understand the case you laid out -- I think it would be easier to comment if / when you have code / tests changes
Reacted by jizezhangWhat I would suggest for this issue is a two part change:
- The PR that adds repartitioning to the RepartitionExec (but don't change the optimization rule)
- A second PR that removes RepartitionExec from the optimizer rule
That way we can better isolate the effects of different changes
Reacted by jizezhangHi @alamb, I was running benchmarks with my PR #19002 (and some variations for comparison), though it seemed that the results somewhat varied a bit across runs, e.g. even with the same main branch, one query can be shown as faster in one run and slower in the next. I just wonder if there are any tips on running benchmarks or interpreting the results, and/or any requirements on the instance itself (e.g. number of cores, memory etc.). Thanks a lot!
Hi @alamb, I was running benchmarks with my PR #19002 (and some variations for comparison), though it seemed that the results somewhat varied a bit across runs, e.g. even with the same main branch, one query can be shown as faster in one run and slower in the next. I just wonder if there are any tips on running benchmarks or interpreting the results, and/or any requirements on the instance itself (e.g. number of cores, memory etc.). Thanks a lot!
Yes I agree there is some variation. I haven't found a good way to handle this other than to make sure the results are reproducable run to run (aka if a performance difference is detected, try rerunning the same benchmarks and see if the same difference is seen)
It would be great to have a more stable benchmarking environment, but I think that would require a more isolated environment than I have available (either a GCP machine or my laptop)
Reacted by jizezhang- added a commit that references this issue
on Dec 9, 2025 Awesome --thank you @jizezhang
It seems like after #19239 all we have left now is to integrate
BatchCoalescerinto theAsyncFuncExecand we really could remove theCoalesceBatchesoptimization rule 🤔Yes, is someone looking into that? or if not yet I can take a stab.
Reacted by Andrew Lamb- changed the title
[-]Integrate `BatchCoalescer` into `RepartitionExec` and remove from `CoalesceBatches` optimization rule[/-][+]Integrate `BatchCoalescer` into `RepartitionExec` [/+]on Dec 15, 2025 It seems like after #19239 all we have left now is to integrate
BatchCoalescerinto theAsyncFuncExecand we really could remove theCoalesceBatchesoptimization rule 🤔Yes, is someone looking into that? or if not yet I can take a stab.
No one I know if is looking at it.
THank you 🙏 I filed this ticket to track the work:
- added a commit that references this issue
on Dec 16, 2025
Is your feature request related to a problem or challenge?
CoalesceBatchesExecoperator #18779The
LimitedBatchCoalescercan now be used to integrate batch coalescing directly inExecutionPlanimplementations, which is better than using an external operator for the reasons described on #18779Currently
RepartitionExecneeds aCoalesceBatchesoperator after for performance, seedatafusion/datafusion/physical-optimizer/src/coalesce_batches.rs
Line 64 in bd30fe2
We should integrate coalescing directly in the operator
Describe the solution you'd like
CoalesceBatchesoptimizer rule:datafusion/datafusion/physical-optimizer/src/coalesce_batches.rs
Line 64 in bd30fe2
Describe alternatives you've considered
@Dandandan did this for
FilterExecin these two PRs, so we can probably follow a similar modelAdditional context
No response