diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs index 1c89927087fd..163696faa0ed 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs @@ -274,9 +274,27 @@ impl ClassicPWMJStream { .output_batches .next_completed_batch() { + // If scanning this stream batch is already done, keep draining the + // remaining completed batches before moving on. `finish_buffered_batch` + // can leave several completed batches queued and only one is returned + // per poll; advancing to `FetchStreamBatch` while any remain (as the + // block below used to do) strands them for the final stream batch and + // drops output rows. + if !self.batch_process_state.continue_process + && self.batch_process_state.output_batches.is_empty() + { + self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; + } return Ok(StatefulStreamResult::Ready(Some(batch))); } + // The scan finished on a previous poll and its output is now fully + // drained; advance without re-scanning (which would re-emit rows). + if !self.batch_process_state.continue_process { + self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; + return Ok(StatefulStreamResult::Continue); + } + // Produce more work let batch = resolve_classic_join( buffered_side, @@ -289,7 +307,11 @@ impl ClassicPWMJStream { )?; if !self.batch_process_state.continue_process { - // We finished scanning this stream batch. + // We finished scanning this stream batch. Finalize the coalescer and + // return the first completed batch; any remaining completed batches are + // drained by the `next_completed_batch` loop at the head of this + // function on subsequent polls (we only advance to `FetchStreamBatch` + // once none remain). self.batch_process_state .output_batches .finish_buffered_batch()?; @@ -298,7 +320,9 @@ impl ClassicPWMJStream { .output_batches .next_completed_batch() { - self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; + if self.batch_process_state.output_batches.is_empty() { + self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; + } return Ok(StatefulStreamResult::Ready(Some(b))); } diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 25c275f4b7e9..f4360034c6b7 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -387,5 +387,44 @@ ORDER BY 1,2; 20 100 NULL NULL +# Regression test for a classic-join batch-draining bug: when a single stream +# batch produced more than one `batch_size` output chunk, `finish_buffered_batch` +# left extra completed batches queued and only one was returned before advancing, +# stranding the rest (dropping output rows) for the final stream batch. A small +# `batch_size` forces multiple output chunks per stream batch to exercise it. +statement ok +set datafusion.execution.batch_size = 4; + +statement ok +CREATE TABLE drain_t1 (id INT, v INT); + +statement ok +CREATE TABLE drain_t2 (id INT, v INT); + +statement ok +INSERT INTO drain_t1 VALUES + (0,0),(1,3),(2,1),(3,-5),(4,NULL),(5,-3),(6,NULL),(7,2),(8,4),(9,NULL), + (10,-1),(11,-2),(12,2),(13,2),(14,-3),(15,NULL),(16,-2),(17,NULL),(18,0),(19,-5), + (20,-1),(21,NULL),(22,-5),(23,0),(24,NULL),(25,NULL),(26,-3),(27,2),(28,-3),(29,NULL); + +statement ok +INSERT INTO drain_t2 VALUES + (0,-5),(1,NULL),(2,1),(3,2),(4,3),(5,-5),(6,NULL),(7,NULL),(8,3),(9,-3), + (10,NULL),(11,-2),(12,NULL),(13,0),(14,NULL),(15,2),(16,-1),(17,NULL),(18,-4),(19,-5), + (20,-3),(21,2),(22,0),(23,-1),(24,-2),(25,NULL),(26,1),(27,-3),(28,-1),(29,-5); + +# All rows must be accounted for: RIGHT JOIN emits every right row (matched or, +# for unmatched right rows including NULL keys, once with NULLs on the left). +query II +SELECT count(*) AS total, count(t1.id) AS matched +FROM drain_t1 t1 +RIGHT JOIN drain_t2 t2 + ON t1.v < t2.v; +---- +198 186 + +statement ok +set datafusion.execution.batch_size = 8192; + statement ok set datafusion.optimizer.enable_piecewise_merge_join = false;