viirya opened a new pull request, #24349:
URL: https://github.com/apache/datafusion/pull/24349
## Which issue does this PR close?
- Closes #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).
```sql
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
