Skip to content

Feat: allow fine tuning output record batch size for repartionionExec - #25176

Merged
kumarUjjawal merged 2 commits into
apache:mainfrom
Rich-T-kid:rich-T-kid/override-buffer-batch-size
Sep 11, 2026
Merged

Feat: allow fine tuning output record batch size for repartionionExec#25176
kumarUjjawal merged 2 commits into
apache:mainfrom
Rich-T-kid:rich-T-kid/override-buffer-batch-size

Conversation

@Rich-T-kid

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

RepartitionExec coalesces output batches to SessionConfig::batch_size (typically 8192). Some callers need a different target size for a single repartition node without altering the global session config. This PR adds a builder method that makes that possible.

What changes are included in this PR?

  • Added optional batch_size: Option field to RepartitionExec
  • Added public builder method with_batch_size(usize) -> Self with a doc comment
  • The override is propagated through consume_input_streams and used when constructing the SharedCoalescer; falls back to session config when None
  • Field is preserved in repartition_target_partitions

What is the testing strategy for this PR?

added 1 test to assert that SharedCoalescer produces the correctly fine-tuned target batch size

Are there any user-facing changes?

yes, user can now change the output batch size for repartionExec streams

@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Sep 10, 2026
@Rich-T-kid
Rich-T-kid force-pushed the rich-T-kid/override-buffer-batch-size branch from a5d07f2 to a45510b Compare September 10, 2026 19:40
@Rich-T-kid
Rich-T-kid force-pushed the rich-T-kid/override-buffer-batch-size branch from a45510b to bf35583 Compare September 10, 2026 19:46
@codecov-commenter

codecov-commenter commented Sep 10, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 84.37500% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.91%. Comparing base (1ec9ede) to head (5e4db24).
⚠️ Report is 3 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/physical-plan/src/repartition/mod.rs 84.37% 1 Missing and 4 partials ⚠️
Additional details and impacted files
@@           Coverage Diff            @@
##             main   #25176    +/-   ##
========================================
  Coverage   81.91%   81.91%            
========================================
  Files        1132     1132            
  Lines      420570   421148   +578     
  Branches   420570   421148   +578     
========================================
+ Hits       344490   344983   +493     
- Misses      55764    55771     +7     
- Partials    20316    20394    +78     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@2010YOUY01 2010YOUY01 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thank you.

/// By default the coalescer targets [`SessionConfig::batch_size`]. Use
/// this method when you need a different batch size for a specific
/// `RepartitionExec` node without changing the global session config.
pub fn with_batch_size(mut self, batch_size: usize) -> Self {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It should error for invalid input (batch_size is 0)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

true, updated the PR

@kumarUjjawal

Copy link
Copy Markdown
Contributor

Thank you @Rich-T-kid and @2010YOUY01 for the review

@kumarUjjawal
kumarUjjawal added this pull request to the merge queue Sep 11, 2026
Merged via the queue into apache:main with commit 517ac6d Sep 11, 2026
41 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

RepartitionExec batch coalescer size is only controllable via session config

4 participants