Skip to content

fix: drain PiecewiseMergeJoin output before state transitions - #24298

Merged
kumarUjjawal merged 5 commits into
apache:mainfrom
buraksenn:pmwj-drain-fix
Aug 15, 2026
Merged

kumarUjjawal merged 5 commits into
apache:mainfrom
buraksenn:pmwj-drain-fix

Conversation

@buraksenn

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

ClassicPWMJStream transitions phases while its output BatchCoalescer still holds completed batches. The Left/Full unmatched pass falls back into its producer when the drained queue is an exact multiple of batch_size — emitting duplicate rows forever — and queued unmatched Right/Full batches are dropped at terminal transitions. A zero-row placeholder batch also escapes as stream output on empty results.

What changes are included in this PR?

  • Add BatchProcessState::next_drained_batch() (finish_buffered_batch + next_completed_batch); None guarantees the coalescer is empty
  • process_stream_batch and process_unmatched_buffered_batch drain one batch per poll and transition only after a confirmed-empty pop, so a finished producer can never re-run and no queued batch is lost
  • Scan completion returns StatefulStreamResult::Continue instead of the empty placeholder batch; empty joins now yield no batches

Are these changes tested?

Yes. The existing unmatched Left/Right tests now run with batch_size = 1 (the Left one bounded with take(6) so the old duplicate loop fails as a snapshot mismatch instead of hanging) — both fail on main and pass here. The empty-join tests assert no batches are emitted.

Are there any user-facing changes?

No (PWMJ is experimental and off by default).

@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Aug 12, 2026
@codecov-commenter

codecov-commenter commented Aug 12, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 74.19355% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.17%. Comparing base (e76f1af) to head (90f333b).

Files with missing lines Patch % Lines
...lan/src/joins/piecewise_merge_join/classic_join.rs 74.19% 0 Missing and 8 partials ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##             main   #24298   +/-   ##
=======================================
  Coverage   81.17%   81.17%           
=======================================
  Files        1109     1109           
  Lines      388033   388001   -32     
  Branches   388033   388001   -32     
=======================================
- Hits       314992   314968   -24     
+ Misses      54501    54489   -12     
- Partials    18540    18544    +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.

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

Thanks @buraksenn looks solid for me. I was thinking to move test to slt, but current implementation is good

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

Looks good!

@kumarUjjawal

Copy link
Copy Markdown
Contributor

Thank you @buraksenn and @comphead

@kumarUjjawal
kumarUjjawal added this pull request to the merge queue Aug 15, 2026
Merged via the queue into apache:main with commit d3c47b4 Aug 15, 2026
40 checks passed
imtherealnaska pushed a commit to imtherealnaska/datafusion that referenced this pull request Aug 16, 2026
…#24298)

## Which issue does this PR close?

- Closes apache#24297

## Rationale for this change

`ClassicPWMJStream` transitions phases while its output `BatchCoalescer`
still holds completed batches. The Left/Full unmatched pass falls back
into its producer when the drained queue is an exact multiple of
`batch_size` — emitting duplicate rows forever — and queued unmatched
Right/Full batches are dropped at terminal transitions. A zero-row
placeholder batch also escapes as stream output on empty results.

## What changes are included in this PR?

- Add `BatchProcessState::next_drained_batch()` (`finish_buffered_batch`
+ `next_completed_batch`); `None` guarantees the coalescer is empty
- `process_stream_batch` and `process_unmatched_buffered_batch` drain
one batch per poll and transition only after a confirmed-empty pop, so a
finished producer can never re-run and no queued batch is lost
- Scan completion returns `StatefulStreamResult::Continue` instead of
the empty placeholder batch; empty joins now yield no batches

## Are these changes tested?

Yes. The existing unmatched Left/Right tests now run with `batch_size =
1` (the Left one bounded with `take(6)` so the old duplicate loop fails
as a snapshot mismatch instead of hanging) — both fail on main and pass
here. The empty-join tests assert no batches are emitted.

## Are there any user-facing changes?

No (PWMJ is experimental and off by default).
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 v56.0.0

Projects

None yet

5 participants