You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "mbutrovich (via GitHub)" <gi...@apache.org> on 2026/03/08 15:51:12 UTC
[PR] feat: specialized SemiAntiSortMergeJoin operator [datafusion]
mbutrovich opened a new pull request, #20806:
URL: https://github.com/apache/datafusion/pull/20806
## Which issue does this PR close?
- N/A.
## Rationale for this change
DataFusion's `SortMergeJoinExec` handles semi/anti joins by materializing `(outer, inner)` row pairs, applying a filter, then deduplicating with a corrected filter mask. Semi/anti joins only need a boolean per outer row — not pairs. The pair-based approach allocates unnecessary intermediate batches and index arrays to materialize output.
Recent PRs have improved SMJ performance within the existing pair-based framework — [#18875](https://github.com/apache/datafusion/pull/18875) (BatchCoalescer to reduce concatenation overhead), [#20463](https://github.com/apache/datafusion/pull/20463) (zero-copy slice instead of take for contiguous indices), [#20478](https://github.com/apache/datafusion/pull/20478) (cached row counts to avoid O(n) recalculation) — but the fundamental mismatch remains: semi/anti joins don't need pairs at all.
## What changes are included in this PR?
A new `SemiAntiSortMergeJoinExec` operator for `LeftSemi`, `LeftAnti`, `RightSemi`, and `RightAnti` sort-merge joins. Instead of materializing row pairs, it maintains a per-outer-batch bitset (`BooleanBufferBuilder`) recording which outer rows have a matching inner row, then emits output via `filter_record_batch`.
**Algorithm:** Merge-scan across two sorted inputs. On key match without filter, set matched bits for the outer key group. With filter, buffer the inner key group and evaluate the filter as outer_slice × inner_scalar, OR-ing results into the bitset with `apply_bitwise_binary_op` (64 bits per iteration). Short-circuit when all outer rows in the group are matched.
**Files:**
- `semi_anti_sort_merge_join/{exec,stream,mod,tests}.rs` — new operator (1035 + 1302 lines)
- `physical_planner.rs` — route semi/anti joins to the new operator when `enable_semi_anti_sort_merge_join` is true (default)
- `config.rs` — session config flag `datafusion.optimizer.enable_semi_anti_sort_merge_join`
- `enforce_distribution.rs` — partition key reordering support for the new operator
- `sort_pushdown.rs` — TODO for future sort pushdown support
- `join_fuzz.rs` — fuzz tests comparing HJ, NLJ, SMJ, and SaSmj
- `smj.rs` — benchmark fix: Q13 filter changed to cross-side predicate
**Benchmark results** (best of 3, `dfbench smj`):
| Query | Type | Old (ms) | New (ms) | Speedup |
|-------|------|----------|----------|---------|
| Q10 | LEFT SEMI, no filter | 4.79 | 4.27 | 1.12x |
| Q11 | LEFT SEMI, 1% filter | 3.00 | 2.30 | 1.31x |
| Q12 | LEFT SEMI, 50% filter | 38.1 | 5.52 | 6.9x |
| Q13 | LEFT SEMI, 90% filter | 66.9 | 10.2 | 6.6x |
| Q14 | LEFT ANTI, no filter | 4.96 | 4.13 | 1.20x |
| Q15 | LEFT ANTI, partial match | 4.80 | 4.22 | 1.14x |
| Q16 | LEFT ANTI, stress | 1.63 | 1.64 | 1.00x |
| Q18 | LEFT SEMI, 2% filter | 7.61 | 5.34 | 1.42x |
| Q19 | LEFT ANTI, partial match | 24.1 | 21.8 | 1.10x |
Non-semi/anti queries are unaffected (same operator).
## Are these changes tested?
- 21 unit tests covering all join types, filters, nulls, batch boundaries, and async re-entry. Compares to current SMJ operator.
- Fuzz tests (`join_fuzz.rs`) compare output against `HashJoinExec`, `NestedLoopJoinExec`, and `SortMergeJoinExec`. Ran thousands of iterations locally (randomly seeded not just repeated the same tests)
- Existing `sort_merge_join.slt` sqllogic tests pass (the operator is transparent to the SQL layer)
## Are there any user-facing changes?
New session config `datafusion.optimizer.enable_semi_anti_sort_merge_join` (default `true`). Set to `false` to use the previous `SortMergeJoinExec` for semi/anti joins. Can be toggled via environment variable `DATAFUSION_OPTIMIZER_ENABLE_SEMI_ANTI_SORT_MERGE_JOIN=false`.
Known limitations to address in follow-up PRs:
- Sort pushdown through the new operator is not yet implemented (TODO in `sort_pushdown.rs`), so some plans may include an unnecessary `SortExec` above the join that the optimizer could otherwise eliminate.
- The inner key group buffer for filtered joins is unbounded (TODO in `stream.rs`). A single highly-skewed join key with many inner rows will buffer all of them in memory. This matches the existing `SortMergeJoinExec` behavior but could benefit from spill support.
Claude wrote the initial draft of the PR description and aspects of the implementation. I understand the implementation though.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027587613
run benchmark tpch
PREFER_HASH_JOIN=false
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2989859081
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see PendingBoundary doc comment.
+ pending_boundary: Option<PendingBoundary>,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ pending_boundary: None,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let saved_inner_keys = slice_keys(&self.inner_key_arrays, num_inner - 1);
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ &self.inner_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let saved_inner_keys = slice_keys(&self.inner_key_arrays, num_inner - 1);
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ &self.inner_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Continue processing an outer key group that spans multiple outer
+ /// batches. Returns `true` if this outer batch was fully consumed
+ /// by the key group and the caller should load another.
+ fn resume_boundary(&mut self) -> Result<bool> {
+ debug_assert!(
+ self.outer_batch.is_some(),
+ "caller must load outer_batch first"
+ );
+ match self.pending_boundary.take() {
+ Some(PendingBoundary::NoFilter { saved_keys }) => {
+ let same_key = keys_match(
+ &saved_keys,
+ &self.outer_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_no_filter()?;
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ self.pending_boundary = Some(PendingBoundary::NoFilter {
+ saved_keys: slice_keys(&self.outer_key_arrays, num_outer - 1),
+ });
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ return Ok(true);
+ }
+ }
+ }
+ Some(PendingBoundary::Filtered { saved_keys }) => {
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "Filtered pending boundary entered but no inner key data exists"
+ );
+ let same_key = keys_match(
+ &saved_keys,
+ &self.outer_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_with_filter()?;
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ self.pending_boundary = Some(PendingBoundary::Filtered {
+ saved_keys: slice_keys(&self.outer_key_arrays, num_outer - 1),
+ });
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ return Ok(true);
+ }
+ }
+ self.clear_inner_key_group();
+ }
+ None => {}
+ }
+ Ok(false)
+ }
+
+ /// Main loop: drive the merge-scan to produce output batches.
+ fn poll_join(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<RecordBatch>>> {
+ let join_time = self.join_time.clone();
+ let _timer = join_time.timer();
+
+ loop {
+ // 1. Ensure we have an outer batch
+ if self.outer_batch.is_none() {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Outer exhausted — flush coalescer
+ self.pending_boundary = None;
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {
+ if self.resume_boundary()? {
+ continue;
+ }
+ }
+ }
+ }
+
+ // 2. Ensure we have an inner batch (unless inner is exhausted).
+ // Skip this when resuming a pending boundary — inner was already
+ // advanced past the key group before the boundary loop started.
+ if self.inner_batch.is_none() && self.pending_boundary.is_none() {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Inner exhausted — emit remaining outer batches.
+ // For semi: no more matches possible.
+ // For anti: all remaining outer rows are unmatched.
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ loop {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => break,
+ Ok(true) => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ }
+ }
+ }
+
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {}
+ }
+ }
+
+ // 3. Main merge-scan loop
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ if self.outer_offset >= num_outer {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ continue;
+ }
+
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ };
+ let num_inner = inner_batch.num_rows();
+
+ if self.inner_offset >= num_inner {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => continue,
+ }
+ }
+
+ // 4. Compare keys at current positions
+ let cmp = compare_join_arrays(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ &self.inner_key_arrays,
+ self.inner_offset,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ match cmp {
+ Ordering::Less => {
+ let group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ self.outer_offset = group_end;
+ }
+ Ordering::Greater => {
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if group_end >= num_inner {
+ let saved_keys =
+ slice_keys(&self.inner_key_arrays, num_inner - 1);
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_keys,
+ &self.inner_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_) => continue,
+ }
+ }
+ continue;
+ }
+ }
+ } else {
+ self.inner_offset = group_end;
+ }
+ }
+ Ordering::Equal => {
+ if self.filter.is_some() {
+ // Buffer inner key group (may span batches)
+ match ready!(self.buffer_inner_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ // Process outer rows against buffered inner group
+ // (may need to handle outer batch boundary)
+ loop {
+ self.process_key_match_with_filter()?;
+
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ if self.outer_offset >= outer_batch.num_rows() {
+ let saved_keys = slice_keys(
+ &self.outer_key_arrays,
+ outer_batch.num_rows() - 1,
+ );
+
+ self.emit_outer_batch()?;
+ debug_assert!(
+ !self.inner_key_buffer.is_empty()
+ || self.inner_key_spill.is_some(),
+ "Filtered pending boundary requires inner key data in buffer or spill"
+ );
+ self.pending_boundary =
+ Some(PendingBoundary::Filtered { saved_keys });
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.pending_boundary = None;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ let Some(PendingBoundary::Filtered {
+ saved_keys,
+ }) = self.pending_boundary.take()
+ else {
+ unreachable!()
+ };
+ let same = keys_match(
+ &saved_keys,
+ &self.outer_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same {
+ continue;
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+
+ self.clear_inner_key_group();
+ } else {
+ // No filter: advance inner past key group, then
+ // mark all outer rows with this key as matched.
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ loop {
+ self.process_key_match_no_filter()?;
+
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let saved_keys =
+ slice_keys(&self.outer_key_arrays, num_outer - 1);
+
+ self.emit_outer_batch()?;
+ self.pending_boundary =
+ Some(PendingBoundary::NoFilter { saved_keys });
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.pending_boundary = None;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ let Some(PendingBoundary::NoFilter {
+ saved_keys,
+ }) = self.pending_boundary.take()
+ else {
+ unreachable!()
+ };
+ let same_key = keys_match(
+ &saved_keys,
+ &self.outer_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ continue;
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+ }
+ }
+ }
+
+ // Check for completed coalescer batch
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ }
+ }
+}
+
+/// Evaluate the filter for all rows in an inner slice against the outer group,
+/// OR-ing results into the matched bitset. Returns the updated matched count.
+/// Extracted as a free function so Rust can split borrows on the stream struct.
+#[expect(clippy::too_many_arguments)]
+fn eval_filter_for_inner_slice(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_slice: &RecordBatch,
+ matched: &mut BooleanBufferBuilder,
+ outer_offset: usize,
+ outer_group_len: usize,
+ mut matched_count: usize,
Review Comment:
can this be derived from `matched` ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2906712552
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1160 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Binary search (not `partition_point`) because `compare_join_arrays`
+/// returns `Result` and `partition_point` requires an infallible predicate
+/// on a `&[T]` — using it here would need either allocation or unsafe
+/// pointer arithmetic to recover the probe index.
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let mut lo = from + 1;
+ let mut hi = len;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
Review Comment:
although binary search is best in theory, we might want to have an algorithm that favor search in joins,
what i was thinking:
1. first check if the next the key in `from + 1` is different
2. if not different: check if the last item is different
3. if not different: do binary search between `from + 2` to `end - 1`
What this saves us, for batch of 8192 rows?
1. when joining on unique key where every item is different, `from + 1` is 1 cmp instead of `log2(8192) = 13`
2. to avoid checking the entire batch, check if different item even exists, so also 1 cmp instead of 13
3. the rest is 13
so in the case where we don't have fall in the 2 optimizations we will waste 2 cmp which will be added to the worst case of 13
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2906790523
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1160 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Binary search (not `partition_point`) because `compare_join_arrays`
+/// returns `Result` and `partition_point` requires an infallible predicate
+/// on a `&[T]` — using it here would need either allocation or unsafe
+/// pointer arithmetic to recover the probe index.
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let mut lo = from + 1;
+ let mut hi = len;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. The `inner_key_buffer` field already holds the buffered
+ /// inner rows needed to resume filter evaluation.
+ FilteredPending,
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Main loop: drive the merge-scan to produce output batches.
+ // Clippy can't see that poll_next_outer_batch guarantees
+ // outer_batch = Some(..) when returning Ok(true).
+ #[expect(clippy::panicking_unwrap)]
+ fn poll_join(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<RecordBatch>>> {
+ let join_time = self.join_time.clone();
+ let _timer = join_time.timer();
+
+ loop {
+ // 1. Ensure we have an outer batch
+ if self.outer_batch.is_none() {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Outer exhausted — flush coalescer
+ self.boundary_state = BoundaryState::Normal;
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {
+ // Check if we're resuming a boundary loop that was
+ // interrupted by Pending. See BoundaryState docs.
+ match std::mem::replace(
+ &mut self.boundary_state,
+ BoundaryState::Normal,
+ ) {
+ BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ } => {
+ let same_key = keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.outer_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_no_filter()?;
+ let num_outer =
+ self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let new_saved = self.outer_key_arrays.clone();
+ let new_idx = num_outer - 1;
+ self.boundary_state =
+ BoundaryState::NoFilterPending {
+ saved_keys: new_saved,
+ saved_idx: new_idx,
+ };
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ }
+ }
+ BoundaryState::FilteredPending => {
+ if !self.inner_key_buffer.is_empty() {
+ let first_inner = &self.inner_key_buffer[0];
+ let inner_keys =
+ evaluate_join_keys(first_inner, &self.on_inner)?;
+ let same_key = keys_match(
+ &self.outer_key_arrays,
+ 0,
+ &inner_keys,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_with_filter()?;
+ let num_outer =
+ self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ self.boundary_state =
+ BoundaryState::FilteredPending;
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ }
+ }
+ self.clear_inner_key_group();
+ }
+ BoundaryState::Normal => {}
+ }
+ }
+ }
+ }
+
+ // 2. Ensure we have an inner batch (unless inner is exhausted).
+ // Skip this when in a boundary state — inner was already
+ // advanced past the key group before the boundary loop started.
+ if self.inner_batch.is_none()
+ && matches!(self.boundary_state, BoundaryState::Normal)
+ {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Inner exhausted — emit remaining outer batches.
+ // For semi: no more matches possible.
+ // For anti: all remaining outer rows are unmatched.
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ loop {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => break,
+ Ok(true) => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ }
+ }
+ }
+
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {}
+ }
+ }
+
+ // 3. Main merge-scan loop
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ if self.outer_offset >= num_outer {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ continue;
+ }
+
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ };
+ let num_inner = inner_batch.num_rows();
+
+ if self.inner_offset >= num_inner {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => continue,
+ }
+ }
+
+ // 4. Compare keys at current positions
+ let cmp = compare_join_arrays(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ &self.inner_key_arrays,
+ self.inner_offset,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ match cmp {
+ Ordering::Less => {
+ let group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ self.outer_offset = group_end;
+ }
+ Ordering::Greater => {
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if group_end >= num_inner {
+ let saved_keys = self.inner_key_arrays.clone();
+ let saved_idx = num_inner - 1;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_) => continue,
+ }
+ }
+ continue;
+ }
+ }
+ } else {
+ self.inner_offset = group_end;
+ }
+ }
+ Ordering::Equal => {
+ if self.filter.is_some() {
+ // Buffer inner key group (may span batches)
+ match ready!(self.buffer_inner_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ // Process outer rows against buffered inner group
+ // (may need to handle outer batch boundary)
+ loop {
+ self.process_key_match_with_filter()?;
+
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ if self.outer_offset >= outer_batch.num_rows() {
+ self.emit_outer_batch()?;
+ self.boundary_state = BoundaryState::FilteredPending;
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.boundary_state = BoundaryState::Normal;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ self.boundary_state = BoundaryState::Normal;
+ if !self.inner_key_buffer.is_empty() {
+ let first_inner = &self.inner_key_buffer[0];
+ let inner_keys = evaluate_join_keys(
+ first_inner,
+ &self.on_inner,
+ )?;
+ let same = keys_match(
+ &self.outer_key_arrays,
+ 0,
+ &inner_keys,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same {
+ continue;
+ }
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+
+ self.clear_inner_key_group();
+ } else {
+ // No filter: advance inner past key group, then
+ // mark all outer rows with this key as matched.
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ loop {
+ self.process_key_match_no_filter()?;
+
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let saved_keys = self.outer_key_arrays.clone();
+ let saved_idx = num_outer - 1;
+
+ self.emit_outer_batch()?;
+ // Save boundary state before polling. If
+ // poll returns Pending, poll_join exits and
+ // re-enters from the top. The saved state
+ // lets step 1 resume the boundary loop.
+ self.boundary_state = BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ };
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.boundary_state = BoundaryState::Normal;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ // Recover saved_keys from boundary state
+ let BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ } = std::mem::replace(
+ &mut self.boundary_state,
+ BoundaryState::Normal,
+ )
+ else {
+ unreachable!()
+ };
+ let same_key = keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.outer_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ continue;
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+ }
+ }
+ }
+
+ // Check for completed coalescer batch
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ }
+ }
+}
+
+/// Evaluate the filter for all rows in an inner slice against the outer group,
+/// OR-ing results into the matched bitset. Returns the updated matched count.
+/// Extracted as a free function so Rust can split borrows on the stream struct.
+#[expect(clippy::too_many_arguments)]
+fn eval_filter_for_inner_slice(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_slice: &RecordBatch,
+ matched: &mut BooleanBufferBuilder,
+ outer_offset: usize,
+ outer_group_len: usize,
+ mut matched_count: usize,
+) -> Result<usize> {
+ for inner_row in 0..inner_slice.num_rows() {
+ if matched_count == outer_group_len {
+ break;
+ }
+
+ let filter_result = evaluate_filter_for_inner_row(
+ outer_is_left,
+ filter,
+ outer_slice,
+ inner_slice,
+ inner_row,
+ )?;
+
+ // OR filter results into the matched bitset. Both sides are
+ // bit-packed [u8] buffers, so apply_bitwise_binary_op
+ // processes 64 bits per loop iteration (not 1 bit at a time).
+ //
+ // The offsets handle alignment: outer_offset is the bit
+ // position within matched where this key group starts,
+ // and filter_buf.offset() is the BooleanBuffer's internal
+ // bit offset (usually 0, but not guaranteed by Arrow).
+ let filter_buf = filter_result.values();
+ apply_bitwise_binary_op(
+ matched.as_slice_mut(),
+ outer_offset,
+ filter_buf.inner().as_slice(),
+ filter_buf.offset(),
+ outer_group_len,
+ |a, b| a | b,
+ );
+
+ // Recount matched bits after the OR. UnalignedBitChunk is
+ // zero-copy — it reads the bytes in place and uses popcnt.
+ matched_count =
+ UnalignedBitChunk::new(matched.as_slice(), outer_offset, outer_group_len)
+ .count_ones();
+ }
+ Ok(matched_count)
+}
+
+/// Compare two key rows using the sort options to determine equality.
+fn keys_match(
+ left_arrays: &[ArrayRef],
+ left_idx: usize,
+ right_arrays: &[ArrayRef],
+ right_idx: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<bool> {
+ let cmp = compare_join_arrays(
+ left_arrays,
+ left_idx,
+ right_arrays,
+ right_idx,
+ sort_options,
+ null_equality,
+ )?;
+ Ok(cmp == Ordering::Equal)
+}
+
+/// Evaluate the join filter for one inner row against a slice of outer rows.
+///
+/// Free function (not a method on SemiAntiSortMergeJoinStream) so that Rust
+/// can split the struct borrow in process_key_match_with_filter: the caller
+/// holds &mut self.matched and &self.inner_key_buffer simultaneously, which
+/// is impossible if this borrows all of &self.
+fn evaluate_filter_for_inner_row(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_batch: &RecordBatch,
+ inner_idx: usize,
+) -> Result<BooleanArray> {
+ let num_outer_rows = outer_slice.num_rows();
+
+ // Build filter input columns in the order the filter expects
+ let mut columns: Vec<ArrayRef> = Vec::with_capacity(filter.column_indices().len());
+ for col_idx in filter.column_indices() {
+ let (side_batch, side_idx) = if outer_is_left {
+ match col_idx.side {
+ JoinSide::Left => (outer_slice, None),
+ JoinSide::Right => (inner_batch, Some(inner_idx)),
+ JoinSide::None => {
+ return internal_err!("Unexpected JoinSide::None in filter");
+ }
+ }
+ } else {
+ match col_idx.side {
+ JoinSide::Left => (inner_batch, Some(inner_idx)),
+ JoinSide::Right => (outer_slice, None),
+ JoinSide::None => {
+ return internal_err!("Unexpected JoinSide::None in filter");
+ }
+ }
+ };
+
+ match side_idx {
+ None => {
+ columns.push(Arc::clone(side_batch.column(col_idx.index)));
+ }
+ Some(idx) => {
+ // Broadcasts inner scalar to N-element array. Arrow's
+ // BinaryExpr handles Scalar×Array natively via the Datum
+ // trait, but Column::evaluate always returns Array, so
+ // we'd need a custom expr to avoid this broadcast.
+ let scalar = ScalarValue::try_from_array(
+ side_batch.column(col_idx.index).as_ref(),
+ idx,
+ )?;
+ columns.push(scalar.to_array_of_size(num_outer_rows)?);
+ }
+ }
+ }
+
+ let filter_batch = RecordBatch::try_new(Arc::clone(filter.schema()), columns)?;
+ let result = filter
+ .expression()
+ .evaluate(&filter_batch)?
+ .into_array(num_outer_rows)?;
Review Comment:
this can be another PR but if the expression returned scalar we can use that optimization later instead of converting to array
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4029436841
🤖 `./gh_compare_branch.sh` [gh_compare_branch.sh](https://github.com/alamb/datafusion-benchmarking/blob/main/scripts/gh_compare_branch.sh) Running
Linux aal-dev 6.14.0-1018-gcp #19~24.04.1-Ubuntu SMP Wed Sep 24 23:23:09 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing semi-anti-smj (0c33ad6b3fbfe2fd0d04421833b8af95e038b959) to 44dfa7b8a5b908d6d82233cc6aaea1137b1d54b5 [diff](https://github.com/apache/datafusion/compare/44dfa7b8a5b908d6d82233cc6aaea1137b1d54b5..0c33ad6b3fbfe2fd0d04421833b8af95e038b959) using: tpch
Results will be posted here when complete
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2906925439
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1160 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Binary search (not `partition_point`) because `compare_join_arrays`
+/// returns `Result` and `partition_point` requires an infallible predicate
+/// on a `&[T]` — using it here would need either allocation or unsafe
+/// pointer arithmetic to recover the probe index.
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let mut lo = from + 1;
+ let mut hi = len;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
Review Comment:
Yeah that makes sense, I'll try that out! Thanks for the feedback!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4063912383
🤖 Criterion benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/20806#issuecomment-4063907063)
`Linux bench-c4063907063-288-xn9pk 6.12.55+ #1 SMP Sun Feb 1 08:59:41 UTC 2026 aarch64 GNU/Linux`
Comparing semi-anti-smj (f14df032c8ec18774dca3bd86017fb85720f8e10) to d2278a9 (merge-base) [diff](https://github.com/apache/datafusion/compare/d2278a90b4543939cefb0f3ffbea8b025fe922f0..f14df032c8ec18774dca3bd86017fb85720f8e10)
BENCH_NAME=smj
BENCH_COMMAND=cargo bench --features=parquet --bench smj
BENCH_FILTER=
Results will be posted here when complete
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2989458647
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
Review Comment:
is it a current cursor position inside relative to `outer_batch`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2989678623
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see PendingBoundary doc comment.
+ pending_boundary: Option<PendingBoundary>,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
Review Comment:
Updated comment. I'd rather not have a bunch of match arms on the hot path, especially inside of a specialized stream where we know some `JoinType` values would be unreachable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027368352
🤖 `./gh_compare_branch.sh` [gh_compare_branch.sh](https://github.com/alamb/datafusion-benchmarking/blob/main/scripts/gh_compare_branch.sh) Running
Linux aal-dev 6.14.0-1018-gcp #19~24.04.1-Ubuntu SMP Wed Sep 24 23:23:09 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing semi-anti-smj (0c33ad6b3fbfe2fd0d04421833b8af95e038b959) to 44dfa7b8a5b908d6d82233cc6aaea1137b1d54b5 [diff](https://github.com/apache/datafusion/compare/44dfa7b8a5b908d6d82233cc6aaea1137b1d54b5..0c33ad6b3fbfe2fd0d04421833b8af95e038b959) using: tpcds
Results will be posted here when complete
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4029190927
run benchmark tpch
PREFER_HASH_JOIN=false
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027404379
🤖: Benchmark completed
<details><summary>Details</summary>
<p>
```
Comparing HEAD and semi-anti-smj
--------------------
Benchmark tpcds_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ semi-anti-smj ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 73.27 ms │ 70.08 ms │ no change │
│ QQuery 2 │ 249.25 ms │ 242.36 ms │ no change │
│ QQuery 3 │ 154.75 ms │ 152.60 ms │ no change │
│ QQuery 4 │ 1852.58 ms │ 1843.21 ms │ no change │
│ QQuery 5 │ 260.42 ms │ 255.45 ms │ no change │
│ QQuery 6 │ 1443.60 ms │ 1427.54 ms │ no change │
│ QQuery 7 │ 496.70 ms │ 492.61 ms │ no change │
│ QQuery 8 │ 174.46 ms │ 168.48 ms │ no change │
│ QQuery 9 │ 293.60 ms │ 272.36 ms │ +1.08x faster │
│ QQuery 10 │ 160.39 ms │ 162.15 ms │ no change │
│ QQuery 11 │ 1248.15 ms │ 1219.25 ms │ no change │
│ QQuery 12 │ 64.31 ms │ 64.58 ms │ no change │
│ QQuery 13 │ 537.58 ms │ 516.68 ms │ no change │
│ QQuery 14 │ 1839.52 ms │ 1819.55 ms │ no change │
│ QQuery 15 │ 28.10 ms │ 28.22 ms │ no change │
│ QQuery 16 │ 64.41 ms │ 63.20 ms │ no change │
│ QQuery 17 │ 353.98 ms │ 339.92 ms │ no change │
│ QQuery 18 │ 188.44 ms │ 188.03 ms │ no change │
│ QQuery 19 │ 215.72 ms │ 208.40 ms │ no change │
│ QQuery 20 │ 24.44 ms │ 23.88 ms │ no change │
│ QQuery 21 │ 35.88 ms │ 35.38 ms │ no change │
│ QQuery 22 │ 716.75 ms │ 708.32 ms │ no change │
│ QQuery 23 │ 1729.75 ms │ 1703.64 ms │ no change │
│ QQuery 24 │ 655.27 ms │ 646.44 ms │ no change │
│ QQuery 25 │ 511.25 ms │ 495.23 ms │ no change │
│ QQuery 26 │ 120.51 ms │ 120.77 ms │ no change │
│ QQuery 27 │ 494.10 ms │ 485.76 ms │ no change │
│ QQuery 28 │ 297.99 ms │ 291.03 ms │ no change │
│ QQuery 29 │ 445.60 ms │ 425.69 ms │ no change │
│ QQuery 30 │ 71.01 ms │ 72.82 ms │ no change │
│ QQuery 31 │ 308.61 ms │ 290.61 ms │ +1.06x faster │
│ QQuery 32 │ 81.79 ms │ 80.99 ms │ no change │
│ QQuery 33 │ 211.76 ms │ 203.33 ms │ no change │
│ QQuery 34 │ 158.86 ms │ 149.36 ms │ +1.06x faster │
│ QQuery 35 │ 170.24 ms │ 161.18 ms │ +1.06x faster │
│ QQuery 36 │ 279.91 ms │ 277.97 ms │ no change │
│ QQuery 37 │ 254.12 ms │ 241.38 ms │ +1.05x faster │
│ QQuery 38 │ 150.65 ms │ 145.03 ms │ no change │
│ QQuery 39 │ 191.80 ms │ 191.03 ms │ no change │
│ QQuery 40 │ 176.22 ms │ 176.14 ms │ no change │
│ QQuery 41 │ 26.88 ms │ 26.97 ms │ no change │
│ QQuery 42 │ 145.01 ms │ 143.12 ms │ no change │
│ QQuery 43 │ 125.57 ms │ 121.06 ms │ no change │
│ QQuery 44 │ 27.84 ms │ 28.64 ms │ no change │
│ QQuery 45 │ 84.10 ms │ 84.19 ms │ no change │
│ QQuery 46 │ 321.15 ms │ 308.41 ms │ no change │
│ QQuery 47 │ 1010.61 ms │ 996.22 ms │ no change │
│ QQuery 48 │ 402.66 ms │ 386.28 ms │ no change │
│ QQuery 49 │ 360.20 ms │ 359.10 ms │ no change │
│ QQuery 50 │ 333.24 ms │ 327.08 ms │ no change │
│ QQuery 51 │ 299.52 ms │ 292.58 ms │ no change │
│ QQuery 52 │ 145.21 ms │ 142.55 ms │ no change │
│ QQuery 53 │ 146.26 ms │ 141.50 ms │ no change │
│ QQuery 54 │ 203.91 ms │ 199.72 ms │ no change │
│ QQuery 55 │ 146.64 ms │ 143.15 ms │ no change │
│ QQuery 56 │ 207.81 ms │ 201.22 ms │ no change │
│ QQuery 57 │ 280.23 ms │ 278.63 ms │ no change │
│ QQuery 58 │ 493.92 ms │ 471.82 ms │ no change │
│ QQuery 59 │ 327.61 ms │ 324.93 ms │ no change │
│ QQuery 60 │ 211.89 ms │ 205.22 ms │ no change │
│ QQuery 61 │ 247.21 ms │ 232.64 ms │ +1.06x faster │
│ QQuery 62 │ 1389.03 ms │ 1360.57 ms │ no change │
│ QQuery 63 │ 147.60 ms │ 142.01 ms │ no change │
│ QQuery 64 │ 1108.76 ms │ 1106.84 ms │ no change │
│ QQuery 65 │ 352.13 ms │ 339.13 ms │ no change │
│ QQuery 66 │ 383.44 ms │ 379.76 ms │ no change │
│ QQuery 67 │ 525.86 ms │ 522.93 ms │ no change │
│ QQuery 68 │ 377.20 ms │ 364.51 ms │ no change │
│ QQuery 69 │ 160.16 ms │ 158.72 ms │ no change │
│ QQuery 70 │ 485.54 ms │ 478.77 ms │ no change │
│ QQuery 71 │ 187.32 ms │ 180.59 ms │ no change │
│ QQuery 72 │ 2031.28 ms │ 2033.70 ms │ no change │
│ QQuery 73 │ 156.60 ms │ 145.48 ms │ +1.08x faster │
│ QQuery 74 │ 805.62 ms │ 798.82 ms │ no change │
│ QQuery 75 │ 410.87 ms │ 393.37 ms │ no change │
│ QQuery 76 │ 184.13 ms │ 177.12 ms │ no change │
│ QQuery 77 │ 277.94 ms │ 274.59 ms │ no change │
│ QQuery 78 │ 576.40 ms │ 572.55 ms │ no change │
│ QQuery 79 │ 321.84 ms │ 310.34 ms │ no change │
│ QQuery 80 │ 498.75 ms │ 486.44 ms │ no change │
│ QQuery 81 │ 51.23 ms │ 50.81 ms │ no change │
│ QQuery 82 │ 278.99 ms │ 263.08 ms │ +1.06x faster │
│ QQuery 83 │ 73.97 ms │ 73.11 ms │ no change │
│ QQuery 84 │ 65.04 ms │ 64.79 ms │ no change │
│ QQuery 85 │ 203.91 ms │ 207.36 ms │ no change │
│ QQuery 86 │ 58.52 ms │ 56.13 ms │ no change │
│ QQuery 87 │ 152.91 ms │ 144.79 ms │ +1.06x faster │
│ QQuery 88 │ 252.37 ms │ 237.09 ms │ +1.06x faster │
│ QQuery 89 │ 163.86 ms │ 162.28 ms │ no change │
│ QQuery 90 │ 45.68 ms │ 43.80 ms │ no change │
│ QQuery 91 │ 93.57 ms │ 99.05 ms │ 1.06x slower │
│ QQuery 92 │ 82.92 ms │ 80.84 ms │ no change │
│ QQuery 93 │ 280.27 ms │ 271.33 ms │ no change │
│ QQuery 94 │ 91.17 ms │ 90.51 ms │ no change │
│ QQuery 95 │ 248.46 ms │ 242.83 ms │ no change │
│ QQuery 96 │ 114.52 ms │ 105.42 ms │ +1.09x faster │
│ QQuery 97 │ 192.13 ms │ 185.32 ms │ no change │
│ QQuery 98 │ 203.97 ms │ 202.49 ms │ no change │
│ QQuery 99 │ 15874.57 ms │ 15887.28 ms │ no change │
└───────────┴─────────────┴───────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ Total Time (HEAD) │ 51241.76ms │
│ Total Time (semi-anti-smj) │ 50570.19ms │
│ Average Time (HEAD) │ 517.59ms │
│ Average Time (semi-anti-smj) │ 510.81ms │
│ Queries Faster │ 11 │
│ Queries Slower │ 1 │
│ Queries with No Change │ 87 │
│ Queries with Failure │ 0 │
└──────────────────────────────┴────────────┘
```
</p>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2913520212
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1164 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
Review Comment:
That's pretty huge refactor since there are a lot of dependent variable names "inner_" and "outer_" that I don't think I want to tackle. SMJ is inherently a nested loop join (where the outer and inner comes from) so I beefed up the comment to hopefully clarify the mapping for readers.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2989458647
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
Review Comment:
is it a position inside current `outer_batch`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2906663047
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1160 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Binary search (not `partition_point`) because `compare_join_arrays`
+/// returns `Result` and `partition_point` requires an infallible predicate
+/// on a `&[T]` — using it here would need either allocation or unsafe
+/// pointer arithmetic to recover the probe index.
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let mut lo = from + 1;
+ let mut hi = len;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. The `inner_key_buffer` field already holds the buffered
+ /// inner rows needed to resume filter evaluation.
+ FilteredPending,
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Main loop: drive the merge-scan to produce output batches.
+ // Clippy can't see that poll_next_outer_batch guarantees
+ // outer_batch = Some(..) when returning Ok(true).
+ #[expect(clippy::panicking_unwrap)]
Review Comment:
Yeah, I can also try to restructure it a bit to make the warning go away. I played with adding a boolean value that flipped in the case when we know it won't panic, but it started to look ugly and was getting late in the evening so I called it for the night.
I will revisit that this week, thank you for the reminder! :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4025013189
Thanks @mbutrovich, IMO it is great to start replacing current SMJ implementation with more performant and easier to maintain structure. Even by pieces. Good sign fuzz tests passed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027281460
🤖 `./gh_compare_branch.sh` [gh_compare_branch.sh](https://github.com/alamb/datafusion-benchmarking/blob/main/scripts/gh_compare_branch.sh) Running
Linux aal-dev 6.14.0-1018-gcp #19~24.04.1-Ubuntu SMP Wed Sep 24 23:23:09 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing semi-anti-smj (0c33ad6b3fbfe2fd0d04421833b8af95e038b959) to 44dfa7b8a5b908d6d82233cc6aaea1137b1d54b5 [diff](https://github.com/apache/datafusion/compare/44dfa7b8a5b908d6d82233cc6aaea1137b1d54b5..0c33ad6b3fbfe2fd0d04421833b8af95e038b959) using: tpch
Results will be posted here when complete
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2989669179
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see PendingBoundary doc comment.
+ pending_boundary: Option<PendingBoundary>,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
Review Comment:
Updated comment.
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
Review Comment:
Updated comment.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2906637232
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1160 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Binary search (not `partition_point`) because `compare_join_arrays`
+/// returns `Result` and `partition_point` requires an infallible predicate
+/// on a `&[T]` — using it here would need either allocation or unsafe
+/// pointer arithmetic to recover the probe index.
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let mut lo = from + 1;
+ let mut hi = len;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. The `inner_key_buffer` field already holds the buffered
+ /// inner rows needed to resume filter evaluation.
+ FilteredPending,
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Main loop: drive the merge-scan to produce output batches.
+ // Clippy can't see that poll_next_outer_batch guarantees
+ // outer_batch = Some(..) when returning Ok(true).
+ #[expect(clippy::panicking_unwrap)]
Review Comment:
can we somehow isolate the expect to the smallest code section that cause the false positive so it won't hide real issues
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4025253220
> One thing to mention: to control the flag `datafusion.optimizer.enable_semi_anti_sort_merge_join` or `DATAFUSION_OPTIMIZER_ENABLE_SEMI_ANTI_SORT_MERGE_JOIN` in `bench.sh`, at least document it
I'd like to understand the long-term implications of having the config. I added it since we're used to having opt outs in Comet from operators in case it is not compatible with a workload, and it makes it easier to benchmark right now. However, would we then be subject to [API health considerations](https://datafusion.apache.org/contributor-guide/api-health.html) on removing the old SMJ's semi/anti support for X releases/months? If so, I might just propose removing the config and saying "this is the new semi/anti SMJ implementation" as if we had changed code internal to `SortMergeJoinExec` like #18875.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027398791
As much as I would like to review and approve, the code is very complex and I'm not THAT familiar with the join code.
what I usually when creating such PRs to help reviewers even when I already have the full implementation:
1. start with PR that does the bare minimum without support for filters and anything that complicate the code and use that when it is possible so tests suites will cover that as well
2. add features/optimizations one by one while also removing from the old code
this will make it much easier to review and understand
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "stuhood (via GitHub)" <gi...@apache.org>.
stuhood commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2913525870
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1164 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
Review Comment:
I think that `next == last` when the array has a single element. But not a big deal obviously.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4048307659
I've opened an epic to track followup work, since I don't want increase the scope of this PR any further.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4048741603
🤖 Hi @Dandandan, thanks for the request (https://github.com/apache/datafusion/pull/20806#issuecomment-4048740287).
[`scrape_comments.py`](https://github.com/alamb/datafusion-benchmarking/blob/main/scripts/scrape_comments.py) only supports whitelisted benchmarks.
- Standard: clickbench_1, clickbench_extended, clickbench_partitioned, clickbench_pushdown, external_aggr, tpcds, tpch, tpch10, tpch_mem, tpch_mem10
- Criterion: aggregate_query_sql, aggregate_vectorized, array_set_ops, case_when, character_length, in_list, left, plan_reuse, range_and_generate_series, replace, reset_plan_states, sort, sql_planner, strpos, substr_index, with_hashes
Please choose one or more of these with `run benchmark <name>` or `run benchmark <name1> <name2>...`
You can also set environment variables on subsequent lines:
```
run benchmark tpch_mem
DATAFUSION_RUNTIME_MEMORY_LIMIT=1G
```
Unsupported benchmarks: smj.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4029571470
🤖: Benchmark completed
<details><summary>Details</summary>
<p>
```
Comparing HEAD and semi-anti-smj
--------------------
Benchmark tpcds_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ semi-anti-smj ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 70.78 ms │ 69.35 ms │ no change │
│ QQuery 2 │ 249.36 ms │ 246.55 ms │ no change │
│ QQuery 3 │ 155.20 ms │ 152.28 ms │ no change │
│ QQuery 4 │ 1826.23 ms │ 1793.10 ms │ no change │
│ QQuery 5 │ 262.63 ms │ 256.74 ms │ no change │
│ QQuery 6 │ 1419.77 ms │ 1412.05 ms │ no change │
│ QQuery 7 │ 493.49 ms │ 484.09 ms │ no change │
│ QQuery 8 │ 172.96 ms │ 165.87 ms │ no change │
│ QQuery 9 │ 293.25 ms │ 268.67 ms │ +1.09x faster │
│ QQuery 10 │ 162.98 ms │ 165.78 ms │ no change │
│ QQuery 11 │ 1197.59 ms │ 1192.61 ms │ no change │
│ QQuery 12 │ 63.40 ms │ 65.42 ms │ no change │
│ QQuery 13 │ 530.69 ms │ 522.41 ms │ no change │
│ QQuery 14 │ 1819.77 ms │ 1803.02 ms │ no change │
│ QQuery 15 │ 27.20 ms │ 27.12 ms │ no change │
│ QQuery 16 │ 63.63 ms │ 62.71 ms │ no change │
│ QQuery 17 │ 351.59 ms │ 336.64 ms │ no change │
│ QQuery 18 │ 182.88 ms │ 185.82 ms │ no change │
│ QQuery 19 │ 214.25 ms │ 205.38 ms │ no change │
│ QQuery 20 │ 23.24 ms │ 22.57 ms │ no change │
│ QQuery 21 │ 35.20 ms │ 34.35 ms │ no change │
│ QQuery 22 │ 708.30 ms │ 715.03 ms │ no change │
│ QQuery 23 │ 1725.45 ms │ 1701.17 ms │ no change │
│ QQuery 24 │ 648.85 ms │ 648.96 ms │ no change │
│ QQuery 25 │ 507.50 ms │ 491.42 ms │ no change │
│ QQuery 26 │ 121.20 ms │ 120.91 ms │ no change │
│ QQuery 27 │ 496.35 ms │ 482.22 ms │ no change │
│ QQuery 28 │ 292.23 ms │ 292.29 ms │ no change │
│ QQuery 29 │ 436.19 ms │ 424.64 ms │ no change │
│ QQuery 30 │ 69.77 ms │ 69.73 ms │ no change │
│ QQuery 31 │ 310.78 ms │ 286.48 ms │ +1.08x faster │
│ QQuery 32 │ 80.33 ms │ 80.04 ms │ no change │
│ QQuery 33 │ 205.04 ms │ 201.66 ms │ no change │
│ QQuery 34 │ 159.36 ms │ 149.55 ms │ +1.07x faster │
│ QQuery 35 │ 168.64 ms │ 162.58 ms │ no change │
│ QQuery 36 │ 280.04 ms │ 277.64 ms │ no change │
│ QQuery 37 │ 258.64 ms │ 238.75 ms │ +1.08x faster │
│ QQuery 38 │ 147.98 ms │ 141.02 ms │ no change │
│ QQuery 39 │ 189.92 ms │ 190.92 ms │ no change │
│ QQuery 40 │ 179.51 ms │ 171.36 ms │ no change │
│ QQuery 41 │ 26.65 ms │ 27.13 ms │ no change │
│ QQuery 42 │ 144.28 ms │ 142.09 ms │ no change │
│ QQuery 43 │ 125.93 ms │ 123.18 ms │ no change │
│ QQuery 44 │ 27.68 ms │ 27.89 ms │ no change │
│ QQuery 45 │ 84.50 ms │ 80.78 ms │ no change │
│ QQuery 46 │ 318.55 ms │ 303.06 ms │ no change │
│ QQuery 47 │ 983.19 ms │ 957.64 ms │ no change │
│ QQuery 48 │ 396.92 ms │ 380.42 ms │ no change │
│ QQuery 49 │ 363.22 ms │ 352.65 ms │ no change │
│ QQuery 50 │ 327.50 ms │ 320.30 ms │ no change │
│ QQuery 51 │ 294.04 ms │ 291.61 ms │ no change │
│ QQuery 52 │ 145.80 ms │ 144.08 ms │ no change │
│ QQuery 53 │ 146.12 ms │ 141.80 ms │ no change │
│ QQuery 54 │ 200.90 ms │ 197.88 ms │ no change │
│ QQuery 55 │ 144.57 ms │ 142.46 ms │ no change │
│ QQuery 56 │ 203.79 ms │ 199.54 ms │ no change │
│ QQuery 57 │ 276.49 ms │ 273.02 ms │ no change │
│ QQuery 58 │ 488.53 ms │ 471.19 ms │ no change │
│ QQuery 59 │ 329.62 ms │ 325.17 ms │ no change │
│ QQuery 60 │ 208.92 ms │ 206.15 ms │ no change │
│ QQuery 61 │ 244.00 ms │ 232.30 ms │ no change │
│ QQuery 62 │ 1345.00 ms │ 1335.60 ms │ no change │
│ QQuery 63 │ 144.41 ms │ 140.05 ms │ no change │
│ QQuery 64 │ 1123.97 ms │ 1113.08 ms │ no change │
│ QQuery 65 │ 350.32 ms │ 341.14 ms │ no change │
│ QQuery 66 │ 399.95 ms │ 373.16 ms │ +1.07x faster │
│ QQuery 67 │ 532.54 ms │ 516.94 ms │ no change │
│ QQuery 68 │ 373.89 ms │ 358.75 ms │ no change │
│ QQuery 69 │ 158.78 ms │ 156.87 ms │ no change │
│ QQuery 70 │ 489.79 ms │ 480.11 ms │ no change │
│ QQuery 71 │ 188.31 ms │ 178.40 ms │ +1.06x faster │
│ QQuery 72 │ 2031.09 ms │ 2010.22 ms │ no change │
│ QQuery 73 │ 153.35 ms │ 144.79 ms │ +1.06x faster │
│ QQuery 74 │ 779.80 ms │ 773.24 ms │ no change │
│ QQuery 75 │ 401.49 ms │ 387.37 ms │ no change │
│ QQuery 76 │ 183.06 ms │ 179.25 ms │ no change │
│ QQuery 77 │ 279.75 ms │ 274.74 ms │ no change │
│ QQuery 78 │ 583.57 ms │ 572.06 ms │ no change │
│ QQuery 79 │ 321.03 ms │ 311.45 ms │ no change │
│ QQuery 80 │ 493.83 ms │ 484.92 ms │ no change │
│ QQuery 81 │ 50.64 ms │ 48.86 ms │ no change │
│ QQuery 82 │ 279.86 ms │ 266.52 ms │ no change │
│ QQuery 83 │ 72.50 ms │ 73.21 ms │ no change │
│ QQuery 84 │ 67.46 ms │ 66.86 ms │ no change │
│ QQuery 85 │ 201.79 ms │ 214.48 ms │ 1.06x slower │
│ QQuery 86 │ 57.20 ms │ 57.54 ms │ no change │
│ QQuery 87 │ 150.36 ms │ 146.11 ms │ no change │
│ QQuery 88 │ 249.76 ms │ 237.40 ms │ no change │
│ QQuery 89 │ 164.15 ms │ 159.76 ms │ no change │
│ QQuery 90 │ 44.40 ms │ 43.94 ms │ no change │
│ QQuery 91 │ 90.97 ms │ 93.82 ms │ no change │
│ QQuery 92 │ 80.57 ms │ 79.47 ms │ no change │
│ QQuery 93 │ 279.55 ms │ 272.39 ms │ no change │
│ QQuery 94 │ 91.20 ms │ 88.72 ms │ no change │
│ QQuery 95 │ 245.93 ms │ 244.96 ms │ no change │
│ QQuery 96 │ 112.95 ms │ 106.87 ms │ +1.06x faster │
│ QQuery 97 │ 188.84 ms │ 181.87 ms │ no change │
│ QQuery 98 │ 202.07 ms │ 204.37 ms │ no change │
│ QQuery 99 │ 15865.68 ms │ 15868.81 ms │ no change │
└───────────┴─────────────┴───────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ Total Time (HEAD) │ 50913.16ms │
│ Total Time (semi-anti-smj) │ 50247.35ms │
│ Average Time (HEAD) │ 514.27ms │
│ Average Time (semi-anti-smj) │ 507.55ms │
│ Queries Faster │ 8 │
│ Queries Slower │ 1 │
│ Queries with No Change │ 90 │
│ Queries with Failure │ 0 │
└──────────────────────────────┴────────────┘
```
</p>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4033508357
@Dandandan could also `un benchmark smj` just as a sanity check that I'm seeing correct speedup here? Thank you!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2990301313
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see PendingBoundary doc comment.
+ pending_boundary: Option<PendingBoundary>,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ pending_boundary: None,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let saved_inner_keys = slice_keys(&self.inner_key_arrays, num_inner - 1);
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ &self.inner_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let saved_inner_keys = slice_keys(&self.inner_key_arrays, num_inner - 1);
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ &self.inner_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Continue processing an outer key group that spans multiple outer
+ /// batches. Returns `true` if this outer batch was fully consumed
+ /// by the key group and the caller should load another.
+ fn resume_boundary(&mut self) -> Result<bool> {
+ debug_assert!(
+ self.outer_batch.is_some(),
+ "caller must load outer_batch first"
+ );
+ match self.pending_boundary.take() {
+ Some(PendingBoundary::NoFilter { saved_keys }) => {
+ let same_key = keys_match(
+ &saved_keys,
+ &self.outer_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_no_filter()?;
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ self.pending_boundary = Some(PendingBoundary::NoFilter {
+ saved_keys: slice_keys(&self.outer_key_arrays, num_outer - 1),
+ });
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ return Ok(true);
+ }
+ }
+ }
+ Some(PendingBoundary::Filtered { saved_keys }) => {
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "Filtered pending boundary entered but no inner key data exists"
+ );
+ let same_key = keys_match(
+ &saved_keys,
+ &self.outer_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_with_filter()?;
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ self.pending_boundary = Some(PendingBoundary::Filtered {
+ saved_keys: slice_keys(&self.outer_key_arrays, num_outer - 1),
+ });
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ return Ok(true);
+ }
+ }
+ self.clear_inner_key_group();
+ }
+ None => {}
+ }
+ Ok(false)
+ }
+
+ /// Main loop: drive the merge-scan to produce output batches.
+ fn poll_join(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<RecordBatch>>> {
+ let join_time = self.join_time.clone();
+ let _timer = join_time.timer();
+
+ loop {
+ // 1. Ensure we have an outer batch
+ if self.outer_batch.is_none() {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Outer exhausted — flush coalescer
+ self.pending_boundary = None;
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {
+ if self.resume_boundary()? {
+ continue;
+ }
+ }
+ }
+ }
+
+ // 2. Ensure we have an inner batch (unless inner is exhausted).
+ // Skip this when resuming a pending boundary — inner was already
+ // advanced past the key group before the boundary loop started.
+ if self.inner_batch.is_none() && self.pending_boundary.is_none() {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Inner exhausted — emit remaining outer batches.
+ // For semi: no more matches possible.
+ // For anti: all remaining outer rows are unmatched.
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ loop {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => break,
+ Ok(true) => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ }
+ }
+ }
+
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {}
+ }
+ }
+
+ // 3. Main merge-scan loop
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ if self.outer_offset >= num_outer {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ continue;
+ }
+
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ };
+ let num_inner = inner_batch.num_rows();
+
+ if self.inner_offset >= num_inner {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => continue,
+ }
+ }
+
+ // 4. Compare keys at current positions
+ let cmp = compare_join_arrays(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ &self.inner_key_arrays,
+ self.inner_offset,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ match cmp {
+ Ordering::Less => {
+ let group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ self.outer_offset = group_end;
+ }
+ Ordering::Greater => {
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if group_end >= num_inner {
+ let saved_keys =
+ slice_keys(&self.inner_key_arrays, num_inner - 1);
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_keys,
+ &self.inner_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_) => continue,
+ }
+ }
+ continue;
+ }
+ }
+ } else {
+ self.inner_offset = group_end;
+ }
+ }
+ Ordering::Equal => {
+ if self.filter.is_some() {
+ // Buffer inner key group (may span batches)
+ match ready!(self.buffer_inner_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ // Process outer rows against buffered inner group
+ // (may need to handle outer batch boundary)
+ loop {
+ self.process_key_match_with_filter()?;
+
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ if self.outer_offset >= outer_batch.num_rows() {
+ let saved_keys = slice_keys(
+ &self.outer_key_arrays,
+ outer_batch.num_rows() - 1,
+ );
+
+ self.emit_outer_batch()?;
+ debug_assert!(
+ !self.inner_key_buffer.is_empty()
+ || self.inner_key_spill.is_some(),
+ "Filtered pending boundary requires inner key data in buffer or spill"
+ );
+ self.pending_boundary =
+ Some(PendingBoundary::Filtered { saved_keys });
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.pending_boundary = None;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ let Some(PendingBoundary::Filtered {
+ saved_keys,
+ }) = self.pending_boundary.take()
+ else {
+ unreachable!()
+ };
+ let same = keys_match(
+ &saved_keys,
+ &self.outer_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same {
+ continue;
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+
+ self.clear_inner_key_group();
+ } else {
+ // No filter: advance inner past key group, then
+ // mark all outer rows with this key as matched.
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ loop {
+ self.process_key_match_no_filter()?;
+
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let saved_keys =
+ slice_keys(&self.outer_key_arrays, num_outer - 1);
+
+ self.emit_outer_batch()?;
+ self.pending_boundary =
+ Some(PendingBoundary::NoFilter { saved_keys });
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.pending_boundary = None;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ let Some(PendingBoundary::NoFilter {
+ saved_keys,
+ }) = self.pending_boundary.take()
+ else {
+ unreachable!()
+ };
+ let same_key = keys_match(
+ &saved_keys,
+ &self.outer_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ continue;
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+ }
+ }
+ }
+
+ // Check for completed coalescer batch
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ }
+ }
+}
+
+/// Evaluate the filter for all rows in an inner slice against the outer group,
+/// OR-ing results into the matched bitset. Returns the updated matched count.
+/// Extracted as a free function so Rust can split borrows on the stream struct.
+#[expect(clippy::too_many_arguments)]
+fn eval_filter_for_inner_slice(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_slice: &RecordBatch,
+ matched: &mut BooleanBufferBuilder,
+ outer_offset: usize,
+ outer_group_len: usize,
+ mut matched_count: usize,
Review Comment:
That requires a recounting when we just did it at the call site. I'll add a `debug_assert` that the arg is correct.
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see PendingBoundary doc comment.
+ pending_boundary: Option<PendingBoundary>,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ pending_boundary: None,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
Review Comment:
Updated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2907668997
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1218 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// outer batch so we can compare with the next batch's first key
+ /// without reading from the inner key buffer (which may have been
+ /// spilled to disk).
+ FilteredPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Continue processing an outer key group that spans multiple outer
+ /// batches. Returns `true` if this outer batch was fully consumed
+ /// by the key group and the caller should load another.
+ fn resume_boundary(&mut self) -> Result<bool> {
+ debug_assert!(
+ self.outer_batch.is_some(),
+ "caller must load outer_batch first"
+ );
+ match std::mem::replace(&mut self.boundary_state, BoundaryState::Normal) {
Review Comment:
Yeah, the `Normal` in `std::mem::replace` is just a placeholder. However, at lines 952 and 1017, `Normal` is also the correct final state (we consumed the pending data). At line 687, it's either correct (for the
`Normal` and non-boundary-continuing cases) or gets overwritten at lines 706/738.
I could add a comment on each replace call to clarify the intent, but I'd avoid adding a `TmpTransition` variant — it adds a match arm everywhere, turns a compile-time guarantee into a runtime panic, and the code paths already ensure `Normal` is correct.
Alternatively, since `boundary_state` is only read via `match`, we could wrap it in `Option<BoundaryState>` and `.take()`. But that adds `Option` unwrapping everywhere else. I want to make sure this is maintainable for other folks, but I don't love an additional state. Especially one that is purely ephemeral and shouldn't ever really match to anything.
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1218 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// outer batch so we can compare with the next batch's first key
+ /// without reading from the inner key buffer (which may have been
+ /// spilled to disk).
+ FilteredPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Continue processing an outer key group that spans multiple outer
+ /// batches. Returns `true` if this outer batch was fully consumed
+ /// by the key group and the caller should load another.
+ fn resume_boundary(&mut self) -> Result<bool> {
+ debug_assert!(
+ self.outer_batch.is_some(),
+ "caller must load outer_batch first"
+ );
+ match std::mem::replace(&mut self.boundary_state, BoundaryState::Normal) {
Review Comment:
Yeah, the `Normal` in `std::mem::replace` is just a placeholder. However, at lines 952 and 1017, `Normal` is also the correct final state (we consumed the pending data). At line 687, it's either correct (for the
`Normal` and non-boundary-continuing cases) or gets overwritten at lines 706/738.
I could add a comment on each replace call to clarify the intent, but I'd avoid adding a `TmpTransition` variant — it adds a match arm everywhere, turns a compile-time guarantee into a runtime panic, and the code paths already ensure `Normal` is correct.
Alternatively, since `boundary_state` is only read via `match`, we could wrap it in `Option<BoundaryState>` and `.take()`. But that adds `Option` unwrapping everywhere else. I want to make sure this is maintainable for other folks, but I don't love an additional state. Especially one that is purely ephemeral and shouldn't ever really `match` to anything.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4019627741
@alamb is it possible to run TPC-H without prefer hash join enabled so we see this operator vs. the old on main?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2907745111
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1218 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// outer batch so we can compare with the next batch's first key
+ /// without reading from the inner key buffer (which may have been
+ /// spilled to disk).
+ FilteredPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Continue processing an outer key group that spans multiple outer
+ /// batches. Returns `true` if this outer batch was fully consumed
+ /// by the key group and the caller should load another.
+ fn resume_boundary(&mut self) -> Result<bool> {
+ debug_assert!(
+ self.outer_batch.is_some(),
+ "caller must load outer_batch first"
+ );
+ match std::mem::replace(&mut self.boundary_state, BoundaryState::Normal) {
Review Comment:
Let me try an `Option` and remove the the `Normal` state. I'll rename it to `PendingBoundary`, where `None` on the `Option` is now obvious, and there will just be two pending states. I'll push that in a little bit.
Thanks for forcing me to think about this again! I appreciate the feedback.
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1218 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// outer batch so we can compare with the next batch's first key
+ /// without reading from the inner key buffer (which may have been
+ /// spilled to disk).
+ FilteredPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Continue processing an outer key group that spans multiple outer
+ /// batches. Returns `true` if this outer batch was fully consumed
+ /// by the key group and the caller should load another.
+ fn resume_boundary(&mut self) -> Result<bool> {
+ debug_assert!(
+ self.outer_batch.is_some(),
+ "caller must load outer_batch first"
+ );
+ match std::mem::replace(&mut self.boundary_state, BoundaryState::Normal) {
Review Comment:
Let me try an `Option` and remove the `Normal` state. I'll rename it to `PendingBoundary`, where `None` on the `Option` is now obvious, and there will just be two pending states. I'll push that in a little bit.
Thanks for forcing me to think about this again! I appreciate the feedback.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027368119
🤖: Benchmark completed
<details><summary>Details</summary>
<p>
```
Comparing HEAD and semi-anti-smj
--------------------
Benchmark tpch_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ semi-anti-smj ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩
│ QQuery 1 │ 178.89 ms │ 178.80 ms │ no change │
│ QQuery 2 │ 82.94 ms │ 83.37 ms │ no change │
│ QQuery 3 │ 117.34 ms │ 116.76 ms │ no change │
│ QQuery 4 │ 74.16 ms │ 74.07 ms │ no change │
│ QQuery 5 │ 163.61 ms │ 158.46 ms │ no change │
│ QQuery 6 │ 58.26 ms │ 64.90 ms │ 1.11x slower │
│ QQuery 7 │ 194.03 ms │ 190.63 ms │ no change │
│ QQuery 8 │ 163.19 ms │ 155.99 ms │ no change │
│ QQuery 9 │ 221.45 ms │ 219.76 ms │ no change │
│ QQuery 10 │ 171.98 ms │ 174.14 ms │ no change │
│ QQuery 11 │ 59.45 ms │ 58.10 ms │ no change │
│ QQuery 12 │ 113.81 ms │ 116.21 ms │ no change │
│ QQuery 13 │ 226.57 ms │ 219.81 ms │ no change │
│ QQuery 14 │ 87.21 ms │ 88.20 ms │ no change │
│ QQuery 15 │ 120.44 ms │ 120.19 ms │ no change │
│ QQuery 16 │ 55.20 ms │ 55.19 ms │ no change │
│ QQuery 17 │ 249.61 ms │ 250.81 ms │ no change │
│ QQuery 18 │ 296.19 ms │ 299.67 ms │ no change │
│ QQuery 19 │ 124.14 ms │ 129.61 ms │ no change │
│ QQuery 20 │ 122.39 ms │ 118.80 ms │ no change │
│ QQuery 21 │ 248.33 ms │ 246.97 ms │ no change │
│ QQuery 22 │ 40.34 ms │ 40.41 ms │ no change │
└───────────┴───────────┴───────────────┴──────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩
│ Total Time (HEAD) │ 3169.55ms │
│ Total Time (semi-anti-smj) │ 3160.86ms │
│ Average Time (HEAD) │ 144.07ms │
│ Average Time (semi-anti-smj) │ 143.68ms │
│ Queries Faster │ 0 │
│ Queries Slower │ 1 │
│ Queries with No Change │ 21 │
│ Queries with Failure │ 0 │
└──────────────────────────────┴───────────┘
```
</p>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2989491999
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see PendingBoundary doc comment.
+ pending_boundary: Option<PendingBoundary>,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
Review Comment:
we still need ON physical expressions to recalculate inner/outer key arrays?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4035709073
> The operator does get swapped in for TPC-H, it's just about a 8% (SF10) to 20% (SF100) improvement just within the SMJ operator on Q21 when running locally, so it doesn't end up as big of a win on query execution time as the microbenchmarks.
right, Q21 is the biggest consumer for ANTI/SEMI
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4031576530
> My only other thought would be to think about the end game on how this will be most maintainable.
>
> Would this be easiest to maintain if it were _only_ a separate `Stream` implementation, with an ExecutionPlan shared with `SortMergeJoin`? We are downstream using a [physical optimizer rule to force `SortMergeJoin` to be used on sorted inputs](https://github.com/paradedb/paradedb/blob/f96213bf7675fbb81936fb52a339d62918f05761/pg_search/src/postgres/customscan/joinscan/planner.rs#L44-L68), and having only one "SortMergeJoin" operator (even if it had multiple backing implementations) would be a bit easier for us.
I didn't look too closely at the rule, but I wonder if you could look for joins that both have `required_input_ordering()` and `maintains_input_order()` rather than the single `SortMergeJoin` operator. That said, the `Stream` implementation is an interesting idea and I'll see what that would look like.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027586477
@Dandandan assuming those run with dfbench, I suspect its default value for `PREFER_HASH_JOIN` is overriding the value you're trying to set.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4029192941
run benchmark tpcds
PREFER_HASH_JOIN =false
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4033417699
@stuhood I pushed a commit that implements your suggestion: just a new `Stream` instead of a whole new operator. I will run some benchmarks but I expect it'll be fine. It also reduced the line count on this PR dramatically since a lot of boilerplate and redundant tests are gone.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4031467906
The operator does get swapped in for TPC-H, it's just about a 8% (SF10) to 20% (SF100) improvement just within the SMJ operator on Q21 when running locally, so it doesn't end up as big of a win on query execution time as the microbenchmarks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich merged PR #20806:
URL: https://github.com/apache/datafusion/pull/20806
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2990302328
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see PendingBoundary doc comment.
+ pending_boundary: Option<PendingBoundary>,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ pending_boundary: None,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
Review Comment:
Updated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2989504604
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see PendingBoundary doc comment.
+ pending_boundary: Option<PendingBoundary>,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
Review Comment:
would it be easier to have a joinType instead of `is_semi` and `outer_is_left`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4048740287
run benchmark smj
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2989838282
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see PendingBoundary doc comment.
+ pending_boundary: Option<PendingBoundary>,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ pending_boundary: None,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
Review Comment:
same
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2989836204
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort-merge join stream specialized for semi/anti joins.
+//!
+//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec)
+//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, or `RightAnti`.
+//!
+//! # Motivation
+//!
+//! The general-purpose `SortMergeJoinStream`
+//! handles semi/anti joins by materializing `(outer, inner)` row pairs,
+//! applying a filter, then using a "corrected filter mask" to deduplicate.
+//! Semi/anti joins only need a boolean per outer row (does a match exist?),
+//! not pairs. The pair-based approach incurs unnecessary memory allocation
+//! and intermediate batches.
+//!
+//! This stream instead tracks matches with a per-outer-batch bitset,
+//! avoiding all pair materialization.
+//!
+//! # "Outer Side" vs "Inner Side"
+//!
+//! For `Left*` join types, left is outer and right is inner.
+//! For `Right*` join types, right is outer and left is inner.
+//! The output schema always equals the outer side's schema.
+//!
+//! # Algorithm
+//!
+//! Both inputs must be sorted by the join keys. The stream performs a merge
+//! scan across the two sorted inputs:
+//!
+//! ```text
+//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7]
+//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7]
+//! ▲
+//! compare keys at cursors
+//! ```
+//!
+//! At each step, the keys at the outer and inner cursors are compared:
+//!
+//! - **outer < inner**: Skip the outer key group (no match exists).
+//! - **outer > inner**: Skip the inner key group.
+//! - **outer == inner**: Process the match (see below).
+//!
+//! Key groups are contiguous runs of equal keys within one side. The scan
+//! advances past entire groups at each step.
+//!
+//! ## Processing a key match
+//!
+//! **Without filter**: All outer rows in the key group are marked as matched.
+//!
+//! **With filter**: The inner key group is buffered (may span multiple inner
+//! batches). For each buffered inner row, the filter is evaluated against the
+//! outer key group as a batch. Results are OR'd into the matched bitset. A
+//! short-circuit exits early when all outer rows in the group are matched.
+//!
+//! ```text
+//! matched bitset: [0, 0, 1, 0, 0, ...]
+//! ▲── one bit per outer row ──▲
+//!
+//! On emit:
+//! Semi → filter_record_batch(outer_batch, &matched)
+//! Anti → filter_record_batch(outer_batch, &NOT(matched))
+//! ```
+//!
+//! ## Batch boundaries
+//!
+//! Key groups can span batch boundaries on either side. The stream handles
+//! this by detecting when a group extends to the end of a batch, loading the
+//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum
+//! preserves loop context across async `Poll::Pending` re-entries.
+//!
+//! # Memory
+//!
+//! Memory usage is bounded and independent of total input size:
+//! - One outer batch at a time (not tracked by reservation — single batch,
+//! cannot be spilled since it's needed for filter evaluation)
+//! - One inner batch at a time (streaming)
+//! - `matched` bitset: one bit per outer row, re-allocated per batch
+//! - Inner key group buffer: only for filtered joins, one key group at a time.
+//! Tracked via `MemoryReservation`; spilled to disk when the memory pool
+//! limit is exceeded.
+//! - `BatchCoalescer`: output buffering to target batch size
+//!
+//! # Degenerate cases
+//!
+//! **Highly skewed key (filtered joins only):** When a filter is present,
+//! the inner key group is buffered so each inner row can be evaluated
+//! against the outer group. If one join key has N inner rows, all N rows
+//! are held in memory simultaneously (or spilled to disk if the memory
+//! pool limit is reached). With uniform key distribution this is small
+//! (inner_rows / num_distinct_keys), but a single hot key can buffer
+//! arbitrarily many rows. The no-filter path does not buffer inner
+//! rows — it only advances the cursor — so it is unaffected.
+//!
+//! **Scalar broadcast during filter evaluation:** Each inner row is
+//! broadcast to match the outer group length for filter evaluation,
+//! allocating one array per inner row × filter column. This is inherent
+//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not
+//! support scalar inputs directly. The total work is
+//! O(inner_group × outer_group) per key, but with much lower constant
+//! factor than the pair-materialization approach.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(crate) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — in the nested-loop model that sort-merge join
+ // implements, "outer" is the driving loop and "inner" is probed for
+ // matches. The existing SortMergeJoinStream calls these "streamed"
+ // and "buffered" respectively. For Left* joins, outer=left; for
+ // Right* joins, outer=right. Output schema equals the outer side.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see PendingBoundary doc comment.
+ pending_boundary: Option<PendingBoundary>,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ pending_boundary: None,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
Review Comment:
```suggestion
let num_rows = batch.num_rows();
if num_rows() == 0 {
continue;
}
self.input_rows.add(num_rows);
let keys = evaluate_join_keys(&batch, &self.on_outer)?;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027289740
run benchmark tpcds
DATAFUSION_OPTIMIZER_PREFER_HASH_JOIN=false
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027275844
run benchmark tpch DATAFUSION_OPTIMIZER_PREFER_HASH_JOIN=false
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "stuhood (via GitHub)" <gi...@apache.org>.
stuhood commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2908102325
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1164 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
Review Comment:
I believe that these cases are equivalent, and so you can just have the "Check if the entire remaining batch shares this key." case?
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1164 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum carries the last key (as single-row sliced arrays) from the
+/// previous batch so we can check whether the next batch continues the same
+/// key group. Stored as `Option<PendingBoundary>`: `None` means normal
+/// processing.
+#[derive(Debug)]
+enum PendingBoundary {
+ /// Resuming a no-filter boundary loop.
+ NoFilter { saved_keys: Vec<ArrayRef> },
+ /// Resuming a filtered boundary loop. Inner key data remains in the
+ /// buffer (or spill file) for the resumed loop.
+ Filtered { saved_keys: Vec<ArrayRef> },
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
Review Comment:
Nit: worth calling them `streamed` vs `buffered` to reduce the cognitive load for readers?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4029527143
🤖: Benchmark completed
<details><summary>Details</summary>
<p>
```
Comparing HEAD and semi-anti-smj
--------------------
Benchmark tpch_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ semi-anti-smj ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 173.98 ms │ 180.18 ms │ no change │
│ QQuery 2 │ 196.59 ms │ 197.54 ms │ no change │
│ QQuery 3 │ 247.97 ms │ 247.25 ms │ no change │
│ QQuery 4 │ 173.42 ms │ 195.71 ms │ 1.13x slower │
│ QQuery 5 │ 390.34 ms │ 379.33 ms │ no change │
│ QQuery 6 │ 63.17 ms │ 58.15 ms │ +1.09x faster │
│ QQuery 7 │ 463.63 ms │ 479.85 ms │ no change │
│ QQuery 8 │ 431.01 ms │ 442.62 ms │ no change │
│ QQuery 9 │ 655.59 ms │ 624.19 ms │ no change │
│ QQuery 10 │ 225.51 ms │ 227.24 ms │ no change │
│ QQuery 11 │ 165.22 ms │ 171.04 ms │ no change │
│ QQuery 12 │ 157.18 ms │ 158.19 ms │ no change │
│ QQuery 13 │ 261.29 ms │ 245.89 ms │ +1.06x faster │
│ QQuery 14 │ 92.20 ms │ 88.01 ms │ no change │
│ QQuery 15 │ 102.65 ms │ 100.99 ms │ no change │
│ QQuery 16 │ 71.06 ms │ 69.70 ms │ no change │
│ QQuery 17 │ 469.17 ms │ 471.08 ms │ no change │
│ QQuery 18 │ 583.84 ms │ 606.05 ms │ no change │
│ QQuery 19 │ 123.10 ms │ 132.66 ms │ 1.08x slower │
│ QQuery 20 │ 180.57 ms │ 191.01 ms │ 1.06x slower │
│ QQuery 21 │ 1034.79 ms │ 956.16 ms │ +1.08x faster │
│ QQuery 22 │ 99.26 ms │ 97.04 ms │ no change │
└───────────┴────────────┴───────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩
│ Total Time (HEAD) │ 6361.56ms │
│ Total Time (semi-anti-smj) │ 6319.88ms │
│ Average Time (HEAD) │ 289.16ms │
│ Average Time (semi-anti-smj) │ 287.27ms │
│ Queries Faster │ 3 │
│ Queries Slower │ 3 │
│ Queries with No Change │ 16 │
│ Queries with Failure │ 0 │
└──────────────────────────────┴───────────┘
```
</p>
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4025063993
One thing to mention: to control the flag `datafusion.optimizer.enable_semi_anti_sort_merge_join` or `DATAFUSION_OPTIMIZER_ENABLE_SEMI_ANTI_SORT_MERGE_JOIN` in `bench.sh`, at least document it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4063912556
Benchmark for [this request](https://github.com/apache/datafusion/pull/20806#issuecomment-4063907063) failed.
Last 20 lines of output:
<details><summary>Click to expand</summary>
```
Cloning into '/workspace/datafusion-branch'...
semi-anti-smj
From https://github.com/apache/datafusion
* [new ref] refs/pull/20806/head -> semi-anti-smj
* branch main -> FETCH_HEAD
Switched to branch 'semi-anti-smj'
d2278a90b4543939cefb0f3ffbea8b025fe922f0
Cloning into '/workspace/datafusion-base'...
HEAD is now at d2278a9 Check sqllogictests for any dangling config settings (#17914) (#20838)
rustc 1.94.0 (4a4ef493e 2026-03-02)
f14df032c8ec18774dca3bd86017fb85720f8e10
d2278a90b4543939cefb0f3ffbea8b025fe922f0
Blocking waiting for file lock on package cache
Blocking waiting for file lock on package cache
Blocking waiting for file lock on package cache
error: no bench target named `smj` in default-run packages
help: a target with a similar name exists: `spm`
```
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4063907063
run benchmark smj
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2907574950
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1218 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// outer batch so we can compare with the next batch's first key
+ /// without reading from the inner key buffer (which may have been
+ /// spilled to disk).
+ FilteredPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Continue processing an outer key group that spans multiple outer
+ /// batches. Returns `true` if this outer batch was fully consumed
+ /// by the key group and the caller should load another.
+ fn resume_boundary(&mut self) -> Result<bool> {
+ debug_assert!(
+ self.outer_batch.is_some(),
+ "caller must load outer_batch first"
+ );
+ match std::mem::replace(&mut self.boundary_state, BoundaryState::Normal) {
Review Comment:
is the replace here to Normal is to get around the compiler due to mutable references?
I also see Normal has some meaning.
because Normal here have 2 meaning, both temp and as actual lifecycle state I want to make it explicit to avoid bugs, so please create another state named `TmpTransition` that whenever you see it anywhere it should panic and use it here
and in places that you actually need Normal set it explicitly.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027588041
🤖 Hi @mbutrovich, thanks for the request (https://github.com/apache/datafusion/pull/20806#issuecomment-4027587613). [`scrape_comments.py`](https://github.com/alamb/datafusion-benchmarking/blob/main/scripts/scrape_comments.py) only responds to whitelisted users. Allowed users: [Dandandan](https://github.com/Dandandan), [Jefffrey](https://github.com/Jefffrey), [Omega359](https://github.com/Omega359), [adriangb](https://github.com/adriangb), [alamb](https://github.com/alamb), [comphead](https://github.com/comphead), [etseidl](https://github.com/etseidl), [gabotechs](https://github.com/gabotechs), [geoffreyclaude](https://github.com/geoffreyclaude), [klion26](https://github.com/klion26), [rluvaton](https://github.com/rluvaton), [xudong963](https://github.com/xudong963), [zhuqi-lucas](https://github.com/zhuqi-lucas).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4031120411
Odd benchmark results. I'll confirm we're actually ending up with the operator in the plan with dfbench. Comet constructs plans manually, and the microbenchmark queries are very targeted.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2907131246
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1160 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Binary search (not `partition_point`) because `compare_join_arrays`
+/// returns `Result` and `partition_point` requires an infallible predicate
+/// on a `&[T]` — using it here would need either allocation or unsafe
+/// pointer arithmetic to recover the probe index.
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let mut lo = from + 1;
+ let mut hi = len;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. The `inner_key_buffer` field already holds the buffered
+ /// inner rows needed to resume filter evaluation.
+ FilteredPending,
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Main loop: drive the merge-scan to produce output batches.
+ // Clippy can't see that poll_next_outer_batch guarantees
+ // outer_batch = Some(..) when returning Ok(true).
+ #[expect(clippy::panicking_unwrap)]
+ fn poll_join(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<RecordBatch>>> {
+ let join_time = self.join_time.clone();
+ let _timer = join_time.timer();
+
+ loop {
+ // 1. Ensure we have an outer batch
+ if self.outer_batch.is_none() {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Outer exhausted — flush coalescer
+ self.boundary_state = BoundaryState::Normal;
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {
+ // Check if we're resuming a boundary loop that was
+ // interrupted by Pending. See BoundaryState docs.
+ match std::mem::replace(
+ &mut self.boundary_state,
+ BoundaryState::Normal,
+ ) {
+ BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ } => {
+ let same_key = keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.outer_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_no_filter()?;
+ let num_outer =
+ self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let new_saved = self.outer_key_arrays.clone();
+ let new_idx = num_outer - 1;
+ self.boundary_state =
+ BoundaryState::NoFilterPending {
+ saved_keys: new_saved,
+ saved_idx: new_idx,
+ };
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ }
+ }
+ BoundaryState::FilteredPending => {
+ if !self.inner_key_buffer.is_empty() {
+ let first_inner = &self.inner_key_buffer[0];
+ let inner_keys =
+ evaluate_join_keys(first_inner, &self.on_inner)?;
+ let same_key = keys_match(
+ &self.outer_key_arrays,
+ 0,
+ &inner_keys,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_with_filter()?;
+ let num_outer =
+ self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ self.boundary_state =
+ BoundaryState::FilteredPending;
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ }
+ }
+ self.clear_inner_key_group();
+ }
+ BoundaryState::Normal => {}
+ }
+ }
+ }
+ }
+
+ // 2. Ensure we have an inner batch (unless inner is exhausted).
+ // Skip this when in a boundary state — inner was already
+ // advanced past the key group before the boundary loop started.
+ if self.inner_batch.is_none()
+ && matches!(self.boundary_state, BoundaryState::Normal)
+ {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Inner exhausted — emit remaining outer batches.
+ // For semi: no more matches possible.
+ // For anti: all remaining outer rows are unmatched.
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ loop {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => break,
+ Ok(true) => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ }
+ }
+ }
+
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {}
+ }
+ }
+
+ // 3. Main merge-scan loop
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ if self.outer_offset >= num_outer {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ continue;
+ }
+
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ };
+ let num_inner = inner_batch.num_rows();
+
+ if self.inner_offset >= num_inner {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => continue,
+ }
+ }
+
+ // 4. Compare keys at current positions
+ let cmp = compare_join_arrays(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ &self.inner_key_arrays,
+ self.inner_offset,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ match cmp {
+ Ordering::Less => {
+ let group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ self.outer_offset = group_end;
+ }
+ Ordering::Greater => {
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if group_end >= num_inner {
+ let saved_keys = self.inner_key_arrays.clone();
+ let saved_idx = num_inner - 1;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_) => continue,
+ }
+ }
+ continue;
+ }
+ }
+ } else {
+ self.inner_offset = group_end;
+ }
+ }
+ Ordering::Equal => {
+ if self.filter.is_some() {
+ // Buffer inner key group (may span batches)
+ match ready!(self.buffer_inner_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ // Process outer rows against buffered inner group
+ // (may need to handle outer batch boundary)
+ loop {
+ self.process_key_match_with_filter()?;
+
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ if self.outer_offset >= outer_batch.num_rows() {
+ self.emit_outer_batch()?;
+ self.boundary_state = BoundaryState::FilteredPending;
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.boundary_state = BoundaryState::Normal;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ self.boundary_state = BoundaryState::Normal;
+ if !self.inner_key_buffer.is_empty() {
+ let first_inner = &self.inner_key_buffer[0];
+ let inner_keys = evaluate_join_keys(
+ first_inner,
+ &self.on_inner,
+ )?;
+ let same = keys_match(
+ &self.outer_key_arrays,
+ 0,
+ &inner_keys,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same {
+ continue;
+ }
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+
+ self.clear_inner_key_group();
+ } else {
+ // No filter: advance inner past key group, then
+ // mark all outer rows with this key as matched.
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ loop {
+ self.process_key_match_no_filter()?;
+
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let saved_keys = self.outer_key_arrays.clone();
+ let saved_idx = num_outer - 1;
+
+ self.emit_outer_batch()?;
+ // Save boundary state before polling. If
+ // poll returns Pending, poll_join exits and
+ // re-enters from the top. The saved state
+ // lets step 1 resume the boundary loop.
+ self.boundary_state = BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ };
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.boundary_state = BoundaryState::Normal;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ // Recover saved_keys from boundary state
+ let BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ } = std::mem::replace(
+ &mut self.boundary_state,
+ BoundaryState::Normal,
+ )
+ else {
+ unreachable!()
+ };
+ let same_key = keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.outer_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ continue;
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+ }
+ }
+ }
+
+ // Check for completed coalescer batch
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ }
+ }
+}
+
+/// Evaluate the filter for all rows in an inner slice against the outer group,
+/// OR-ing results into the matched bitset. Returns the updated matched count.
+/// Extracted as a free function so Rust can split borrows on the stream struct.
+#[expect(clippy::too_many_arguments)]
+fn eval_filter_for_inner_slice(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_slice: &RecordBatch,
+ matched: &mut BooleanBufferBuilder,
+ outer_offset: usize,
+ outer_group_len: usize,
+ mut matched_count: usize,
+) -> Result<usize> {
+ for inner_row in 0..inner_slice.num_rows() {
+ if matched_count == outer_group_len {
+ break;
+ }
+
+ let filter_result = evaluate_filter_for_inner_row(
+ outer_is_left,
+ filter,
+ outer_slice,
+ inner_slice,
+ inner_row,
+ )?;
+
+ // OR filter results into the matched bitset. Both sides are
+ // bit-packed [u8] buffers, so apply_bitwise_binary_op
+ // processes 64 bits per loop iteration (not 1 bit at a time).
+ //
+ // The offsets handle alignment: outer_offset is the bit
+ // position within matched where this key group starts,
+ // and filter_buf.offset() is the BooleanBuffer's internal
+ // bit offset (usually 0, but not guaranteed by Arrow).
+ let filter_buf = filter_result.values();
+ apply_bitwise_binary_op(
+ matched.as_slice_mut(),
+ outer_offset,
+ filter_buf.inner().as_slice(),
+ filter_buf.offset(),
+ outer_group_len,
+ |a, b| a | b,
+ );
+
+ // Recount matched bits after the OR. UnalignedBitChunk is
+ // zero-copy — it reads the bytes in place and uses popcnt.
+ matched_count =
+ UnalignedBitChunk::new(matched.as_slice(), outer_offset, outer_group_len)
+ .count_ones();
+ }
+ Ok(matched_count)
+}
+
+/// Compare two key rows using the sort options to determine equality.
+fn keys_match(
+ left_arrays: &[ArrayRef],
+ left_idx: usize,
+ right_arrays: &[ArrayRef],
+ right_idx: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<bool> {
+ let cmp = compare_join_arrays(
+ left_arrays,
+ left_idx,
+ right_arrays,
+ right_idx,
+ sort_options,
+ null_equality,
+ )?;
+ Ok(cmp == Ordering::Equal)
+}
+
+/// Evaluate the join filter for one inner row against a slice of outer rows.
+///
+/// Free function (not a method on SemiAntiSortMergeJoinStream) so that Rust
+/// can split the struct borrow in process_key_match_with_filter: the caller
+/// holds &mut self.matched and &self.inner_key_buffer simultaneously, which
+/// is impossible if this borrows all of &self.
+fn evaluate_filter_for_inner_row(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_batch: &RecordBatch,
+ inner_idx: usize,
+) -> Result<BooleanArray> {
+ let num_outer_rows = outer_slice.num_rows();
+
+ // Build filter input columns in the order the filter expects
+ let mut columns: Vec<ArrayRef> = Vec::with_capacity(filter.column_indices().len());
+ for col_idx in filter.column_indices() {
+ let (side_batch, side_idx) = if outer_is_left {
+ match col_idx.side {
+ JoinSide::Left => (outer_slice, None),
+ JoinSide::Right => (inner_batch, Some(inner_idx)),
+ JoinSide::None => {
+ return internal_err!("Unexpected JoinSide::None in filter");
+ }
+ }
+ } else {
+ match col_idx.side {
+ JoinSide::Left => (inner_batch, Some(inner_idx)),
+ JoinSide::Right => (outer_slice, None),
+ JoinSide::None => {
+ return internal_err!("Unexpected JoinSide::None in filter");
+ }
+ }
+ };
+
+ match side_idx {
+ None => {
+ columns.push(Arc::clone(side_batch.column(col_idx.index)));
+ }
+ Some(idx) => {
+ // Broadcasts inner scalar to N-element array. Arrow's
+ // BinaryExpr handles Scalar×Array natively via the Datum
+ // trait, but Column::evaluate always returns Array, so
+ // we'd need a custom expr to avoid this broadcast.
+ let scalar = ScalarValue::try_from_array(
+ side_batch.column(col_idx.index).as_ref(),
+ idx,
+ )?;
+ columns.push(scalar.to_array_of_size(num_outer_rows)?);
+ }
+ }
+ }
+
+ let filter_batch = RecordBatch::try_new(Arc::clone(filter.schema()), columns)?;
+ let result = filter
+ .expression()
+ .evaluate(&filter_batch)?
+ .into_array(num_outer_rows)?;
Review Comment:
I'm trying to think through if the expression ever _could_ return a Scalar in this context. Doesn't that only happen if the leaves of the expression are `Literal`?
Let me think about it.
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1160 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Binary search (not `partition_point`) because `compare_join_arrays`
+/// returns `Result` and `partition_point` requires an infallible predicate
+/// on a `&[T]` — using it here would need either allocation or unsafe
+/// pointer arithmetic to recover the probe index.
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let mut lo = from + 1;
+ let mut hi = len;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. The `inner_key_buffer` field already holds the buffered
+ /// inner rows needed to resume filter evaluation.
+ FilteredPending,
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Main loop: drive the merge-scan to produce output batches.
+ // Clippy can't see that poll_next_outer_batch guarantees
+ // outer_batch = Some(..) when returning Ok(true).
+ #[expect(clippy::panicking_unwrap)]
+ fn poll_join(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<RecordBatch>>> {
+ let join_time = self.join_time.clone();
+ let _timer = join_time.timer();
+
+ loop {
+ // 1. Ensure we have an outer batch
+ if self.outer_batch.is_none() {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Outer exhausted — flush coalescer
+ self.boundary_state = BoundaryState::Normal;
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {
+ // Check if we're resuming a boundary loop that was
+ // interrupted by Pending. See BoundaryState docs.
+ match std::mem::replace(
+ &mut self.boundary_state,
+ BoundaryState::Normal,
+ ) {
+ BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ } => {
+ let same_key = keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.outer_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_no_filter()?;
+ let num_outer =
+ self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let new_saved = self.outer_key_arrays.clone();
+ let new_idx = num_outer - 1;
+ self.boundary_state =
+ BoundaryState::NoFilterPending {
+ saved_keys: new_saved,
+ saved_idx: new_idx,
+ };
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ }
+ }
+ BoundaryState::FilteredPending => {
+ if !self.inner_key_buffer.is_empty() {
+ let first_inner = &self.inner_key_buffer[0];
+ let inner_keys =
+ evaluate_join_keys(first_inner, &self.on_inner)?;
+ let same_key = keys_match(
+ &self.outer_key_arrays,
+ 0,
+ &inner_keys,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_with_filter()?;
+ let num_outer =
+ self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ self.boundary_state =
+ BoundaryState::FilteredPending;
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ }
+ }
+ self.clear_inner_key_group();
+ }
+ BoundaryState::Normal => {}
+ }
+ }
+ }
+ }
+
+ // 2. Ensure we have an inner batch (unless inner is exhausted).
+ // Skip this when in a boundary state — inner was already
+ // advanced past the key group before the boundary loop started.
+ if self.inner_batch.is_none()
+ && matches!(self.boundary_state, BoundaryState::Normal)
+ {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Inner exhausted — emit remaining outer batches.
+ // For semi: no more matches possible.
+ // For anti: all remaining outer rows are unmatched.
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ loop {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => break,
+ Ok(true) => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ }
+ }
+ }
+
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {}
+ }
+ }
+
+ // 3. Main merge-scan loop
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ if self.outer_offset >= num_outer {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ continue;
+ }
+
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ };
+ let num_inner = inner_batch.num_rows();
+
+ if self.inner_offset >= num_inner {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => continue,
+ }
+ }
+
+ // 4. Compare keys at current positions
+ let cmp = compare_join_arrays(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ &self.inner_key_arrays,
+ self.inner_offset,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ match cmp {
+ Ordering::Less => {
+ let group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ self.outer_offset = group_end;
+ }
+ Ordering::Greater => {
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if group_end >= num_inner {
+ let saved_keys = self.inner_key_arrays.clone();
+ let saved_idx = num_inner - 1;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_) => continue,
+ }
+ }
+ continue;
+ }
+ }
+ } else {
+ self.inner_offset = group_end;
+ }
+ }
+ Ordering::Equal => {
+ if self.filter.is_some() {
+ // Buffer inner key group (may span batches)
+ match ready!(self.buffer_inner_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ // Process outer rows against buffered inner group
+ // (may need to handle outer batch boundary)
+ loop {
+ self.process_key_match_with_filter()?;
+
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ if self.outer_offset >= outer_batch.num_rows() {
+ self.emit_outer_batch()?;
+ self.boundary_state = BoundaryState::FilteredPending;
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.boundary_state = BoundaryState::Normal;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ self.boundary_state = BoundaryState::Normal;
+ if !self.inner_key_buffer.is_empty() {
+ let first_inner = &self.inner_key_buffer[0];
+ let inner_keys = evaluate_join_keys(
+ first_inner,
+ &self.on_inner,
+ )?;
+ let same = keys_match(
+ &self.outer_key_arrays,
+ 0,
+ &inner_keys,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same {
+ continue;
+ }
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+
+ self.clear_inner_key_group();
+ } else {
+ // No filter: advance inner past key group, then
+ // mark all outer rows with this key as matched.
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ loop {
+ self.process_key_match_no_filter()?;
+
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let saved_keys = self.outer_key_arrays.clone();
+ let saved_idx = num_outer - 1;
+
+ self.emit_outer_batch()?;
+ // Save boundary state before polling. If
+ // poll returns Pending, poll_join exits and
+ // re-enters from the top. The saved state
+ // lets step 1 resume the boundary loop.
+ self.boundary_state = BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ };
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.boundary_state = BoundaryState::Normal;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ // Recover saved_keys from boundary state
+ let BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ } = std::mem::replace(
+ &mut self.boundary_state,
+ BoundaryState::Normal,
+ )
+ else {
+ unreachable!()
+ };
+ let same_key = keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.outer_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ continue;
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+ }
+ }
+ }
+
+ // Check for completed coalescer batch
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ }
+ }
+}
+
+/// Evaluate the filter for all rows in an inner slice against the outer group,
+/// OR-ing results into the matched bitset. Returns the updated matched count.
+/// Extracted as a free function so Rust can split borrows on the stream struct.
+#[expect(clippy::too_many_arguments)]
+fn eval_filter_for_inner_slice(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_slice: &RecordBatch,
+ matched: &mut BooleanBufferBuilder,
+ outer_offset: usize,
+ outer_group_len: usize,
+ mut matched_count: usize,
+) -> Result<usize> {
+ for inner_row in 0..inner_slice.num_rows() {
+ if matched_count == outer_group_len {
+ break;
+ }
+
+ let filter_result = evaluate_filter_for_inner_row(
+ outer_is_left,
+ filter,
+ outer_slice,
+ inner_slice,
+ inner_row,
+ )?;
+
+ // OR filter results into the matched bitset. Both sides are
+ // bit-packed [u8] buffers, so apply_bitwise_binary_op
+ // processes 64 bits per loop iteration (not 1 bit at a time).
+ //
+ // The offsets handle alignment: outer_offset is the bit
+ // position within matched where this key group starts,
+ // and filter_buf.offset() is the BooleanBuffer's internal
+ // bit offset (usually 0, but not guaranteed by Arrow).
+ let filter_buf = filter_result.values();
+ apply_bitwise_binary_op(
+ matched.as_slice_mut(),
+ outer_offset,
+ filter_buf.inner().as_slice(),
+ filter_buf.offset(),
+ outer_group_len,
+ |a, b| a | b,
+ );
+
+ // Recount matched bits after the OR. UnalignedBitChunk is
+ // zero-copy — it reads the bytes in place and uses popcnt.
+ matched_count =
+ UnalignedBitChunk::new(matched.as_slice(), outer_offset, outer_group_len)
+ .count_ones();
+ }
+ Ok(matched_count)
+}
+
+/// Compare two key rows using the sort options to determine equality.
+fn keys_match(
+ left_arrays: &[ArrayRef],
+ left_idx: usize,
+ right_arrays: &[ArrayRef],
+ right_idx: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<bool> {
+ let cmp = compare_join_arrays(
+ left_arrays,
+ left_idx,
+ right_arrays,
+ right_idx,
+ sort_options,
+ null_equality,
+ )?;
+ Ok(cmp == Ordering::Equal)
+}
+
+/// Evaluate the join filter for one inner row against a slice of outer rows.
+///
+/// Free function (not a method on SemiAntiSortMergeJoinStream) so that Rust
+/// can split the struct borrow in process_key_match_with_filter: the caller
+/// holds &mut self.matched and &self.inner_key_buffer simultaneously, which
+/// is impossible if this borrows all of &self.
+fn evaluate_filter_for_inner_row(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_batch: &RecordBatch,
+ inner_idx: usize,
+) -> Result<BooleanArray> {
+ let num_outer_rows = outer_slice.num_rows();
+
+ // Build filter input columns in the order the filter expects
+ let mut columns: Vec<ArrayRef> = Vec::with_capacity(filter.column_indices().len());
+ for col_idx in filter.column_indices() {
+ let (side_batch, side_idx) = if outer_is_left {
+ match col_idx.side {
+ JoinSide::Left => (outer_slice, None),
+ JoinSide::Right => (inner_batch, Some(inner_idx)),
+ JoinSide::None => {
+ return internal_err!("Unexpected JoinSide::None in filter");
+ }
+ }
+ } else {
+ match col_idx.side {
+ JoinSide::Left => (inner_batch, Some(inner_idx)),
+ JoinSide::Right => (outer_slice, None),
+ JoinSide::None => {
+ return internal_err!("Unexpected JoinSide::None in filter");
+ }
+ }
+ };
+
+ match side_idx {
+ None => {
+ columns.push(Arc::clone(side_batch.column(col_idx.index)));
+ }
+ Some(idx) => {
+ // Broadcasts inner scalar to N-element array. Arrow's
+ // BinaryExpr handles Scalar×Array natively via the Datum
+ // trait, but Column::evaluate always returns Array, so
+ // we'd need a custom expr to avoid this broadcast.
+ let scalar = ScalarValue::try_from_array(
+ side_batch.column(col_idx.index).as_ref(),
+ idx,
+ )?;
+ columns.push(scalar.to_array_of_size(num_outer_rows)?);
+ }
+ }
+ }
+
+ let filter_batch = RecordBatch::try_new(Arc::clone(filter.schema()), columns)?;
+ let result = filter
+ .expression()
+ .evaluate(&filter_batch)?
+ .into_array(num_outer_rows)?;
Review Comment:
I'm trying to think through if the expression ever _could_ return a scalar in this context. Doesn't that only happen if the leaves of the expression are `Literal`?
Let me think about it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2907131246
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1160 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Binary search (not `partition_point`) because `compare_join_arrays`
+/// returns `Result` and `partition_point` requires an infallible predicate
+/// on a `&[T]` — using it here would need either allocation or unsafe
+/// pointer arithmetic to recover the probe index.
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let mut lo = from + 1;
+ let mut hi = len;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. The `inner_key_buffer` field already holds the buffered
+ /// inner rows needed to resume filter evaluation.
+ FilteredPending,
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Main loop: drive the merge-scan to produce output batches.
+ // Clippy can't see that poll_next_outer_batch guarantees
+ // outer_batch = Some(..) when returning Ok(true).
+ #[expect(clippy::panicking_unwrap)]
+ fn poll_join(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<RecordBatch>>> {
+ let join_time = self.join_time.clone();
+ let _timer = join_time.timer();
+
+ loop {
+ // 1. Ensure we have an outer batch
+ if self.outer_batch.is_none() {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Outer exhausted — flush coalescer
+ self.boundary_state = BoundaryState::Normal;
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {
+ // Check if we're resuming a boundary loop that was
+ // interrupted by Pending. See BoundaryState docs.
+ match std::mem::replace(
+ &mut self.boundary_state,
+ BoundaryState::Normal,
+ ) {
+ BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ } => {
+ let same_key = keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.outer_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_no_filter()?;
+ let num_outer =
+ self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let new_saved = self.outer_key_arrays.clone();
+ let new_idx = num_outer - 1;
+ self.boundary_state =
+ BoundaryState::NoFilterPending {
+ saved_keys: new_saved,
+ saved_idx: new_idx,
+ };
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ }
+ }
+ BoundaryState::FilteredPending => {
+ if !self.inner_key_buffer.is_empty() {
+ let first_inner = &self.inner_key_buffer[0];
+ let inner_keys =
+ evaluate_join_keys(first_inner, &self.on_inner)?;
+ let same_key = keys_match(
+ &self.outer_key_arrays,
+ 0,
+ &inner_keys,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_with_filter()?;
+ let num_outer =
+ self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ self.boundary_state =
+ BoundaryState::FilteredPending;
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ }
+ }
+ self.clear_inner_key_group();
+ }
+ BoundaryState::Normal => {}
+ }
+ }
+ }
+ }
+
+ // 2. Ensure we have an inner batch (unless inner is exhausted).
+ // Skip this when in a boundary state — inner was already
+ // advanced past the key group before the boundary loop started.
+ if self.inner_batch.is_none()
+ && matches!(self.boundary_state, BoundaryState::Normal)
+ {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Inner exhausted — emit remaining outer batches.
+ // For semi: no more matches possible.
+ // For anti: all remaining outer rows are unmatched.
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ loop {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => break,
+ Ok(true) => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ }
+ }
+ }
+
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {}
+ }
+ }
+
+ // 3. Main merge-scan loop
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ if self.outer_offset >= num_outer {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ continue;
+ }
+
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ };
+ let num_inner = inner_batch.num_rows();
+
+ if self.inner_offset >= num_inner {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => continue,
+ }
+ }
+
+ // 4. Compare keys at current positions
+ let cmp = compare_join_arrays(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ &self.inner_key_arrays,
+ self.inner_offset,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ match cmp {
+ Ordering::Less => {
+ let group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ self.outer_offset = group_end;
+ }
+ Ordering::Greater => {
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if group_end >= num_inner {
+ let saved_keys = self.inner_key_arrays.clone();
+ let saved_idx = num_inner - 1;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_) => continue,
+ }
+ }
+ continue;
+ }
+ }
+ } else {
+ self.inner_offset = group_end;
+ }
+ }
+ Ordering::Equal => {
+ if self.filter.is_some() {
+ // Buffer inner key group (may span batches)
+ match ready!(self.buffer_inner_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ // Process outer rows against buffered inner group
+ // (may need to handle outer batch boundary)
+ loop {
+ self.process_key_match_with_filter()?;
+
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ if self.outer_offset >= outer_batch.num_rows() {
+ self.emit_outer_batch()?;
+ self.boundary_state = BoundaryState::FilteredPending;
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.boundary_state = BoundaryState::Normal;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ self.boundary_state = BoundaryState::Normal;
+ if !self.inner_key_buffer.is_empty() {
+ let first_inner = &self.inner_key_buffer[0];
+ let inner_keys = evaluate_join_keys(
+ first_inner,
+ &self.on_inner,
+ )?;
+ let same = keys_match(
+ &self.outer_key_arrays,
+ 0,
+ &inner_keys,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same {
+ continue;
+ }
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+
+ self.clear_inner_key_group();
+ } else {
+ // No filter: advance inner past key group, then
+ // mark all outer rows with this key as matched.
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ loop {
+ self.process_key_match_no_filter()?;
+
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let saved_keys = self.outer_key_arrays.clone();
+ let saved_idx = num_outer - 1;
+
+ self.emit_outer_batch()?;
+ // Save boundary state before polling. If
+ // poll returns Pending, poll_join exits and
+ // re-enters from the top. The saved state
+ // lets step 1 resume the boundary loop.
+ self.boundary_state = BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ };
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.boundary_state = BoundaryState::Normal;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ // Recover saved_keys from boundary state
+ let BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ } = std::mem::replace(
+ &mut self.boundary_state,
+ BoundaryState::Normal,
+ )
+ else {
+ unreachable!()
+ };
+ let same_key = keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.outer_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ continue;
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+ }
+ }
+ }
+
+ // Check for completed coalescer batch
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ }
+ }
+}
+
+/// Evaluate the filter for all rows in an inner slice against the outer group,
+/// OR-ing results into the matched bitset. Returns the updated matched count.
+/// Extracted as a free function so Rust can split borrows on the stream struct.
+#[expect(clippy::too_many_arguments)]
+fn eval_filter_for_inner_slice(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_slice: &RecordBatch,
+ matched: &mut BooleanBufferBuilder,
+ outer_offset: usize,
+ outer_group_len: usize,
+ mut matched_count: usize,
+) -> Result<usize> {
+ for inner_row in 0..inner_slice.num_rows() {
+ if matched_count == outer_group_len {
+ break;
+ }
+
+ let filter_result = evaluate_filter_for_inner_row(
+ outer_is_left,
+ filter,
+ outer_slice,
+ inner_slice,
+ inner_row,
+ )?;
+
+ // OR filter results into the matched bitset. Both sides are
+ // bit-packed [u8] buffers, so apply_bitwise_binary_op
+ // processes 64 bits per loop iteration (not 1 bit at a time).
+ //
+ // The offsets handle alignment: outer_offset is the bit
+ // position within matched where this key group starts,
+ // and filter_buf.offset() is the BooleanBuffer's internal
+ // bit offset (usually 0, but not guaranteed by Arrow).
+ let filter_buf = filter_result.values();
+ apply_bitwise_binary_op(
+ matched.as_slice_mut(),
+ outer_offset,
+ filter_buf.inner().as_slice(),
+ filter_buf.offset(),
+ outer_group_len,
+ |a, b| a | b,
+ );
+
+ // Recount matched bits after the OR. UnalignedBitChunk is
+ // zero-copy — it reads the bytes in place and uses popcnt.
+ matched_count =
+ UnalignedBitChunk::new(matched.as_slice(), outer_offset, outer_group_len)
+ .count_ones();
+ }
+ Ok(matched_count)
+}
+
+/// Compare two key rows using the sort options to determine equality.
+fn keys_match(
+ left_arrays: &[ArrayRef],
+ left_idx: usize,
+ right_arrays: &[ArrayRef],
+ right_idx: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<bool> {
+ let cmp = compare_join_arrays(
+ left_arrays,
+ left_idx,
+ right_arrays,
+ right_idx,
+ sort_options,
+ null_equality,
+ )?;
+ Ok(cmp == Ordering::Equal)
+}
+
+/// Evaluate the join filter for one inner row against a slice of outer rows.
+///
+/// Free function (not a method on SemiAntiSortMergeJoinStream) so that Rust
+/// can split the struct borrow in process_key_match_with_filter: the caller
+/// holds &mut self.matched and &self.inner_key_buffer simultaneously, which
+/// is impossible if this borrows all of &self.
+fn evaluate_filter_for_inner_row(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_batch: &RecordBatch,
+ inner_idx: usize,
+) -> Result<BooleanArray> {
+ let num_outer_rows = outer_slice.num_rows();
+
+ // Build filter input columns in the order the filter expects
+ let mut columns: Vec<ArrayRef> = Vec::with_capacity(filter.column_indices().len());
+ for col_idx in filter.column_indices() {
+ let (side_batch, side_idx) = if outer_is_left {
+ match col_idx.side {
+ JoinSide::Left => (outer_slice, None),
+ JoinSide::Right => (inner_batch, Some(inner_idx)),
+ JoinSide::None => {
+ return internal_err!("Unexpected JoinSide::None in filter");
+ }
+ }
+ } else {
+ match col_idx.side {
+ JoinSide::Left => (inner_batch, Some(inner_idx)),
+ JoinSide::Right => (outer_slice, None),
+ JoinSide::None => {
+ return internal_err!("Unexpected JoinSide::None in filter");
+ }
+ }
+ };
+
+ match side_idx {
+ None => {
+ columns.push(Arc::clone(side_batch.column(col_idx.index)));
+ }
+ Some(idx) => {
+ // Broadcasts inner scalar to N-element array. Arrow's
+ // BinaryExpr handles Scalar×Array natively via the Datum
+ // trait, but Column::evaluate always returns Array, so
+ // we'd need a custom expr to avoid this broadcast.
+ let scalar = ScalarValue::try_from_array(
+ side_batch.column(col_idx.index).as_ref(),
+ idx,
+ )?;
+ columns.push(scalar.to_array_of_size(num_outer_rows)?);
+ }
+ }
+ }
+
+ let filter_batch = RecordBatch::try_new(Arc::clone(filter.schema()), columns)?;
+ let result = filter
+ .expression()
+ .evaluate(&filter_batch)?
+ .into_array(num_outer_rows)?;
Review Comment:
I'm trying to think through if the expression ever _could_ return a scalar in this context.
Let me think about it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2906537029
##########
benchmarks/src/smj.rs:
##########
@@ -277,6 +277,7 @@ const SMJ_QUERIES: &[&str] = &[
WHERE EXISTS (
SELECT 1 FROM t2_sorted
WHERE t2_sorted.key = t1_sorted.key
+ AND t2_sorted.data <> t1_sorted.data
Review Comment:
I added this benchmark in #18875 but neglected to apply the filter correctly. Now this 90% scenario has the same logic as the 10% and 50% filters above it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2906790523
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1160 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Binary search (not `partition_point`) because `compare_join_arrays`
+/// returns `Result` and `partition_point` requires an infallible predicate
+/// on a `&[T]` — using it here would need either allocation or unsafe
+/// pointer arithmetic to recover the probe index.
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let mut lo = from + 1;
+ let mut hi = len;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. The `inner_key_buffer` field already holds the buffered
+ /// inner rows needed to resume filter evaluation.
+ FilteredPending,
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Main loop: drive the merge-scan to produce output batches.
+ // Clippy can't see that poll_next_outer_batch guarantees
+ // outer_batch = Some(..) when returning Ok(true).
+ #[expect(clippy::panicking_unwrap)]
+ fn poll_join(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<RecordBatch>>> {
+ let join_time = self.join_time.clone();
+ let _timer = join_time.timer();
+
+ loop {
+ // 1. Ensure we have an outer batch
+ if self.outer_batch.is_none() {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Outer exhausted — flush coalescer
+ self.boundary_state = BoundaryState::Normal;
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {
+ // Check if we're resuming a boundary loop that was
+ // interrupted by Pending. See BoundaryState docs.
+ match std::mem::replace(
+ &mut self.boundary_state,
+ BoundaryState::Normal,
+ ) {
+ BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ } => {
+ let same_key = keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.outer_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_no_filter()?;
+ let num_outer =
+ self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let new_saved = self.outer_key_arrays.clone();
+ let new_idx = num_outer - 1;
+ self.boundary_state =
+ BoundaryState::NoFilterPending {
+ saved_keys: new_saved,
+ saved_idx: new_idx,
+ };
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ }
+ }
+ BoundaryState::FilteredPending => {
+ if !self.inner_key_buffer.is_empty() {
+ let first_inner = &self.inner_key_buffer[0];
+ let inner_keys =
+ evaluate_join_keys(first_inner, &self.on_inner)?;
+ let same_key = keys_match(
+ &self.outer_key_arrays,
+ 0,
+ &inner_keys,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ self.process_key_match_with_filter()?;
+ let num_outer =
+ self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ self.boundary_state =
+ BoundaryState::FilteredPending;
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ }
+ }
+ self.clear_inner_key_group();
+ }
+ BoundaryState::Normal => {}
+ }
+ }
+ }
+ }
+
+ // 2. Ensure we have an inner batch (unless inner is exhausted).
+ // Skip this when in a boundary state — inner was already
+ // advanced past the key group before the boundary loop started.
+ if self.inner_batch.is_none()
+ && matches!(self.boundary_state, BoundaryState::Normal)
+ {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ // Inner exhausted — emit remaining outer batches.
+ // For semi: no more matches possible.
+ // For anti: all remaining outer rows are unmatched.
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ loop {
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => break,
+ Ok(true) => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ }
+ }
+ }
+
+ self.coalescer.finish_buffered_batch()?;
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ return Poll::Ready(Ok(None));
+ }
+ Ok(true) => {}
+ }
+ }
+
+ // 3. Main merge-scan loop
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ if self.outer_offset >= num_outer {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ continue;
+ }
+
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => {
+ self.emit_outer_batch()?;
+ self.outer_batch = None;
+ continue;
+ }
+ };
+ let num_inner = inner_batch.num_rows();
+
+ if self.inner_offset >= num_inner {
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => continue,
+ }
+ }
+
+ // 4. Compare keys at current positions
+ let cmp = compare_join_arrays(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ &self.inner_key_arrays,
+ self.inner_offset,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ match cmp {
+ Ordering::Less => {
+ let group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ self.outer_offset = group_end;
+ }
+ Ordering::Greater => {
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if group_end >= num_inner {
+ let saved_keys = self.inner_key_arrays.clone();
+ let saved_idx = num_inner - 1;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.inner_batch = None;
+ continue;
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_) => continue,
+ }
+ }
+ continue;
+ }
+ }
+ } else {
+ self.inner_offset = group_end;
+ }
+ }
+ Ordering::Equal => {
+ if self.filter.is_some() {
+ // Buffer inner key group (may span batches)
+ match ready!(self.buffer_inner_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ // Process outer rows against buffered inner group
+ // (may need to handle outer batch boundary)
+ loop {
+ self.process_key_match_with_filter()?;
+
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ if self.outer_offset >= outer_batch.num_rows() {
+ self.emit_outer_batch()?;
+ self.boundary_state = BoundaryState::FilteredPending;
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.boundary_state = BoundaryState::Normal;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ self.boundary_state = BoundaryState::Normal;
+ if !self.inner_key_buffer.is_empty() {
+ let first_inner = &self.inner_key_buffer[0];
+ let inner_keys = evaluate_join_keys(
+ first_inner,
+ &self.on_inner,
+ )?;
+ let same = keys_match(
+ &self.outer_key_arrays,
+ 0,
+ &inner_keys,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same {
+ continue;
+ }
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+
+ self.clear_inner_key_group();
+ } else {
+ // No filter: advance inner past key group, then
+ // mark all outer rows with this key as matched.
+ match ready!(self.advance_inner_past_key_group(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(_inner_exhausted) => {}
+ }
+
+ loop {
+ self.process_key_match_no_filter()?;
+
+ let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
+ if self.outer_offset >= num_outer {
+ let saved_keys = self.outer_key_arrays.clone();
+ let saved_idx = num_outer - 1;
+
+ self.emit_outer_batch()?;
+ // Save boundary state before polling. If
+ // poll returns Pending, poll_join exits and
+ // re-enters from the top. The saved state
+ // lets step 1 resume the boundary loop.
+ self.boundary_state = BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ };
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.boundary_state = BoundaryState::Normal;
+ self.outer_batch = None;
+ break;
+ }
+ Ok(true) => {
+ // Recover saved_keys from boundary state
+ let BoundaryState::NoFilterPending {
+ saved_keys,
+ saved_idx,
+ } = std::mem::replace(
+ &mut self.boundary_state,
+ BoundaryState::Normal,
+ )
+ else {
+ unreachable!()
+ };
+ let same_key = keys_match(
+ &saved_keys,
+ saved_idx,
+ &self.outer_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same_key {
+ continue;
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+ }
+ }
+ }
+
+ // Check for completed coalescer batch
+ if let Some(batch) = self.coalescer.next_completed_batch() {
+ return Poll::Ready(Ok(Some(batch)));
+ }
+ }
+ }
+}
+
+/// Evaluate the filter for all rows in an inner slice against the outer group,
+/// OR-ing results into the matched bitset. Returns the updated matched count.
+/// Extracted as a free function so Rust can split borrows on the stream struct.
+#[expect(clippy::too_many_arguments)]
+fn eval_filter_for_inner_slice(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_slice: &RecordBatch,
+ matched: &mut BooleanBufferBuilder,
+ outer_offset: usize,
+ outer_group_len: usize,
+ mut matched_count: usize,
+) -> Result<usize> {
+ for inner_row in 0..inner_slice.num_rows() {
+ if matched_count == outer_group_len {
+ break;
+ }
+
+ let filter_result = evaluate_filter_for_inner_row(
+ outer_is_left,
+ filter,
+ outer_slice,
+ inner_slice,
+ inner_row,
+ )?;
+
+ // OR filter results into the matched bitset. Both sides are
+ // bit-packed [u8] buffers, so apply_bitwise_binary_op
+ // processes 64 bits per loop iteration (not 1 bit at a time).
+ //
+ // The offsets handle alignment: outer_offset is the bit
+ // position within matched where this key group starts,
+ // and filter_buf.offset() is the BooleanBuffer's internal
+ // bit offset (usually 0, but not guaranteed by Arrow).
+ let filter_buf = filter_result.values();
+ apply_bitwise_binary_op(
+ matched.as_slice_mut(),
+ outer_offset,
+ filter_buf.inner().as_slice(),
+ filter_buf.offset(),
+ outer_group_len,
+ |a, b| a | b,
+ );
+
+ // Recount matched bits after the OR. UnalignedBitChunk is
+ // zero-copy — it reads the bytes in place and uses popcnt.
+ matched_count =
+ UnalignedBitChunk::new(matched.as_slice(), outer_offset, outer_group_len)
+ .count_ones();
+ }
+ Ok(matched_count)
+}
+
+/// Compare two key rows using the sort options to determine equality.
+fn keys_match(
+ left_arrays: &[ArrayRef],
+ left_idx: usize,
+ right_arrays: &[ArrayRef],
+ right_idx: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<bool> {
+ let cmp = compare_join_arrays(
+ left_arrays,
+ left_idx,
+ right_arrays,
+ right_idx,
+ sort_options,
+ null_equality,
+ )?;
+ Ok(cmp == Ordering::Equal)
+}
+
+/// Evaluate the join filter for one inner row against a slice of outer rows.
+///
+/// Free function (not a method on SemiAntiSortMergeJoinStream) so that Rust
+/// can split the struct borrow in process_key_match_with_filter: the caller
+/// holds &mut self.matched and &self.inner_key_buffer simultaneously, which
+/// is impossible if this borrows all of &self.
+fn evaluate_filter_for_inner_row(
+ outer_is_left: bool,
+ filter: &JoinFilter,
+ outer_slice: &RecordBatch,
+ inner_batch: &RecordBatch,
+ inner_idx: usize,
+) -> Result<BooleanArray> {
+ let num_outer_rows = outer_slice.num_rows();
+
+ // Build filter input columns in the order the filter expects
+ let mut columns: Vec<ArrayRef> = Vec::with_capacity(filter.column_indices().len());
+ for col_idx in filter.column_indices() {
+ let (side_batch, side_idx) = if outer_is_left {
+ match col_idx.side {
+ JoinSide::Left => (outer_slice, None),
+ JoinSide::Right => (inner_batch, Some(inner_idx)),
+ JoinSide::None => {
+ return internal_err!("Unexpected JoinSide::None in filter");
+ }
+ }
+ } else {
+ match col_idx.side {
+ JoinSide::Left => (inner_batch, Some(inner_idx)),
+ JoinSide::Right => (outer_slice, None),
+ JoinSide::None => {
+ return internal_err!("Unexpected JoinSide::None in filter");
+ }
+ }
+ };
+
+ match side_idx {
+ None => {
+ columns.push(Arc::clone(side_batch.column(col_idx.index)));
+ }
+ Some(idx) => {
+ // Broadcasts inner scalar to N-element array. Arrow's
+ // BinaryExpr handles Scalar×Array natively via the Datum
+ // trait, but Column::evaluate always returns Array, so
+ // we'd need a custom expr to avoid this broadcast.
+ let scalar = ScalarValue::try_from_array(
+ side_batch.column(col_idx.index).as_ref(),
+ idx,
+ )?;
+ columns.push(scalar.to_array_of_size(num_outer_rows)?);
+ }
+ }
+ }
+
+ let filter_batch = RecordBatch::try_new(Arc::clone(filter.schema()), columns)?;
+ let result = filter
+ .expression()
+ .evaluate(&filter_batch)?
+ .into_array(num_outer_rows)?;
Review Comment:
this can be another PR but if the expression returned scalar we can use that optimization later...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2907668997
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1218 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// outer batch so we can compare with the next batch's first key
+ /// without reading from the inner key buffer (which may have been
+ /// spilled to disk).
+ FilteredPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Continue processing an outer key group that spans multiple outer
+ /// batches. Returns `true` if this outer batch was fully consumed
+ /// by the key group and the caller should load another.
+ fn resume_boundary(&mut self) -> Result<bool> {
+ debug_assert!(
+ self.outer_batch.is_some(),
+ "caller must load outer_batch first"
+ );
+ match std::mem::replace(&mut self.boundary_state, BoundaryState::Normal) {
Review Comment:
Yeah, the `Normal` in `std::mem::replace` is just a placeholder. However, at lines 952 and 1017, `Normal` is also the correct final state (we consumed the pending data). At line 687, it's either correct (for the
`Normal` and non-boundary-continuing cases) or gets overwritten at lines 706/738.
I could add a comment on each `replace` call to clarify the intent, but I'd avoid adding a `TmpTransition` variant — it adds a match arm everywhere, turns a compile-time guarantee into a runtime panic, and the code paths already ensure `Normal` is correct.
Alternatively, since `boundary_state` is only read via `match`, we could wrap it in `Option<BoundaryState>` and `.take()`. But that adds `Option` unwrapping everywhere else. I want to make sure this is maintainable for other folks, but I don't love an additional state. Especially one that is purely ephemeral and shouldn't ever really `match` to anything.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2907668997
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1218 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
+
+ // Binary search the interior: key at `next` matches, key at `last` doesn't.
+ let mut lo = next + 1;
+ let mut hi = last;
+ while lo < hi {
+ let mid = lo + (hi - lo) / 2;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ mid,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+ Ok(lo)
+}
+
+/// Tracks whether we're mid-key-group when `poll_next_outer_batch` returns
+/// `Poll::Pending` inside the Equal branch's boundary loop.
+///
+/// When an outer key group spans a batch boundary, the boundary loop emits
+/// the current batch, then polls for the next. If that poll returns Pending,
+/// `ready!` exits `poll_join` and we re-enter from the top on the next call.
+/// Without this state, the new batch would be processed fresh by the
+/// merge-scan — but inner already advanced past this key, so the matching
+/// outer rows would be skipped via `Ordering::Less` and never marked.
+///
+/// This enum saves the context needed to resume the boundary loop on
+/// re-entry: compare the new batch's first key with the saved key to
+/// decide whether to continue marking or move on.
+#[derive(Debug)]
+enum BoundaryState {
+ /// Normal processing — not inside a boundary poll.
+ Normal,
+ /// The no-filter boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// batch so we can compare with the next batch's first key.
+ NoFilterPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+ /// The filtered boundary loop's `poll_next_outer_batch` returned
+ /// Pending. Carries the key arrays and index from the last emitted
+ /// outer batch so we can compare with the next batch's first key
+ /// without reading from the inner key buffer (which may have been
+ /// spilled to disk).
+ FilteredPending {
+ saved_keys: Vec<ArrayRef>,
+ saved_idx: usize,
+ },
+}
+
+pub(super) struct SemiAntiSortMergeJoinStream {
+ /// true for semi (emit matched), false for anti (emit unmatched)
+ is_semi: bool,
+
+ // Input streams — "outer" is the streamed side whose rows we output,
+ // "inner" is the buffered side we match against.
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+
+ // Current batches and cursor positions
+ outer_batch: Option<RecordBatch>,
+ outer_offset: usize,
+ outer_key_arrays: Vec<ArrayRef>,
+ inner_batch: Option<RecordBatch>,
+ inner_offset: usize,
+ inner_key_arrays: Vec<ArrayRef>,
+
+ // Per-outer-batch match tracking, reused across batches.
+ // Bit-packed (not Vec<bool>) so that:
+ // - emit: finish() yields a BooleanBuffer directly (no packing iteration)
+ // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks
+ // - count: UnalignedBitChunk::count_ones uses popcnt
+ matched: BooleanBufferBuilder,
+
+ // Inner key group buffer: all inner rows sharing the current join key.
+ // Only populated when a filter is present. Unbounded — a single key
+ // with many inner rows will buffer them all. See "Degenerate cases"
+ // in exec.rs. Spilled to disk when memory reservation fails.
+ inner_key_buffer: Vec<RecordBatch>,
+ inner_key_spill: Option<RefCountedTempFile>,
+
+ // True when buffer_inner_key_group returned Pending after partially
+ // filling inner_key_buffer. On re-entry, buffer_inner_key_group
+ // must skip clear() and resume from poll_next_inner_batch (the
+ // current inner_batch was already sliced and pushed before Pending).
+ buffering_inner_pending: bool,
+
+ // Boundary re-entry state — see BoundaryState doc comment.
+ boundary_state: BoundaryState,
+
+ // Join condition
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ // When join_type is RightSemi/RightAnti, outer=right, inner=left,
+ // so we need to swap sides when building the filter batch.
+ outer_is_left: bool,
+
+ // Output
+ coalescer: BatchCoalescer,
+ schema: SchemaRef,
+
+ // Metrics
+ join_time: crate::metrics::Time,
+ input_batches: Count,
+ input_rows: Count,
+ baseline_metrics: BaselineMetrics,
+ peak_mem_used: Gauge,
+
+ // Memory / spill — only the inner key buffer is tracked via reservation,
+ // matching existing SMJ (which tracks only the buffered side). The outer
+ // batch is a single batch at a time and cannot be spilled.
+ reservation: MemoryReservation,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ inner_buffer_size: usize,
+
+ // True once the current outer batch has been emitted. The Equal
+ // branch's inner loops call emit then `ready!(poll_next_outer_batch)`.
+ // If that poll returns Pending, poll_join re-enters from the top
+ // on the next poll — with outer_batch still Some and outer_offset
+ // past the end. The main loop's step 3 would re-emit without this
+ // guard. Cleared when poll_next_outer_batch loads a new batch.
+ batch_emitted: bool,
+}
+
+impl SemiAntiSortMergeJoinStream {
+ #[expect(clippy::too_many_arguments)]
+ pub fn try_new(
+ schema: SchemaRef,
+ sort_options: Vec<SortOptions>,
+ null_equality: NullEquality,
+ outer: SendableRecordBatchStream,
+ inner: SendableRecordBatchStream,
+ on_outer: Vec<PhysicalExprRef>,
+ on_inner: Vec<PhysicalExprRef>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ batch_size: usize,
+ partition: usize,
+ metrics: &ExecutionPlanMetricsSet,
+ reservation: MemoryReservation,
+ peak_mem_used: Gauge,
+ spill_manager: SpillManager,
+ runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>,
+ ) -> Result<Self> {
+ let is_semi = matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi);
+ let outer_is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti);
+
+ let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition);
+ let input_batches =
+ MetricBuilder::new(metrics).counter("input_batches", partition);
+ let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition);
+ let baseline_metrics = BaselineMetrics::new(metrics, partition);
+
+ Ok(Self {
+ is_semi,
+ outer,
+ inner,
+ outer_batch: None,
+ outer_offset: 0,
+ outer_key_arrays: vec![],
+ inner_batch: None,
+ inner_offset: 0,
+ inner_key_arrays: vec![],
+ matched: BooleanBufferBuilder::new(0),
+ inner_key_buffer: vec![],
+ inner_key_spill: None,
+ buffering_inner_pending: false,
+ boundary_state: BoundaryState::Normal,
+ on_outer,
+ on_inner,
+ filter,
+ sort_options,
+ null_equality,
+ outer_is_left,
+ coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size)
+ .with_biggest_coalesce_batch_size(Some(batch_size / 2)),
+ schema,
+ join_time,
+ input_batches,
+ input_rows,
+ baseline_metrics,
+ peak_mem_used,
+ reservation,
+ spill_manager,
+ runtime_env,
+ inner_buffer_size: 0,
+ batch_emitted: false,
+ })
+ }
+
+ /// Resize the memory reservation to match current tracked usage.
+ fn try_resize_reservation(&mut self) -> Result<()> {
+ let needed = self.inner_buffer_size;
+ self.reservation.try_resize(needed)?;
+ self.peak_mem_used.set_max(self.reservation.size());
+ Ok(())
+ }
+
+ /// Spill the in-memory inner key buffer to disk and clear it.
+ fn spill_inner_key_buffer(&mut self) -> Result<()> {
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_and_finish(
+ &self.inner_key_buffer,
+ "semi_anti_smj_inner_key_spill",
+ )?
+ .expect("inner_key_buffer is non-empty when spilling");
+ self.inner_key_buffer.clear();
+ self.inner_buffer_size = 0;
+ self.inner_key_spill = Some(spill_file);
+ // Should succeed now — inner buffer has been spilled.
+ self.try_resize_reservation()
+ }
+
+ /// Clear inner key group state after processing. Does not resize the
+ /// reservation — the next key group will resize when buffering, or
+ /// the stream's Drop will free it. This avoids unnecessary memory
+ /// pool interactions (see apache/datafusion#20729).
+ fn clear_inner_key_group(&mut self) {
+ self.inner_key_buffer.clear();
+ self.inner_key_spill = None;
+ self.inner_buffer_size = 0;
+ }
+
+ /// Poll for the next outer batch. Returns true if a batch was loaded.
+ fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.outer.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_outer)?;
+ let num_rows = batch.num_rows();
+ self.outer_batch = Some(batch);
+ self.outer_offset = 0;
+ self.outer_key_arrays = keys;
+ self.batch_emitted = false;
+ self.matched = BooleanBufferBuilder::new(num_rows);
+ self.matched.append_n(num_rows, false);
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Poll for the next inner batch. Returns true if a batch was loaded.
+ fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ loop {
+ match ready!(self.inner.poll_next_unpin(cx)) {
+ None => return Poll::Ready(Ok(false)),
+ Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let keys = evaluate_join_keys(&batch, &self.on_inner)?;
+ self.inner_batch = Some(batch);
+ self.inner_offset = 0;
+ self.inner_key_arrays = keys;
+ return Poll::Ready(Ok(true));
+ }
+ }
+ }
+ }
+
+ /// Emit the current outer batch through the coalescer, applying the
+ /// matched bitset as a selection mask. No-op if already emitted
+ /// (see `batch_emitted` field).
+ fn emit_outer_batch(&mut self) -> Result<()> {
+ if self.batch_emitted {
+ return Ok(());
+ }
+ self.batch_emitted = true;
+
+ let batch = self.outer_batch.as_ref().unwrap();
+
+ // finish() converts the bit-packed builder directly to a
+ // BooleanBuffer — no iteration or repacking needed.
+ let selection = BooleanArray::new(self.matched.finish(), None);
+
+ let selection = if self.is_semi {
+ selection
+ } else {
+ not(&selection)?
+ };
+
+ let filtered = filter_record_batch(batch, &selection)?;
+ if filtered.num_rows() > 0 {
+ self.coalescer.push_batch(filtered)?;
+ }
+ Ok(())
+ }
+
+ /// Process a key match between outer and inner sides (no filter).
+ /// Sets matched bits for all outer rows sharing the current key.
+ fn process_key_match_no_filter(&mut self) -> Result<()> {
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ for i in self.outer_offset..outer_group_end {
+ self.matched.set_bit(i, true);
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Advance inner past the current key group. Returns Ok(true) if inner
+ /// is exhausted.
+ fn advance_inner_past_key_group(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<bool>> {
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+
+ // Key group extends to end of batch — need to check next batch
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Buffer inner key group for filter evaluation. Collects all inner rows
+ /// with the current key across batch boundaries.
+ ///
+ /// If poll_next_inner_batch returns Pending, we save progress via
+ /// buffering_inner_pending. On re-entry (from the Equal branch in
+ /// poll_join), we skip clear() and the slice+push for the current
+ /// batch (which was already buffered before Pending), and go directly
+ /// to polling for the next inner batch.
+ fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
+ // On re-entry after Pending: don't clear the partially-filled
+ // buffer. The current inner_batch was already sliced and pushed
+ // before Pending, so jump to polling for the next batch.
+ let mut resume_from_poll = false;
+ if self.buffering_inner_pending {
+ self.buffering_inner_pending = false;
+ resume_from_poll = true;
+ } else {
+ self.clear_inner_key_group();
+ }
+
+ loop {
+ let inner_batch = match &self.inner_batch {
+ Some(b) => b,
+ None => return Poll::Ready(Ok(true)),
+ };
+ let num_inner = inner_batch.num_rows();
+ let group_end = find_key_group_end(
+ &self.inner_key_arrays,
+ self.inner_offset,
+ num_inner,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+
+ if !resume_from_poll {
+ let slice =
+ inner_batch.slice(self.inner_offset, group_end - self.inner_offset);
+ self.inner_buffer_size += slice.get_array_memory_size();
+ self.inner_key_buffer.push(slice);
+
+ // Reserve memory for the newly buffered slice. If the pool
+ // is exhausted, spill the entire buffer to disk.
+ if self.try_resize_reservation().is_err() {
+ if self.runtime_env.disk_manager.tmp_files_enabled() {
+ self.spill_inner_key_buffer()?;
+ } else {
+ // Re-attempt to get the error message
+ self.try_resize_reservation().map_err(|e| {
+ datafusion_common::DataFusionError::Execution(format!(
+ "{e}. Disk spilling disabled."
+ ))
+ })?;
+ }
+ }
+
+ if group_end < num_inner {
+ self.inner_offset = group_end;
+ return Poll::Ready(Ok(false));
+ }
+ }
+ resume_from_poll = false;
+
+ // Key group extends to end of batch — check next
+ let last_key_idx = num_inner - 1;
+ let saved_inner_keys = self.inner_key_arrays.clone();
+
+ // If poll returns Pending, the current batch is already
+ // in inner_key_buffer.
+ self.buffering_inner_pending = true;
+ match ready!(self.poll_next_inner_batch(cx)) {
+ Err(e) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Err(e));
+ }
+ Ok(false) => {
+ self.buffering_inner_pending = false;
+ return Poll::Ready(Ok(true));
+ }
+ Ok(true) => {
+ self.buffering_inner_pending = false;
+ if keys_match(
+ &saved_inner_keys,
+ last_key_idx,
+ &self.inner_key_arrays,
+ 0,
+ &self.sort_options,
+ self.null_equality,
+ )? {
+ continue;
+ } else {
+ return Poll::Ready(Ok(false));
+ }
+ }
+ }
+ }
+ }
+
+ /// Process a key match with a filter. For each inner row in the buffered
+ /// key group, evaluates the filter against the outer key group and ORs
+ /// the results into the matched bitset using u64-chunked bitwise ops.
+ fn process_key_match_with_filter(&mut self) -> Result<()> {
+ let filter = self.filter.as_ref().unwrap();
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ let num_outer = outer_batch.num_rows();
+
+ // buffer_inner_key_group must be called before this function
+ debug_assert!(
+ !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+ "process_key_match_with_filter called with no inner key data"
+ );
+ debug_assert!(
+ self.outer_offset < num_outer,
+ "outer_offset must be within the current batch"
+ );
+ debug_assert!(
+ self.matched.len() == num_outer,
+ "matched vector must be sized for the current outer batch"
+ );
+
+ let outer_group_end = find_key_group_end(
+ &self.outer_key_arrays,
+ self.outer_offset,
+ num_outer,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ let outer_group_len = outer_group_end - self.outer_offset;
+ let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len);
+
+ // Count already-matched bits using popcnt on u64 chunks (zero-copy).
+ let mut matched_count = UnalignedBitChunk::new(
+ self.matched.as_slice(),
+ self.outer_offset,
+ outer_group_len,
+ )
+ .count_ones();
+
+ // Process spilled inner batches first (read back from disk).
+ if let Some(spill_file) = &self.inner_key_spill {
+ let file = BufReader::new(File::open(spill_file.path())?);
+ let reader = StreamReader::try_new(file, None)?;
+ for batch_result in reader {
+ let inner_slice = batch_result?;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break;
+ }
+ }
+ }
+
+ // Then process in-memory inner batches.
+ // evaluate_filter_for_inner_row is a free function (not &self method)
+ // so that Rust can split the struct borrow: &mut self.matched coexists
+ // with &self.inner_key_buffer and &self.filter inside this loop.
+ if matched_count < outer_group_len {
+ 'outer: for inner_slice in &self.inner_key_buffer {
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ if matched_count == outer_group_len {
+ break 'outer;
+ }
+ }
+ }
+
+ self.outer_offset = outer_group_end;
+ Ok(())
+ }
+
+ /// Continue processing an outer key group that spans multiple outer
+ /// batches. Returns `true` if this outer batch was fully consumed
+ /// by the key group and the caller should load another.
+ fn resume_boundary(&mut self) -> Result<bool> {
+ debug_assert!(
+ self.outer_batch.is_some(),
+ "caller must load outer_batch first"
+ );
+ match std::mem::replace(&mut self.boundary_state, BoundaryState::Normal) {
Review Comment:
Yeah, the `Normal` in `std::mem::replace` is just a placeholder. However, at lines 952 and 1017, `Normal` is also the correct final state (we consumed the pending data). At line 687, it's either correct (for the
`Normal` and non-boundary-continuing cases) or gets overwritten at lines 706/738.
I could add a comment on each replace call to clarify the intent, but I'd avoid adding a `TmpTransition` variant — it adds a match arm everywhere, turns a compile-time guarantee into a runtime panic, and the code paths already ensure `Normal` is correct.
Alternatively, since `boundary_state` is only read via `match`, we could wrap it in `Option<BoundaryState>` and `.take()`. But that adds `Option` unwrapping everywhere else. I want to make sure this is maintainable for other folks, but I don't love the addition state.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4027276009
🤖 Hi @Dandandan, thanks for the request (https://github.com/apache/datafusion/pull/20806#issuecomment-4027275844).
[`scrape_comments.py`](https://github.com/alamb/datafusion-benchmarking/blob/main/scripts/scrape_comments.py) only supports whitelisted benchmarks.
- Standard: clickbench_1, clickbench_extended, clickbench_partitioned, clickbench_pushdown, external_aggr, tpcds, tpch, tpch10, tpch_mem, tpch_mem10
- Criterion: aggregate_query_sql, aggregate_vectorized, array_set_ops, case_when, character_length, in_list, left, plan_reuse, range_and_generate_series, replace, reset_plan_states, sort, sql_planner, strpos, substr_index, with_hashes
Please choose one or more of these with `run benchmark <name>` or `run benchmark <name1> <name2>...`
You can also set environment variables on subsequent lines:
```
run benchmark tpch_mem
DATAFUSION_RUNTIME_MEMORY_LIMIT=1G
```
Unsupported benchmarks: DATAFUSION_OPTIMIZER_PREFER_HASH_JOIN=false.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#discussion_r2913355108
##########
datafusion/physical-plan/src/joins/semi_anti_sort_merge_join/stream.rs:
##########
@@ -0,0 +1,1164 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Stream implementation for semi/anti sort-merge joins.
+
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::BufReader;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use crate::RecordBatchStream;
+use crate::joins::utils::{JoinFilter, compare_join_arrays};
+use crate::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
+};
+use crate::spill::spill_manager::SpillManager;
+use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
+use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
+use arrow::util::bit_util::apply_bitwise_binary_op;
+use datafusion_common::{
+ JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err,
+};
+use datafusion_execution::SendableRecordBatchStream;
+use datafusion_execution::disk_manager::RefCountedTempFile;
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
+
+use futures::{Stream, StreamExt, ready};
+
+/// Evaluates join key expressions against a batch, returning one array per key.
+fn evaluate_join_keys(
+ batch: &RecordBatch,
+ on: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+ on.iter()
+ .map(|expr| {
+ let num_rows = batch.num_rows();
+ let val = expr.evaluate(batch)?;
+ val.into_array(num_rows)
+ })
+ .collect()
+}
+
+/// Find the first index in `key_arrays` starting from `from` where the key
+/// differs from the key at `from`. Uses `compare_join_arrays` for zero-alloc
+/// ordinal comparison.
+///
+/// Optimized for join workloads: checks adjacent and boundary keys before
+/// falling back to binary search, since most key groups are small (often 1).
+fn find_key_group_end(
+ key_arrays: &[ArrayRef],
+ from: usize,
+ len: usize,
+ sort_options: &[SortOptions],
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let next = from + 1;
+ if next >= len {
+ return Ok(len);
+ }
+
+ // Fast path: single-row group (common with unique keys).
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ next,
+ sort_options,
+ null_equality,
+ )? != Ordering::Equal
+ {
+ return Ok(next);
+ }
+
+ // Check if the entire remaining batch shares this key.
+ let last = len - 1;
+ if compare_join_arrays(
+ key_arrays,
+ from,
+ key_arrays,
+ last,
+ sort_options,
+ null_equality,
+ )? == Ordering::Equal
+ {
+ return Ok(len);
+ }
Review Comment:
One is checking `next` (is the next row the boundary), the other is checking `last` (is the last row the boundary), and then finally fall back to binary search.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoin operator [datafusion]
Posted by "alamb-ghbot (via GitHub)" <gi...@apache.org>.
alamb-ghbot commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4029527329
🤖 `./gh_compare_branch.sh` [gh_compare_branch.sh](https://github.com/alamb/datafusion-benchmarking/blob/main/scripts/gh_compare_branch.sh) Running
Linux aal-dev 6.14.0-1018-gcp #19~24.04.1-Ubuntu SMP Wed Sep 24 23:23:09 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing semi-anti-smj (0c33ad6b3fbfe2fd0d04421833b8af95e038b959) to 44dfa7b8a5b908d6d82233cc6aaea1137b1d54b5 [diff](https://github.com/apache/datafusion/compare/44dfa7b8a5b908d6d82233cc6aaea1137b1d54b5..0c33ad6b3fbfe2fd0d04421833b8af95e038b959) using: tpcds
Results will be posted here when complete
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: specialized SemiAntiSortMergeJoinStream [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #20806:
URL: https://github.com/apache/datafusion/pull/20806#issuecomment-4129466370
Thanks for the reviews and feedback @rluvaton @comphead @stuhood!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org