Skip to content

fix: drain all completed batches in classic PiecewiseMergeJoin - #24349

Open
viirya wants to merge 1 commit into
apache:mainfrom
viirya:pwmj-classic-batch-drain-fix
Open

fix: drain all completed batches in classic PiecewiseMergeJoin#24349
viirya wants to merge 1 commit into
apache:mainfrom
viirya:pwmj-classic-batch-drain-fix

Conversation

@viirya

@viirya viirya commented Aug 13, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Rationale for this change

A classic PiecewiseMergeJoin can drop output rows when a single stream batch produces more than one batch_size output chunk.

ClassicPWMJStream::process_stream_batch returns one completed batch from the coalescer and then advances to fetch the next stream batch. But finish_buffered_batch can leave several completed batches queued; advancing while some remain strands them for the final stream batch, so those rows are never emitted. It only shows up at small batch_size (the default 8192 usually keeps a stream batch's output within a single chunk).

set datafusion.optimizer.enable_piecewise_merge_join = true;
set datafusion.execution.batch_size = 4;
-- (see #24348 for the full data)
select count(*) from l right join r on l.v < r.v;   -- 196 with PWMJ, 198 with NestedLoopJoin

What changes are included in this PR?

  • In process_stream_batch, only advance to FetchStreamBatch once all completed batches are drained (the head-of-function drain returns one per poll). Once the scan is finished and the output is fully drained, advance without re-scanning (re-scanning would re-emit rows).

Are these changes tested?

Yes.

  • Regression test in pwmj.slt: a RIGHT JOIN with SET batch_size = 4 over a 30-row dataset, asserting count(*), count(l.id) = 198 186. The test fails without this change (196) and passes with it.
  • Verified more broadly with a differential fuzz against NestedLoopJoin (enable_piecewise_merge_join on vs off): all join types, operators </<=/>/>=, batch_size in {8,16,32,64,100,8192} and 1/4 partitions, random data with NULLs — 0 mismatches.

Are there any user-facing changes?

PiecewiseMergeJoin (behind enable_piecewise_merge_join, default off) no longer drops rows at small batch_size, matching NestedLoopJoin. No API changes.

`ClassicPWMJStream::process_stream_batch` returned a single completed
batch and then advanced to `FetchStreamBatch`. When a stream batch
produced more than one `batch_size` output chunk, `finish_buffered_batch`
left additional completed batches queued; advancing while they remained
stranded them for the final stream batch, dropping output rows. This only
manifested at small `batch_size` values (the default 8192 hides it).

Only advance to `FetchStreamBatch` once all completed batches are drained
(the head-of-function drain returns one per poll). Once the scan is done
and the output is drained, advance without re-scanning.

Closes apache#24348.

Co-authored-by: Claude Code
@github-actions github-actions Bot added sqllogictest SQL Logic Tests (.slt) physical-plan Changes to the physical-plan crate labels Aug 13, 2026
@viirya
viirya requested a review from comphead August 13, 2026 22:30
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 83.33333% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.17%. Comparing base (0a429a3) to head (45108d9).

Files with missing lines Patch % Lines
...lan/src/joins/piecewise_merge_join/classic_join.rs 83.33% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #24349      +/-   ##
==========================================
- Coverage   81.17%   81.17%   -0.01%     
==========================================
  Files        1109     1109              
  Lines      388038   388049      +11     
  Branches   388038   388049      +11     
==========================================
+ Hits       314992   314998       +6     
- Misses      54507    54508       +1     
- Partials    18539    18543       +4     

☔ 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.

@saadtajwar saadtajwar 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.

Not a maintainer but this makes sense and LGTM!

@viirya

viirya commented Aug 14, 2026

Copy link
Copy Markdown
Member Author

Thanks @saadtajwar

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 sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

PiecewiseMergeJoin drops rows when a stream batch spans multiple output chunks (small batch_size)

3 participants