You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/15 14:19:05 UTC
(datafusion) branch gh-readonly-queue/main/pr-24298-4403c5d84d25adee68cd2f8a581d2003dad52bb4 created (now d3c47b4a80)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24298-4403c5d84d25adee68cd2f8a581d2003dad52bb4
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at d3c47b4a80 fix: drain PiecewiseMergeJoin output before state transitions (#24298)
This branch includes the following new commits:
new d3c47b4a80 fix: drain PiecewiseMergeJoin output before state transitions (#24298)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix: drain PiecewiseMergeJoin output before state transitions (#24298)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24298-4403c5d84d25adee68cd2f8a581d2003dad52bb4
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit d3c47b4a80baa3ca941ca0fc2cc80f96687c5a27
Author: Burak Εen <bu...@gmail.com>
AuthorDate: Sat Aug 15 14:18:41 2026 +0000
fix: drain PiecewiseMergeJoin output before state transitions (#24298)
## Which issue does this PR close?
- Closes #24297
## Rationale for this change
`ClassicPWMJStream` transitions phases while its output `BatchCoalescer`
still holds completed batches. The Left/Full unmatched pass falls back
into its producer when the drained queue is an exact multiple of
`batch_size` β emitting duplicate rows forever β and queued unmatched
Right/Full batches are dropped at terminal transitions. A zero-row
placeholder batch also escapes as stream output on empty results.
## What changes are included in this PR?
- Add `BatchProcessState::next_drained_batch()` (`finish_buffered_batch`
+ `next_completed_batch`); `None` guarantees the coalescer is empty
- `process_stream_batch` and `process_unmatched_buffered_batch` drain
one batch per poll and transition only after a confirmed-empty pop, so a
finished producer can never re-run and no queued batch is lost
- Scan completion returns `StatefulStreamResult::Continue` instead of
the empty placeholder batch; empty joins now yield no batches
## Are these changes tested?
Yes. The existing unmatched Left/Right tests now run with `batch_size =
1` (the Left one bounded with `take(6)` so the old duplicate loop fails
as a snapshot mismatch instead of hanging) β both fail on main and pass
here. The empty-join tests assert no batches are emitted.
## Are there any user-facing changes?
No (PWMJ is experimental and off by default).
---
.../src/joins/piecewise_merge_join/classic_join.rs | 123 ++++++++-------------
1 file changed, 49 insertions(+), 74 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 1c89927087..bc69ae5140 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,18 @@ 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)));
+ }
+
+ self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch;
+ return Ok(StatefulStreamResult::Continue);
+ }
+
// Produce more work
let batch = resolve_classic_join(
buffered_side,
@@ -289,25 +301,8 @@ impl ClassicPWMJStream {
)?;
if !self.batch_process_state.continue_process {
- // We finished scanning this stream batch.
- self.batch_process_state
- .output_batches
- .finish_buffered_batch()?;
- if let Some(b) = 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)));
- }
+ // Scan finished; re-enter through the drain guard above.
+ return Ok(StatefulStreamResult::Continue);
}
Ok(StatefulStreamResult::Ready(Some(batch)))
@@ -324,25 +319,13 @@ impl ClassicPWMJStream {
}
if !self.batch_process_state.continue_process {
- 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.batch_process_state
- .output_batches
- .finish_buffered_batch()?;
- if let Some(batch) = self
- .batch_process_state
- .output_batches
- .next_completed_batch()
- {
- self.state = PiecewiseMergeJoinStreamState::Completed;
- return Ok(StatefulStreamResult::Ready(Some(batch)));
- }
+ // Fully drained; finish instead of re-running the pass.
+ self.state = PiecewiseMergeJoinStreamState::Completed;
+ return Ok(StatefulStreamResult::Continue);
}
let buffered_data =
@@ -372,29 +355,8 @@ 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()
- {
- self.state = PiecewiseMergeJoinStreamState::Completed;
- return Ok(StatefulStreamResult::Ready(Some(batch)));
- }
-
- self.state = PiecewiseMergeJoinStreamState::Completed;
- self.batch_process_state.reset();
- Ok(StatefulStreamResult::Ready(None))
+ // Re-enter through the drain guard above.
+ Ok(StatefulStreamResult::Continue)
}
}
@@ -438,6 +400,13 @@ impl BatchProcessState {
self.continue_process = true;
self.processed_null_count = false;
}
+
+ // `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<Option<RecordBatch>> {
+ self.output_batches.finish_buffered_batch()?;
+ Ok(self.output_batches.next_completed_batch())
+ }
}
impl Stream for ClassicPWMJStream {
@@ -674,7 +643,9 @@ mod tests {
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 futures::TryStreamExt;
use insta::assert_snapshot;
use std::sync::Arc;
@@ -964,12 +935,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(())
}
@@ -1298,8 +1265,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::<Vec<_>>()
+ .await?;
assert_snapshot!(batches_to_string(&batches), @r"
+----+----+----+----+----+----+
@@ -1344,8 +1319,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"
+----+----+----+----+----+----+
@@ -1408,12 +1387,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(())
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org