finalize pwmj drain fix

Author: buraksenn

datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs

@@ -278,14 +278,7 @@ impl ClassicPWMJStream {
         }
 
         if !self.batch_process_state.continue_process {
-            self.batch_process_state
-                .output_batches
-                .finish_buffered_batch()?;
-            if let Some(batch) = self
-                .batch_process_state
-                .output_batches
-                .next_completed_batch()
-            {
+            if let Some(batch) = self.batch_process_state.next_drained_batch()? {
                 return Ok(StatefulStreamResult::Ready(Some(batch)));
             }
 
@@ -294,7 +287,7 @@ impl ClassicPWMJStream {
         }
 
         // Produce more work
-        let batch = resolve_classic_join(
+        let scan_batch = resolve_classic_join(
             buffered_side,
             stream_batch,
             &self.schema,
@@ -305,23 +298,18 @@ impl ClassicPWMJStream {
         )?;
 
         if !self.batch_process_state.continue_process {
-            // A flush can queue multiple batches, so transition only after draining.
-            self.batch_process_state
-                .output_batches
-                .finish_buffered_batch()?;
-            if let Some(batch) = self
-                .batch_process_state
-                .output_batches
-                .next_completed_batch()
-            {
+            // The queue can hold several completed batches, so transition only
+            // after draining all of them.
+            if let Some(batch) = self.batch_process_state.next_drained_batch()? {
                 return Ok(StatefulStreamResult::Ready(Some(batch)));
             }
 
+            // Fully drained; emit the scan's empty tail batch and move on.
             self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch;
-            return Ok(StatefulStreamResult::Ready(Some(batch)));
+            return Ok(StatefulStreamResult::Ready(Some(scan_batch)));
         }
 
-        Ok(StatefulStreamResult::Ready(Some(batch)))
+        Ok(StatefulStreamResult::Ready(Some(scan_batch)))
     }
 
     // Process remaining unmatched rows
@@ -335,22 +323,7 @@ impl ClassicPWMJStream {
         }
 
         if !self.batch_process_state.continue_process {
-            if let Some(batch) = self
-                .batch_process_state
-                .output_batches
-                .next_completed_batch()
-            {
-                return Ok(StatefulStreamResult::Ready(Some(batch)));
-            }
-
-            self.batch_process_state
-                .output_batches
-                .finish_buffered_batch()?;
-            if let Some(batch) = self
-                .batch_process_state
-                .output_batches
-                .next_completed_batch()
-            {
+            if let Some(batch) = self.batch_process_state.next_drained_batch()? {
                 return Ok(StatefulStreamResult::Ready(Some(batch)));
             }
 
@@ -386,27 +359,11 @@ impl ClassicPWMJStream {
         self.batch_process_state.output_batches.push_batch(batch)?;
 
         self.batch_process_state.continue_process = false;
-        if let Some(batch) = self
-            .batch_process_state
-            .output_batches
-            .next_completed_batch()
-        {
-            return Ok(StatefulStreamResult::Ready(Some(batch)));
-        }
-
-        self.batch_process_state
-            .output_batches
-            .finish_buffered_batch()?;
-        if let Some(batch) = self
-            .batch_process_state
-            .output_batches
-            .next_completed_batch()
-        {
+        if let Some(batch) = self.batch_process_state.next_drained_batch()? {
             return Ok(StatefulStreamResult::Ready(Some(batch)));
         }
 
         self.state = PiecewiseMergeJoinStreamState::Completed;
-        self.batch_process_state.reset();
         Ok(StatefulStreamResult::Continue)
     }
 }
@@ -451,6 +408,17 @@ impl BatchProcessState {
         self.continue_process = true;
         self.processed_null_count = false;
     }
+
+    // Pops the next completed batch, flushing the partial remainder once the
+    // queue is empty. `None` guarantees the coalescer holds no pending rows,
+    // so the stream may leave its current state without losing output.
+    fn next_drained_batch(&mut self) -> Result<Option<RecordBatch>> {
+        if let Some(batch) = self.output_batches.next_completed_batch() {
+            return Ok(Some(batch));
+        }
+        self.output_batches.finish_buffered_batch()?;
+        Ok(self.output_batches.next_completed_batch())
+    }
 }
 
 impl Stream for ClassicPWMJStream {
@@ -836,7 +804,9 @@ mod tests {
 
     #[tokio::test]
     async fn join_right_unmatched_rows_exceeding_batch_size() -> Result<()> {
-        // 100 < {1, 2, 3} is false, making every streamed row unmatched.
+        // 100 < {1, 2, 3} is false, making every streamed row unmatched; with
+        // batch_size 2 the scan finishes with more than one completed batch
+        // still queued, and all of them must be emitted before the stream ends.
         let left = build_table(("a1", &vec![0]), ("b1", &vec![100]), ("c1", &vec![0]));
         let right = build_table(
             ("a2", &vec![10, 20, 30]),
@@ -878,6 +848,9 @@ mod tests {
 
     #[tokio::test]
     async fn join_left_unmatched_rows_exact_batch_multiple() -> Result<()> {
+        // Four unmatched buffered rows with batch_size 2 flush with no
+        // remainder (an exact multiple); the unmatched pass must terminate
+        // once the queue drains instead of recomputing the same rows.
         let left = build_table(
             ("a1", &vec![10, 20, 30, 40]),
             ("b1", &vec![100, 101, 102, 103]),
@@ -896,6 +869,9 @@ mod tests {
         let mut stream = join.execute(0, task_ctx)?;
 
         let mut batches = Vec::with_capacity(2);
+        // Bounded polls (ignoring the empty batch the scan phase emits): a
+        // non-terminating stream fails the count assertion below instead of
+        // hanging the test.
         for _ in 0..3 {
             match stream.next().await.transpose()? {
                 Some(batch) if batch.num_rows() > 0 => batches.push(batch),

datafusion/sqllogictest/test_files/pwmj.slt

@@ -342,5 +342,122 @@ ORDER BY 1,2;
 1 3
 2 3
 
+# The queries below shrink the output batch size so the join finishes with
+# several completed output batches still queued; every queued batch must be
+# emitted before the operator changes state.
+
+statement ok
+CREATE TABLE drain_t1 (t1_v INT);
+
+statement ok
+CREATE TABLE drain_t2 (t2_v INT);
+
+# No `t1_v < t2_v` pair matches, so every row of both tables is unmatched.
+statement ok
+INSERT INTO drain_t1 VALUES (100), (101), (102), (103);
+
+statement ok
+INSERT INTO drain_t2 VALUES (1), (2), (3);
+
+# Shrink the batch size only for query execution so each table above stays a
+# single input batch larger than one output batch.
+statement ok
+set datafusion.execution.batch_size = 2;
+
+# Right join: every unmatched streamed row must be emitted null-extended.
+query II
+SELECT t1.t1_v, t2.t2_v
+FROM drain_t1 t1
+RIGHT JOIN drain_t2 t2
+  ON t1.t1_v < t2.t2_v
+ORDER BY 2;
+----
+NULL 1
+NULL 2
+NULL 3
+
+query TT
+EXPLAIN
+SELECT t1.t1_v, t2.t2_v
+FROM drain_t1 t1
+RIGHT JOIN drain_t2 t2
+  ON t1.t1_v < t2.t2_v
+ORDER BY 2;
+----
+logical_plan
+01)Sort: t2.t2_v ASC NULLS LAST
+02)--Right Join:  Filter: t1.t1_v < t2.t2_v
+03)----SubqueryAlias: t1
+04)------TableScan: drain_t1 projection=[t1_v]
+05)----SubqueryAlias: t2
+06)------TableScan: drain_t2 projection=[t2_v]
+physical_plan
+01)SortPreservingMergeExec: [t2_v@1 ASC NULLS LAST]
+02)--SortExec: expr=[t2_v@1 ASC NULLS LAST], preserve_partitioning=[true]
+03)----PiecewiseMergeJoin: operator=Lt, join_type=Right, on=(t1_v < t2_v)
+04)------SortExec: expr=[t1_v@0 DESC], preserve_partitioning=[false]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+06)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+07)--------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Left join: the four unmatched buffered rows are an exact multiple of the
+# batch size. The LIMIT keeps the query bounded even if the unmatched pass
+# fails to terminate, so a regression fails with extra rows instead of
+# hanging the runner.
+query II rowsort
+SELECT t1.t1_v, t2.t2_v
+FROM drain_t1 t1
+LEFT JOIN drain_t2 t2
+  ON t1.t1_v < t2.t2_v
+LIMIT 10;
+----
+100 NULL
+101 NULL
+102 NULL
+103 NULL
+
+query TT
+EXPLAIN
+SELECT t1.t1_v, t2.t2_v
+FROM drain_t1 t1
+LEFT JOIN drain_t2 t2
+  ON t1.t1_v < t2.t2_v
+LIMIT 10;
+----
+logical_plan
+01)Limit: skip=0, fetch=10
+02)--Left Join:  Filter: t1.t1_v < t2.t2_v
+03)----SubqueryAlias: t1
+04)------Limit: skip=0, fetch=10
+05)--------TableScan: drain_t1 projection=[t1_v], fetch=10
+06)----SubqueryAlias: t2
+07)------TableScan: drain_t2 projection=[t2_v]
+physical_plan
+01)CoalescePartitionsExec: fetch=10
+02)--PiecewiseMergeJoin: operator=Lt, join_type=Left, on=(t1_v < t2_v)
+03)----SortExec: TopK(fetch=10), expr=[t1_v@0 DESC], preserve_partitioning=[false]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+05)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+06)------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Full join: both unmatched sides must drain across the state transitions.
+query II rowsort
+SELECT t1.t1_v, t2.t2_v
+FROM drain_t1 t1
+FULL JOIN drain_t2 t2
+  ON t1.t1_v < t2.t2_v
+LIMIT 20;
+----
+100 NULL
+101 NULL
+102 NULL
+103 NULL
+NULL 1
+NULL 2
+NULL 3
+
+statement ok
+reset datafusion.execution.batch_size;
+
 statement ok
 set datafusion.optimizer.enable_piecewise_merge_join = false;