fix: drain all completed batches in classic PiecewiseMergeJoin

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)));
             }
 

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;