Integrate `BatchCoalescer` into `RepartitionExec` and remove from `CoalesceBatches` optimization rule

[alamb]

Is your feature request related to a problem or challenge?

• Part of [EPIC] Remove CoalesceBatchesExec operator #18779

The LimitedBatchCoalescer can now be used to integrate batch coalescing directly in ExecutionPlan implementations, which is better than using an external operator for the reasons described on #18779

Currently RepartitionExec needs a CoalesceBatches operator after for performance, see

datafusion/datafusion/physical-optimizer/src/coalesce_batches.rs

Line 64 in bd30fe2

.downcast_ref::<RepartitionExec>()

We should integrate coalescing directly in the operator

Describe the solution you'd like

1. Add batch coalescing into RepartitionExec
2. Remove RepartitionExec from the CoalesceBatches optimizer rule:

   datafusion/datafusion/physical-optimizer/src/coalesce_batches.rs

   Line 64 in bd30fe2

   .downcast_ref::<RepartitionExec>()

Describe alternatives you've considered

@Dandandan did this for FilterExec in these two PRs, so we can probably follow a similar model

• Coalesce batches inside FilterExec #18604
• Remove FilterExec from CoalesceBatches optimization rule, add fetch support #18630

Additional context

No response

[alamb]

Not sure about this one, closing for now

[jizezhang]

take

[jizezhang]

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!

[alamb]

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

[jizezhang]

Hi @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_input before channel.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_next is called, e.g.
  • For non order-preserving case, output stream is a PerPartitionStream; for order-preserving case, output stream is a SortPreservingMergeStream whose input is vector of PerPartitionStreams.
  • 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 PerPartitionStream that coalesces while calling PerPartitionStream::poll_next. This would look very much like a CoalesceBatches but not as a separate node.
  • For order preserving case, maybe coalesce inside SortPreservingMergeStream::poll_next_inner.

Thanks a lot!