feat: integrate batch coalescer with repartition exec #19002
+87
−27
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
github-actions
bot
added
core
jizezhang
changed the title
jizezhang
changed the title
jizezhang
commented
Dec 1, 2025
| Ok(Box::pin(RecordBatchStreamAdapter::new( | ||
| self.schema(), | ||
| futures::stream::iter(self.batches.clone().into_iter().map(Ok)), | ||
| futures::stream::iter(self.batches.clone().into_iter().map(move |batch| { |
Contributor
Author
There was a problem hiding this comment.
Please see #18782 (comment) for my thoughts/reason on updating this test.
Contributor
|
run benchmark tpch tpch_mem tpch10 |
Contributor
|
🤖 Hi @Dandandan, thanks for the request (#19002 (comment)). |
Contributor
|
run benchmark tpch tpch_mem tpch10 |
Contributor
|
🤖 Hi @Dandandan, thanks for the request (#19002 (comment)). |
Contributor
|
run benchmark tpch |
Contributor
|
🤖 |
Contributor
|
run benchmark tpch10 |
Contributor
|
run benchmark clickbench_partitioned |
Contributor
I will update my scraper to support this syntax FYI |
Contributor
|
🤖: Benchmark completed Details
|
martin-g
reviewed
Dec 1, 2025
| break; | ||
| } | ||
| let inner_poll = self.poll_next_inner(cx); | ||
| let _timer = cloned_time.timer(); |
Member
There was a problem hiding this comment.
Shouldn't this be before the poll_next_inner() call ?
| } | ||
| Poll::Ready(None) => { | ||
| completed = true; | ||
| coalescer.finish()?; |
Member
There was a problem hiding this comment.
Suggested change
| coalescer.finish()?; | |
| if let Err(e) = coalescer.finish() { | |
| self.batch_coalescer = Some(coalescer); | |
| return self.baseline_metrics.record_poll(Poll::Ready(Some(Err(e)))); | |
| } |
Otherwise in case of an error the self.batch_coalescer won't be restored.
| } | ||
|
|
||
| impl PerPartitionStream { | ||
| #[allow(clippy::too_many_arguments)] |
Member
There was a problem hiding this comment.
Suggested change
| #[allow(clippy::too_many_arguments)] | |
| #[expect(clippy::too_many_arguments)] |
By using expect instead of allow the Clippy rule will fail too once it is no more needed and the developer will have to remove it. Otherwise it may become obsolete.
| coalescer.finish()?; | ||
| } | ||
| Poll::Ready(Some(Ok(batch))) => { | ||
| coalescer.push_batch(batch)?; |
Member
There was a problem hiding this comment.
Suggested change
| coalescer.push_batch(batch)?; | |
| coalescer.?; | |
| if let Err(e) = coalescer.push_batch(batch) { | |
| self.batch_coalescer = Some(coalescer); | |
| return self.baseline_metrics.record_poll(Poll::Ready(Some(Err(e)))); | |
| } |
Contributor
|
🤖 |
Contributor
|
🤖: Benchmark completed Details
|
Contributor
|
🤖 |
Contributor
|
🤖: Benchmark completed Details
|
Contributor
|
run benchmark clickbench_partitioned |
Contributor
|
🤖 |
Contributor
|
🤖: Benchmark completed Details
|
This was referenced Dec 1, 2025
Closed
Contributor
|
Thank you for this PR @jizezhang 🙏 My reading of the benchmark results so far is that this PR may be slightly faster, and doesn't cause any regressions. I'll try and patch up my ability to run tpch SF10 tests and then give this one a good review |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
BatchCoalescerintoRepartitionExecand remove fromCoalesceBatchesoptimization rule #18782.Rationale for this change
RepartitionExechas two cases: sort-preserving and non sort-preserving. This change integratesLimitedBatchCoalescerwith the latter. For the former, it seems thatSortPreservingMergeStreamthat builds on top ofPerPartitionStreamshas batching logic built in:datafusion/datafusion/physical-plan/src/sorts/merge.rs
Lines 279 to 289 in e4dcf0c
What changes are included in this PR?
Are these changes tested?
Yes
Are there any user-facing changes?
No