From c66d9fc9a8b9edd4e8547f3222ba0996ae11b41b Mon Sep 17 00:00:00 2001 From: Burak Sen Date: Thu, 6 Aug 2026 00:23:25 +0300 Subject: [PATCH 1/4] fix pmwj drain issue encountered --- .../piecewise_merge_join/classic_join.rs | 130 ++++++++++++++++-- 1 file changed, 116 insertions(+), 14 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 50ef78f18bf6..5645e205962e 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 @@ -277,6 +277,22 @@ impl ClassicPWMJStream { return Ok(StatefulStreamResult::Ready(Some(batch))); } + 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() + { + return Ok(StatefulStreamResult::Ready(Some(batch))); + } + + self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; + return Ok(StatefulStreamResult::Continue); + } + // Produce more work let batch = resolve_classic_join( buffered_side, @@ -289,25 +305,20 @@ impl ClassicPWMJStream { )?; if !self.batch_process_state.continue_process { - // We finished scanning this stream batch. + // A flush can queue multiple batches, so transition only after draining. self.batch_process_state .output_batches .finish_buffered_batch()?; - if let Some(b) = self + if let Some(batch) = self .batch_process_state .output_batches .next_completed_batch() { - self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; - return Ok(StatefulStreamResult::Ready(Some(b))); - } - - // Nothing pending; hand back whatever `resolve` returned (often empty) and move on. - if self.batch_process_state.output_batches.is_empty() { - self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; - return Ok(StatefulStreamResult::Ready(Some(batch))); } + + self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; + return Ok(StatefulStreamResult::Ready(Some(batch))); } Ok(StatefulStreamResult::Ready(Some(batch))) @@ -340,9 +351,12 @@ impl ClassicPWMJStream { .output_batches .next_completed_batch() { - self.state = PiecewiseMergeJoinStreamState::Completed; return Ok(StatefulStreamResult::Ready(Some(batch))); } + + // Avoid restarting the unmatched pass when there is no remainder. + self.state = PiecewiseMergeJoinStreamState::Completed; + return Ok(StatefulStreamResult::Continue); } let buffered_data = @@ -388,13 +402,12 @@ impl ClassicPWMJStream { .output_batches .next_completed_batch() { - self.state = PiecewiseMergeJoinStreamState::Completed; return Ok(StatefulStreamResult::Ready(Some(batch))); } self.state = PiecewiseMergeJoinStreamState::Completed; self.batch_process_state.reset(); - Ok(StatefulStreamResult::Ready(None)) + Ok(StatefulStreamResult::Continue) } } @@ -660,10 +673,12 @@ mod tests { joins::PiecewiseMergeJoinExec, test::{TestMemoryExec, build_table_i32}, }; - use arrow::array::{Date32Array, Date64Array}; + use arrow::array::{Date32Array, Date64Array, Int32Array}; + use arrow::compute::concat_batches; use arrow_schema::{DataType, Field}; use datafusion_common::test_util::batches_to_string; use datafusion_execution::TaskContext; + use datafusion_execution::config::SessionConfig; use datafusion_physical_expr::{PhysicalExpr, expressions::Column}; use insta::assert_snapshot; use std::sync::Arc; @@ -819,6 +834,93 @@ mod tests { Ok(()) } + #[tokio::test] + async fn join_right_unmatched_rows_exceeding_batch_size() -> Result<()> { + // 100 < {1, 2, 3} is false, making every streamed row unmatched. + let left = build_table(("a1", &vec![0]), ("b1", &vec![100]), ("c1", &vec![0])); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![1, 2, 3]), + ("c2", &vec![70, 80, 90]), + ); + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let session_config = SessionConfig::new().with_batch_size(2); + let task_ctx = + Arc::new(TaskContext::default().with_session_config(session_config)); + let join = join(left, right, on, Operator::Lt, JoinType::Right)?; + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + let total_rows: usize = batches.iter().map(RecordBatch::num_rows).sum(); + assert_eq!( + total_rows, 3, + "every unmatched streamed row must be emitted" + ); + + let combined = concat_batches(&join.schema(), &batches)?; + for buffered_column in 0..3 { + assert_eq!(combined.column(buffered_column).null_count(), 3); + } + let a2 = combined + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + let mut streamed_values: Vec = a2.iter().flatten().collect(); + streamed_values.sort_unstable(); + assert_eq!(streamed_values, vec![10, 20, 30]); + Ok(()) + } + + #[tokio::test] + async fn join_left_unmatched_rows_exact_batch_multiple() -> Result<()> { + let left = build_table( + ("a1", &vec![10, 20, 30, 40]), + ("b1", &vec![100, 101, 102, 103]), + ("c1", &vec![70, 80, 90, 100]), + ); + let right = build_table(("a2", &vec![0]), ("b1", &vec![1]), ("c2", &vec![0])); + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let session_config = SessionConfig::new().with_batch_size(2); + let task_ctx = + Arc::new(TaskContext::default().with_session_config(session_config)); + let join = join(left, right, on, Operator::Lt, JoinType::Left)?; + let mut stream = join.execute(0, task_ctx)?; + + let mut batches = Vec::with_capacity(2); + for _ in 0..3 { + match stream.next().await.transpose()? { + Some(batch) if batch.num_rows() > 0 => batches.push(batch), + Some(_) => {} + None => break, + } + } + assert_eq!(batches.len(), 2); + assert!(stream.next().await.is_none()); + + let combined = concat_batches(&join.schema(), &batches)?; + for streamed_column in 3..6 { + assert_eq!(combined.column(streamed_column).null_count(), 4); + } + let a1 = combined + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let mut buffered_values: Vec = a1.iter().flatten().collect(); + buffered_values.sort_unstable(); + assert_eq!(buffered_values, vec![10, 20, 30, 40]); + Ok(()) + } + #[tokio::test] async fn join_inner_less_than_unsorted() -> Result<()> { // +----+----+----+ From b54e39a66087d9581c8316ed60bf14d654fb0524 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Fri, 7 Aug 2026 00:47:40 +0300 Subject: [PATCH 2/4] finalize pwmj drain fix --- .../piecewise_merge_join/classic_join.rs | 84 +++++-------- datafusion/sqllogictest/test_files/pwmj.slt | 117 ++++++++++++++++++ 2 files changed, 147 insertions(+), 54 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 5645e205962e..075f0a6bb5dc 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 @@ -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> { + 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), diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 9789c0e4e539..408cb4dfca9f 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/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; From 4b38954e777763a8794e248f7f7209d4327615ce Mon Sep 17 00:00:00 2001 From: Burak Sen Date: Mon, 10 Aug 2026 14:08:29 +0300 Subject: [PATCH 3/4] pmj --- .../piecewise_merge_join/classic_join.rs | 137 +++++++----------- datafusion/sqllogictest/test_files/pwmj.slt | 28 ++-- 2 files changed, 64 insertions(+), 101 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 075f0a6bb5dc..c52b66396ec4 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 @@ -277,6 +277,9 @@ impl ClassicPWMJStream { return Ok(StatefulStreamResult::Ready(Some(batch))); } + // A finished scan can leave several completed batches queued; emit + // them one per poll and transition only once the queue is empty, so + // no output is lost. if !self.batch_process_state.continue_process { if let Some(batch) = self.batch_process_state.next_drained_batch()? { return Ok(StatefulStreamResult::Ready(Some(batch))); @@ -287,7 +290,7 @@ impl ClassicPWMJStream { } // Produce more work - let scan_batch = resolve_classic_join( + let batch = resolve_classic_join( buffered_side, stream_batch, &self.schema, @@ -298,18 +301,11 @@ impl ClassicPWMJStream { )?; if !self.batch_process_state.continue_process { - // 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(scan_batch))); + // Scan finished; re-enter through the drain guard above. + return Ok(StatefulStreamResult::Continue); } - Ok(StatefulStreamResult::Ready(Some(scan_batch))) + Ok(StatefulStreamResult::Ready(Some(batch))) } // Process remaining unmatched rows @@ -327,7 +323,7 @@ impl ClassicPWMJStream { return Ok(StatefulStreamResult::Ready(Some(batch))); } - // Avoid restarting the unmatched pass when there is no remainder. + // Fully drained; finish instead of re-running the pass. self.state = PiecewiseMergeJoinStreamState::Completed; return Ok(StatefulStreamResult::Continue); } @@ -359,11 +355,7 @@ 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.next_drained_batch()? { - return Ok(StatefulStreamResult::Ready(Some(batch))); - } - - self.state = PiecewiseMergeJoinStreamState::Completed; + // Re-enter through the drain guard above. Ok(StatefulStreamResult::Continue) } } @@ -409,9 +401,8 @@ impl BatchProcessState { 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. + // `None` guarantees the coalescer holds no pending rows, so the caller + // may safely transition state without losing output. fn next_drained_batch(&mut self) -> Result> { if let Some(batch) = self.output_batches.next_completed_batch() { return Ok(Some(batch)); @@ -641,10 +632,9 @@ mod tests { joins::PiecewiseMergeJoinExec, test::{TestMemoryExec, build_table_i32}, }; - use arrow::array::{Date32Array, Date64Array, Int32Array}; - use arrow::compute::concat_batches; + use arrow::array::{Date32Array, Date64Array}; use arrow_schema::{DataType, Field}; - use datafusion_common::test_util::batches_to_string; + use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; use datafusion_physical_expr::{PhysicalExpr, expressions::Column}; @@ -804,9 +794,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; 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. + // 100 < {1, 2, 3} never matches, so with batch_size 2 the scan ends + // with several completed batches still queued; all of them must be + // emitted (regression: only one was, silently dropping rows). let left = build_table(("a1", &vec![0]), ("b1", &vec![100]), ("c1", &vec![0])); let right = build_table( ("a2", &vec![10, 20, 30]), @@ -825,32 +815,24 @@ mod tests { let stream = join.execute(0, task_ctx)?; let batches = common::collect(stream).await?; - let total_rows: usize = batches.iter().map(RecordBatch::num_rows).sum(); - assert_eq!( - total_rows, 3, - "every unmatched streamed row must be emitted" - ); - - let combined = concat_batches(&join.schema(), &batches)?; - for buffered_column in 0..3 { - assert_eq!(combined.column(buffered_column).null_count(), 3); - } - let a2 = combined - .column(3) - .as_any() - .downcast_ref::() - .unwrap(); - let mut streamed_values: Vec = a2.iter().flatten().collect(); - streamed_values.sort_unstable(); - assert_eq!(streamed_values, vec![10, 20, 30]); + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | | | | 10 | 1 | 70 | + | | | | 20 | 2 | 80 | + | | | | 30 | 3 | 90 | + +----+----+----+----+----+----+ + "); Ok(()) } #[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. + // Four unmatched buffered rows flush as an exact multiple of + // batch_size 2, leaving no partial remainder; the unmatched pass must + // then terminate (regression: it restarted and emitted duplicates + // forever). The timeout turns that hang into a test failure. let left = build_table( ("a1", &vec![10, 20, 30, 40]), ("b1", &vec![100, 101, 102, 103]), @@ -866,34 +848,25 @@ mod tests { let task_ctx = Arc::new(TaskContext::default().with_session_config(session_config)); let join = join(left, right, on, Operator::Lt, JoinType::Left)?; - 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), - Some(_) => {} - None => break, - } - } - assert_eq!(batches.len(), 2); - assert!(stream.next().await.is_none()); + let stream = join.execute(0, task_ctx)?; - let combined = concat_batches(&join.schema(), &batches)?; - for streamed_column in 3..6 { - assert_eq!(combined.column(streamed_column).null_count(), 4); - } - let a1 = combined - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let mut buffered_values: Vec = a1.iter().flatten().collect(); - buffered_values.sort_unstable(); - assert_eq!(buffered_values, vec![10, 20, 30, 40]); + let batches = tokio::time::timeout( + std::time::Duration::from_secs(5), + common::collect(stream), + ) + .await + .expect("left join unmatched pass should terminate")?; + + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+-----+-----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+-----+-----+----+----+----+ + | 10 | 100 | 70 | | | | + | 20 | 101 | 80 | | | | + | 30 | 102 | 90 | | | | + | 40 | 103 | 100 | | | | + +----+-----+-----+----+----+----+ + "); Ok(()) } @@ -1032,12 +1005,8 @@ mod tests { ); let (_, batches) = join_collect(left, right, on, Operator::LtEq, JoinType::Inner).await?; - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+----+----+----+ - | a1 | b1 | c1 | a2 | b1 | c2 | - +----+----+----+----+----+----+ - +----+----+----+----+----+----+ - "); + // An empty join result produces no batches at all, not an empty batch. + assert!(batches.is_empty()); Ok(()) } @@ -1476,12 +1445,8 @@ mod tests { let (_, batches) = join_collect(left, right, on, Operator::Gt, JoinType::Inner).await?; - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+----+----+----+ - | a1 | b1 | c1 | a2 | b1 | c2 | - +----+----+----+----+----+----+ - +----+----+----+----+----+----+ - "); + // An empty join result produces no batches at all, not an empty batch. + assert!(batches.is_empty()); Ok(()) } diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 408cb4dfca9f..06f4f9c81b91 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -416,29 +416,27 @@ LIMIT 10; 102 NULL 103 NULL +# The EXPLAIN omits the LIMIT so it only pins plan selection, not the +# limit-pushdown shape. 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; + ON t1.t1_v < t2.t2_v; ---- 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] +01)Left Join: Filter: t1.t1_v < t2.t2_v +02)--SubqueryAlias: t1 +03)----TableScan: drain_t1 projection=[t1_v] +04)--SubqueryAlias: t2 +05)----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] +01)PiecewiseMergeJoin: operator=Lt, join_type=Left, on=(t1_v < t2_v) +02)--SortExec: expr=[t1_v@0 DESC], preserve_partitioning=[false] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)----DataSourceExec: partitions=1, partition_sizes=[1] # Full join: both unmatched sides must drain across the state transitions. query II rowsort From 107b9fd695acf2d0b99ebf8bf1cf67835101fb01 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Wed, 12 Aug 2026 21:17:07 +0300 Subject: [PATCH 4/4] get rid of unnecessary test --- .../piecewise_merge_join/classic_join.rs | 104 +++------------- datafusion/sqllogictest/test_files/pwmj.slt | 115 ------------------ 2 files changed, 18 insertions(+), 201 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 c52b66396ec4..26432972bc88 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 @@ -404,9 +404,6 @@ impl BatchProcessState { // `None` guarantees the coalescer holds no pending rows, so the caller // may safely transition state without losing output. fn next_drained_batch(&mut self) -> Result> { - 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()) } @@ -634,10 +631,11 @@ mod tests { }; use arrow::array::{Date32Array, Date64Array}; use arrow_schema::{DataType, Field}; - use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; + use datafusion_common::test_util::batches_to_string; use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; use datafusion_physical_expr::{PhysicalExpr, expressions::Column}; + use futures::TryStreamExt; use insta::assert_snapshot; use std::sync::Arc; @@ -792,84 +790,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn join_right_unmatched_rows_exceeding_batch_size() -> Result<()> { - // 100 < {1, 2, 3} never matches, so with batch_size 2 the scan ends - // with several completed batches still queued; all of them must be - // emitted (regression: only one was, silently dropping rows). - let left = build_table(("a1", &vec![0]), ("b1", &vec![100]), ("c1", &vec![0])); - let right = build_table( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![1, 2, 3]), - ("c2", &vec![70, 80, 90]), - ); - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let session_config = SessionConfig::new().with_batch_size(2); - let task_ctx = - Arc::new(TaskContext::default().with_session_config(session_config)); - let join = join(left, right, on, Operator::Lt, JoinType::Right)?; - let stream = join.execute(0, task_ctx)?; - let batches = common::collect(stream).await?; - - assert_snapshot!(batches_to_sort_string(&batches), @r" - +----+----+----+----+----+----+ - | a1 | b1 | c1 | a2 | b1 | c2 | - +----+----+----+----+----+----+ - | | | | 10 | 1 | 70 | - | | | | 20 | 2 | 80 | - | | | | 30 | 3 | 90 | - +----+----+----+----+----+----+ - "); - Ok(()) - } - - #[tokio::test] - async fn join_left_unmatched_rows_exact_batch_multiple() -> Result<()> { - // Four unmatched buffered rows flush as an exact multiple of - // batch_size 2, leaving no partial remainder; the unmatched pass must - // then terminate (regression: it restarted and emitted duplicates - // forever). The timeout turns that hang into a test failure. - let left = build_table( - ("a1", &vec![10, 20, 30, 40]), - ("b1", &vec![100, 101, 102, 103]), - ("c1", &vec![70, 80, 90, 100]), - ); - let right = build_table(("a2", &vec![0]), ("b1", &vec![1]), ("c2", &vec![0])); - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let session_config = SessionConfig::new().with_batch_size(2); - let task_ctx = - Arc::new(TaskContext::default().with_session_config(session_config)); - let join = join(left, right, on, Operator::Lt, JoinType::Left)?; - let stream = join.execute(0, task_ctx)?; - - let batches = tokio::time::timeout( - std::time::Duration::from_secs(5), - common::collect(stream), - ) - .await - .expect("left join unmatched pass should terminate")?; - - assert_snapshot!(batches_to_sort_string(&batches), @r" - +----+-----+-----+----+----+----+ - | a1 | b1 | c1 | a2 | b1 | c2 | - +----+-----+-----+----+----+----+ - | 10 | 100 | 70 | | | | - | 20 | 101 | 80 | | | | - | 30 | 102 | 90 | | | | - | 40 | 103 | 100 | | | | - +----+-----+-----+----+----+----+ - "); - Ok(()) - } - #[tokio::test] async fn join_inner_less_than_unsorted() -> Result<()> { // +----+----+----+ @@ -1335,8 +1255,16 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = - join_collect(left, right, on, Operator::LtEq, JoinType::Left).await?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(1)), + ); + // Bound collection so the old duplicate loop becomes a snapshot mismatch. + let batches = join(left, right, on, Operator::LtEq, JoinType::Left)? + .execute(0, task_ctx)? + .take(6) + .try_collect::>() + .await?; assert_snapshot!(batches_to_string(&batches), @r" +----+----+----+----+----+----+ @@ -1381,8 +1309,12 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = - join_collect(left, right, on, Operator::GtEq, JoinType::Right).await?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(1)), + ); + let join = join(left, right, on, Operator::GtEq, JoinType::Right)?; + let batches = common::collect(join.execute(0, task_ctx)?).await?; assert_snapshot!(batches_to_string(&batches), @r" +----+----+----+----+----+----+ diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 06f4f9c81b91..9789c0e4e539 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -342,120 +342,5 @@ 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 - -# The EXPLAIN omits the LIMIT so it only pins plan selection, not the -# limit-pushdown shape. -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; ----- -logical_plan -01)Left Join: Filter: t1.t1_v < t2.t2_v -02)--SubqueryAlias: t1 -03)----TableScan: drain_t1 projection=[t1_v] -04)--SubqueryAlias: t2 -05)----TableScan: drain_t2 projection=[t2_v] -physical_plan -01)PiecewiseMergeJoin: operator=Lt, join_type=Left, on=(t1_v < t2_v) -02)--SortExec: expr=[t1_v@0 DESC], preserve_partitioning=[false] -03)----DataSourceExec: partitions=1, partition_sizes=[1] -04)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -05)----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;