fix: drain all completed batches in classic PiecewiseMergeJoin

[viirya]

Which issue does this PR close?

• Closes PiecewiseMergeJoin drops rows when a stream batch spans multiple output chunks (small batch_size) #24348.

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.

[codecov-commenter]

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]

Not a maintainer but this makes sense and LGTM!

[viirya]

Thanks @saadtajwar

[buraksenn]

Hi @viirya, I've also encountered this issue and filed #24297 then opened the PR #24298. As far as I can see, they overlap but with minor differences. I can rebase my PR on top of yours or close it completely. Can you check which one to proceed?

[viirya]

Thanks @buraksenn — I checked both and found that this PR is a duplicate. Same root cause (transitioning before the output coalescer is drained), but yours fixes it at both phase boundaries via next_drained_batch() and also covers the Left/Full unmatched-pass infinite-loop (bug 1 in #24297) and the empty-placeholder batch — this PR only fixed the process_stream_batch right/full drop (your bug 2). I confirmed your PR resolves my repro too (RIGHT JOIN at small batch_size → correct count).

Please go ahead with #24298 — no need to rebase onto it. I'll close #24349 and issue #24348 as duplicates of #24297.

[viirya]

Closing as a duplicate of #24298 (which is more complete — see my comment above).