From 45108d9291ffc4ea7cc05780b495e6352c54d7aa Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Thu, 13 Aug 2026 15:16:03 -0700 Subject: [PATCH] fix: drain all completed batches in classic PiecewiseMergeJoin `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 #24348. Co-authored-by: Claude Code --- .../piecewise_merge_join/classic_join.rs | 28 ++++++++++++- datafusion/sqllogictest/test_files/pwmj.slt | 39 +++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) 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;