You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "2010YOUY01 (via GitHub)" <gi...@apache.org> on 2025/07/31 15:20:10 UTC
[PR] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
2010YOUY01 opened a new pull request, #16996:
URL: https://github.com/apache/datafusion/pull/16996
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
- Closes #.
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
# Summary
This PR rewrites the NLJ operator from scratch with a different approach, to limit the extra intermediate data overhead to only 1 `RecordBatch`, and eliminate other redundant conversions in the old implementation.
Using the NLJ micro-bench introduced in https://github.com/apache/datafusion/pull/16819, this PR can introduce up to **3.5X speed-up**, and **uses only 1% memory** in extreme cases.
### Speed benchmark
Note: have to cherry-pick the micro-bench PR https://github.com/apache/datafusion/pull/16819 that has not merged.
```
┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ pr-16819 ┃ nlj-rewrite ┃ Change ┃
┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 218.38 ms │ 86.36 ms │ +2.53x faster │
│ QQuery 2 │ 258.65 ms │ 116.88 ms │ +2.21x faster │
│ QQuery 3 │ 335.23 ms │ 146.01 ms │ +2.30x faster │
│ QQuery 4 │ 900.49 ms │ 371.38 ms │ +2.42x faster │
│ QQuery 5 │ 634.22 ms │ 240.42 ms │ +2.64x faster │
│ QQuery 6 │ 5897.75 ms │ 1694.10 ms │ +3.48x faster │
│ QQuery 7 │ 617.83 ms │ 248.83 ms │ +2.48x faster │
│ QQuery 8 │ 5712.21 ms │ 1693.09 ms │ +3.37x faster │
│ QQuery 9 │ 671.34 ms │ 269.08 ms │ +2.49x faster │
│ QQuery 10 │ 1731.51 ms │ 488.54 ms │ +3.54x faster │
└──────────────┴────────────┴─────────────┴───────────────┘
```
### Memory Usage Benchmark
TODO
@ding-young could you help to re-run it? I have pushed one change to the micro-bench PR to fix the memory issue in https://github.com/apache/datafusion/pull/16889#issuecomment-3121126110
### Next Steps
I think it is ready for review, the major potential optimizations have all been done. Though there is one minor chore to be done:
- [ ] Add join metrics tracking
# Why a Rewrite?
(TLDR: it's the easiest way to address the existing problem)
The original implementation performs a Cartesian product of (all-left-batches x right-batch), materializes that intermediate result for predicate evaluation, and then materializes the (potentially very large) final result all at once. This design is inherently inefficient, and although many patches have attempted to alleviate the problem, the fundamental issue remains.
A key challenge is that the original design and the ideal design (i.e., one that produces small intermediates during execution) are fundamentally different. As a result, it's practically impossible to make small incremental changes that fully address the inefficiency. These patches may also increase code complexity, making long-term maintenance more difficult.
### Example of Prior Work
Here's a recent example of a small patch intended to improve the situation:
https://github.com/apache/datafusion/pull/16443
Even with careful engineering, I still feel the entropy in the code increases.
Since NLJ is a relatively straightforward operator, a full rewrite seemed worthwhile. This allows for a clean, simplified design focused on current goals—performance and memory efficiency—without being constrained by the legacy implementation.
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
### Implementation
The implementation/design doc can be found in the source code.
A brief comparison between old implementation and this PR:
#### Old implementation
1. Do a Cartesian produce of (`all_buffered_left_batch x one_right_batch`) to calculate indices
2. Construct intermediate batch and output it chunk by chunk
#### This PR
- For the inner loop, only perform (`one_left_row x one_right_batch`), and do filtering, and output construction directly on this small intermediate.
- Eagerly yield output when the output buffer has reached `batch_size` threshold
#### Old v.s. PR
- The old implementation requires multiple conversions between `indices <--> batch`, and this PR can use the right batch directly. This avoids unnecessary transformations and make the implementation more cache-friendly
- The old implementation has extra memory overhead of `left_buffered_batches_total_rows * right_batch_size (default 8192) * 12 Bytes (left and right indices are represented by uint64 and uint32)`, can be significant for large dataset.
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
4. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
Existing tests
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
The old implementation can maintain right input order in certain cases, this PR's new design is not able to maintain that property -- preserving it would require a different design, and it has significant performance and memory overhead.
If this property is important to some user, we can keep the old implementation (maybe rename to `RightOrderPreservingNLJ` and use a configuration to control 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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3184292918
Thanks @2010YOUY01 @mbutrovich @jonathanc-n @UBarney @ding-young for the fantastic team work. I'm going to merge this PR in the end of the day
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2251126428
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1048 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
-
- // Produces 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
- *left_indices_cache = UInt64Array::from_iter_values(
- (0..output_row_count as u64).map(|i| i % left_row_count as u64),
- );
-
- // Produces 0, 0, 0, 1, 1, 1, 2, 2, 2, ...
- *right_indices_cache = UInt32Array::from_iter_values(
- (0..output_row_count as u32).map(|i| i / left_row_count as u32),
- );
-
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- };
-
- if let Some(filter) = filter {
- apply_join_filter_to_indices(
- left_batch,
- right_batch,
- left_indices,
- right_indices,
- filter,
- JoinSide::Left,
- Some(max_intermediate_batch_size),
- )
- } else {
- Ok((left_indices, right_indices))
- }
-}
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
-impl NestedLoopJoinStream {
- fn poll_next_impl(
- &mut self,
+ /// See the comments [`NestedLoopJoinExec`] for high-level design ideas.
+ ///
+ /// # Implementation
+ ///
+ /// This function is the entry point of NLJ operator's state machine
+ /// transitions. The rough state transition graph is as follow, for more
+ /// details see the comment in each state's matching arm.
+ ///
+ /// ============================
+ /// State transition graph:
+ /// ============================
+ ///
+ /// (start) --> BufferingLeft
+ /// ----------------------------
+ /// BufferingLeft → FetchingRight
+ ///
+ /// FetchingRight → ProbeRight (if right batch available)
+ /// FetchingRight → EmitLeftUnmatched (if right exhausted)
+ ///
+ /// ProbeRight → ProbeRight (next left row or after yielding output)
+ /// ProbeRight → EmitRightUnmatched (for special join types like right join)
+ /// ProbeRight → FetchingRight (done with the current right batch)
+ ///
+ /// EmitRightUnmatched → FetchingRight
+ ///
+ /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each
+ /// iteration)
+ /// EmitLeftUnmatched → Done (if finished)
+ /// ----------------------------
+ /// Done → (end)
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
+ ) -> Poll<Option<Self::Item>> {
loop {
- return match self.state {
- NestedLoopJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
+ match self.state {
+ // # NLJState transitions
+ // --> FetchingRight
+ // This state will prepare the left side batches, next state
+ // `FetchingRight` is responsible for preparing a single probe
+ // side batch, before start joining.
+ NLJState::BufferingLeft => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ match ready!(self.inner_table.get_shared(cx)) {
+ Ok(left_data) => {
+ self.buffered_left_data = Some(left_data);
+ // TODO: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
+
+ // # NLJState transitions:
+ // 1. --> ProbeRight
+ // Start processing the join for the newly fetched right
+ // batch.
+ // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit
+ // unmatched left side rows.
+ //
+ // After fetching a new batch from the right side, it will
+ // process all rows from the buffered left data:
+ // ```text
+ // for batch in right_side:
+ // for row in left_buffer:
+ // join(batch, row)
+ // ```
+ // Note: the implementation does this step incrementally,
+ // instead of materializing all intermediate Cartesian products
+ // at once in memory.
+ //
+ // So after the right side input is exhausted, the join phase
+ // for the current buffered left data is finished. We can go to
+ // the next `EmitLeftUnmatched` phase to check if there is any
+ // special handling (e.g., in cases like left join).
+ NLJState::FetchingRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ match ready!(self.outer_table.poll_next_unpin(cx)) {
+ Some(Ok(right_batch)) => {
+ let right_batch_size = right_batch.num_rows();
+
+ // Skip the empty batch
+ if right_batch_size == 0 {
+ continue;
+ }
+
+ self.current_right_batch = Some(right_batch);
+
+ // Prepare right bitmap
+ if self.should_track_unmatched_right {
+ let zeroed_buf =
+ BooleanBuffer::new_unset(right_batch_size);
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(zeroed_buf, None));
+ }
+
+ self.l_probe_idx = 0;
+ self.state = NLJState::ProbeRight;
+ continue;
+ }
+ Some(Err(e)) => return Poll::Ready(Some(Err(e))),
+ None => {
+ // Right stream exhausted/as
+ self.state = NLJState::EmitLeftUnmatched;
+ continue;
+ }
+ }
}
- NestedLoopJoinStreamState::ProcessProbeBatch(_) => {
- let poll = handle_state!(self.process_probe_batch());
- self.join_metrics.baseline.record_poll(poll)
+
+ // NLJState transitions:
+ // 1. --> ProbeRight(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> ProbeRight(2)
+ // After probing one right batch, and evaluating the
+ // join filter on (left-row x right-batch), it will advance
+ // to the next left row, then re-enter the current state and
+ // continue joining.
+ // 3. --> FetchRight
+ // After it has done with the current right batch (to join
+ // with all rows in the left buffer), it will go to
+ // FetchRight state to check what to do next.
+ NLJState::ProbeRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // Process current probe state
+ match self.process_probe_batch() {
+ // State unchanged (ProbeRight)
+ // Continue probing until we have done joining the
+ // current right batch with all buffered left rows.
+ Ok(true) => continue,
+ // To next FetchRightState
+ // We have finished joining
+ // (cur_right_batch x buffered_left_batches)
+ Ok(false) => {
+ // Left exhausted, transition to FetchingRight
+ self.l_probe_idx = 0;
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ self.state = NLJState::EmitRightUnmatched;
+ } else {
+ self.current_right_batch = None;
+ self.state = NLJState::FetchingRight;
+ }
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::PrepareUnmatchedBuildRows => {
- handle_state!(self.prepare_unmatched_output_indices())
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been output by the join. In this state we have to
+ // output unmatched rows for current right batch (with null
+ // padding for left relation)
+ // Precondition: we have checked the join type so that it's
+ // possible to output right unmatched (e.g. it's right join)
+ NLJState::EmitRightUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ debug_assert!(self.current_right_batch.is_some());
+ debug_assert!(self.current_right_batch_matched.is_some());
+
+ // Construct the result batch for unmatched right rows using a utility function
+ if let Some(batch) = self.process_right_unmatched()? {
+ self.output_buffer.push_batch(batch)?;
+ }
+
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
}
- NestedLoopJoinStreamState::OutputUnmatchedBuildRows(_) => {
- let poll = handle_state!(self.build_unmatched_output());
- self.join_metrics.baseline.record_poll(poll)
+
+ // NLJState transitions:
+ // 1. --> EmitLeftUnmatched(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> EmitLeftUnmatched(2)
+ // After processing some unmatched rows, it will re-enter
+ // the same state, to check if there are any more final
+ // results to output.
+ // 3. --> Done
+ // It has processed all data, go to the final state and ready
+ // to exit.
+ //
+ // TODO: For memory-limited case, go back to `BufferingLeft`
+ // state again.
+ NLJState::EmitLeftUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // Process current unmatched state
+ match self.process_left_unmatched() {
+ // State unchanged (EmitLeftUnmatched)
+ // Continue processing until we have processed all unmatched rows
+ Ok(true) => continue,
+ // To Done state
+ // We have finished processing all unmatched rows
+ Ok(false) => {
+ self.output_buffer.finish_buffered_batch()?;
+ self.state = NLJState::Done;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::Completed => Poll::Ready(None),
- };
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any remaining completed batches before final termination
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // HACK for the doc test in https://github.com/apache/datafusion/blob/main/datafusion/core/src/dataframe/mod.rs#L1265
+ // If this operator directly return `Poll::Ready(None)`
+ // for empty result, the final result will become an empty
+ // batch with empty schema, however the expected result
+ // should be with the expected schema for this operator
+ if !self.handled_empty_output {
+ let zero_count = Count::new();
+ if *self.join_metrics.baseline.output_rows() == zero_count {
+ let empty_batch =
+ RecordBatch::new_empty(Arc::clone(&self.output_schema));
+ self.handled_empty_output = true;
+ return Poll::Ready(Some(Ok(empty_batch)));
+ }
+ }
+
+ return Poll::Ready(None);
+ }
+ }
}
}
+}
- // This function's main job is to construct an output `RecordBatch` based on pre-calculated join indices.
- // It operates in a chunk-based manner, meaning it processes a portion of the results in each call,
- // making it suitable for streaming large datasets without high memory consumption.
- // This function behaves like an iterator. It returns `Ok(None)`
- // to signal that the result stream is exhausted and there is no more data.
- fn get_next_join_result(&mut self) -> Result<Option<RecordBatch>> {
- let status = self.join_result_status.as_mut().ok_or_else(|| {
- internal_datafusion_err!(
- "get_next_join_result called without initializing join_result_status"
- )
- })?;
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
+ }
+}
- let (left_indices, right_indices, current_start) = (
- &status.build_indices,
- &status.probe_indices,
- status.processed_count,
+impl NLJStream {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new(
+ schema: Arc<Schema>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ outer_table: SendableRecordBatchStream,
+ inner_table: OnceFut<JoinLeftData>,
+ column_indices: Vec<ColumnIndex>,
+ join_metrics: BuildProbeJoinMetrics,
+ cfg_batch_size: usize,
+ ) -> Self {
+ let should_track_unmatched_right = matches!(
+ join_type,
+ JoinType::Full
+ | JoinType::Right
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
);
- let left_batch = self
- .left_data
- .as_ref()
- .ok_or_else(|| internal_datafusion_err!("should have left_batch"))?
- .batch();
+ Self {
+ output_schema: Arc::clone(&schema),
+ join_filter: filter,
+ join_type,
+ outer_table,
+ column_indices,
+ inner_table,
+ join_metrics,
+ buffered_left_data: None,
+ output_buffer: Box::new(BatchCoalescer::new(schema, cfg_batch_size)),
+ cfg_batch_size,
+ current_right_batch: None,
+ current_right_batch_matched: None,
+ state: NLJState::BufferingLeft,
+ l_probe_idx: 0,
+ l_emit_idx: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
- let right_batch = match &self.state {
- NestedLoopJoinStreamState::ProcessProbeBatch(record_batch)
- | NestedLoopJoinStreamState::OutputUnmatchedBuildRows(record_batch) => {
- record_batch
- }
- _ => {
- return internal_err!(
- "State should be ProcessProbeBatch or OutputUnmatchedBuildRows"
- )
- }
- };
+ // ==== Core logic handling for each state ====
- if left_indices.is_empty() && right_indices.is_empty() && current_start == 0 {
- // To match the behavior of the previous implementation, return an empty RecordBatch.
- let res = RecordBatch::new_empty(Arc::clone(&self.schema));
- status.processed_count = 1;
- return Ok(Some(res));
- }
+ /// Returns bool to indicate should it continue probing
+ /// true -> continue in the same ProbeRight state
+ /// false -> It has done with the (buffered_left x cur_right_batch), go to
+ /// next state (ProbeRight)
+ fn process_probe_batch(&mut self) -> Result<bool> {
+ let left_data = Arc::clone(self.get_left_data()?);
+ let right_batch = self
+ .current_right_batch
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("Right batch should be available"))?
+ .clone();
- if matches!(self.join_type, JoinType::RightSemi | JoinType::RightAnti) {
- // in this case left_indices.num_rows() == 0
- let end = min(
- current_start + self.intermediate_batch_size,
- right_indices.len(),
- );
+ // stop probing, the caller will go to the next state
+ if self.l_probe_idx >= left_data.batch().num_rows() {
+ return Ok(false);
+ }
- if current_start >= end {
- return Ok(None);
- }
+ // ========
+ // Join (l_row x right_batch)
+ // and push the result into output_buffer
+ // ========
- let res = Some(build_batch_from_indices(
- &self.schema,
- left_batch,
- right_batch,
- left_indices,
- &right_indices.slice(current_start, end - current_start),
- &self.column_indices,
- JoinSide::Left,
- )?);
+ let l_idx = self.l_probe_idx;
+ let join_batch =
+ self.process_single_left_row_join(&left_data, &right_batch, l_idx)?;
- status.processed_count = end;
- return Ok(res);
+ if let Some(batch) = join_batch {
+ self.output_buffer.push_batch(batch)?;
}
- if current_start >= left_indices.len() {
+ // ==== Prepare for the next iteration ====
+
+ // Advance left cursor
+ self.l_probe_idx += 1;
+
+ // Return true to continue probing
+ Ok(true)
+ }
+
+ /// Process a single left row join with the current right batch.
+ /// Returns a RecordBatch containing the join results (None if empty)
+ fn process_single_left_row_join(
+ &mut self,
+ left_data: &JoinLeftData,
+ right_batch: &RecordBatch,
+ l_index: usize,
+ ) -> Result<Option<RecordBatch>> {
+ let right_row_count = right_batch.num_rows();
+ if right_row_count == 0 {
return Ok(None);
}
- let end = min(
- current_start + self.intermediate_batch_size,
- left_indices.len(),
- );
+ let cur_right_bitmap = if let Some(filter) = &self.join_filter {
+ apply_filter_to_row_join_batch(
+ left_data.batch(),
+ l_index,
+ right_batch,
+ filter,
+ )?
+ } else {
+ BooleanArray::from(vec![true; right_row_count])
+ };
- let left_indices = &left_indices.slice(current_start, end - current_start);
- let right_indices = &right_indices.slice(current_start, end - current_start);
+ self.update_matched_bitmap(l_index, &cur_right_bitmap)?;
- // Switch around the build side and probe side for `JoinType::RightMark`
- // because in a RightMark join, we want to mark rows on the right table
- // by looking for matches in the left.
- let res = if self.join_type == JoinType::RightMark {
- build_batch_from_indices(
- &self.schema,
- right_batch,
- left_batch,
- left_indices,
- right_indices,
- &self.column_indices,
- JoinSide::Right,
- )
+ // For the following join types: here we only have to set the left/right
+ // bitmap, and no need to output result
+ if matches!(
+ self.join_type,
+ JoinType::LeftAnti
+ | JoinType::LeftSemi
+ | JoinType::LeftMark
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ ) {
+ return Ok(None);
+ }
+
+ if cur_right_bitmap.true_count() == 0 {
+ // If none of the pairs has passed the join predicate/filter
+ Ok(None)
} else {
- build_batch_from_indices(
- &self.schema,
- left_batch,
+ // Use the optimized approach similar to build_intermediate_batch_for_single_left_row
+ let join_batch = build_row_join_batch(
+ &self.output_schema,
+ left_data.batch(),
+ l_index,
right_batch,
- left_indices,
- right_indices,
+ Some(cur_right_bitmap),
&self.column_indices,
JoinSide::Left,
- )
- }?;
+ )?;
+ Ok(join_batch)
+ }
+ }
- status.processed_count = end;
+ /// Returns bool to indicate should it continue processing unmatched rows
+ /// true -> continue in the same EmitLeftUnmatched state
+ /// false -> next state (Done)
+ fn process_left_unmatched(&mut self) -> Result<bool> {
+ let left_data = self.get_left_data()?;
+ let left_batch = left_data.batch();
- Ok(Some(res))
- }
+ // Early return if join type can't have unmatched rows
+ if !need_produce_result_in_final(self.join_type) {
+ return Ok(false);
+ }
- fn collect_build_side(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- let build_timer = self.join_metrics.build_time.timer();
- // build hash table from left (build) side, if not yet done
- self.left_data = Some(ready!(self.inner_table.get_shared(cx))?);
- build_timer.done();
+ // Early return if another thread is already processing unmatched rows
+ if self.l_emit_idx == 0 && !left_data.report_probe_completed() {
+ return Ok(false);
+ }
- self.state = NestedLoopJoinStreamState::FetchProbeBatch;
+ // Stop processing unmatched rows, the caller will go to the next state
+ if self.l_emit_idx >= left_batch.num_rows() {
+ return Ok(false);
+ }
- Poll::Ready(Ok(StatefulStreamResult::Continue))
- }
+ // ========
+ // Process unmatched rows and push the result into output_buffer
+ // Each time, the number to process is up to batch size
+ // ========
+ let start_idx = self.l_emit_idx;
+ let end_idx =
+ std::cmp::min(start_idx + self.cfg_batch_size, left_batch.num_rows());
- /// Fetches next batch from probe-side
- ///
- /// If a non-empty batch has been fetched, updates state to
- /// `ProcessProbeBatchState`, otherwise updates state to `ExhaustedProbeSide`.
- fn fetch_probe_batch(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- match ready!(self.outer_table.poll_next_unpin(cx)) {
- None => {
- self.state = NestedLoopJoinStreamState::PrepareUnmatchedBuildRows;
- }
- Some(Ok(right_batch)) => {
- self.join_metrics.input_batches.add(1);
- self.join_metrics.input_rows.add(right_batch.num_rows());
+ if let Some(batch) =
+ self.process_left_unmatched_range(left_data, start_idx, end_idx)?
+ {
+ self.output_buffer.push_batch(batch)?;
+ }
- self.state = NestedLoopJoinStreamState::ProcessProbeBatch(right_batch);
- }
- Some(Err(err)) => return Poll::Ready(Err(err)),
- };
+ // ==== Prepare for the next iteration ====
+ self.l_emit_idx = end_idx;
- Poll::Ready(Ok(StatefulStreamResult::Continue))
+ // Return true to continue processing unmatched rows
+ Ok(true)
}
- /// Joins current probe batch with build-side data and produces batch with
- /// matched output, updates state to `FetchProbeBatch`.
- fn process_probe_batch(
- &mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let Some(left_data) = self.left_data.clone() else {
- return internal_err!(
- "Expected left_data to be Some in ProcessProbeBatch state"
+ /// Process unmatched rows from the left data within the specified range.
+ /// Returns a RecordBatch containing the unmatched rows (None if empty).
+ ///
+ /// # Arguments
+ /// * `left_data` - The left side data containing the batch and bitmap
+ /// * `start_idx` - Start index (inclusive) of the range to process
+ /// * `end_idx` - End index (exclusive) of the range to process
+ ///
+ /// # Safety
+ /// The caller is responsible for ensuring that `start_idx` and `end_idx` are
+ /// within valid bounds of the left batch. This function does not perform
+ /// bounds checking.
+ fn process_left_unmatched_range(
+ &self,
+ left_data: &JoinLeftData,
+ start_idx: usize,
+ end_idx: usize,
+ ) -> Result<Option<RecordBatch>> {
+ if start_idx == end_idx {
+ return Ok(None);
+ }
+
+ // Slice both left batch, and bitmap to range [start_idx, end_idx)
+ // The range is bit index (not byte)
+ let left_batch = left_data.batch();
+ let left_batch_sliced = left_batch.slice(start_idx, end_idx - start_idx);
+
+ // Can this be more efficient?
+ let mut bitmap_sliced = BooleanBufferBuilder::new(end_idx - start_idx);
+ bitmap_sliced.append_n(end_idx - start_idx, false);
+ let bitmap = left_data.bitmap().lock();
+ for i in start_idx..end_idx {
+ assert!(
+ i - start_idx < bitmap_sliced.capacity(),
+ "DBG: {start_idx}, {end_idx}"
);
- };
- let visited_left_side = left_data.bitmap();
- let batch = self.state.try_as_process_probe_batch()?;
+ bitmap_sliced.set_bit(i - start_idx, bitmap.get_bit(i));
+ }
+ let bitmap_sliced = BooleanArray::new(bitmap_sliced.finish(), None);
+
+ build_unmatched_batch(
+ Arc::clone(&self.output_schema),
+ &left_batch_sliced,
+ bitmap_sliced,
+ self.outer_table.schema(),
+ &self.column_indices,
+ self.join_type,
+ JoinSide::Left,
+ )
+ }
- let binding = self.join_metrics.join_time.clone();
- let _timer = binding.timer();
+ /// Process unmatched rows from the current right batch and reset the bitmap.
+ /// Returns a RecordBatch containing the unmatched right rows (None if empty).
+ fn process_right_unmatched(&mut self) -> Result<Option<RecordBatch>> {
+ // ==== Take current right batch and its bitmap ====
+ let right_batch_bitmap: BooleanArray =
+ std::mem::take(&mut self.current_right_batch_matched).ok_or_else(|| {
+ internal_datafusion_err!("right bitmap should be available")
+ })?;
+
+ let right_batch = self.current_right_batch.take();
+ let cur_right_batch = unwrap_or_internal_err!(right_batch);
+
+ let left_data = self.get_left_data()?;
+ let left_schema = left_data.batch().schema();
+
+ let res = build_unmatched_batch(
+ Arc::clone(&self.output_schema),
+ &cur_right_batch,
+ right_batch_bitmap,
+ left_schema,
+ &self.column_indices,
+ self.join_type,
+ JoinSide::Right,
+ );
- if self.join_result_status.is_none() {
- let (left_side_indices, right_side_indices) = join_left_and_right_batch(
- left_data.batch(),
- batch,
- self.join_type,
- self.filter.as_ref(),
- visited_left_side,
- &mut self.indices_cache,
- self.right_side_ordered,
- self.intermediate_batch_size,
- )?;
- self.join_result_status = Some(JoinResultProgress {
- build_indices: left_side_indices,
- probe_indices: right_side_indices,
- processed_count: 0,
- })
- }
+ // ==== Clean-up ====
+ self.current_right_batch_matched = None;
- let join_result = self.get_next_join_result()?;
+ res
+ }
- match join_result {
- Some(res) => {
- self.join_metrics.output_batches.add(1);
- Ok(StatefulStreamResult::Ready(Some(res)))
- }
- None => {
- self.state = NestedLoopJoinStreamState::FetchProbeBatch;
- self.join_result_status = None;
- Ok(StatefulStreamResult::Continue)
- }
- }
+ // ==== Utilities ====
+
+ /// Get the inner data of left data, errors if it's None
+ fn get_left_data(&self) -> Result<&Arc<JoinLeftData>> {
+ self.buffered_left_data
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("LeftData should be available"))
}
- fn build_unmatched_output(
- &mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let binding = self.join_metrics.join_time.clone();
- let _timer = binding.timer();
-
- let res = self.get_next_join_result()?;
- match res {
- Some(res) => {
- self.join_metrics.output_batches.add(1);
- Ok(StatefulStreamResult::Ready(Some(res)))
- }
- None => {
- self.state = NestedLoopJoinStreamState::Completed;
- Ok(StatefulStreamResult::Ready(None))
+ /// Flush the `output_buffer` if there are batches ready to output
+ /// None if no result batch ready.
+ fn maybe_flush_ready_batch(&mut self) -> Option<Poll<Option<Result<RecordBatch>>>> {
+ if self.output_buffer.has_completed_batch() {
+ if let Some(batch) = self.output_buffer.next_completed_batch() {
+ let poll = Poll::Ready(Some(Ok(batch)));
+ return Some(self.join_metrics.baseline.record_poll(poll));
}
}
+
+ None
}
- /// This function's primary purpose is to handle the final output stage required by specific join types after all right-side (probe) data has been exhausted.
- /// It is critically important for LEFT*/FULL joins, which must emit left-side (build) rows that found no match. For these cases, it identifies the unmatched rows and prepares the necessary state to output them.
- fn prepare_unmatched_output_indices(
+ /// After joining (l_index@left_buffer x current_right_batch), it will result
+ /// in a bitmap (the same length as current_right_batch) as the join match
+ /// result. Use this bitmap to update the global bitmap, for special join
+ /// types like full joins.
+ ///
+ /// Example:
+ /// After joining l_index=1 (1-indexed row in the left buffer), and the
+ /// current right batch with 3 elements, this function will be called with
+ /// arguments: l_index = 1, r_matched = [false, false, true]
+ /// - If the join type is FullJoin, the 1-index in the left bitmap will be
+ /// set to true, and also the right bitmap will be bitwise-ORed with the
+ /// input r_matched bitmap.
+ /// - For join types that don't require output unmatched rows, this
+ /// function can be a no-op. For inner joins, this function is a no-op; for left
+ /// joins, only the left bitmap may be updated.
+ fn update_matched_bitmap(
&mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let Some(left_data) = self.left_data.clone() else {
- return internal_err!(
- "Expected left_data to be Some in ExhaustedProbeSide state"
- );
- };
- let visited_left_side = left_data.bitmap();
- if need_produce_result_in_final(self.join_type) {
- // At this stage `visited_left_side` won't be updated, so it's
- // safe to report about probe completion.
- //
- // Setting `is_exhausted` / returning None will prevent from
- // multiple calls of `report_probe_completed()`
- if !left_data.report_probe_completed() {
- self.state = NestedLoopJoinStreamState::Completed;
- return Ok(StatefulStreamResult::Ready(None));
- };
+ l_index: usize,
+ r_matched: &BooleanArray,
+ ) -> Result<()> {
+ let left_data = self.get_left_data()?;
- // Only setting up timer, input is exhausted
- let _timer = self.join_metrics.join_time.timer();
- // use the global left bitmap to produce the left indices and right indices
- let (left_side, right_side) =
- get_final_indices_from_shared_bitmap(visited_left_side, self.join_type);
-
- self.join_result_status = Some(JoinResultProgress {
- build_indices: left_side,
- probe_indices: right_side,
- processed_count: 0,
- });
- self.state = NestedLoopJoinStreamState::OutputUnmatchedBuildRows(
- RecordBatch::new_empty(self.outer_table.schema()),
- );
+ // number of successfully joined pairs from (l_index x cur_right_batch)
+ let joined_len = r_matched.true_count();
- Ok(StatefulStreamResult::Continue)
- } else {
- // end of the join loop
- self.state = NestedLoopJoinStreamState::Completed;
- Ok(StatefulStreamResult::Ready(None))
+ // 1. Maybe update the left bitmap
+ if need_produce_result_in_final(self.join_type) && (joined_len > 0) {
+ let mut bitmap = left_data.bitmap().lock();
+ bitmap.set_bit(l_index, true);
+ }
+
+ // 2. Maybe updateh the right bitmap
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ // after bit-wise or, it will be put back
+ let right_bitmap = std::mem::take(&mut self.current_right_batch_matched)
+ .ok_or_else(|| {
+ internal_datafusion_err!("right batch's bitmap should be present")
+ })?;
+ let (buf, nulls) = right_bitmap.into_parts();
+ debug_assert!(nulls.is_none());
+ let updated_right_bitmap = buf.bitor(r_matched.values());
+
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(updated_right_bitmap, None));
}
+
+ Ok(())
}
}
-#[allow(clippy::too_many_arguments)]
-fn join_left_and_right_batch(
+// ==== Utilities ====
+
+/// Apply the join filter between:
+/// (l_index th row in left buffer) x (right batch)
+/// Returns a bitmap, with successfully joined indices set to true
+fn apply_filter_to_row_join_batch(
left_batch: &RecordBatch,
+ l_index: usize,
right_batch: &RecordBatch,
- join_type: JoinType,
- filter: Option<&JoinFilter>,
- visited_left_side: &SharedBitmapBuilder,
- indices_cache: &mut (UInt64Array, UInt32Array),
- right_side_ordered: bool,
- max_intermediate_batch_size: usize,
-) -> Result<(PrimitiveArray<UInt64Type>, PrimitiveArray<UInt32Type>)> {
- let (left_side, right_side) = build_join_indices(
- left_batch,
- right_batch,
- filter,
- indices_cache,
- max_intermediate_batch_size,
- )
- .map_err(|e| {
- exec_datafusion_err!(
- "Fail to build join indices in NestedLoopJoinExec, error: {e}"
- )
- })?;
-
- // set the left bitmap
- // and only full join need the left bitmap
- if need_produce_result_in_final(join_type) {
- let mut bitmap = visited_left_side.lock();
- left_side.values().iter().for_each(|x| {
- bitmap.set_bit(*x as usize, true);
- });
- }
- // adjust the two side indices base on the join type
- let (left_side, right_side) = adjust_indices_by_join_type(
- left_side,
- right_side,
- 0..right_batch.num_rows(),
- join_type,
- right_side_ordered,
- )?;
+ filter: &JoinFilter,
+) -> Result<BooleanArray> {
+ debug_assert!(left_batch.num_rows() != 0 && right_batch.num_rows() != 0);
+
+ let intermediate_batch = if filter.schema.fields().is_empty() {
+ // If filter is constant (e.g. literal `true`), empty batch can be used
+ // in the later filter step.
+ create_record_batch_with_empty_schema(
+ Arc::new((*filter.schema).clone()),
+ right_batch.num_rows(),
+ )?
+ } else {
+ build_row_join_batch(
+ &filter.schema,
+ left_batch,
+ l_index,
+ right_batch,
+ None,
+ &filter.column_indices,
+ JoinSide::Left,
+ )?
+ .ok_or_else(|| internal_datafusion_err!("This function assume input batch is not empty, so the intermediate batch can't be empty too"))?
+ };
+
+ let filter_result = filter
+ .expression()
+ .evaluate(&intermediate_batch)?
+ .into_array(intermediate_batch.num_rows())?;
- Ok((left_side, right_side))
+ Ok(as_boolean_array(&filter_result)?.clone())
}
-impl Stream for NestedLoopJoinStream {
- type Item = Result<RecordBatch>;
+/// This function performs the following steps:
+/// 1. Apply filter to probe-side batch
+/// 2. Broadcast the left row (build_side_batch\[build_side_index\]) to the
+/// filtered probe-side batch
+/// 3. Concat them together according to `col_indices`, and return the result
+/// (None if the result is empty)
+///
+/// Example:
+/// build_side_batch:
+/// a
+/// ----
+/// 1
+/// 2
+/// 3
+///
+/// # 0 index element in the build_side_batch (that is `1`) will be used
+/// build_side_index: 0
+///
+/// probe_side_batch:
+/// b
+/// ----
+/// 10
+/// 20
+/// 30
+/// 40
+///
+/// # After applying it, only index 1 and 3 elemnt in probe_side_batch will be
+/// # kept
+/// probe_side_filter:
+/// false
+/// true
+/// false
+/// true
+///
+///
+/// # Projections to the build/probe side batch, to construct the output batch
+/// col_indices:
+/// [(left, 0), (right, 0)]
+///
+/// build_side: left
+///
+/// ====
+/// Result batch:
+/// a b
+/// ----
+/// 1 20
+/// 1 40
+fn build_row_join_batch(
+ output_schema: &Schema,
+ build_side_batch: &RecordBatch,
+ build_side_index: usize,
+ probe_side_batch: &RecordBatch,
+ probe_side_filter: Option<BooleanArray>,
+ // See [`NLJStream`] struct's `column_indices` field for more detail
+ col_indices: &[ColumnIndex],
+ // If the build side is left or right, used to interpret the side information
+ // in `col_indices`
+ build_side: JoinSide,
+) -> Result<Option<RecordBatch>> {
+ debug_assert!(build_side != JoinSide::None);
+
+ // TODO(perf): since the output might be projection of right batch, this
+ // filtering step is more efficient to be done inside the column_index loop
+ let filtered_probe_batch = if let Some(filter) = probe_side_filter {
+ &filter_record_batch(probe_side_batch, &filter)?
+ } else {
+ probe_side_batch
+ };
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
+ if filtered_probe_batch.num_rows() == 0 {
+ return Ok(None);
+ }
+
+ let mut columns: Vec<Arc<dyn Array>> =
+ Vec::with_capacity(output_schema.fields().len());
+
+ for column_index in col_indices {
+ let array = if column_index.side == build_side {
+ // Broadcast the single build-side row to match the filtered
+ // probe-side batch length
+ let original_left_array = build_side_batch.column(column_index.index);
+ let scalar_value = ScalarValue::try_from_array(
+ original_left_array.as_ref(),
+ build_side_index,
+ )?;
+ scalar_value.to_array_of_size(filtered_probe_batch.num_rows())?
Review Comment:
> I tried but it only gets slightly faster, I think this optimization might not worth the extra complexity.
>
> The reason for not getting significantly faster:
>
> 1. This step is not the bottleneck (for this simple NLJ micro-bench, now the actual bottleneck is evaluating the expression, especially modulo)
> 2. I used a utility in the optimizer to rewrite the join filter, this step is slow (6% in total time)
>
> Benchmark result:
>
> ```
> ┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
> ┃ Query ┃ nlj-rewrite ┃ optimize-join-filter ┃ Change ┃
> ┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
> │ QQuery 1 │ 87.14 ms │ 83.12 ms │ no change │
> │ QQuery 2 │ 118.03 ms │ 110.51 ms │ +1.07x faster │
> │ QQuery 3 │ 164.04 ms │ 160.60 ms │ no change │
> │ QQuery 4 │ 373.67 ms │ 358.43 ms │ no change │
> │ QQuery 5 │ 232.66 ms │ 259.75 ms │ 1.12x slower │
> │ QQuery 6 │ 1708.19 ms │ 1633.82 ms │ no change │
> │ QQuery 7 │ 236.27 ms │ 257.24 ms │ 1.09x slower │
> │ QQuery 8 │ 1710.80 ms │ 1634.83 ms │ no change │
> │ QQuery 9 │ 270.86 ms │ 261.67 ms │ no change │
> │ QQuery 10 │ 499.70 ms │ 493.96 ms │ no change │
> └──────────────┴─────────────┴──────────────────────┴───────────────┘
> ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓
> ┃ Benchmark Summary ┃ ┃
> ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩
> │ Total Time (nlj-rewrite) │ 5401.37ms │
> │ Total Time (optimize-join-filter) │ 5253.94ms │
> │ Average Time (nlj-rewrite) │ 540.14ms │
> │ Average Time (optimize-join-filter) │ 525.39ms │
> │ Queries Faster │ 1 │
> │ Queries Slower │ 2 │
> │ Queries with No Change │ 7 │
> │ Queries with Failure │ 0 │
> └─────────────────────────────────────┴───────────┘
> ```
>
> My experiment branch: https://github.com/2010YOUY01/arrow-datafusion/tree/optimize-join-filter
Very nice, thanks for trying.
I wonder if we can try with a different expression, e.g. greater than?
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "jonathanc-n (via GitHub)" <gi...@apache.org>.
jonathanc-n commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3176847893
> Thanks @2010YOUY01 I think this PR is good the way it is. Given the change I would probably ask another pair of eyes, like @jonathanc-n who is also working on join in #16660
I'll try to take a look today
--
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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3188444016
I have tested the sqlite test suit locally, and it's passing now.
```
yongting@Yongtings-MacBook-Pro-2 ~/C/datafusion (pr-16996)> INCLUDE_SQLITE=true cargo test --profile release-nonlto --test sqllogictests
Completed 943 test files in 4 minutes
```
This PR is ready to merge again.
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2245710943
##########
datafusion/sqllogictest/test_files/joins.slt:
##########
@@ -4127,9 +4127,9 @@ logical_plan
03)----TableScan: left_table projection=[a, b, c]
04)----TableScan: right_table projection=[x, y, z]
physical_plan
-01)NestedLoopJoinExec: join_type=Inner, filter=a@0 < x@1
-02)--DataSourceExec: partitions=1, partition_sizes=[0]
-03)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false]
+01)SortExec: expr=[x@3 ASC NULLS LAST], preserve_partitioning=[false]
Review Comment:
See the PR write-up's 'breaking change' section for the reason.
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2269761709
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +691,1168 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NestedLoopJoinStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
Review Comment:
I believe Snowflake flips the convention and calls the left table "inner" which confuses the issue, regardless of physical join implementation. I think left/build and right/probe is consistent and fine, even though build and probe don't have as much meaning in a NLJ. I just was concerned about inner and outer getting confused particularly for a NLJ. Thanks!
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2266071623
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
Review Comment:
Makes sense, let's use `batch_size` for now.
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "jonathanc-n (via GitHub)" <gi...@apache.org>.
jonathanc-n commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2268648557
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +691,1168 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NestedLoopJoinStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
Review Comment:
Yes I think we should just keep it as left/build + right/probe. I can see the use of referring as it to inner/outer but it gets confusing because in academics (as @mbutrovich says), the term 'outer table' is the table which drives the outer loop. despite this, I've noticed that people do associate the left table to be the inner one, so just removing that ambiguity is probably the best idea
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2272059908
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +691,1168 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NestedLoopJoinStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
Review Comment:
Updated in [8cc3654](https://github.com/apache/datafusion/pull/16996/commits/8cc365491cdeab7dec8aac5d3da458191b2f3e16). Thanks!
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2253750647
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1048 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
-
- // Produces 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
- *left_indices_cache = UInt64Array::from_iter_values(
- (0..output_row_count as u64).map(|i| i % left_row_count as u64),
- );
-
- // Produces 0, 0, 0, 1, 1, 1, 2, 2, 2, ...
- *right_indices_cache = UInt32Array::from_iter_values(
- (0..output_row_count as u32).map(|i| i / left_row_count as u32),
- );
-
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- };
-
- if let Some(filter) = filter {
- apply_join_filter_to_indices(
- left_batch,
- right_batch,
- left_indices,
- right_indices,
- filter,
- JoinSide::Left,
- Some(max_intermediate_batch_size),
- )
- } else {
- Ok((left_indices, right_indices))
- }
-}
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
-impl NestedLoopJoinStream {
- fn poll_next_impl(
- &mut self,
+ /// See the comments [`NestedLoopJoinExec`] for high-level design ideas.
+ ///
+ /// # Implementation
+ ///
+ /// This function is the entry point of NLJ operator's state machine
+ /// transitions. The rough state transition graph is as follow, for more
+ /// details see the comment in each state's matching arm.
+ ///
+ /// ============================
+ /// State transition graph:
+ /// ============================
+ ///
+ /// (start) --> BufferingLeft
+ /// ----------------------------
+ /// BufferingLeft → FetchingRight
+ ///
+ /// FetchingRight → ProbeRight (if right batch available)
+ /// FetchingRight → EmitLeftUnmatched (if right exhausted)
+ ///
+ /// ProbeRight → ProbeRight (next left row or after yielding output)
+ /// ProbeRight → EmitRightUnmatched (for special join types like right join)
+ /// ProbeRight → FetchingRight (done with the current right batch)
+ ///
+ /// EmitRightUnmatched → FetchingRight
+ ///
+ /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each
+ /// iteration)
+ /// EmitLeftUnmatched → Done (if finished)
+ /// ----------------------------
+ /// Done → (end)
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
+ ) -> Poll<Option<Self::Item>> {
loop {
- return match self.state {
- NestedLoopJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
+ match self.state {
+ // # NLJState transitions
+ // --> FetchingRight
+ // This state will prepare the left side batches, next state
+ // `FetchingRight` is responsible for preparing a single probe
+ // side batch, before start joining.
+ NLJState::BufferingLeft => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ match ready!(self.inner_table.get_shared(cx)) {
+ Ok(left_data) => {
+ self.buffered_left_data = Some(left_data);
+ // TODO: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
+
+ // # NLJState transitions:
+ // 1. --> ProbeRight
+ // Start processing the join for the newly fetched right
+ // batch.
+ // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit
+ // unmatched left side rows.
+ //
+ // After fetching a new batch from the right side, it will
+ // process all rows from the buffered left data:
+ // ```text
+ // for batch in right_side:
+ // for row in left_buffer:
+ // join(batch, row)
+ // ```
+ // Note: the implementation does this step incrementally,
+ // instead of materializing all intermediate Cartesian products
+ // at once in memory.
+ //
+ // So after the right side input is exhausted, the join phase
+ // for the current buffered left data is finished. We can go to
+ // the next `EmitLeftUnmatched` phase to check if there is any
+ // special handling (e.g., in cases like left join).
+ NLJState::FetchingRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ match ready!(self.outer_table.poll_next_unpin(cx)) {
+ Some(Ok(right_batch)) => {
+ let right_batch_size = right_batch.num_rows();
+
+ // Skip the empty batch
+ if right_batch_size == 0 {
+ continue;
+ }
+
+ self.current_right_batch = Some(right_batch);
+
+ // Prepare right bitmap
+ if self.should_track_unmatched_right {
+ let zeroed_buf =
+ BooleanBuffer::new_unset(right_batch_size);
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(zeroed_buf, None));
+ }
+
+ self.l_probe_idx = 0;
+ self.state = NLJState::ProbeRight;
+ continue;
+ }
+ Some(Err(e)) => return Poll::Ready(Some(Err(e))),
+ None => {
+ // Right stream exhausted/as
+ self.state = NLJState::EmitLeftUnmatched;
+ continue;
+ }
+ }
}
- NestedLoopJoinStreamState::ProcessProbeBatch(_) => {
- let poll = handle_state!(self.process_probe_batch());
- self.join_metrics.baseline.record_poll(poll)
+
+ // NLJState transitions:
+ // 1. --> ProbeRight(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> ProbeRight(2)
+ // After probing one right batch, and evaluating the
+ // join filter on (left-row x right-batch), it will advance
+ // to the next left row, then re-enter the current state and
+ // continue joining.
+ // 3. --> FetchRight
+ // After it has done with the current right batch (to join
+ // with all rows in the left buffer), it will go to
+ // FetchRight state to check what to do next.
+ NLJState::ProbeRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // Process current probe state
+ match self.process_probe_batch() {
+ // State unchanged (ProbeRight)
+ // Continue probing until we have done joining the
+ // current right batch with all buffered left rows.
+ Ok(true) => continue,
+ // To next FetchRightState
+ // We have finished joining
+ // (cur_right_batch x buffered_left_batches)
+ Ok(false) => {
+ // Left exhausted, transition to FetchingRight
+ self.l_probe_idx = 0;
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ self.state = NLJState::EmitRightUnmatched;
+ } else {
+ self.current_right_batch = None;
+ self.state = NLJState::FetchingRight;
+ }
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::PrepareUnmatchedBuildRows => {
- handle_state!(self.prepare_unmatched_output_indices())
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been output by the join. In this state we have to
+ // output unmatched rows for current right batch (with null
+ // padding for left relation)
+ // Precondition: we have checked the join type so that it's
+ // possible to output right unmatched (e.g. it's right join)
+ NLJState::EmitRightUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ debug_assert!(self.current_right_batch.is_some());
+ debug_assert!(self.current_right_batch_matched.is_some());
+
+ // Construct the result batch for unmatched right rows using a utility function
+ if let Some(batch) = self.process_right_unmatched()? {
+ self.output_buffer.push_batch(batch)?;
+ }
+
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
}
- NestedLoopJoinStreamState::OutputUnmatchedBuildRows(_) => {
- let poll = handle_state!(self.build_unmatched_output());
- self.join_metrics.baseline.record_poll(poll)
+
+ // NLJState transitions:
+ // 1. --> EmitLeftUnmatched(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> EmitLeftUnmatched(2)
+ // After processing some unmatched rows, it will re-enter
+ // the same state, to check if there are any more final
+ // results to output.
+ // 3. --> Done
+ // It has processed all data, go to the final state and ready
+ // to exit.
+ //
+ // TODO: For memory-limited case, go back to `BufferingLeft`
+ // state again.
+ NLJState::EmitLeftUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // Process current unmatched state
+ match self.process_left_unmatched() {
+ // State unchanged (EmitLeftUnmatched)
+ // Continue processing until we have processed all unmatched rows
+ Ok(true) => continue,
+ // To Done state
+ // We have finished processing all unmatched rows
+ Ok(false) => {
+ self.output_buffer.finish_buffered_batch()?;
+ self.state = NLJState::Done;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::Completed => Poll::Ready(None),
- };
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any remaining completed batches before final termination
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // HACK for the doc test in https://github.com/apache/datafusion/blob/main/datafusion/core/src/dataframe/mod.rs#L1265
+ // If this operator directly return `Poll::Ready(None)`
+ // for empty result, the final result will become an empty
+ // batch with empty schema, however the expected result
+ // should be with the expected schema for this operator
+ if !self.handled_empty_output {
+ let zero_count = Count::new();
+ if *self.join_metrics.baseline.output_rows() == zero_count {
+ let empty_batch =
+ RecordBatch::new_empty(Arc::clone(&self.output_schema));
+ self.handled_empty_output = true;
+ return Poll::Ready(Some(Ok(empty_batch)));
+ }
+ }
+
+ return Poll::Ready(None);
+ }
+ }
}
}
+}
- // This function's main job is to construct an output `RecordBatch` based on pre-calculated join indices.
- // It operates in a chunk-based manner, meaning it processes a portion of the results in each call,
- // making it suitable for streaming large datasets without high memory consumption.
- // This function behaves like an iterator. It returns `Ok(None)`
- // to signal that the result stream is exhausted and there is no more data.
- fn get_next_join_result(&mut self) -> Result<Option<RecordBatch>> {
- let status = self.join_result_status.as_mut().ok_or_else(|| {
- internal_datafusion_err!(
- "get_next_join_result called without initializing join_result_status"
- )
- })?;
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
+ }
+}
- let (left_indices, right_indices, current_start) = (
- &status.build_indices,
- &status.probe_indices,
- status.processed_count,
+impl NLJStream {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new(
+ schema: Arc<Schema>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ outer_table: SendableRecordBatchStream,
+ inner_table: OnceFut<JoinLeftData>,
+ column_indices: Vec<ColumnIndex>,
+ join_metrics: BuildProbeJoinMetrics,
+ cfg_batch_size: usize,
+ ) -> Self {
+ let should_track_unmatched_right = matches!(
+ join_type,
+ JoinType::Full
+ | JoinType::Right
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
);
- let left_batch = self
- .left_data
- .as_ref()
- .ok_or_else(|| internal_datafusion_err!("should have left_batch"))?
- .batch();
+ Self {
+ output_schema: Arc::clone(&schema),
+ join_filter: filter,
+ join_type,
+ outer_table,
+ column_indices,
+ inner_table,
+ join_metrics,
+ buffered_left_data: None,
+ output_buffer: Box::new(BatchCoalescer::new(schema, cfg_batch_size)),
+ cfg_batch_size,
+ current_right_batch: None,
+ current_right_batch_matched: None,
+ state: NLJState::BufferingLeft,
+ l_probe_idx: 0,
+ l_emit_idx: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
- let right_batch = match &self.state {
- NestedLoopJoinStreamState::ProcessProbeBatch(record_batch)
- | NestedLoopJoinStreamState::OutputUnmatchedBuildRows(record_batch) => {
- record_batch
- }
- _ => {
- return internal_err!(
- "State should be ProcessProbeBatch or OutputUnmatchedBuildRows"
- )
- }
- };
+ // ==== Core logic handling for each state ====
- if left_indices.is_empty() && right_indices.is_empty() && current_start == 0 {
- // To match the behavior of the previous implementation, return an empty RecordBatch.
- let res = RecordBatch::new_empty(Arc::clone(&self.schema));
- status.processed_count = 1;
- return Ok(Some(res));
- }
+ /// Returns bool to indicate should it continue probing
+ /// true -> continue in the same ProbeRight state
+ /// false -> It has done with the (buffered_left x cur_right_batch), go to
+ /// next state (ProbeRight)
+ fn process_probe_batch(&mut self) -> Result<bool> {
+ let left_data = Arc::clone(self.get_left_data()?);
+ let right_batch = self
+ .current_right_batch
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("Right batch should be available"))?
+ .clone();
- if matches!(self.join_type, JoinType::RightSemi | JoinType::RightAnti) {
- // in this case left_indices.num_rows() == 0
- let end = min(
- current_start + self.intermediate_batch_size,
- right_indices.len(),
- );
+ // stop probing, the caller will go to the next state
+ if self.l_probe_idx >= left_data.batch().num_rows() {
+ return Ok(false);
+ }
- if current_start >= end {
- return Ok(None);
- }
+ // ========
+ // Join (l_row x right_batch)
+ // and push the result into output_buffer
+ // ========
- let res = Some(build_batch_from_indices(
- &self.schema,
- left_batch,
- right_batch,
- left_indices,
- &right_indices.slice(current_start, end - current_start),
- &self.column_indices,
- JoinSide::Left,
- )?);
+ let l_idx = self.l_probe_idx;
+ let join_batch =
+ self.process_single_left_row_join(&left_data, &right_batch, l_idx)?;
- status.processed_count = end;
- return Ok(res);
+ if let Some(batch) = join_batch {
+ self.output_buffer.push_batch(batch)?;
}
- if current_start >= left_indices.len() {
+ // ==== Prepare for the next iteration ====
+
+ // Advance left cursor
+ self.l_probe_idx += 1;
+
+ // Return true to continue probing
+ Ok(true)
+ }
+
+ /// Process a single left row join with the current right batch.
+ /// Returns a RecordBatch containing the join results (None if empty)
+ fn process_single_left_row_join(
+ &mut self,
+ left_data: &JoinLeftData,
+ right_batch: &RecordBatch,
+ l_index: usize,
+ ) -> Result<Option<RecordBatch>> {
+ let right_row_count = right_batch.num_rows();
+ if right_row_count == 0 {
return Ok(None);
}
- let end = min(
- current_start + self.intermediate_batch_size,
- left_indices.len(),
- );
+ let cur_right_bitmap = if let Some(filter) = &self.join_filter {
+ apply_filter_to_row_join_batch(
+ left_data.batch(),
+ l_index,
+ right_batch,
+ filter,
+ )?
+ } else {
+ BooleanArray::from(vec![true; right_row_count])
+ };
- let left_indices = &left_indices.slice(current_start, end - current_start);
- let right_indices = &right_indices.slice(current_start, end - current_start);
+ self.update_matched_bitmap(l_index, &cur_right_bitmap)?;
- // Switch around the build side and probe side for `JoinType::RightMark`
- // because in a RightMark join, we want to mark rows on the right table
- // by looking for matches in the left.
- let res = if self.join_type == JoinType::RightMark {
- build_batch_from_indices(
- &self.schema,
- right_batch,
- left_batch,
- left_indices,
- right_indices,
- &self.column_indices,
- JoinSide::Right,
- )
+ // For the following join types: here we only have to set the left/right
+ // bitmap, and no need to output result
+ if matches!(
+ self.join_type,
+ JoinType::LeftAnti
+ | JoinType::LeftSemi
+ | JoinType::LeftMark
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ ) {
+ return Ok(None);
+ }
+
+ if cur_right_bitmap.true_count() == 0 {
+ // If none of the pairs has passed the join predicate/filter
+ Ok(None)
} else {
- build_batch_from_indices(
- &self.schema,
- left_batch,
+ // Use the optimized approach similar to build_intermediate_batch_for_single_left_row
+ let join_batch = build_row_join_batch(
+ &self.output_schema,
+ left_data.batch(),
+ l_index,
right_batch,
- left_indices,
- right_indices,
+ Some(cur_right_bitmap),
&self.column_indices,
JoinSide::Left,
- )
- }?;
+ )?;
+ Ok(join_batch)
+ }
+ }
- status.processed_count = end;
+ /// Returns bool to indicate should it continue processing unmatched rows
+ /// true -> continue in the same EmitLeftUnmatched state
+ /// false -> next state (Done)
+ fn process_left_unmatched(&mut self) -> Result<bool> {
+ let left_data = self.get_left_data()?;
+ let left_batch = left_data.batch();
- Ok(Some(res))
- }
+ // Early return if join type can't have unmatched rows
+ if !need_produce_result_in_final(self.join_type) {
+ return Ok(false);
+ }
- fn collect_build_side(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- let build_timer = self.join_metrics.build_time.timer();
- // build hash table from left (build) side, if not yet done
- self.left_data = Some(ready!(self.inner_table.get_shared(cx))?);
- build_timer.done();
+ // Early return if another thread is already processing unmatched rows
+ if self.l_emit_idx == 0 && !left_data.report_probe_completed() {
+ return Ok(false);
+ }
- self.state = NestedLoopJoinStreamState::FetchProbeBatch;
+ // Stop processing unmatched rows, the caller will go to the next state
+ if self.l_emit_idx >= left_batch.num_rows() {
+ return Ok(false);
+ }
- Poll::Ready(Ok(StatefulStreamResult::Continue))
- }
+ // ========
+ // Process unmatched rows and push the result into output_buffer
+ // Each time, the number to process is up to batch size
+ // ========
+ let start_idx = self.l_emit_idx;
+ let end_idx =
+ std::cmp::min(start_idx + self.cfg_batch_size, left_batch.num_rows());
- /// Fetches next batch from probe-side
- ///
- /// If a non-empty batch has been fetched, updates state to
- /// `ProcessProbeBatchState`, otherwise updates state to `ExhaustedProbeSide`.
- fn fetch_probe_batch(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- match ready!(self.outer_table.poll_next_unpin(cx)) {
- None => {
- self.state = NestedLoopJoinStreamState::PrepareUnmatchedBuildRows;
- }
- Some(Ok(right_batch)) => {
- self.join_metrics.input_batches.add(1);
- self.join_metrics.input_rows.add(right_batch.num_rows());
+ if let Some(batch) =
+ self.process_left_unmatched_range(left_data, start_idx, end_idx)?
+ {
+ self.output_buffer.push_batch(batch)?;
+ }
- self.state = NestedLoopJoinStreamState::ProcessProbeBatch(right_batch);
- }
- Some(Err(err)) => return Poll::Ready(Err(err)),
- };
+ // ==== Prepare for the next iteration ====
+ self.l_emit_idx = end_idx;
- Poll::Ready(Ok(StatefulStreamResult::Continue))
+ // Return true to continue processing unmatched rows
+ Ok(true)
}
- /// Joins current probe batch with build-side data and produces batch with
- /// matched output, updates state to `FetchProbeBatch`.
- fn process_probe_batch(
- &mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let Some(left_data) = self.left_data.clone() else {
- return internal_err!(
- "Expected left_data to be Some in ProcessProbeBatch state"
+ /// Process unmatched rows from the left data within the specified range.
+ /// Returns a RecordBatch containing the unmatched rows (None if empty).
+ ///
+ /// # Arguments
+ /// * `left_data` - The left side data containing the batch and bitmap
+ /// * `start_idx` - Start index (inclusive) of the range to process
+ /// * `end_idx` - End index (exclusive) of the range to process
+ ///
+ /// # Safety
+ /// The caller is responsible for ensuring that `start_idx` and `end_idx` are
+ /// within valid bounds of the left batch. This function does not perform
+ /// bounds checking.
+ fn process_left_unmatched_range(
+ &self,
+ left_data: &JoinLeftData,
+ start_idx: usize,
+ end_idx: usize,
+ ) -> Result<Option<RecordBatch>> {
+ if start_idx == end_idx {
+ return Ok(None);
+ }
+
+ // Slice both left batch, and bitmap to range [start_idx, end_idx)
+ // The range is bit index (not byte)
+ let left_batch = left_data.batch();
+ let left_batch_sliced = left_batch.slice(start_idx, end_idx - start_idx);
+
+ // Can this be more efficient?
+ let mut bitmap_sliced = BooleanBufferBuilder::new(end_idx - start_idx);
+ bitmap_sliced.append_n(end_idx - start_idx, false);
+ let bitmap = left_data.bitmap().lock();
+ for i in start_idx..end_idx {
+ assert!(
+ i - start_idx < bitmap_sliced.capacity(),
+ "DBG: {start_idx}, {end_idx}"
);
- };
- let visited_left_side = left_data.bitmap();
- let batch = self.state.try_as_process_probe_batch()?;
+ bitmap_sliced.set_bit(i - start_idx, bitmap.get_bit(i));
+ }
+ let bitmap_sliced = BooleanArray::new(bitmap_sliced.finish(), None);
+
+ build_unmatched_batch(
+ Arc::clone(&self.output_schema),
+ &left_batch_sliced,
+ bitmap_sliced,
+ self.outer_table.schema(),
+ &self.column_indices,
+ self.join_type,
+ JoinSide::Left,
+ )
+ }
- let binding = self.join_metrics.join_time.clone();
- let _timer = binding.timer();
+ /// Process unmatched rows from the current right batch and reset the bitmap.
+ /// Returns a RecordBatch containing the unmatched right rows (None if empty).
+ fn process_right_unmatched(&mut self) -> Result<Option<RecordBatch>> {
+ // ==== Take current right batch and its bitmap ====
+ let right_batch_bitmap: BooleanArray =
+ std::mem::take(&mut self.current_right_batch_matched).ok_or_else(|| {
+ internal_datafusion_err!("right bitmap should be available")
+ })?;
+
+ let right_batch = self.current_right_batch.take();
+ let cur_right_batch = unwrap_or_internal_err!(right_batch);
+
+ let left_data = self.get_left_data()?;
+ let left_schema = left_data.batch().schema();
+
+ let res = build_unmatched_batch(
+ Arc::clone(&self.output_schema),
+ &cur_right_batch,
+ right_batch_bitmap,
+ left_schema,
+ &self.column_indices,
+ self.join_type,
+ JoinSide::Right,
+ );
- if self.join_result_status.is_none() {
- let (left_side_indices, right_side_indices) = join_left_and_right_batch(
- left_data.batch(),
- batch,
- self.join_type,
- self.filter.as_ref(),
- visited_left_side,
- &mut self.indices_cache,
- self.right_side_ordered,
- self.intermediate_batch_size,
- )?;
- self.join_result_status = Some(JoinResultProgress {
- build_indices: left_side_indices,
- probe_indices: right_side_indices,
- processed_count: 0,
- })
- }
+ // ==== Clean-up ====
+ self.current_right_batch_matched = None;
- let join_result = self.get_next_join_result()?;
+ res
+ }
- match join_result {
- Some(res) => {
- self.join_metrics.output_batches.add(1);
- Ok(StatefulStreamResult::Ready(Some(res)))
- }
- None => {
- self.state = NestedLoopJoinStreamState::FetchProbeBatch;
- self.join_result_status = None;
- Ok(StatefulStreamResult::Continue)
- }
- }
+ // ==== Utilities ====
+
+ /// Get the inner data of left data, errors if it's None
+ fn get_left_data(&self) -> Result<&Arc<JoinLeftData>> {
+ self.buffered_left_data
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("LeftData should be available"))
}
- fn build_unmatched_output(
- &mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let binding = self.join_metrics.join_time.clone();
- let _timer = binding.timer();
-
- let res = self.get_next_join_result()?;
- match res {
- Some(res) => {
- self.join_metrics.output_batches.add(1);
- Ok(StatefulStreamResult::Ready(Some(res)))
- }
- None => {
- self.state = NestedLoopJoinStreamState::Completed;
- Ok(StatefulStreamResult::Ready(None))
+ /// Flush the `output_buffer` if there are batches ready to output
+ /// None if no result batch ready.
+ fn maybe_flush_ready_batch(&mut self) -> Option<Poll<Option<Result<RecordBatch>>>> {
+ if self.output_buffer.has_completed_batch() {
+ if let Some(batch) = self.output_buffer.next_completed_batch() {
+ let poll = Poll::Ready(Some(Ok(batch)));
+ return Some(self.join_metrics.baseline.record_poll(poll));
}
}
+
+ None
}
- /// This function's primary purpose is to handle the final output stage required by specific join types after all right-side (probe) data has been exhausted.
- /// It is critically important for LEFT*/FULL joins, which must emit left-side (build) rows that found no match. For these cases, it identifies the unmatched rows and prepares the necessary state to output them.
- fn prepare_unmatched_output_indices(
+ /// After joining (l_index@left_buffer x current_right_batch), it will result
+ /// in a bitmap (the same length as current_right_batch) as the join match
+ /// result. Use this bitmap to update the global bitmap, for special join
+ /// types like full joins.
+ ///
+ /// Example:
+ /// After joining l_index=1 (1-indexed row in the left buffer), and the
+ /// current right batch with 3 elements, this function will be called with
+ /// arguments: l_index = 1, r_matched = [false, false, true]
+ /// - If the join type is FullJoin, the 1-index in the left bitmap will be
+ /// set to true, and also the right bitmap will be bitwise-ORed with the
+ /// input r_matched bitmap.
+ /// - For join types that don't require output unmatched rows, this
+ /// function can be a no-op. For inner joins, this function is a no-op; for left
+ /// joins, only the left bitmap may be updated.
+ fn update_matched_bitmap(
&mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let Some(left_data) = self.left_data.clone() else {
- return internal_err!(
- "Expected left_data to be Some in ExhaustedProbeSide state"
- );
- };
- let visited_left_side = left_data.bitmap();
- if need_produce_result_in_final(self.join_type) {
- // At this stage `visited_left_side` won't be updated, so it's
- // safe to report about probe completion.
- //
- // Setting `is_exhausted` / returning None will prevent from
- // multiple calls of `report_probe_completed()`
- if !left_data.report_probe_completed() {
- self.state = NestedLoopJoinStreamState::Completed;
- return Ok(StatefulStreamResult::Ready(None));
- };
+ l_index: usize,
+ r_matched: &BooleanArray,
+ ) -> Result<()> {
+ let left_data = self.get_left_data()?;
- // Only setting up timer, input is exhausted
- let _timer = self.join_metrics.join_time.timer();
- // use the global left bitmap to produce the left indices and right indices
- let (left_side, right_side) =
- get_final_indices_from_shared_bitmap(visited_left_side, self.join_type);
-
- self.join_result_status = Some(JoinResultProgress {
- build_indices: left_side,
- probe_indices: right_side,
- processed_count: 0,
- });
- self.state = NestedLoopJoinStreamState::OutputUnmatchedBuildRows(
- RecordBatch::new_empty(self.outer_table.schema()),
- );
+ // number of successfully joined pairs from (l_index x cur_right_batch)
+ let joined_len = r_matched.true_count();
- Ok(StatefulStreamResult::Continue)
- } else {
- // end of the join loop
- self.state = NestedLoopJoinStreamState::Completed;
- Ok(StatefulStreamResult::Ready(None))
+ // 1. Maybe update the left bitmap
+ if need_produce_result_in_final(self.join_type) && (joined_len > 0) {
+ let mut bitmap = left_data.bitmap().lock();
+ bitmap.set_bit(l_index, true);
+ }
+
+ // 2. Maybe updateh the right bitmap
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ // after bit-wise or, it will be put back
+ let right_bitmap = std::mem::take(&mut self.current_right_batch_matched)
+ .ok_or_else(|| {
+ internal_datafusion_err!("right batch's bitmap should be present")
+ })?;
+ let (buf, nulls) = right_bitmap.into_parts();
+ debug_assert!(nulls.is_none());
+ let updated_right_bitmap = buf.bitor(r_matched.values());
+
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(updated_right_bitmap, None));
}
+
+ Ok(())
}
}
-#[allow(clippy::too_many_arguments)]
-fn join_left_and_right_batch(
+// ==== Utilities ====
+
+/// Apply the join filter between:
+/// (l_index th row in left buffer) x (right batch)
+/// Returns a bitmap, with successfully joined indices set to true
+fn apply_filter_to_row_join_batch(
left_batch: &RecordBatch,
+ l_index: usize,
right_batch: &RecordBatch,
- join_type: JoinType,
- filter: Option<&JoinFilter>,
- visited_left_side: &SharedBitmapBuilder,
- indices_cache: &mut (UInt64Array, UInt32Array),
- right_side_ordered: bool,
- max_intermediate_batch_size: usize,
-) -> Result<(PrimitiveArray<UInt64Type>, PrimitiveArray<UInt32Type>)> {
- let (left_side, right_side) = build_join_indices(
- left_batch,
- right_batch,
- filter,
- indices_cache,
- max_intermediate_batch_size,
- )
- .map_err(|e| {
- exec_datafusion_err!(
- "Fail to build join indices in NestedLoopJoinExec, error: {e}"
- )
- })?;
-
- // set the left bitmap
- // and only full join need the left bitmap
- if need_produce_result_in_final(join_type) {
- let mut bitmap = visited_left_side.lock();
- left_side.values().iter().for_each(|x| {
- bitmap.set_bit(*x as usize, true);
- });
- }
- // adjust the two side indices base on the join type
- let (left_side, right_side) = adjust_indices_by_join_type(
- left_side,
- right_side,
- 0..right_batch.num_rows(),
- join_type,
- right_side_ordered,
- )?;
+ filter: &JoinFilter,
+) -> Result<BooleanArray> {
+ debug_assert!(left_batch.num_rows() != 0 && right_batch.num_rows() != 0);
+
+ let intermediate_batch = if filter.schema.fields().is_empty() {
+ // If filter is constant (e.g. literal `true`), empty batch can be used
+ // in the later filter step.
+ create_record_batch_with_empty_schema(
+ Arc::new((*filter.schema).clone()),
+ right_batch.num_rows(),
+ )?
+ } else {
+ build_row_join_batch(
+ &filter.schema,
+ left_batch,
+ l_index,
+ right_batch,
+ None,
+ &filter.column_indices,
+ JoinSide::Left,
+ )?
+ .ok_or_else(|| internal_datafusion_err!("This function assume input batch is not empty, so the intermediate batch can't be empty too"))?
+ };
+
+ let filter_result = filter
+ .expression()
+ .evaluate(&intermediate_batch)?
+ .into_array(intermediate_batch.num_rows())?;
- Ok((left_side, right_side))
+ Ok(as_boolean_array(&filter_result)?.clone())
}
-impl Stream for NestedLoopJoinStream {
- type Item = Result<RecordBatch>;
+/// This function performs the following steps:
+/// 1. Apply filter to probe-side batch
+/// 2. Broadcast the left row (build_side_batch\[build_side_index\]) to the
+/// filtered probe-side batch
+/// 3. Concat them together according to `col_indices`, and return the result
+/// (None if the result is empty)
+///
+/// Example:
+/// build_side_batch:
+/// a
+/// ----
+/// 1
+/// 2
+/// 3
+///
+/// # 0 index element in the build_side_batch (that is `1`) will be used
+/// build_side_index: 0
+///
+/// probe_side_batch:
+/// b
+/// ----
+/// 10
+/// 20
+/// 30
+/// 40
+///
+/// # After applying it, only index 1 and 3 elemnt in probe_side_batch will be
+/// # kept
+/// probe_side_filter:
+/// false
+/// true
+/// false
+/// true
+///
+///
+/// # Projections to the build/probe side batch, to construct the output batch
+/// col_indices:
+/// [(left, 0), (right, 0)]
+///
+/// build_side: left
+///
+/// ====
+/// Result batch:
+/// a b
+/// ----
+/// 1 20
+/// 1 40
+fn build_row_join_batch(
+ output_schema: &Schema,
+ build_side_batch: &RecordBatch,
+ build_side_index: usize,
+ probe_side_batch: &RecordBatch,
+ probe_side_filter: Option<BooleanArray>,
+ // See [`NLJStream`] struct's `column_indices` field for more detail
+ col_indices: &[ColumnIndex],
+ // If the build side is left or right, used to interpret the side information
+ // in `col_indices`
+ build_side: JoinSide,
+) -> Result<Option<RecordBatch>> {
+ debug_assert!(build_side != JoinSide::None);
+
+ // TODO(perf): since the output might be projection of right batch, this
+ // filtering step is more efficient to be done inside the column_index loop
+ let filtered_probe_batch = if let Some(filter) = probe_side_filter {
+ &filter_record_batch(probe_side_batch, &filter)?
+ } else {
+ probe_side_batch
+ };
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
+ if filtered_probe_batch.num_rows() == 0 {
+ return Ok(None);
+ }
+
+ let mut columns: Vec<Arc<dyn Array>> =
+ Vec::with_capacity(output_schema.fields().len());
+
+ for column_index in col_indices {
+ let array = if column_index.side == build_side {
+ // Broadcast the single build-side row to match the filtered
+ // probe-side batch length
+ let original_left_array = build_side_batch.column(column_index.index);
+ let scalar_value = ScalarValue::try_from_array(
+ original_left_array.as_ref(),
+ build_side_index,
+ )?;
+ scalar_value.to_array_of_size(filtered_probe_batch.num_rows())?
Review Comment:
For such simple expressions, it's 20% faster on top of this PR. I think we should include this optimization as a follow-up.
After it, the bottleneck shifts to filter + concating the final output, this TODO might help
https://github.com/apache/arrow-rs/blob/a9b6077b2d8d5b5aee0e97e7d4335878e8da1876/arrow-select/src/coalesce.rs#L206
--
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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "Omega359 (via GitHub)" <gi...@apache.org>.
Omega359 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3188362336
> Run extended tests
That was removed recently unfortunately.
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2249090965
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1048 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
-
- // Produces 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
- *left_indices_cache = UInt64Array::from_iter_values(
- (0..output_row_count as u64).map(|i| i % left_row_count as u64),
- );
-
- // Produces 0, 0, 0, 1, 1, 1, 2, 2, 2, ...
- *right_indices_cache = UInt32Array::from_iter_values(
- (0..output_row_count as u32).map(|i| i / left_row_count as u32),
- );
-
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- };
-
- if let Some(filter) = filter {
- apply_join_filter_to_indices(
- left_batch,
- right_batch,
- left_indices,
- right_indices,
- filter,
- JoinSide::Left,
- Some(max_intermediate_batch_size),
- )
- } else {
- Ok((left_indices, right_indices))
- }
-}
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
-impl NestedLoopJoinStream {
- fn poll_next_impl(
- &mut self,
+ /// See the comments [`NestedLoopJoinExec`] for high-level design ideas.
+ ///
+ /// # Implementation
+ ///
+ /// This function is the entry point of NLJ operator's state machine
+ /// transitions. The rough state transition graph is as follow, for more
+ /// details see the comment in each state's matching arm.
+ ///
+ /// ============================
+ /// State transition graph:
+ /// ============================
+ ///
+ /// (start) --> BufferingLeft
+ /// ----------------------------
+ /// BufferingLeft → FetchingRight
+ ///
+ /// FetchingRight → ProbeRight (if right batch available)
+ /// FetchingRight → EmitLeftUnmatched (if right exhausted)
+ ///
+ /// ProbeRight → ProbeRight (next left row or after yielding output)
+ /// ProbeRight → EmitRightUnmatched (for special join types like right join)
+ /// ProbeRight → FetchingRight (done with the current right batch)
+ ///
+ /// EmitRightUnmatched → FetchingRight
+ ///
+ /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each
+ /// iteration)
+ /// EmitLeftUnmatched → Done (if finished)
+ /// ----------------------------
+ /// Done → (end)
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
+ ) -> Poll<Option<Self::Item>> {
loop {
- return match self.state {
- NestedLoopJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
+ match self.state {
+ // # NLJState transitions
+ // --> FetchingRight
+ // This state will prepare the left side batches, next state
+ // `FetchingRight` is responsible for preparing a single probe
+ // side batch, before start joining.
+ NLJState::BufferingLeft => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ match ready!(self.inner_table.get_shared(cx)) {
+ Ok(left_data) => {
+ self.buffered_left_data = Some(left_data);
+ // TODO: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
+
+ // # NLJState transitions:
+ // 1. --> ProbeRight
+ // Start processing the join for the newly fetched right
+ // batch.
+ // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit
+ // unmatched left side rows.
+ //
+ // After fetching a new batch from the right side, it will
+ // process all rows from the buffered left data:
+ // ```text
+ // for batch in right_side:
+ // for row in left_buffer:
+ // join(batch, row)
+ // ```
+ // Note: the implementation does this step incrementally,
+ // instead of materializing all intermediate Cartesian products
+ // at once in memory.
+ //
+ // So after the right side input is exhausted, the join phase
+ // for the current buffered left data is finished. We can go to
+ // the next `EmitLeftUnmatched` phase to check if there is any
+ // special handling (e.g., in cases like left join).
+ NLJState::FetchingRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ match ready!(self.outer_table.poll_next_unpin(cx)) {
+ Some(Ok(right_batch)) => {
+ let right_batch_size = right_batch.num_rows();
+
+ // Skip the empty batch
+ if right_batch_size == 0 {
+ continue;
+ }
+
+ self.current_right_batch = Some(right_batch);
+
+ // Prepare right bitmap
+ if self.should_track_unmatched_right {
+ let zeroed_buf =
+ BooleanBuffer::new_unset(right_batch_size);
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(zeroed_buf, None));
+ }
+
+ self.l_probe_idx = 0;
+ self.state = NLJState::ProbeRight;
+ continue;
+ }
+ Some(Err(e)) => return Poll::Ready(Some(Err(e))),
+ None => {
+ // Right stream exhausted/as
+ self.state = NLJState::EmitLeftUnmatched;
+ continue;
+ }
+ }
}
- NestedLoopJoinStreamState::ProcessProbeBatch(_) => {
- let poll = handle_state!(self.process_probe_batch());
- self.join_metrics.baseline.record_poll(poll)
+
+ // NLJState transitions:
+ // 1. --> ProbeRight(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> ProbeRight(2)
+ // After probing one right batch, and evaluating the
+ // join filter on (left-row x right-batch), it will advance
+ // to the next left row, then re-enter the current state and
+ // continue joining.
+ // 3. --> FetchRight
+ // After it has done with the current right batch (to join
+ // with all rows in the left buffer), it will go to
+ // FetchRight state to check what to do next.
+ NLJState::ProbeRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // Process current probe state
+ match self.process_probe_batch() {
+ // State unchanged (ProbeRight)
+ // Continue probing until we have done joining the
+ // current right batch with all buffered left rows.
+ Ok(true) => continue,
+ // To next FetchRightState
+ // We have finished joining
+ // (cur_right_batch x buffered_left_batches)
+ Ok(false) => {
+ // Left exhausted, transition to FetchingRight
+ self.l_probe_idx = 0;
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ self.state = NLJState::EmitRightUnmatched;
+ } else {
+ self.current_right_batch = None;
+ self.state = NLJState::FetchingRight;
+ }
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::PrepareUnmatchedBuildRows => {
- handle_state!(self.prepare_unmatched_output_indices())
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been output by the join. In this state we have to
+ // output unmatched rows for current right batch (with null
+ // padding for left relation)
+ // Precondition: we have checked the join type so that it's
+ // possible to output right unmatched (e.g. it's right join)
+ NLJState::EmitRightUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ debug_assert!(self.current_right_batch.is_some());
+ debug_assert!(self.current_right_batch_matched.is_some());
+
+ // Construct the result batch for unmatched right rows using a utility function
+ if let Some(batch) = self.process_right_unmatched()? {
+ self.output_buffer.push_batch(batch)?;
+ }
+
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
}
- NestedLoopJoinStreamState::OutputUnmatchedBuildRows(_) => {
- let poll = handle_state!(self.build_unmatched_output());
- self.join_metrics.baseline.record_poll(poll)
+
+ // NLJState transitions:
+ // 1. --> EmitLeftUnmatched(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> EmitLeftUnmatched(2)
+ // After processing some unmatched rows, it will re-enter
+ // the same state, to check if there are any more final
+ // results to output.
+ // 3. --> Done
+ // It has processed all data, go to the final state and ready
+ // to exit.
+ //
+ // TODO: For memory-limited case, go back to `BufferingLeft`
+ // state again.
+ NLJState::EmitLeftUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // Process current unmatched state
+ match self.process_left_unmatched() {
+ // State unchanged (EmitLeftUnmatched)
+ // Continue processing until we have processed all unmatched rows
+ Ok(true) => continue,
+ // To Done state
+ // We have finished processing all unmatched rows
+ Ok(false) => {
+ self.output_buffer.finish_buffered_batch()?;
+ self.state = NLJState::Done;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::Completed => Poll::Ready(None),
- };
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any remaining completed batches before final termination
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // HACK for the doc test in https://github.com/apache/datafusion/blob/main/datafusion/core/src/dataframe/mod.rs#L1265
+ // If this operator directly return `Poll::Ready(None)`
+ // for empty result, the final result will become an empty
+ // batch with empty schema, however the expected result
+ // should be with the expected schema for this operator
+ if !self.handled_empty_output {
+ let zero_count = Count::new();
+ if *self.join_metrics.baseline.output_rows() == zero_count {
+ let empty_batch =
+ RecordBatch::new_empty(Arc::clone(&self.output_schema));
+ self.handled_empty_output = true;
+ return Poll::Ready(Some(Ok(empty_batch)));
+ }
+ }
+
+ return Poll::Ready(None);
+ }
+ }
}
}
+}
- // This function's main job is to construct an output `RecordBatch` based on pre-calculated join indices.
- // It operates in a chunk-based manner, meaning it processes a portion of the results in each call,
- // making it suitable for streaming large datasets without high memory consumption.
- // This function behaves like an iterator. It returns `Ok(None)`
- // to signal that the result stream is exhausted and there is no more data.
- fn get_next_join_result(&mut self) -> Result<Option<RecordBatch>> {
- let status = self.join_result_status.as_mut().ok_or_else(|| {
- internal_datafusion_err!(
- "get_next_join_result called without initializing join_result_status"
- )
- })?;
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
+ }
+}
- let (left_indices, right_indices, current_start) = (
- &status.build_indices,
- &status.probe_indices,
- status.processed_count,
+impl NLJStream {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new(
+ schema: Arc<Schema>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ outer_table: SendableRecordBatchStream,
+ inner_table: OnceFut<JoinLeftData>,
+ column_indices: Vec<ColumnIndex>,
+ join_metrics: BuildProbeJoinMetrics,
+ cfg_batch_size: usize,
+ ) -> Self {
+ let should_track_unmatched_right = matches!(
+ join_type,
+ JoinType::Full
+ | JoinType::Right
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
);
- let left_batch = self
- .left_data
- .as_ref()
- .ok_or_else(|| internal_datafusion_err!("should have left_batch"))?
- .batch();
+ Self {
+ output_schema: Arc::clone(&schema),
+ join_filter: filter,
+ join_type,
+ outer_table,
+ column_indices,
+ inner_table,
+ join_metrics,
+ buffered_left_data: None,
+ output_buffer: Box::new(BatchCoalescer::new(schema, cfg_batch_size)),
+ cfg_batch_size,
+ current_right_batch: None,
+ current_right_batch_matched: None,
+ state: NLJState::BufferingLeft,
+ l_probe_idx: 0,
+ l_emit_idx: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
- let right_batch = match &self.state {
- NestedLoopJoinStreamState::ProcessProbeBatch(record_batch)
- | NestedLoopJoinStreamState::OutputUnmatchedBuildRows(record_batch) => {
- record_batch
- }
- _ => {
- return internal_err!(
- "State should be ProcessProbeBatch or OutputUnmatchedBuildRows"
- )
- }
- };
+ // ==== Core logic handling for each state ====
- if left_indices.is_empty() && right_indices.is_empty() && current_start == 0 {
- // To match the behavior of the previous implementation, return an empty RecordBatch.
- let res = RecordBatch::new_empty(Arc::clone(&self.schema));
- status.processed_count = 1;
- return Ok(Some(res));
- }
+ /// Returns bool to indicate should it continue probing
+ /// true -> continue in the same ProbeRight state
+ /// false -> It has done with the (buffered_left x cur_right_batch), go to
+ /// next state (ProbeRight)
+ fn process_probe_batch(&mut self) -> Result<bool> {
+ let left_data = Arc::clone(self.get_left_data()?);
+ let right_batch = self
+ .current_right_batch
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("Right batch should be available"))?
+ .clone();
- if matches!(self.join_type, JoinType::RightSemi | JoinType::RightAnti) {
- // in this case left_indices.num_rows() == 0
- let end = min(
- current_start + self.intermediate_batch_size,
- right_indices.len(),
- );
+ // stop probing, the caller will go to the next state
+ if self.l_probe_idx >= left_data.batch().num_rows() {
+ return Ok(false);
+ }
- if current_start >= end {
- return Ok(None);
- }
+ // ========
+ // Join (l_row x right_batch)
+ // and push the result into output_buffer
+ // ========
- let res = Some(build_batch_from_indices(
- &self.schema,
- left_batch,
- right_batch,
- left_indices,
- &right_indices.slice(current_start, end - current_start),
- &self.column_indices,
- JoinSide::Left,
- )?);
+ let l_idx = self.l_probe_idx;
+ let join_batch =
+ self.process_single_left_row_join(&left_data, &right_batch, l_idx)?;
- status.processed_count = end;
- return Ok(res);
+ if let Some(batch) = join_batch {
+ self.output_buffer.push_batch(batch)?;
}
- if current_start >= left_indices.len() {
+ // ==== Prepare for the next iteration ====
+
+ // Advance left cursor
+ self.l_probe_idx += 1;
+
+ // Return true to continue probing
+ Ok(true)
+ }
+
+ /// Process a single left row join with the current right batch.
+ /// Returns a RecordBatch containing the join results (None if empty)
+ fn process_single_left_row_join(
+ &mut self,
+ left_data: &JoinLeftData,
+ right_batch: &RecordBatch,
+ l_index: usize,
+ ) -> Result<Option<RecordBatch>> {
+ let right_row_count = right_batch.num_rows();
+ if right_row_count == 0 {
return Ok(None);
}
- let end = min(
- current_start + self.intermediate_batch_size,
- left_indices.len(),
- );
+ let cur_right_bitmap = if let Some(filter) = &self.join_filter {
+ apply_filter_to_row_join_batch(
+ left_data.batch(),
+ l_index,
+ right_batch,
+ filter,
+ )?
+ } else {
+ BooleanArray::from(vec![true; right_row_count])
+ };
- let left_indices = &left_indices.slice(current_start, end - current_start);
- let right_indices = &right_indices.slice(current_start, end - current_start);
+ self.update_matched_bitmap(l_index, &cur_right_bitmap)?;
- // Switch around the build side and probe side for `JoinType::RightMark`
- // because in a RightMark join, we want to mark rows on the right table
- // by looking for matches in the left.
- let res = if self.join_type == JoinType::RightMark {
- build_batch_from_indices(
- &self.schema,
- right_batch,
- left_batch,
- left_indices,
- right_indices,
- &self.column_indices,
- JoinSide::Right,
- )
+ // For the following join types: here we only have to set the left/right
+ // bitmap, and no need to output result
+ if matches!(
+ self.join_type,
+ JoinType::LeftAnti
+ | JoinType::LeftSemi
+ | JoinType::LeftMark
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ ) {
+ return Ok(None);
+ }
+
+ if cur_right_bitmap.true_count() == 0 {
+ // If none of the pairs has passed the join predicate/filter
+ Ok(None)
} else {
- build_batch_from_indices(
- &self.schema,
- left_batch,
+ // Use the optimized approach similar to build_intermediate_batch_for_single_left_row
+ let join_batch = build_row_join_batch(
+ &self.output_schema,
+ left_data.batch(),
+ l_index,
right_batch,
- left_indices,
- right_indices,
+ Some(cur_right_bitmap),
&self.column_indices,
JoinSide::Left,
- )
- }?;
+ )?;
+ Ok(join_batch)
+ }
+ }
- status.processed_count = end;
+ /// Returns bool to indicate should it continue processing unmatched rows
+ /// true -> continue in the same EmitLeftUnmatched state
+ /// false -> next state (Done)
+ fn process_left_unmatched(&mut self) -> Result<bool> {
+ let left_data = self.get_left_data()?;
+ let left_batch = left_data.batch();
- Ok(Some(res))
- }
+ // Early return if join type can't have unmatched rows
+ if !need_produce_result_in_final(self.join_type) {
+ return Ok(false);
+ }
- fn collect_build_side(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- let build_timer = self.join_metrics.build_time.timer();
- // build hash table from left (build) side, if not yet done
- self.left_data = Some(ready!(self.inner_table.get_shared(cx))?);
- build_timer.done();
+ // Early return if another thread is already processing unmatched rows
+ if self.l_emit_idx == 0 && !left_data.report_probe_completed() {
+ return Ok(false);
+ }
- self.state = NestedLoopJoinStreamState::FetchProbeBatch;
+ // Stop processing unmatched rows, the caller will go to the next state
+ if self.l_emit_idx >= left_batch.num_rows() {
+ return Ok(false);
+ }
- Poll::Ready(Ok(StatefulStreamResult::Continue))
- }
+ // ========
+ // Process unmatched rows and push the result into output_buffer
+ // Each time, the number to process is up to batch size
+ // ========
+ let start_idx = self.l_emit_idx;
+ let end_idx =
+ std::cmp::min(start_idx + self.cfg_batch_size, left_batch.num_rows());
- /// Fetches next batch from probe-side
- ///
- /// If a non-empty batch has been fetched, updates state to
- /// `ProcessProbeBatchState`, otherwise updates state to `ExhaustedProbeSide`.
- fn fetch_probe_batch(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- match ready!(self.outer_table.poll_next_unpin(cx)) {
- None => {
- self.state = NestedLoopJoinStreamState::PrepareUnmatchedBuildRows;
- }
- Some(Ok(right_batch)) => {
- self.join_metrics.input_batches.add(1);
- self.join_metrics.input_rows.add(right_batch.num_rows());
+ if let Some(batch) =
+ self.process_left_unmatched_range(left_data, start_idx, end_idx)?
+ {
+ self.output_buffer.push_batch(batch)?;
+ }
- self.state = NestedLoopJoinStreamState::ProcessProbeBatch(right_batch);
- }
- Some(Err(err)) => return Poll::Ready(Err(err)),
- };
+ // ==== Prepare for the next iteration ====
+ self.l_emit_idx = end_idx;
- Poll::Ready(Ok(StatefulStreamResult::Continue))
+ // Return true to continue processing unmatched rows
+ Ok(true)
}
- /// Joins current probe batch with build-side data and produces batch with
- /// matched output, updates state to `FetchProbeBatch`.
- fn process_probe_batch(
- &mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let Some(left_data) = self.left_data.clone() else {
- return internal_err!(
- "Expected left_data to be Some in ProcessProbeBatch state"
+ /// Process unmatched rows from the left data within the specified range.
+ /// Returns a RecordBatch containing the unmatched rows (None if empty).
+ ///
+ /// # Arguments
+ /// * `left_data` - The left side data containing the batch and bitmap
+ /// * `start_idx` - Start index (inclusive) of the range to process
+ /// * `end_idx` - End index (exclusive) of the range to process
+ ///
+ /// # Safety
+ /// The caller is responsible for ensuring that `start_idx` and `end_idx` are
+ /// within valid bounds of the left batch. This function does not perform
+ /// bounds checking.
+ fn process_left_unmatched_range(
+ &self,
+ left_data: &JoinLeftData,
+ start_idx: usize,
+ end_idx: usize,
+ ) -> Result<Option<RecordBatch>> {
+ if start_idx == end_idx {
+ return Ok(None);
+ }
+
+ // Slice both left batch, and bitmap to range [start_idx, end_idx)
+ // The range is bit index (not byte)
+ let left_batch = left_data.batch();
+ let left_batch_sliced = left_batch.slice(start_idx, end_idx - start_idx);
+
+ // Can this be more efficient?
+ let mut bitmap_sliced = BooleanBufferBuilder::new(end_idx - start_idx);
+ bitmap_sliced.append_n(end_idx - start_idx, false);
+ let bitmap = left_data.bitmap().lock();
+ for i in start_idx..end_idx {
+ assert!(
+ i - start_idx < bitmap_sliced.capacity(),
+ "DBG: {start_idx}, {end_idx}"
);
- };
- let visited_left_side = left_data.bitmap();
- let batch = self.state.try_as_process_probe_batch()?;
+ bitmap_sliced.set_bit(i - start_idx, bitmap.get_bit(i));
+ }
+ let bitmap_sliced = BooleanArray::new(bitmap_sliced.finish(), None);
+
+ build_unmatched_batch(
+ Arc::clone(&self.output_schema),
+ &left_batch_sliced,
+ bitmap_sliced,
+ self.outer_table.schema(),
+ &self.column_indices,
+ self.join_type,
+ JoinSide::Left,
+ )
+ }
- let binding = self.join_metrics.join_time.clone();
- let _timer = binding.timer();
+ /// Process unmatched rows from the current right batch and reset the bitmap.
+ /// Returns a RecordBatch containing the unmatched right rows (None if empty).
+ fn process_right_unmatched(&mut self) -> Result<Option<RecordBatch>> {
+ // ==== Take current right batch and its bitmap ====
+ let right_batch_bitmap: BooleanArray =
+ std::mem::take(&mut self.current_right_batch_matched).ok_or_else(|| {
+ internal_datafusion_err!("right bitmap should be available")
+ })?;
+
+ let right_batch = self.current_right_batch.take();
+ let cur_right_batch = unwrap_or_internal_err!(right_batch);
+
+ let left_data = self.get_left_data()?;
+ let left_schema = left_data.batch().schema();
+
+ let res = build_unmatched_batch(
+ Arc::clone(&self.output_schema),
+ &cur_right_batch,
+ right_batch_bitmap,
+ left_schema,
+ &self.column_indices,
+ self.join_type,
+ JoinSide::Right,
+ );
- if self.join_result_status.is_none() {
- let (left_side_indices, right_side_indices) = join_left_and_right_batch(
- left_data.batch(),
- batch,
- self.join_type,
- self.filter.as_ref(),
- visited_left_side,
- &mut self.indices_cache,
- self.right_side_ordered,
- self.intermediate_batch_size,
- )?;
- self.join_result_status = Some(JoinResultProgress {
- build_indices: left_side_indices,
- probe_indices: right_side_indices,
- processed_count: 0,
- })
- }
+ // ==== Clean-up ====
+ self.current_right_batch_matched = None;
- let join_result = self.get_next_join_result()?;
+ res
+ }
- match join_result {
- Some(res) => {
- self.join_metrics.output_batches.add(1);
- Ok(StatefulStreamResult::Ready(Some(res)))
- }
- None => {
- self.state = NestedLoopJoinStreamState::FetchProbeBatch;
- self.join_result_status = None;
- Ok(StatefulStreamResult::Continue)
- }
- }
+ // ==== Utilities ====
+
+ /// Get the inner data of left data, errors if it's None
+ fn get_left_data(&self) -> Result<&Arc<JoinLeftData>> {
+ self.buffered_left_data
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("LeftData should be available"))
}
- fn build_unmatched_output(
- &mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let binding = self.join_metrics.join_time.clone();
- let _timer = binding.timer();
-
- let res = self.get_next_join_result()?;
- match res {
- Some(res) => {
- self.join_metrics.output_batches.add(1);
- Ok(StatefulStreamResult::Ready(Some(res)))
- }
- None => {
- self.state = NestedLoopJoinStreamState::Completed;
- Ok(StatefulStreamResult::Ready(None))
+ /// Flush the `output_buffer` if there are batches ready to output
+ /// None if no result batch ready.
+ fn maybe_flush_ready_batch(&mut self) -> Option<Poll<Option<Result<RecordBatch>>>> {
+ if self.output_buffer.has_completed_batch() {
+ if let Some(batch) = self.output_buffer.next_completed_batch() {
+ let poll = Poll::Ready(Some(Ok(batch)));
+ return Some(self.join_metrics.baseline.record_poll(poll));
}
}
+
+ None
}
- /// This function's primary purpose is to handle the final output stage required by specific join types after all right-side (probe) data has been exhausted.
- /// It is critically important for LEFT*/FULL joins, which must emit left-side (build) rows that found no match. For these cases, it identifies the unmatched rows and prepares the necessary state to output them.
- fn prepare_unmatched_output_indices(
+ /// After joining (l_index@left_buffer x current_right_batch), it will result
+ /// in a bitmap (the same length as current_right_batch) as the join match
+ /// result. Use this bitmap to update the global bitmap, for special join
+ /// types like full joins.
+ ///
+ /// Example:
+ /// After joining l_index=1 (1-indexed row in the left buffer), and the
+ /// current right batch with 3 elements, this function will be called with
+ /// arguments: l_index = 1, r_matched = [false, false, true]
+ /// - If the join type is FullJoin, the 1-index in the left bitmap will be
+ /// set to true, and also the right bitmap will be bitwise-ORed with the
+ /// input r_matched bitmap.
+ /// - For join types that don't require output unmatched rows, this
+ /// function can be a no-op. For inner joins, this function is a no-op; for left
+ /// joins, only the left bitmap may be updated.
+ fn update_matched_bitmap(
&mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let Some(left_data) = self.left_data.clone() else {
- return internal_err!(
- "Expected left_data to be Some in ExhaustedProbeSide state"
- );
- };
- let visited_left_side = left_data.bitmap();
- if need_produce_result_in_final(self.join_type) {
- // At this stage `visited_left_side` won't be updated, so it's
- // safe to report about probe completion.
- //
- // Setting `is_exhausted` / returning None will prevent from
- // multiple calls of `report_probe_completed()`
- if !left_data.report_probe_completed() {
- self.state = NestedLoopJoinStreamState::Completed;
- return Ok(StatefulStreamResult::Ready(None));
- };
+ l_index: usize,
+ r_matched: &BooleanArray,
+ ) -> Result<()> {
+ let left_data = self.get_left_data()?;
- // Only setting up timer, input is exhausted
- let _timer = self.join_metrics.join_time.timer();
- // use the global left bitmap to produce the left indices and right indices
- let (left_side, right_side) =
- get_final_indices_from_shared_bitmap(visited_left_side, self.join_type);
-
- self.join_result_status = Some(JoinResultProgress {
- build_indices: left_side,
- probe_indices: right_side,
- processed_count: 0,
- });
- self.state = NestedLoopJoinStreamState::OutputUnmatchedBuildRows(
- RecordBatch::new_empty(self.outer_table.schema()),
- );
+ // number of successfully joined pairs from (l_index x cur_right_batch)
+ let joined_len = r_matched.true_count();
- Ok(StatefulStreamResult::Continue)
- } else {
- // end of the join loop
- self.state = NestedLoopJoinStreamState::Completed;
- Ok(StatefulStreamResult::Ready(None))
+ // 1. Maybe update the left bitmap
+ if need_produce_result_in_final(self.join_type) && (joined_len > 0) {
+ let mut bitmap = left_data.bitmap().lock();
+ bitmap.set_bit(l_index, true);
+ }
+
+ // 2. Maybe updateh the right bitmap
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ // after bit-wise or, it will be put back
+ let right_bitmap = std::mem::take(&mut self.current_right_batch_matched)
+ .ok_or_else(|| {
+ internal_datafusion_err!("right batch's bitmap should be present")
+ })?;
+ let (buf, nulls) = right_bitmap.into_parts();
+ debug_assert!(nulls.is_none());
+ let updated_right_bitmap = buf.bitor(r_matched.values());
+
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(updated_right_bitmap, None));
}
+
+ Ok(())
}
}
-#[allow(clippy::too_many_arguments)]
-fn join_left_and_right_batch(
+// ==== Utilities ====
+
+/// Apply the join filter between:
+/// (l_index th row in left buffer) x (right batch)
+/// Returns a bitmap, with successfully joined indices set to true
+fn apply_filter_to_row_join_batch(
left_batch: &RecordBatch,
+ l_index: usize,
right_batch: &RecordBatch,
- join_type: JoinType,
- filter: Option<&JoinFilter>,
- visited_left_side: &SharedBitmapBuilder,
- indices_cache: &mut (UInt64Array, UInt32Array),
- right_side_ordered: bool,
- max_intermediate_batch_size: usize,
-) -> Result<(PrimitiveArray<UInt64Type>, PrimitiveArray<UInt32Type>)> {
- let (left_side, right_side) = build_join_indices(
- left_batch,
- right_batch,
- filter,
- indices_cache,
- max_intermediate_batch_size,
- )
- .map_err(|e| {
- exec_datafusion_err!(
- "Fail to build join indices in NestedLoopJoinExec, error: {e}"
- )
- })?;
-
- // set the left bitmap
- // and only full join need the left bitmap
- if need_produce_result_in_final(join_type) {
- let mut bitmap = visited_left_side.lock();
- left_side.values().iter().for_each(|x| {
- bitmap.set_bit(*x as usize, true);
- });
- }
- // adjust the two side indices base on the join type
- let (left_side, right_side) = adjust_indices_by_join_type(
- left_side,
- right_side,
- 0..right_batch.num_rows(),
- join_type,
- right_side_ordered,
- )?;
+ filter: &JoinFilter,
+) -> Result<BooleanArray> {
+ debug_assert!(left_batch.num_rows() != 0 && right_batch.num_rows() != 0);
+
+ let intermediate_batch = if filter.schema.fields().is_empty() {
+ // If filter is constant (e.g. literal `true`), empty batch can be used
+ // in the later filter step.
+ create_record_batch_with_empty_schema(
+ Arc::new((*filter.schema).clone()),
+ right_batch.num_rows(),
+ )?
+ } else {
+ build_row_join_batch(
+ &filter.schema,
+ left_batch,
+ l_index,
+ right_batch,
+ None,
+ &filter.column_indices,
+ JoinSide::Left,
+ )?
+ .ok_or_else(|| internal_datafusion_err!("This function assume input batch is not empty, so the intermediate batch can't be empty too"))?
+ };
+
+ let filter_result = filter
+ .expression()
+ .evaluate(&intermediate_batch)?
+ .into_array(intermediate_batch.num_rows())?;
- Ok((left_side, right_side))
+ Ok(as_boolean_array(&filter_result)?.clone())
}
-impl Stream for NestedLoopJoinStream {
- type Item = Result<RecordBatch>;
+/// This function performs the following steps:
+/// 1. Apply filter to probe-side batch
+/// 2. Broadcast the left row (build_side_batch\[build_side_index\]) to the
+/// filtered probe-side batch
+/// 3. Concat them together according to `col_indices`, and return the result
+/// (None if the result is empty)
+///
+/// Example:
+/// build_side_batch:
+/// a
+/// ----
+/// 1
+/// 2
+/// 3
+///
+/// # 0 index element in the build_side_batch (that is `1`) will be used
+/// build_side_index: 0
+///
+/// probe_side_batch:
+/// b
+/// ----
+/// 10
+/// 20
+/// 30
+/// 40
+///
+/// # After applying it, only index 1 and 3 elemnt in probe_side_batch will be
+/// # kept
+/// probe_side_filter:
+/// false
+/// true
+/// false
+/// true
+///
+///
+/// # Projections to the build/probe side batch, to construct the output batch
+/// col_indices:
+/// [(left, 0), (right, 0)]
+///
+/// build_side: left
+///
+/// ====
+/// Result batch:
+/// a b
+/// ----
+/// 1 20
+/// 1 40
+fn build_row_join_batch(
+ output_schema: &Schema,
+ build_side_batch: &RecordBatch,
+ build_side_index: usize,
+ probe_side_batch: &RecordBatch,
+ probe_side_filter: Option<BooleanArray>,
+ // See [`NLJStream`] struct's `column_indices` field for more detail
+ col_indices: &[ColumnIndex],
+ // If the build side is left or right, used to interpret the side information
+ // in `col_indices`
+ build_side: JoinSide,
+) -> Result<Option<RecordBatch>> {
+ debug_assert!(build_side != JoinSide::None);
+
+ // TODO(perf): since the output might be projection of right batch, this
+ // filtering step is more efficient to be done inside the column_index loop
+ let filtered_probe_batch = if let Some(filter) = probe_side_filter {
+ &filter_record_batch(probe_side_batch, &filter)?
+ } else {
+ probe_side_batch
+ };
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
+ if filtered_probe_batch.num_rows() == 0 {
+ return Ok(None);
+ }
+
+ let mut columns: Vec<Arc<dyn Array>> =
+ Vec::with_capacity(output_schema.fields().len());
+
+ for column_index in col_indices {
+ let array = if column_index.side == build_side {
+ // Broadcast the single build-side row to match the filtered
+ // probe-side batch length
+ let original_left_array = build_side_batch.column(column_index.index);
+ let scalar_value = ScalarValue::try_from_array(
+ original_left_array.as_ref(),
+ build_side_index,
+ )?;
+ scalar_value.to_array_of_size(filtered_probe_batch.num_rows())?
Review Comment:
I was expecting this optimization to bring > 10% speedup to the NLJ micro-bench, based on the flamegraphs.
Though I think this change requires some transformation to `JoinFilter`, and we don't have such utility function yet. I'll open an issue to track this idea 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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "Omega359 (via GitHub)" <gi...@apache.org>.
Omega359 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3186369426
Can someone please run the extended tests on this branch? I am pretty sure that the extended tests are failing against this branch but I am having issues with git locally so I'd like confirmation.
--
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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "Omega359 (via GitHub)" <gi...@apache.org>.
Omega359 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3186490312
Here is an example of a query that passes in main and fails on this branch:
```sql
CREATE TABLE tab0(col0 INTEGER, col1 INTEGER, col2 INTEGER);
CREATE TABLE tab1(col0 INTEGER, col1 INTEGER, col2 INTEGER);
CREATE TABLE tab2(col0 INTEGER, col1 INTEGER, col2 INTEGER);
INSERT INTO tab0 VALUES(97,1,99);
INSERT INTO tab0 VALUES(15,81,47);
INSERT INTO tab0 VALUES(87,21,10);
INSERT INTO tab1 VALUES(51,14,96);
INSERT INTO tab1 VALUES(85,5,59);
INSERT INTO tab1 VALUES(91,47,68);
INSERT INTO tab2 VALUES(64,77,40);
INSERT INTO tab2 VALUES(75,67,58);
INSERT INTO tab2 VALUES(46,51,23);
SELECT DISTINCT 32 AS col2 FROM tab0 AS cor0 LEFT OUTER JOIN tab2 AS cor1 ON ( NULL ) IS NULL;
Arrow error: Invalid argument error: must either specify a row count or at least one column
```
There are a bunch of other failures where the result count between main and this branch do not match. I verified a bunch of the queries do have NestedLoopJoinExec in the explain result.
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2262497844
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -63,92 +64,95 @@ use datafusion_physical_expr::equivalence::{
join_equivalence_properties, ProjectionMapping,
};
-use futures::{ready, Stream, StreamExt, TryStreamExt};
+use futures::{Stream, StreamExt, TryStreamExt};
+use log::debug;
use parking_lot::Mutex;
-/// Left (build-side) data
-struct JoinLeftData {
- /// Build-side data collected to single batch
- batch: RecordBatch,
- /// Shared bitmap builder for visited left indices
- bitmap: SharedBitmapBuilder,
- /// Counter of running probe-threads, potentially able to update `bitmap`
- probe_threads_counter: AtomicUsize,
- /// Memory reservation for tracking batch and bitmap
- /// Cleared on `JoinLeftData` drop
- /// reservation is cleared on Drop
- #[expect(dead_code)]
- reservation: MemoryReservation,
-}
-
-impl JoinLeftData {
- fn new(
- batch: RecordBatch,
- bitmap: SharedBitmapBuilder,
- probe_threads_counter: AtomicUsize,
- reservation: MemoryReservation,
- ) -> Self {
- Self {
- batch,
- bitmap,
- probe_threads_counter,
- reservation,
- }
- }
-
- fn batch(&self) -> &RecordBatch {
- &self.batch
- }
-
- fn bitmap(&self) -> &SharedBitmapBuilder {
- &self.bitmap
- }
-
- /// Decrements counter of running threads, and returns `true`
- /// if caller is the last running thread
- fn report_probe_completed(&self) -> bool {
- self.probe_threads_counter.fetch_sub(1, Ordering::Relaxed) == 1
- }
-}
-
#[allow(rustdoc::private_intra_doc_links)]
-/// NestedLoopJoinExec is build-probe join operator, whose main task is to
-/// perform joins without any equijoin conditions in `ON` clause.
+/// NestedLoopJoinExec is a build-probe join operator designed for joins that
+/// do not have equijoin keys in their `ON` clause.
///
-/// Execution consists of following phases:
+/// # Execution Flow
///
-/// #### 1. Build phase
-/// Collecting build-side data in memory, by polling all available data from build-side input.
-/// Due to the absence of equijoin conditions, it's not possible to partition build-side data
-/// across multiple threads of the operator, so build-side is always collected in a single
-/// batch shared across all threads.
-/// The operator always considers LEFT input as build-side input, so it's crucial to adjust
-/// smaller input to be the LEFT one. Normally this selection is handled by physical optimizer.
+/// ```text
+/// Incoming right batch
+/// Left Side Buffered Batches
+/// ┌───────────┐ ┌───────────────┐
+/// │ ┌───────┐ │ │ │
+/// │ │ │ │ │ │
+/// Current Left Row ───▶│ ├───────├─┤──────────┐ │ │
+/// │ │ │ │ │ └───────────────┘
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ └───────┘ │ │ │
+/// │ ┌───────┐ │ │ │
+/// │ │ │ │ │ ┌─────┘
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ └───────┘ │ ▼ ▼
+/// │ ...... │ ┌──────────────────────┐
+/// │ │ │X (Cartesian Product) │
+/// │ │ └──────────┬───────────┘
+/// └───────────┘ │
+/// │
+/// ▼
+/// ┌───────┬───────────────┐
+/// │ │ │
+/// │ │ │
+/// │ │ │
+/// └───────┴───────────────┘
+/// Intermediate Batch
+/// (For join predicate evaluation)
+/// ```
///
-/// #### 2. Probe phase
-/// Sequentially polling batches from the probe-side input and processing them according to the
-/// following logic:
-/// - apply join filter (`ON` clause) to Cartesian product of probe batch and build side data
-/// -- filter evaluation is executed once per build-side data row
-/// - update shared bitmap of joined ("visited") build-side row indices, if required -- allows
-/// to produce unmatched build-side data in case of e.g. LEFT/FULL JOIN after probing phase
-/// completed
-/// - perform join index alignment is required -- depending on `JoinType`
-/// - produce output join batch
+/// The execution follows a two-phase design:
///
-/// Probing phase is executed in parallel, according to probe-side input partitioning -- one
-/// thread per partition. After probe input is exhausted, each thread **ATTEMPTS** to produce
-/// unmatched build-side data.
+/// ## 1. Buffering Left Input
+/// - The operator eagerly buffers all left-side input batches into memory,
+/// unless a memory limit is reached (see 'Memory-limited Execution').
Review Comment:
I have updated the comment to explain 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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2251087268
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1048 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
-
- // Produces 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
- *left_indices_cache = UInt64Array::from_iter_values(
- (0..output_row_count as u64).map(|i| i % left_row_count as u64),
- );
-
- // Produces 0, 0, 0, 1, 1, 1, 2, 2, 2, ...
- *right_indices_cache = UInt32Array::from_iter_values(
- (0..output_row_count as u32).map(|i| i / left_row_count as u32),
- );
-
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- };
-
- if let Some(filter) = filter {
- apply_join_filter_to_indices(
- left_batch,
- right_batch,
- left_indices,
- right_indices,
- filter,
- JoinSide::Left,
- Some(max_intermediate_batch_size),
- )
- } else {
- Ok((left_indices, right_indices))
- }
-}
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
-impl NestedLoopJoinStream {
- fn poll_next_impl(
- &mut self,
+ /// See the comments [`NestedLoopJoinExec`] for high-level design ideas.
+ ///
+ /// # Implementation
+ ///
+ /// This function is the entry point of NLJ operator's state machine
+ /// transitions. The rough state transition graph is as follow, for more
+ /// details see the comment in each state's matching arm.
+ ///
+ /// ============================
+ /// State transition graph:
+ /// ============================
+ ///
+ /// (start) --> BufferingLeft
+ /// ----------------------------
+ /// BufferingLeft → FetchingRight
+ ///
+ /// FetchingRight → ProbeRight (if right batch available)
+ /// FetchingRight → EmitLeftUnmatched (if right exhausted)
+ ///
+ /// ProbeRight → ProbeRight (next left row or after yielding output)
+ /// ProbeRight → EmitRightUnmatched (for special join types like right join)
+ /// ProbeRight → FetchingRight (done with the current right batch)
+ ///
+ /// EmitRightUnmatched → FetchingRight
+ ///
+ /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each
+ /// iteration)
+ /// EmitLeftUnmatched → Done (if finished)
+ /// ----------------------------
+ /// Done → (end)
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
+ ) -> Poll<Option<Self::Item>> {
loop {
- return match self.state {
- NestedLoopJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
+ match self.state {
+ // # NLJState transitions
+ // --> FetchingRight
+ // This state will prepare the left side batches, next state
+ // `FetchingRight` is responsible for preparing a single probe
+ // side batch, before start joining.
+ NLJState::BufferingLeft => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ match ready!(self.inner_table.get_shared(cx)) {
+ Ok(left_data) => {
+ self.buffered_left_data = Some(left_data);
+ // TODO: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
+
+ // # NLJState transitions:
+ // 1. --> ProbeRight
+ // Start processing the join for the newly fetched right
+ // batch.
+ // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit
+ // unmatched left side rows.
+ //
+ // After fetching a new batch from the right side, it will
+ // process all rows from the buffered left data:
+ // ```text
+ // for batch in right_side:
+ // for row in left_buffer:
+ // join(batch, row)
+ // ```
+ // Note: the implementation does this step incrementally,
+ // instead of materializing all intermediate Cartesian products
+ // at once in memory.
+ //
+ // So after the right side input is exhausted, the join phase
+ // for the current buffered left data is finished. We can go to
+ // the next `EmitLeftUnmatched` phase to check if there is any
+ // special handling (e.g., in cases like left join).
+ NLJState::FetchingRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ match ready!(self.outer_table.poll_next_unpin(cx)) {
+ Some(Ok(right_batch)) => {
+ let right_batch_size = right_batch.num_rows();
+
+ // Skip the empty batch
+ if right_batch_size == 0 {
+ continue;
+ }
+
+ self.current_right_batch = Some(right_batch);
+
+ // Prepare right bitmap
+ if self.should_track_unmatched_right {
+ let zeroed_buf =
+ BooleanBuffer::new_unset(right_batch_size);
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(zeroed_buf, None));
+ }
+
+ self.l_probe_idx = 0;
+ self.state = NLJState::ProbeRight;
+ continue;
+ }
+ Some(Err(e)) => return Poll::Ready(Some(Err(e))),
+ None => {
+ // Right stream exhausted/as
+ self.state = NLJState::EmitLeftUnmatched;
+ continue;
+ }
+ }
}
- NestedLoopJoinStreamState::ProcessProbeBatch(_) => {
- let poll = handle_state!(self.process_probe_batch());
- self.join_metrics.baseline.record_poll(poll)
+
+ // NLJState transitions:
+ // 1. --> ProbeRight(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> ProbeRight(2)
+ // After probing one right batch, and evaluating the
+ // join filter on (left-row x right-batch), it will advance
+ // to the next left row, then re-enter the current state and
+ // continue joining.
+ // 3. --> FetchRight
+ // After it has done with the current right batch (to join
+ // with all rows in the left buffer), it will go to
+ // FetchRight state to check what to do next.
+ NLJState::ProbeRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // Process current probe state
+ match self.process_probe_batch() {
+ // State unchanged (ProbeRight)
+ // Continue probing until we have done joining the
+ // current right batch with all buffered left rows.
+ Ok(true) => continue,
+ // To next FetchRightState
+ // We have finished joining
+ // (cur_right_batch x buffered_left_batches)
+ Ok(false) => {
+ // Left exhausted, transition to FetchingRight
+ self.l_probe_idx = 0;
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ self.state = NLJState::EmitRightUnmatched;
+ } else {
+ self.current_right_batch = None;
+ self.state = NLJState::FetchingRight;
+ }
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::PrepareUnmatchedBuildRows => {
- handle_state!(self.prepare_unmatched_output_indices())
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been output by the join. In this state we have to
+ // output unmatched rows for current right batch (with null
+ // padding for left relation)
+ // Precondition: we have checked the join type so that it's
+ // possible to output right unmatched (e.g. it's right join)
+ NLJState::EmitRightUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ debug_assert!(self.current_right_batch.is_some());
+ debug_assert!(self.current_right_batch_matched.is_some());
+
+ // Construct the result batch for unmatched right rows using a utility function
+ if let Some(batch) = self.process_right_unmatched()? {
+ self.output_buffer.push_batch(batch)?;
+ }
+
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
}
- NestedLoopJoinStreamState::OutputUnmatchedBuildRows(_) => {
- let poll = handle_state!(self.build_unmatched_output());
- self.join_metrics.baseline.record_poll(poll)
+
+ // NLJState transitions:
+ // 1. --> EmitLeftUnmatched(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> EmitLeftUnmatched(2)
+ // After processing some unmatched rows, it will re-enter
+ // the same state, to check if there are any more final
+ // results to output.
+ // 3. --> Done
+ // It has processed all data, go to the final state and ready
+ // to exit.
+ //
+ // TODO: For memory-limited case, go back to `BufferingLeft`
+ // state again.
+ NLJState::EmitLeftUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // Process current unmatched state
+ match self.process_left_unmatched() {
+ // State unchanged (EmitLeftUnmatched)
+ // Continue processing until we have processed all unmatched rows
+ Ok(true) => continue,
+ // To Done state
+ // We have finished processing all unmatched rows
+ Ok(false) => {
+ self.output_buffer.finish_buffered_batch()?;
+ self.state = NLJState::Done;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::Completed => Poll::Ready(None),
- };
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any remaining completed batches before final termination
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // HACK for the doc test in https://github.com/apache/datafusion/blob/main/datafusion/core/src/dataframe/mod.rs#L1265
+ // If this operator directly return `Poll::Ready(None)`
+ // for empty result, the final result will become an empty
+ // batch with empty schema, however the expected result
+ // should be with the expected schema for this operator
+ if !self.handled_empty_output {
+ let zero_count = Count::new();
+ if *self.join_metrics.baseline.output_rows() == zero_count {
+ let empty_batch =
+ RecordBatch::new_empty(Arc::clone(&self.output_schema));
+ self.handled_empty_output = true;
+ return Poll::Ready(Some(Ok(empty_batch)));
+ }
+ }
+
+ return Poll::Ready(None);
+ }
+ }
}
}
+}
- // This function's main job is to construct an output `RecordBatch` based on pre-calculated join indices.
- // It operates in a chunk-based manner, meaning it processes a portion of the results in each call,
- // making it suitable for streaming large datasets without high memory consumption.
- // This function behaves like an iterator. It returns `Ok(None)`
- // to signal that the result stream is exhausted and there is no more data.
- fn get_next_join_result(&mut self) -> Result<Option<RecordBatch>> {
- let status = self.join_result_status.as_mut().ok_or_else(|| {
- internal_datafusion_err!(
- "get_next_join_result called without initializing join_result_status"
- )
- })?;
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
+ }
+}
- let (left_indices, right_indices, current_start) = (
- &status.build_indices,
- &status.probe_indices,
- status.processed_count,
+impl NLJStream {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new(
+ schema: Arc<Schema>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ outer_table: SendableRecordBatchStream,
+ inner_table: OnceFut<JoinLeftData>,
+ column_indices: Vec<ColumnIndex>,
+ join_metrics: BuildProbeJoinMetrics,
+ cfg_batch_size: usize,
+ ) -> Self {
+ let should_track_unmatched_right = matches!(
+ join_type,
+ JoinType::Full
+ | JoinType::Right
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
);
- let left_batch = self
- .left_data
- .as_ref()
- .ok_or_else(|| internal_datafusion_err!("should have left_batch"))?
- .batch();
+ Self {
+ output_schema: Arc::clone(&schema),
+ join_filter: filter,
+ join_type,
+ outer_table,
+ column_indices,
+ inner_table,
+ join_metrics,
+ buffered_left_data: None,
+ output_buffer: Box::new(BatchCoalescer::new(schema, cfg_batch_size)),
+ cfg_batch_size,
+ current_right_batch: None,
+ current_right_batch_matched: None,
+ state: NLJState::BufferingLeft,
+ l_probe_idx: 0,
+ l_emit_idx: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
- let right_batch = match &self.state {
- NestedLoopJoinStreamState::ProcessProbeBatch(record_batch)
- | NestedLoopJoinStreamState::OutputUnmatchedBuildRows(record_batch) => {
- record_batch
- }
- _ => {
- return internal_err!(
- "State should be ProcessProbeBatch or OutputUnmatchedBuildRows"
- )
- }
- };
+ // ==== Core logic handling for each state ====
- if left_indices.is_empty() && right_indices.is_empty() && current_start == 0 {
- // To match the behavior of the previous implementation, return an empty RecordBatch.
- let res = RecordBatch::new_empty(Arc::clone(&self.schema));
- status.processed_count = 1;
- return Ok(Some(res));
- }
+ /// Returns bool to indicate should it continue probing
+ /// true -> continue in the same ProbeRight state
+ /// false -> It has done with the (buffered_left x cur_right_batch), go to
+ /// next state (ProbeRight)
+ fn process_probe_batch(&mut self) -> Result<bool> {
+ let left_data = Arc::clone(self.get_left_data()?);
+ let right_batch = self
+ .current_right_batch
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("Right batch should be available"))?
+ .clone();
- if matches!(self.join_type, JoinType::RightSemi | JoinType::RightAnti) {
- // in this case left_indices.num_rows() == 0
- let end = min(
- current_start + self.intermediate_batch_size,
- right_indices.len(),
- );
+ // stop probing, the caller will go to the next state
+ if self.l_probe_idx >= left_data.batch().num_rows() {
+ return Ok(false);
+ }
- if current_start >= end {
- return Ok(None);
- }
+ // ========
+ // Join (l_row x right_batch)
+ // and push the result into output_buffer
+ // ========
- let res = Some(build_batch_from_indices(
- &self.schema,
- left_batch,
- right_batch,
- left_indices,
- &right_indices.slice(current_start, end - current_start),
- &self.column_indices,
- JoinSide::Left,
- )?);
+ let l_idx = self.l_probe_idx;
+ let join_batch =
+ self.process_single_left_row_join(&left_data, &right_batch, l_idx)?;
- status.processed_count = end;
- return Ok(res);
+ if let Some(batch) = join_batch {
+ self.output_buffer.push_batch(batch)?;
}
- if current_start >= left_indices.len() {
+ // ==== Prepare for the next iteration ====
+
+ // Advance left cursor
+ self.l_probe_idx += 1;
+
+ // Return true to continue probing
+ Ok(true)
+ }
+
+ /// Process a single left row join with the current right batch.
+ /// Returns a RecordBatch containing the join results (None if empty)
+ fn process_single_left_row_join(
+ &mut self,
+ left_data: &JoinLeftData,
+ right_batch: &RecordBatch,
+ l_index: usize,
+ ) -> Result<Option<RecordBatch>> {
+ let right_row_count = right_batch.num_rows();
+ if right_row_count == 0 {
return Ok(None);
}
- let end = min(
- current_start + self.intermediate_batch_size,
- left_indices.len(),
- );
+ let cur_right_bitmap = if let Some(filter) = &self.join_filter {
+ apply_filter_to_row_join_batch(
+ left_data.batch(),
+ l_index,
+ right_batch,
+ filter,
+ )?
+ } else {
+ BooleanArray::from(vec![true; right_row_count])
+ };
- let left_indices = &left_indices.slice(current_start, end - current_start);
- let right_indices = &right_indices.slice(current_start, end - current_start);
+ self.update_matched_bitmap(l_index, &cur_right_bitmap)?;
- // Switch around the build side and probe side for `JoinType::RightMark`
- // because in a RightMark join, we want to mark rows on the right table
- // by looking for matches in the left.
- let res = if self.join_type == JoinType::RightMark {
- build_batch_from_indices(
- &self.schema,
- right_batch,
- left_batch,
- left_indices,
- right_indices,
- &self.column_indices,
- JoinSide::Right,
- )
+ // For the following join types: here we only have to set the left/right
+ // bitmap, and no need to output result
+ if matches!(
+ self.join_type,
+ JoinType::LeftAnti
+ | JoinType::LeftSemi
+ | JoinType::LeftMark
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ ) {
+ return Ok(None);
+ }
+
+ if cur_right_bitmap.true_count() == 0 {
+ // If none of the pairs has passed the join predicate/filter
+ Ok(None)
} else {
- build_batch_from_indices(
- &self.schema,
- left_batch,
+ // Use the optimized approach similar to build_intermediate_batch_for_single_left_row
+ let join_batch = build_row_join_batch(
+ &self.output_schema,
+ left_data.batch(),
+ l_index,
right_batch,
- left_indices,
- right_indices,
+ Some(cur_right_bitmap),
&self.column_indices,
JoinSide::Left,
- )
- }?;
+ )?;
+ Ok(join_batch)
+ }
+ }
- status.processed_count = end;
+ /// Returns bool to indicate should it continue processing unmatched rows
+ /// true -> continue in the same EmitLeftUnmatched state
+ /// false -> next state (Done)
+ fn process_left_unmatched(&mut self) -> Result<bool> {
+ let left_data = self.get_left_data()?;
+ let left_batch = left_data.batch();
- Ok(Some(res))
- }
+ // Early return if join type can't have unmatched rows
+ if !need_produce_result_in_final(self.join_type) {
+ return Ok(false);
+ }
- fn collect_build_side(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- let build_timer = self.join_metrics.build_time.timer();
- // build hash table from left (build) side, if not yet done
- self.left_data = Some(ready!(self.inner_table.get_shared(cx))?);
- build_timer.done();
+ // Early return if another thread is already processing unmatched rows
+ if self.l_emit_idx == 0 && !left_data.report_probe_completed() {
+ return Ok(false);
+ }
- self.state = NestedLoopJoinStreamState::FetchProbeBatch;
+ // Stop processing unmatched rows, the caller will go to the next state
+ if self.l_emit_idx >= left_batch.num_rows() {
+ return Ok(false);
+ }
- Poll::Ready(Ok(StatefulStreamResult::Continue))
- }
+ // ========
+ // Process unmatched rows and push the result into output_buffer
+ // Each time, the number to process is up to batch size
+ // ========
+ let start_idx = self.l_emit_idx;
+ let end_idx =
+ std::cmp::min(start_idx + self.cfg_batch_size, left_batch.num_rows());
- /// Fetches next batch from probe-side
- ///
- /// If a non-empty batch has been fetched, updates state to
- /// `ProcessProbeBatchState`, otherwise updates state to `ExhaustedProbeSide`.
- fn fetch_probe_batch(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- match ready!(self.outer_table.poll_next_unpin(cx)) {
- None => {
- self.state = NestedLoopJoinStreamState::PrepareUnmatchedBuildRows;
- }
- Some(Ok(right_batch)) => {
- self.join_metrics.input_batches.add(1);
- self.join_metrics.input_rows.add(right_batch.num_rows());
+ if let Some(batch) =
+ self.process_left_unmatched_range(left_data, start_idx, end_idx)?
+ {
+ self.output_buffer.push_batch(batch)?;
+ }
- self.state = NestedLoopJoinStreamState::ProcessProbeBatch(right_batch);
- }
- Some(Err(err)) => return Poll::Ready(Err(err)),
- };
+ // ==== Prepare for the next iteration ====
+ self.l_emit_idx = end_idx;
- Poll::Ready(Ok(StatefulStreamResult::Continue))
+ // Return true to continue processing unmatched rows
+ Ok(true)
}
- /// Joins current probe batch with build-side data and produces batch with
- /// matched output, updates state to `FetchProbeBatch`.
- fn process_probe_batch(
- &mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let Some(left_data) = self.left_data.clone() else {
- return internal_err!(
- "Expected left_data to be Some in ProcessProbeBatch state"
+ /// Process unmatched rows from the left data within the specified range.
+ /// Returns a RecordBatch containing the unmatched rows (None if empty).
+ ///
+ /// # Arguments
+ /// * `left_data` - The left side data containing the batch and bitmap
+ /// * `start_idx` - Start index (inclusive) of the range to process
+ /// * `end_idx` - End index (exclusive) of the range to process
+ ///
+ /// # Safety
+ /// The caller is responsible for ensuring that `start_idx` and `end_idx` are
+ /// within valid bounds of the left batch. This function does not perform
+ /// bounds checking.
+ fn process_left_unmatched_range(
+ &self,
+ left_data: &JoinLeftData,
+ start_idx: usize,
+ end_idx: usize,
+ ) -> Result<Option<RecordBatch>> {
+ if start_idx == end_idx {
+ return Ok(None);
+ }
+
+ // Slice both left batch, and bitmap to range [start_idx, end_idx)
+ // The range is bit index (not byte)
+ let left_batch = left_data.batch();
+ let left_batch_sliced = left_batch.slice(start_idx, end_idx - start_idx);
+
+ // Can this be more efficient?
+ let mut bitmap_sliced = BooleanBufferBuilder::new(end_idx - start_idx);
+ bitmap_sliced.append_n(end_idx - start_idx, false);
+ let bitmap = left_data.bitmap().lock();
+ for i in start_idx..end_idx {
+ assert!(
+ i - start_idx < bitmap_sliced.capacity(),
+ "DBG: {start_idx}, {end_idx}"
);
- };
- let visited_left_side = left_data.bitmap();
- let batch = self.state.try_as_process_probe_batch()?;
+ bitmap_sliced.set_bit(i - start_idx, bitmap.get_bit(i));
+ }
+ let bitmap_sliced = BooleanArray::new(bitmap_sliced.finish(), None);
+
+ build_unmatched_batch(
+ Arc::clone(&self.output_schema),
+ &left_batch_sliced,
+ bitmap_sliced,
+ self.outer_table.schema(),
+ &self.column_indices,
+ self.join_type,
+ JoinSide::Left,
+ )
+ }
- let binding = self.join_metrics.join_time.clone();
- let _timer = binding.timer();
+ /// Process unmatched rows from the current right batch and reset the bitmap.
+ /// Returns a RecordBatch containing the unmatched right rows (None if empty).
+ fn process_right_unmatched(&mut self) -> Result<Option<RecordBatch>> {
+ // ==== Take current right batch and its bitmap ====
+ let right_batch_bitmap: BooleanArray =
+ std::mem::take(&mut self.current_right_batch_matched).ok_or_else(|| {
+ internal_datafusion_err!("right bitmap should be available")
+ })?;
+
+ let right_batch = self.current_right_batch.take();
+ let cur_right_batch = unwrap_or_internal_err!(right_batch);
+
+ let left_data = self.get_left_data()?;
+ let left_schema = left_data.batch().schema();
+
+ let res = build_unmatched_batch(
+ Arc::clone(&self.output_schema),
+ &cur_right_batch,
+ right_batch_bitmap,
+ left_schema,
+ &self.column_indices,
+ self.join_type,
+ JoinSide::Right,
+ );
- if self.join_result_status.is_none() {
- let (left_side_indices, right_side_indices) = join_left_and_right_batch(
- left_data.batch(),
- batch,
- self.join_type,
- self.filter.as_ref(),
- visited_left_side,
- &mut self.indices_cache,
- self.right_side_ordered,
- self.intermediate_batch_size,
- )?;
- self.join_result_status = Some(JoinResultProgress {
- build_indices: left_side_indices,
- probe_indices: right_side_indices,
- processed_count: 0,
- })
- }
+ // ==== Clean-up ====
+ self.current_right_batch_matched = None;
- let join_result = self.get_next_join_result()?;
+ res
+ }
- match join_result {
- Some(res) => {
- self.join_metrics.output_batches.add(1);
- Ok(StatefulStreamResult::Ready(Some(res)))
- }
- None => {
- self.state = NestedLoopJoinStreamState::FetchProbeBatch;
- self.join_result_status = None;
- Ok(StatefulStreamResult::Continue)
- }
- }
+ // ==== Utilities ====
+
+ /// Get the inner data of left data, errors if it's None
+ fn get_left_data(&self) -> Result<&Arc<JoinLeftData>> {
+ self.buffered_left_data
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("LeftData should be available"))
}
- fn build_unmatched_output(
- &mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let binding = self.join_metrics.join_time.clone();
- let _timer = binding.timer();
-
- let res = self.get_next_join_result()?;
- match res {
- Some(res) => {
- self.join_metrics.output_batches.add(1);
- Ok(StatefulStreamResult::Ready(Some(res)))
- }
- None => {
- self.state = NestedLoopJoinStreamState::Completed;
- Ok(StatefulStreamResult::Ready(None))
+ /// Flush the `output_buffer` if there are batches ready to output
+ /// None if no result batch ready.
+ fn maybe_flush_ready_batch(&mut self) -> Option<Poll<Option<Result<RecordBatch>>>> {
+ if self.output_buffer.has_completed_batch() {
+ if let Some(batch) = self.output_buffer.next_completed_batch() {
+ let poll = Poll::Ready(Some(Ok(batch)));
+ return Some(self.join_metrics.baseline.record_poll(poll));
}
}
+
+ None
}
- /// This function's primary purpose is to handle the final output stage required by specific join types after all right-side (probe) data has been exhausted.
- /// It is critically important for LEFT*/FULL joins, which must emit left-side (build) rows that found no match. For these cases, it identifies the unmatched rows and prepares the necessary state to output them.
- fn prepare_unmatched_output_indices(
+ /// After joining (l_index@left_buffer x current_right_batch), it will result
+ /// in a bitmap (the same length as current_right_batch) as the join match
+ /// result. Use this bitmap to update the global bitmap, for special join
+ /// types like full joins.
+ ///
+ /// Example:
+ /// After joining l_index=1 (1-indexed row in the left buffer), and the
+ /// current right batch with 3 elements, this function will be called with
+ /// arguments: l_index = 1, r_matched = [false, false, true]
+ /// - If the join type is FullJoin, the 1-index in the left bitmap will be
+ /// set to true, and also the right bitmap will be bitwise-ORed with the
+ /// input r_matched bitmap.
+ /// - For join types that don't require output unmatched rows, this
+ /// function can be a no-op. For inner joins, this function is a no-op; for left
+ /// joins, only the left bitmap may be updated.
+ fn update_matched_bitmap(
&mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let Some(left_data) = self.left_data.clone() else {
- return internal_err!(
- "Expected left_data to be Some in ExhaustedProbeSide state"
- );
- };
- let visited_left_side = left_data.bitmap();
- if need_produce_result_in_final(self.join_type) {
- // At this stage `visited_left_side` won't be updated, so it's
- // safe to report about probe completion.
- //
- // Setting `is_exhausted` / returning None will prevent from
- // multiple calls of `report_probe_completed()`
- if !left_data.report_probe_completed() {
- self.state = NestedLoopJoinStreamState::Completed;
- return Ok(StatefulStreamResult::Ready(None));
- };
+ l_index: usize,
+ r_matched: &BooleanArray,
+ ) -> Result<()> {
+ let left_data = self.get_left_data()?;
- // Only setting up timer, input is exhausted
- let _timer = self.join_metrics.join_time.timer();
- // use the global left bitmap to produce the left indices and right indices
- let (left_side, right_side) =
- get_final_indices_from_shared_bitmap(visited_left_side, self.join_type);
-
- self.join_result_status = Some(JoinResultProgress {
- build_indices: left_side,
- probe_indices: right_side,
- processed_count: 0,
- });
- self.state = NestedLoopJoinStreamState::OutputUnmatchedBuildRows(
- RecordBatch::new_empty(self.outer_table.schema()),
- );
+ // number of successfully joined pairs from (l_index x cur_right_batch)
+ let joined_len = r_matched.true_count();
- Ok(StatefulStreamResult::Continue)
- } else {
- // end of the join loop
- self.state = NestedLoopJoinStreamState::Completed;
- Ok(StatefulStreamResult::Ready(None))
+ // 1. Maybe update the left bitmap
+ if need_produce_result_in_final(self.join_type) && (joined_len > 0) {
+ let mut bitmap = left_data.bitmap().lock();
+ bitmap.set_bit(l_index, true);
+ }
+
+ // 2. Maybe updateh the right bitmap
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ // after bit-wise or, it will be put back
+ let right_bitmap = std::mem::take(&mut self.current_right_batch_matched)
+ .ok_or_else(|| {
+ internal_datafusion_err!("right batch's bitmap should be present")
+ })?;
+ let (buf, nulls) = right_bitmap.into_parts();
+ debug_assert!(nulls.is_none());
+ let updated_right_bitmap = buf.bitor(r_matched.values());
+
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(updated_right_bitmap, None));
}
+
+ Ok(())
}
}
-#[allow(clippy::too_many_arguments)]
-fn join_left_and_right_batch(
+// ==== Utilities ====
+
+/// Apply the join filter between:
+/// (l_index th row in left buffer) x (right batch)
+/// Returns a bitmap, with successfully joined indices set to true
+fn apply_filter_to_row_join_batch(
left_batch: &RecordBatch,
+ l_index: usize,
right_batch: &RecordBatch,
- join_type: JoinType,
- filter: Option<&JoinFilter>,
- visited_left_side: &SharedBitmapBuilder,
- indices_cache: &mut (UInt64Array, UInt32Array),
- right_side_ordered: bool,
- max_intermediate_batch_size: usize,
-) -> Result<(PrimitiveArray<UInt64Type>, PrimitiveArray<UInt32Type>)> {
- let (left_side, right_side) = build_join_indices(
- left_batch,
- right_batch,
- filter,
- indices_cache,
- max_intermediate_batch_size,
- )
- .map_err(|e| {
- exec_datafusion_err!(
- "Fail to build join indices in NestedLoopJoinExec, error: {e}"
- )
- })?;
-
- // set the left bitmap
- // and only full join need the left bitmap
- if need_produce_result_in_final(join_type) {
- let mut bitmap = visited_left_side.lock();
- left_side.values().iter().for_each(|x| {
- bitmap.set_bit(*x as usize, true);
- });
- }
- // adjust the two side indices base on the join type
- let (left_side, right_side) = adjust_indices_by_join_type(
- left_side,
- right_side,
- 0..right_batch.num_rows(),
- join_type,
- right_side_ordered,
- )?;
+ filter: &JoinFilter,
+) -> Result<BooleanArray> {
+ debug_assert!(left_batch.num_rows() != 0 && right_batch.num_rows() != 0);
+
+ let intermediate_batch = if filter.schema.fields().is_empty() {
+ // If filter is constant (e.g. literal `true`), empty batch can be used
+ // in the later filter step.
+ create_record_batch_with_empty_schema(
+ Arc::new((*filter.schema).clone()),
+ right_batch.num_rows(),
+ )?
+ } else {
+ build_row_join_batch(
+ &filter.schema,
+ left_batch,
+ l_index,
+ right_batch,
+ None,
+ &filter.column_indices,
+ JoinSide::Left,
+ )?
+ .ok_or_else(|| internal_datafusion_err!("This function assume input batch is not empty, so the intermediate batch can't be empty too"))?
+ };
+
+ let filter_result = filter
+ .expression()
+ .evaluate(&intermediate_batch)?
+ .into_array(intermediate_batch.num_rows())?;
- Ok((left_side, right_side))
+ Ok(as_boolean_array(&filter_result)?.clone())
}
-impl Stream for NestedLoopJoinStream {
- type Item = Result<RecordBatch>;
+/// This function performs the following steps:
+/// 1. Apply filter to probe-side batch
+/// 2. Broadcast the left row (build_side_batch\[build_side_index\]) to the
+/// filtered probe-side batch
+/// 3. Concat them together according to `col_indices`, and return the result
+/// (None if the result is empty)
+///
+/// Example:
+/// build_side_batch:
+/// a
+/// ----
+/// 1
+/// 2
+/// 3
+///
+/// # 0 index element in the build_side_batch (that is `1`) will be used
+/// build_side_index: 0
+///
+/// probe_side_batch:
+/// b
+/// ----
+/// 10
+/// 20
+/// 30
+/// 40
+///
+/// # After applying it, only index 1 and 3 elemnt in probe_side_batch will be
+/// # kept
+/// probe_side_filter:
+/// false
+/// true
+/// false
+/// true
+///
+///
+/// # Projections to the build/probe side batch, to construct the output batch
+/// col_indices:
+/// [(left, 0), (right, 0)]
+///
+/// build_side: left
+///
+/// ====
+/// Result batch:
+/// a b
+/// ----
+/// 1 20
+/// 1 40
+fn build_row_join_batch(
+ output_schema: &Schema,
+ build_side_batch: &RecordBatch,
+ build_side_index: usize,
+ probe_side_batch: &RecordBatch,
+ probe_side_filter: Option<BooleanArray>,
+ // See [`NLJStream`] struct's `column_indices` field for more detail
+ col_indices: &[ColumnIndex],
+ // If the build side is left or right, used to interpret the side information
+ // in `col_indices`
+ build_side: JoinSide,
+) -> Result<Option<RecordBatch>> {
+ debug_assert!(build_side != JoinSide::None);
+
+ // TODO(perf): since the output might be projection of right batch, this
+ // filtering step is more efficient to be done inside the column_index loop
+ let filtered_probe_batch = if let Some(filter) = probe_side_filter {
+ &filter_record_batch(probe_side_batch, &filter)?
+ } else {
+ probe_side_batch
+ };
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
+ if filtered_probe_batch.num_rows() == 0 {
+ return Ok(None);
+ }
+
+ let mut columns: Vec<Arc<dyn Array>> =
+ Vec::with_capacity(output_schema.fields().len());
+
+ for column_index in col_indices {
+ let array = if column_index.side == build_side {
+ // Broadcast the single build-side row to match the filtered
+ // probe-side batch length
+ let original_left_array = build_side_batch.column(column_index.index);
+ let scalar_value = ScalarValue::try_from_array(
+ original_left_array.as_ref(),
+ build_side_index,
+ )?;
+ scalar_value.to_array_of_size(filtered_probe_batch.num_rows())?
Review Comment:
I tried but it only gets slightly faster, I think this optimization might not worth the extra complexity.
The reason for not getting significantly faster:
1. This step is not the bottleneck (for this simple NLJ micro-bench, now the actual bottleneck is evaluating the expression, especially modulo)
2. I used a utility in the optimizer to rewrite the join filter, this step is slow (6% in total time)
Benchmark result:
```
┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ nlj-rewrite ┃ optimize-join-filter ┃ Change ┃
┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 87.14 ms │ 83.12 ms │ no change │
│ QQuery 2 │ 118.03 ms │ 110.51 ms │ +1.07x faster │
│ QQuery 3 │ 164.04 ms │ 160.60 ms │ no change │
│ QQuery 4 │ 373.67 ms │ 358.43 ms │ no change │
│ QQuery 5 │ 232.66 ms │ 259.75 ms │ 1.12x slower │
│ QQuery 6 │ 1708.19 ms │ 1633.82 ms │ no change │
│ QQuery 7 │ 236.27 ms │ 257.24 ms │ 1.09x slower │
│ QQuery 8 │ 1710.80 ms │ 1634.83 ms │ no change │
│ QQuery 9 │ 270.86 ms │ 261.67 ms │ no change │
│ QQuery 10 │ 499.70 ms │ 493.96 ms │ no change │
└──────────────┴─────────────┴──────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩
│ Total Time (nlj-rewrite) │ 5401.37ms │
│ Total Time (optimize-join-filter) │ 5253.94ms │
│ Average Time (nlj-rewrite) │ 540.14ms │
│ Average Time (optimize-join-filter) │ 525.39ms │
│ Queries Faster │ 1 │
│ Queries Slower │ 2 │
│ Queries with No Change │ 7 │
│ Queries with Failure │ 0 │
└─────────────────────────────────────┴───────────┘
```
My experiment branch: https://github.com/2010YOUY01/arrow-datafusion/tree/optimize-join-filter
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2260641740
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -63,92 +64,95 @@ use datafusion_physical_expr::equivalence::{
join_equivalence_properties, ProjectionMapping,
};
-use futures::{ready, Stream, StreamExt, TryStreamExt};
+use futures::{Stream, StreamExt, TryStreamExt};
+use log::debug;
use parking_lot::Mutex;
-/// Left (build-side) data
-struct JoinLeftData {
- /// Build-side data collected to single batch
- batch: RecordBatch,
- /// Shared bitmap builder for visited left indices
- bitmap: SharedBitmapBuilder,
- /// Counter of running probe-threads, potentially able to update `bitmap`
- probe_threads_counter: AtomicUsize,
- /// Memory reservation for tracking batch and bitmap
- /// Cleared on `JoinLeftData` drop
- /// reservation is cleared on Drop
- #[expect(dead_code)]
- reservation: MemoryReservation,
-}
-
-impl JoinLeftData {
- fn new(
- batch: RecordBatch,
- bitmap: SharedBitmapBuilder,
- probe_threads_counter: AtomicUsize,
- reservation: MemoryReservation,
- ) -> Self {
- Self {
- batch,
- bitmap,
- probe_threads_counter,
- reservation,
- }
- }
-
- fn batch(&self) -> &RecordBatch {
- &self.batch
- }
-
- fn bitmap(&self) -> &SharedBitmapBuilder {
- &self.bitmap
- }
-
- /// Decrements counter of running threads, and returns `true`
- /// if caller is the last running thread
- fn report_probe_completed(&self) -> bool {
- self.probe_threads_counter.fetch_sub(1, Ordering::Relaxed) == 1
- }
-}
-
#[allow(rustdoc::private_intra_doc_links)]
-/// NestedLoopJoinExec is build-probe join operator, whose main task is to
-/// perform joins without any equijoin conditions in `ON` clause.
+/// NestedLoopJoinExec is a build-probe join operator designed for joins that
+/// do not have equijoin keys in their `ON` clause.
///
-/// Execution consists of following phases:
+/// # Execution Flow
///
-/// #### 1. Build phase
-/// Collecting build-side data in memory, by polling all available data from build-side input.
-/// Due to the absence of equijoin conditions, it's not possible to partition build-side data
-/// across multiple threads of the operator, so build-side is always collected in a single
-/// batch shared across all threads.
-/// The operator always considers LEFT input as build-side input, so it's crucial to adjust
-/// smaller input to be the LEFT one. Normally this selection is handled by physical optimizer.
+/// ```text
+/// Incoming right batch
+/// Left Side Buffered Batches
+/// ┌───────────┐ ┌───────────────┐
+/// │ ┌───────┐ │ │ │
+/// │ │ │ │ │ │
+/// Current Left Row ───▶│ ├───────├─┤──────────┐ │ │
+/// │ │ │ │ │ └───────────────┘
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ └───────┘ │ │ │
+/// │ ┌───────┐ │ │ │
+/// │ │ │ │ │ ┌─────┘
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ └───────┘ │ ▼ ▼
+/// │ ...... │ ┌──────────────────────┐
+/// │ │ │X (Cartesian Product) │
+/// │ │ └──────────┬───────────┘
+/// └───────────┘ │
+/// │
+/// ▼
+/// ┌───────┬───────────────┐
+/// │ │ │
+/// │ │ │
+/// │ │ │
+/// └───────┴───────────────┘
+/// Intermediate Batch
+/// (For join predicate evaluation)
+/// ```
///
-/// #### 2. Probe phase
-/// Sequentially polling batches from the probe-side input and processing them according to the
-/// following logic:
-/// - apply join filter (`ON` clause) to Cartesian product of probe batch and build side data
-/// -- filter evaluation is executed once per build-side data row
-/// - update shared bitmap of joined ("visited") build-side row indices, if required -- allows
-/// to produce unmatched build-side data in case of e.g. LEFT/FULL JOIN after probing phase
-/// completed
-/// - perform join index alignment is required -- depending on `JoinType`
-/// - produce output join batch
+/// The execution follows a two-phase design:
///
-/// Probing phase is executed in parallel, according to probe-side input partitioning -- one
-/// thread per partition. After probe input is exhausted, each thread **ATTEMPTS** to produce
-/// unmatched build-side data.
+/// ## 1. Buffering Left Input
+/// - The operator eagerly buffers all left-side input batches into memory,
+/// unless a memory limit is reached (see 'Memory-limited Execution').
+/// - The rationale for buffering the left side is that scanning the right side
+/// can be expensive (e.g., decoding Parquet files), so buffering more left
+/// rows reduces the number of right-side scan passes required.
///
-/// #### 3. Producing unmatched build-side data
-/// Producing unmatched build-side data as an output batch, after probe input is exhausted.
-/// This step is also executed in parallel (once per probe input partition), and to avoid
-/// duplicate output of unmatched data (due to shared nature build-side data), each thread
-/// "reports" about probe phase completion (which means that "visited" bitmap won't be
-/// updated anymore), and only the last thread, reporting about completion, will return output.
+/// ## 2. Probing Right Input
+/// - Right-side input is streamed batch by batch.
+/// - For each right-side batch:
+/// - It evaluates the join filter against the full buffered left input.
+/// This results in a Cartesian product between the right batch and each
+/// left row -- with the join predicate/filter applied -- for each inner
+/// loop iteration.
+/// - Matched results are accumulated into an output buffer. (see more in
+/// `Output Buffering Strategy` section)
+/// - This process continues until all right-side input is consumed.
///
-/// # Clone / Shared State
+/// # Producing unmatched build-side data
+/// - For special join types like left/full joins, it's required to also output
+/// unmatched pairs. During execution, bitmaps are kept for both left and right
+/// sides of the input; they'll be handled by dedicated states in `NLJStream`.
+/// - The final output of the left side unmatched rows is handled by a single
+/// partition for simplicity, since it only counts a small portion of the
+/// execution time. (e.g. if probe side has 10k rows, the final output of
+/// unmatched build side only roughly counts for 1/10k of the total time)
+///
+/// # Output Buffering Strategy
+/// The operator uses an intermediate output buffer to accumulate results. Once
+/// the output threshold is reached (currently set to the same value as
+/// `batch_size` in the configuration), the results will be eagerly output.
///
+/// # Extra Notes
+/// - The operator always considers the **left** side as the build (buffered) side.
+/// Therefore, the physical optimizer should assign the smaller input to the left.
+/// - The design try to minimize the intermediate data size to approximately
+/// 1 batch, for better cache locality and memory efficiency.
+///
+/// # TODO: Memory-limited Execution
Review Comment:
it would be good to create a followup ticket and refer it.
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
Review Comment:
inner_table, outer_table, perhaps left_data right_data
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -63,92 +64,95 @@ use datafusion_physical_expr::equivalence::{
join_equivalence_properties, ProjectionMapping,
};
-use futures::{ready, Stream, StreamExt, TryStreamExt};
+use futures::{Stream, StreamExt, TryStreamExt};
+use log::debug;
use parking_lot::Mutex;
-/// Left (build-side) data
-struct JoinLeftData {
- /// Build-side data collected to single batch
- batch: RecordBatch,
- /// Shared bitmap builder for visited left indices
- bitmap: SharedBitmapBuilder,
- /// Counter of running probe-threads, potentially able to update `bitmap`
- probe_threads_counter: AtomicUsize,
- /// Memory reservation for tracking batch and bitmap
- /// Cleared on `JoinLeftData` drop
- /// reservation is cleared on Drop
- #[expect(dead_code)]
- reservation: MemoryReservation,
-}
-
-impl JoinLeftData {
- fn new(
- batch: RecordBatch,
- bitmap: SharedBitmapBuilder,
- probe_threads_counter: AtomicUsize,
- reservation: MemoryReservation,
- ) -> Self {
- Self {
- batch,
- bitmap,
- probe_threads_counter,
- reservation,
- }
- }
-
- fn batch(&self) -> &RecordBatch {
- &self.batch
- }
-
- fn bitmap(&self) -> &SharedBitmapBuilder {
- &self.bitmap
- }
-
- /// Decrements counter of running threads, and returns `true`
- /// if caller is the last running thread
- fn report_probe_completed(&self) -> bool {
- self.probe_threads_counter.fetch_sub(1, Ordering::Relaxed) == 1
- }
-}
-
#[allow(rustdoc::private_intra_doc_links)]
-/// NestedLoopJoinExec is build-probe join operator, whose main task is to
-/// perform joins without any equijoin conditions in `ON` clause.
+/// NestedLoopJoinExec is a build-probe join operator designed for joins that
+/// do not have equijoin keys in their `ON` clause.
///
-/// Execution consists of following phases:
+/// # Execution Flow
///
-/// #### 1. Build phase
-/// Collecting build-side data in memory, by polling all available data from build-side input.
-/// Due to the absence of equijoin conditions, it's not possible to partition build-side data
-/// across multiple threads of the operator, so build-side is always collected in a single
-/// batch shared across all threads.
-/// The operator always considers LEFT input as build-side input, so it's crucial to adjust
-/// smaller input to be the LEFT one. Normally this selection is handled by physical optimizer.
+/// ```text
+/// Incoming right batch
+/// Left Side Buffered Batches
+/// ┌───────────┐ ┌───────────────┐
+/// │ ┌───────┐ │ │ │
+/// │ │ │ │ │ │
+/// Current Left Row ───▶│ ├───────├─┤──────────┐ │ │
+/// │ │ │ │ │ └───────────────┘
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ └───────┘ │ │ │
+/// │ ┌───────┐ │ │ │
+/// │ │ │ │ │ ┌─────┘
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ │ │ │ │ │
+/// │ └───────┘ │ ▼ ▼
+/// │ ...... │ ┌──────────────────────┐
+/// │ │ │X (Cartesian Product) │
+/// │ │ └──────────┬───────────┘
+/// └───────────┘ │
+/// │
+/// ▼
+/// ┌───────┬───────────────┐
+/// │ │ │
+/// │ │ │
+/// │ │ │
+/// └───────┴───────────────┘
+/// Intermediate Batch
+/// (For join predicate evaluation)
+/// ```
///
-/// #### 2. Probe phase
-/// Sequentially polling batches from the probe-side input and processing them according to the
-/// following logic:
-/// - apply join filter (`ON` clause) to Cartesian product of probe batch and build side data
-/// -- filter evaluation is executed once per build-side data row
-/// - update shared bitmap of joined ("visited") build-side row indices, if required -- allows
-/// to produce unmatched build-side data in case of e.g. LEFT/FULL JOIN after probing phase
-/// completed
-/// - perform join index alignment is required -- depending on `JoinType`
-/// - produce output join batch
+/// The execution follows a two-phase design:
///
-/// Probing phase is executed in parallel, according to probe-side input partitioning -- one
-/// thread per partition. After probe input is exhausted, each thread **ATTEMPTS** to produce
-/// unmatched build-side data.
+/// ## 1. Buffering Left Input
+/// - The operator eagerly buffers all left-side input batches into memory,
+/// unless a memory limit is reached (see 'Memory-limited Execution').
Review Comment:
what would happen if memory limit is reached?
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
+impl Stream for NLJStream {
Review Comment:
perhaps we can come up with another naming? recently we renamed SMJStream to user friendly name. I got the fact the proper name already taken, so perhaps we can use NestedLoopJoinStream_v2 or some postfix showing implementation specifics?
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
Review Comment:
when it should be false or true, is it for specific join types?
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
Review Comment:
lets use just `batch_size`?
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
Review Comment:
`l` means left? if so lets call it `left` ?
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
- // Produces 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
- *left_indices_cache = UInt64Array::from_iter_values(
- (0..output_row_count as u64).map(|i| i % left_row_count as u64),
- );
+ /// See the comments [`NestedLoopJoinExec`] for high-level design ideas.
+ ///
+ /// # Implementation
+ ///
+ /// This function is the entry point of NLJ operator's state machine
+ /// transitions. The rough state transition graph is as follow, for more
+ /// details see the comment in each state's matching arm.
+ ///
+ /// ============================
+ /// State transition graph:
+ /// ============================
+ ///
+ /// (start) --> BufferingLeft
+ /// ----------------------------
+ /// BufferingLeft → FetchingRight
+ ///
+ /// FetchingRight → ProbeRight (if right batch available)
+ /// FetchingRight → EmitLeftUnmatched (if right exhausted)
+ ///
+ /// ProbeRight → ProbeRight (next left row or after yielding output)
+ /// ProbeRight → EmitRightUnmatched (for special join types like right join)
+ /// ProbeRight → FetchingRight (done with the current right batch)
+ ///
+ /// EmitRightUnmatched → FetchingRight
+ ///
+ /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each
+ /// iteration)
+ /// EmitLeftUnmatched → Done (if finished)
+ /// ----------------------------
+ /// Done → (end)
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ loop {
+ match self.state {
+ // # NLJState transitions
+ // --> FetchingRight
+ // This state will prepare the left side batches, next state
+ // `FetchingRight` is responsible for preparing a single probe
+ // side batch, before start joining.
+ NLJState::BufferingLeft => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // inside `collect_left_input` (the rountine to buffer build
+ // -side batches), related metrics except build time will be
+ // updated.
+ // stop on drop
+ let build_metric = self.join_metrics.build_time.clone();
+ let _build_timer = build_metric.timer();
+
+ match self.handle_buffering_left(cx) {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => return poll,
+ }
+ }
- // Produces 0, 0, 0, 1, 1, 1, 2, 2, 2, ...
- *right_indices_cache = UInt32Array::from_iter_values(
- (0..output_row_count as u32).map(|i| i / left_row_count as u32),
- );
+ // # NLJState transitions:
+ // 1. --> ProbeRight
+ // Start processing the join for the newly fetched right
+ // batch.
+ // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit
+ // unmatched left side rows.
+ //
+ // After fetching a new batch from the right side, it will
+ // process all rows from the buffered left data:
+ // ```text
+ // for batch in right_side:
+ // for row in left_buffer:
+ // join(batch, row)
+ // ```
+ // Note: the implementation does this step incrementally,
+ // instead of materializing all intermediate Cartesian products
+ // at once in memory.
+ //
+ // So after the right side input is exhausted, the join phase
+ // for the current buffered left data is finished. We can go to
+ // the next `EmitLeftUnmatched` phase to check if there is any
+ // special handling (e.g., in cases like left join).
+ NLJState::FetchingRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_fetching_right(cx) {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => return poll,
+ }
+ }
+
+ // NLJState transitions:
+ // 1. --> ProbeRight(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> ProbeRight(2)
+ // After probing one right batch, and evaluating the
+ // join filter on (left-row x right-batch), it will advance
+ // to the next left row, then re-enter the current state and
+ // continue joining.
+ // 3. --> FetchRight
+ // After it has done with the current right batch (to join
+ // with all rows in the left buffer), it will go to
+ // FetchRight state to check what to do next.
+ NLJState::ProbeRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_probe_right() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been output by the join. In this state we have to
+ // output unmatched rows for current right batch (with null
+ // padding for left relation)
+ // Precondition: we have checked the join type so that it's
+ // possible to output right unmatched (e.g. it's right join)
+ NLJState::EmitRightUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_emit_right_unmatched() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // NLJState transitions:
+ // 1. --> EmitLeftUnmatched(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> EmitLeftUnmatched(2)
+ // After processing some unmatched rows, it will re-enter
+ // the same state, to check if there are any more final
+ // results to output.
+ // 3. --> Done
+ // It has processed all data, go to the final state and ready
+ // to exit.
+ //
+ // TODO: For memory-limited case, go back to `BufferingLeft`
+ // state again.
+ NLJState::EmitLeftUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_emit_left_unmatched() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
- (left_indices_cache.clone(), right_indices_cache.clone())
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+ // counting it in join timer due to there might be some
+ // final resout batches to output in this state
+
+ let poll = self.handle_done();
+ return self.join_metrics.baseline.record_poll(poll);
+ }
}
- };
+ }
+ }
+}
- if let Some(filter) = filter {
- apply_join_filter_to_indices(
- left_batch,
- right_batch,
- left_indices,
- right_indices,
- filter,
- JoinSide::Left,
- Some(max_intermediate_batch_size),
- )
- } else {
- Ok((left_indices, right_indices))
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
}
}
-impl NestedLoopJoinStream {
- fn poll_next_impl(
+impl NLJStream {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new(
+ schema: Arc<Schema>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ outer_table: SendableRecordBatchStream,
+ inner_table: OnceFut<JoinLeftData>,
+ column_indices: Vec<ColumnIndex>,
+ join_metrics: BuildProbeJoinMetrics,
+ cfg_batch_size: usize,
+ ) -> Self {
+ let should_track_unmatched_right = matches!(
+ join_type,
+ JoinType::Full
+ | JoinType::Right
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ );
+
+ Self {
+ output_schema: Arc::clone(&schema),
+ join_filter: filter,
+ join_type,
+ outer_table,
+ column_indices,
+ inner_table,
+ join_metrics,
+ buffered_left_data: None,
+ output_buffer: Box::new(BatchCoalescer::new(schema, cfg_batch_size)),
+ cfg_batch_size,
+ current_right_batch: None,
+ current_right_batch_matched: None,
+ state: NLJState::BufferingLeft,
+ l_probe_idx: 0,
+ l_emit_idx: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
+
+ // ==== State handler functions ====
+
+ /// Handle BufferingLeft state - prepare left side batches
+ fn handle_buffering_left(
&mut self,
cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- NestedLoopJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- NestedLoopJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- NestedLoopJoinStreamState::ProcessProbeBatch(_) => {
- let poll = handle_state!(self.process_probe_batch());
- self.join_metrics.baseline.record_poll(poll)
- }
- NestedLoopJoinStreamState::PrepareUnmatchedBuildRows => {
- handle_state!(self.prepare_unmatched_output_indices())
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ match self.inner_table.get_shared(cx) {
+ Poll::Ready(Ok(left_data)) => {
+ self.buffered_left_data = Some(left_data);
+ // TODO: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ // Continue to next state immediately
+ ControlFlow::Continue(())
+ }
+ Poll::Ready(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ Poll::Pending => ControlFlow::Break(Poll::Pending),
+ }
+ }
+
+ /// Handle FetchingRight state - fetch next right batch and prepare for processing
+ fn handle_fetching_right(
+ &mut self,
+ cx: &mut std::task::Context<'_>,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ match self.outer_table.poll_next_unpin(cx) {
+ Poll::Ready(result) => match result {
+ Some(Ok(right_batch)) => {
+ // Update metrics
+ self.join_metrics.input_rows.add(right_batch.num_rows());
+ self.join_metrics.input_batches.add(1);
+
+ let right_batch_size = right_batch.num_rows();
Review Comment:
lets move right_batch_size up and reuse it instead of calling `num_rows`
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
- // Produces 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
- *left_indices_cache = UInt64Array::from_iter_values(
- (0..output_row_count as u64).map(|i| i % left_row_count as u64),
- );
+ /// See the comments [`NestedLoopJoinExec`] for high-level design ideas.
+ ///
+ /// # Implementation
+ ///
+ /// This function is the entry point of NLJ operator's state machine
+ /// transitions. The rough state transition graph is as follow, for more
+ /// details see the comment in each state's matching arm.
+ ///
+ /// ============================
+ /// State transition graph:
+ /// ============================
+ ///
+ /// (start) --> BufferingLeft
+ /// ----------------------------
+ /// BufferingLeft → FetchingRight
+ ///
+ /// FetchingRight → ProbeRight (if right batch available)
+ /// FetchingRight → EmitLeftUnmatched (if right exhausted)
+ ///
+ /// ProbeRight → ProbeRight (next left row or after yielding output)
+ /// ProbeRight → EmitRightUnmatched (for special join types like right join)
+ /// ProbeRight → FetchingRight (done with the current right batch)
+ ///
+ /// EmitRightUnmatched → FetchingRight
+ ///
+ /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each
+ /// iteration)
+ /// EmitLeftUnmatched → Done (if finished)
+ /// ----------------------------
+ /// Done → (end)
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ loop {
+ match self.state {
+ // # NLJState transitions
+ // --> FetchingRight
+ // This state will prepare the left side batches, next state
+ // `FetchingRight` is responsible for preparing a single probe
+ // side batch, before start joining.
+ NLJState::BufferingLeft => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // inside `collect_left_input` (the rountine to buffer build
+ // -side batches), related metrics except build time will be
+ // updated.
+ // stop on drop
+ let build_metric = self.join_metrics.build_time.clone();
+ let _build_timer = build_metric.timer();
+
+ match self.handle_buffering_left(cx) {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => return poll,
+ }
+ }
- // Produces 0, 0, 0, 1, 1, 1, 2, 2, 2, ...
- *right_indices_cache = UInt32Array::from_iter_values(
- (0..output_row_count as u32).map(|i| i / left_row_count as u32),
- );
+ // # NLJState transitions:
+ // 1. --> ProbeRight
+ // Start processing the join for the newly fetched right
+ // batch.
+ // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit
+ // unmatched left side rows.
+ //
+ // After fetching a new batch from the right side, it will
+ // process all rows from the buffered left data:
+ // ```text
+ // for batch in right_side:
+ // for row in left_buffer:
+ // join(batch, row)
+ // ```
+ // Note: the implementation does this step incrementally,
+ // instead of materializing all intermediate Cartesian products
+ // at once in memory.
+ //
+ // So after the right side input is exhausted, the join phase
+ // for the current buffered left data is finished. We can go to
+ // the next `EmitLeftUnmatched` phase to check if there is any
+ // special handling (e.g., in cases like left join).
+ NLJState::FetchingRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_fetching_right(cx) {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => return poll,
+ }
+ }
+
+ // NLJState transitions:
+ // 1. --> ProbeRight(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> ProbeRight(2)
+ // After probing one right batch, and evaluating the
+ // join filter on (left-row x right-batch), it will advance
+ // to the next left row, then re-enter the current state and
+ // continue joining.
+ // 3. --> FetchRight
+ // After it has done with the current right batch (to join
+ // with all rows in the left buffer), it will go to
+ // FetchRight state to check what to do next.
+ NLJState::ProbeRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_probe_right() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been output by the join. In this state we have to
+ // output unmatched rows for current right batch (with null
+ // padding for left relation)
+ // Precondition: we have checked the join type so that it's
+ // possible to output right unmatched (e.g. it's right join)
+ NLJState::EmitRightUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_emit_right_unmatched() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // NLJState transitions:
+ // 1. --> EmitLeftUnmatched(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> EmitLeftUnmatched(2)
+ // After processing some unmatched rows, it will re-enter
+ // the same state, to check if there are any more final
+ // results to output.
+ // 3. --> Done
+ // It has processed all data, go to the final state and ready
+ // to exit.
+ //
+ // TODO: For memory-limited case, go back to `BufferingLeft`
+ // state again.
+ NLJState::EmitLeftUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_emit_left_unmatched() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
- (left_indices_cache.clone(), right_indices_cache.clone())
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+ // counting it in join timer due to there might be some
+ // final resout batches to output in this state
+
+ let poll = self.handle_done();
+ return self.join_metrics.baseline.record_poll(poll);
+ }
}
- };
+ }
+ }
+}
- if let Some(filter) = filter {
- apply_join_filter_to_indices(
- left_batch,
- right_batch,
- left_indices,
- right_indices,
- filter,
- JoinSide::Left,
- Some(max_intermediate_batch_size),
- )
- } else {
- Ok((left_indices, right_indices))
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
}
}
-impl NestedLoopJoinStream {
- fn poll_next_impl(
+impl NLJStream {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new(
+ schema: Arc<Schema>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ outer_table: SendableRecordBatchStream,
+ inner_table: OnceFut<JoinLeftData>,
+ column_indices: Vec<ColumnIndex>,
+ join_metrics: BuildProbeJoinMetrics,
+ cfg_batch_size: usize,
+ ) -> Self {
+ let should_track_unmatched_right = matches!(
+ join_type,
+ JoinType::Full
+ | JoinType::Right
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ );
+
+ Self {
+ output_schema: Arc::clone(&schema),
+ join_filter: filter,
+ join_type,
+ outer_table,
+ column_indices,
+ inner_table,
+ join_metrics,
+ buffered_left_data: None,
+ output_buffer: Box::new(BatchCoalescer::new(schema, cfg_batch_size)),
+ cfg_batch_size,
+ current_right_batch: None,
+ current_right_batch_matched: None,
+ state: NLJState::BufferingLeft,
+ l_probe_idx: 0,
+ l_emit_idx: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
+
+ // ==== State handler functions ====
+
+ /// Handle BufferingLeft state - prepare left side batches
+ fn handle_buffering_left(
&mut self,
cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- NestedLoopJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- NestedLoopJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- NestedLoopJoinStreamState::ProcessProbeBatch(_) => {
- let poll = handle_state!(self.process_probe_batch());
- self.join_metrics.baseline.record_poll(poll)
- }
- NestedLoopJoinStreamState::PrepareUnmatchedBuildRows => {
- handle_state!(self.prepare_unmatched_output_indices())
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ match self.inner_table.get_shared(cx) {
+ Poll::Ready(Ok(left_data)) => {
+ self.buffered_left_data = Some(left_data);
+ // TODO: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ // Continue to next state immediately
+ ControlFlow::Continue(())
+ }
+ Poll::Ready(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ Poll::Pending => ControlFlow::Break(Poll::Pending),
+ }
+ }
+
+ /// Handle FetchingRight state - fetch next right batch and prepare for processing
+ fn handle_fetching_right(
+ &mut self,
+ cx: &mut std::task::Context<'_>,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ match self.outer_table.poll_next_unpin(cx) {
+ Poll::Ready(result) => match result {
+ Some(Ok(right_batch)) => {
+ // Update metrics
+ self.join_metrics.input_rows.add(right_batch.num_rows());
+ self.join_metrics.input_batches.add(1);
+
+ let right_batch_size = right_batch.num_rows();
+
+ // Skip the empty batch
+ if right_batch_size == 0 {
+ return ControlFlow::Continue(());
+ }
+
+ self.current_right_batch = Some(right_batch);
+
+ // Prepare right bitmap
+ if self.should_track_unmatched_right {
+ let zeroed_buf = BooleanBuffer::new_unset(right_batch_size);
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(zeroed_buf, None));
+ }
+
+ self.l_probe_idx = 0;
+ self.state = NLJState::ProbeRight;
+ ControlFlow::Continue(())
}
- NestedLoopJoinStreamState::OutputUnmatchedBuildRows(_) => {
- let poll = handle_state!(self.build_unmatched_output());
- self.join_metrics.baseline.record_poll(poll)
+ Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ None => {
+ // Right stream exhausted
+ self.state = NLJState::EmitLeftUnmatched;
+ ControlFlow::Continue(())
}
- NestedLoopJoinStreamState::Completed => Poll::Ready(None),
- };
+ },
+ Poll::Pending => ControlFlow::Break(Poll::Pending),
}
}
- // This function's main job is to construct an output `RecordBatch` based on pre-calculated join indices.
- // It operates in a chunk-based manner, meaning it processes a portion of the results in each call,
- // making it suitable for streaming large datasets without high memory consumption.
- // This function behaves like an iterator. It returns `Ok(None)`
- // to signal that the result stream is exhausted and there is no more data.
- fn get_next_join_result(&mut self) -> Result<Option<RecordBatch>> {
- let status = self.join_result_status.as_mut().ok_or_else(|| {
- internal_datafusion_err!(
- "get_next_join_result called without initializing join_result_status"
- )
- })?;
+ /// Handle ProbeRight state - process current probe batch
+ fn handle_probe_right(&mut self) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return ControlFlow::Break(poll);
+ }
- let (left_indices, right_indices, current_start) = (
- &status.build_indices,
- &status.probe_indices,
- status.processed_count,
- );
+ // Process current probe state
+ match self.process_probe_batch() {
+ // State unchanged (ProbeRight)
+ // Continue probing until we have done joining the
+ // current right batch with all buffered left rows.
+ Ok(true) => ControlFlow::Continue(()),
+ // To next FetchRightState
+ // We have finished joining
+ // (cur_right_batch x buffered_left_batches)
+ Ok(false) => {
+ // Left exhausted, transition to FetchingRight
+ self.l_probe_idx = 0;
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
Review Comment:
Please add meaningful message
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
- // Produces 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
- *left_indices_cache = UInt64Array::from_iter_values(
- (0..output_row_count as u64).map(|i| i % left_row_count as u64),
- );
+ /// See the comments [`NestedLoopJoinExec`] for high-level design ideas.
+ ///
+ /// # Implementation
+ ///
+ /// This function is the entry point of NLJ operator's state machine
+ /// transitions. The rough state transition graph is as follow, for more
+ /// details see the comment in each state's matching arm.
+ ///
+ /// ============================
+ /// State transition graph:
+ /// ============================
+ ///
+ /// (start) --> BufferingLeft
+ /// ----------------------------
+ /// BufferingLeft → FetchingRight
+ ///
+ /// FetchingRight → ProbeRight (if right batch available)
+ /// FetchingRight → EmitLeftUnmatched (if right exhausted)
+ ///
+ /// ProbeRight → ProbeRight (next left row or after yielding output)
+ /// ProbeRight → EmitRightUnmatched (for special join types like right join)
+ /// ProbeRight → FetchingRight (done with the current right batch)
+ ///
+ /// EmitRightUnmatched → FetchingRight
+ ///
+ /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each
+ /// iteration)
+ /// EmitLeftUnmatched → Done (if finished)
+ /// ----------------------------
+ /// Done → (end)
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ loop {
+ match self.state {
+ // # NLJState transitions
+ // --> FetchingRight
+ // This state will prepare the left side batches, next state
+ // `FetchingRight` is responsible for preparing a single probe
+ // side batch, before start joining.
+ NLJState::BufferingLeft => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // inside `collect_left_input` (the rountine to buffer build
+ // -side batches), related metrics except build time will be
+ // updated.
+ // stop on drop
+ let build_metric = self.join_metrics.build_time.clone();
+ let _build_timer = build_metric.timer();
+
+ match self.handle_buffering_left(cx) {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => return poll,
+ }
+ }
- // Produces 0, 0, 0, 1, 1, 1, 2, 2, 2, ...
- *right_indices_cache = UInt32Array::from_iter_values(
- (0..output_row_count as u32).map(|i| i / left_row_count as u32),
- );
+ // # NLJState transitions:
+ // 1. --> ProbeRight
+ // Start processing the join for the newly fetched right
+ // batch.
+ // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit
+ // unmatched left side rows.
+ //
+ // After fetching a new batch from the right side, it will
+ // process all rows from the buffered left data:
+ // ```text
+ // for batch in right_side:
+ // for row in left_buffer:
+ // join(batch, row)
+ // ```
+ // Note: the implementation does this step incrementally,
+ // instead of materializing all intermediate Cartesian products
+ // at once in memory.
+ //
+ // So after the right side input is exhausted, the join phase
+ // for the current buffered left data is finished. We can go to
+ // the next `EmitLeftUnmatched` phase to check if there is any
+ // special handling (e.g., in cases like left join).
+ NLJState::FetchingRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_fetching_right(cx) {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => return poll,
+ }
+ }
+
+ // NLJState transitions:
+ // 1. --> ProbeRight(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> ProbeRight(2)
+ // After probing one right batch, and evaluating the
+ // join filter on (left-row x right-batch), it will advance
+ // to the next left row, then re-enter the current state and
+ // continue joining.
+ // 3. --> FetchRight
+ // After it has done with the current right batch (to join
+ // with all rows in the left buffer), it will go to
+ // FetchRight state to check what to do next.
+ NLJState::ProbeRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_probe_right() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been output by the join. In this state we have to
+ // output unmatched rows for current right batch (with null
+ // padding for left relation)
+ // Precondition: we have checked the join type so that it's
+ // possible to output right unmatched (e.g. it's right join)
+ NLJState::EmitRightUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_emit_right_unmatched() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // NLJState transitions:
+ // 1. --> EmitLeftUnmatched(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> EmitLeftUnmatched(2)
+ // After processing some unmatched rows, it will re-enter
+ // the same state, to check if there are any more final
+ // results to output.
+ // 3. --> Done
+ // It has processed all data, go to the final state and ready
+ // to exit.
+ //
+ // TODO: For memory-limited case, go back to `BufferingLeft`
+ // state again.
+ NLJState::EmitLeftUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_emit_left_unmatched() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
- (left_indices_cache.clone(), right_indices_cache.clone())
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+ // counting it in join timer due to there might be some
+ // final resout batches to output in this state
+
+ let poll = self.handle_done();
+ return self.join_metrics.baseline.record_poll(poll);
+ }
}
- };
+ }
+ }
+}
- if let Some(filter) = filter {
- apply_join_filter_to_indices(
- left_batch,
- right_batch,
- left_indices,
- right_indices,
- filter,
- JoinSide::Left,
- Some(max_intermediate_batch_size),
- )
- } else {
- Ok((left_indices, right_indices))
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
}
}
-impl NestedLoopJoinStream {
- fn poll_next_impl(
+impl NLJStream {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new(
+ schema: Arc<Schema>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ outer_table: SendableRecordBatchStream,
+ inner_table: OnceFut<JoinLeftData>,
+ column_indices: Vec<ColumnIndex>,
+ join_metrics: BuildProbeJoinMetrics,
+ cfg_batch_size: usize,
+ ) -> Self {
+ let should_track_unmatched_right = matches!(
+ join_type,
+ JoinType::Full
+ | JoinType::Right
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ );
+
+ Self {
+ output_schema: Arc::clone(&schema),
+ join_filter: filter,
+ join_type,
+ outer_table,
+ column_indices,
+ inner_table,
+ join_metrics,
+ buffered_left_data: None,
+ output_buffer: Box::new(BatchCoalescer::new(schema, cfg_batch_size)),
+ cfg_batch_size,
+ current_right_batch: None,
+ current_right_batch_matched: None,
+ state: NLJState::BufferingLeft,
+ l_probe_idx: 0,
+ l_emit_idx: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
+
+ // ==== State handler functions ====
+
+ /// Handle BufferingLeft state - prepare left side batches
+ fn handle_buffering_left(
&mut self,
cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- NestedLoopJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- NestedLoopJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- NestedLoopJoinStreamState::ProcessProbeBatch(_) => {
- let poll = handle_state!(self.process_probe_batch());
- self.join_metrics.baseline.record_poll(poll)
- }
- NestedLoopJoinStreamState::PrepareUnmatchedBuildRows => {
- handle_state!(self.prepare_unmatched_output_indices())
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ match self.inner_table.get_shared(cx) {
+ Poll::Ready(Ok(left_data)) => {
+ self.buffered_left_data = Some(left_data);
+ // TODO: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ // Continue to next state immediately
+ ControlFlow::Continue(())
+ }
+ Poll::Ready(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ Poll::Pending => ControlFlow::Break(Poll::Pending),
+ }
+ }
+
+ /// Handle FetchingRight state - fetch next right batch and prepare for processing
+ fn handle_fetching_right(
+ &mut self,
+ cx: &mut std::task::Context<'_>,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ match self.outer_table.poll_next_unpin(cx) {
+ Poll::Ready(result) => match result {
+ Some(Ok(right_batch)) => {
+ // Update metrics
+ self.join_metrics.input_rows.add(right_batch.num_rows());
+ self.join_metrics.input_batches.add(1);
+
+ let right_batch_size = right_batch.num_rows();
+
+ // Skip the empty batch
+ if right_batch_size == 0 {
+ return ControlFlow::Continue(());
+ }
+
+ self.current_right_batch = Some(right_batch);
+
+ // Prepare right bitmap
+ if self.should_track_unmatched_right {
+ let zeroed_buf = BooleanBuffer::new_unset(right_batch_size);
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(zeroed_buf, None));
+ }
+
+ self.l_probe_idx = 0;
+ self.state = NLJState::ProbeRight;
+ ControlFlow::Continue(())
}
- NestedLoopJoinStreamState::OutputUnmatchedBuildRows(_) => {
- let poll = handle_state!(self.build_unmatched_output());
- self.join_metrics.baseline.record_poll(poll)
+ Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ None => {
+ // Right stream exhausted
+ self.state = NLJState::EmitLeftUnmatched;
+ ControlFlow::Continue(())
}
- NestedLoopJoinStreamState::Completed => Poll::Ready(None),
- };
+ },
+ Poll::Pending => ControlFlow::Break(Poll::Pending),
}
}
- // This function's main job is to construct an output `RecordBatch` based on pre-calculated join indices.
- // It operates in a chunk-based manner, meaning it processes a portion of the results in each call,
- // making it suitable for streaming large datasets without high memory consumption.
- // This function behaves like an iterator. It returns `Ok(None)`
- // to signal that the result stream is exhausted and there is no more data.
- fn get_next_join_result(&mut self) -> Result<Option<RecordBatch>> {
- let status = self.join_result_status.as_mut().ok_or_else(|| {
- internal_datafusion_err!(
- "get_next_join_result called without initializing join_result_status"
- )
- })?;
+ /// Handle ProbeRight state - process current probe batch
+ fn handle_probe_right(&mut self) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return ControlFlow::Break(poll);
+ }
- let (left_indices, right_indices, current_start) = (
- &status.build_indices,
- &status.probe_indices,
- status.processed_count,
- );
+ // Process current probe state
+ match self.process_probe_batch() {
+ // State unchanged (ProbeRight)
+ // Continue probing until we have done joining the
+ // current right batch with all buffered left rows.
+ Ok(true) => ControlFlow::Continue(()),
+ // To next FetchRightState
+ // We have finished joining
+ // (cur_right_batch x buffered_left_batches)
+ Ok(false) => {
+ // Left exhausted, transition to FetchingRight
+ self.l_probe_idx = 0;
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ self.state = NLJState::EmitRightUnmatched;
+ } else {
+ self.current_right_batch = None;
+ self.state = NLJState::FetchingRight;
+ }
+ ControlFlow::Continue(())
+ }
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ }
+ }
- let left_batch = self
- .left_data
- .as_ref()
- .ok_or_else(|| internal_datafusion_err!("should have left_batch"))?
- .batch();
+ /// Handle EmitRightUnmatched state - emit unmatched right rows
+ fn handle_emit_right_unmatched(
+ &mut self,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return ControlFlow::Break(poll);
+ }
- let right_batch = match &self.state {
- NestedLoopJoinStreamState::ProcessProbeBatch(record_batch)
- | NestedLoopJoinStreamState::OutputUnmatchedBuildRows(record_batch) => {
- record_batch
+ debug_assert!(self.current_right_batch.is_some());
+ debug_assert!(self.current_right_batch_matched.is_some());
+
+ // Construct the result batch for unmatched right rows using a utility function
+ match self.process_right_unmatched() {
+ Ok(Some(batch)) => {
+ match self.output_buffer.push_batch(batch) {
+ Ok(()) => {
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
+ ControlFlow::Continue(())
+ }
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(
+ DataFusionError::ArrowError(Box::new(e), None),
+ )))),
+ }
}
- _ => {
- return internal_err!(
- "State should be ProcessProbeBatch or OutputUnmatchedBuildRows"
- )
+ Ok(None) => {
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
+ ControlFlow::Continue(())
}
- };
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ }
+ }
- if left_indices.is_empty() && right_indices.is_empty() && current_start == 0 {
- // To match the behavior of the previous implementation, return an empty RecordBatch.
- let res = RecordBatch::new_empty(Arc::clone(&self.schema));
- status.processed_count = 1;
- return Ok(Some(res));
+ /// Handle EmitLeftUnmatched state - emit unmatched left rows
+ fn handle_emit_left_unmatched(
+ &mut self,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return ControlFlow::Break(poll);
}
- if matches!(self.join_type, JoinType::RightSemi | JoinType::RightAnti) {
- // in this case left_indices.num_rows() == 0
- let end = min(
- current_start + self.intermediate_batch_size,
- right_indices.len(),
- );
+ // Process current unmatched state
+ match self.process_left_unmatched() {
+ // State unchanged (EmitLeftUnmatched)
+ // Continue processing until we have processed all unmatched rows
+ Ok(true) => ControlFlow::Continue(()),
+ // To Done state
+ // We have finished processing all unmatched rows
+ Ok(false) => match self.output_buffer.finish_buffered_batch() {
+ Ok(()) => {
+ self.state = NLJState::Done;
+ ControlFlow::Continue(())
+ }
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(
+ DataFusionError::ArrowError(Box::new(e), None),
+ )))),
+ },
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ }
+ }
- if current_start >= end {
- return Ok(None);
+ /// Handle Done state - final state processing
+ fn handle_done(&mut self) -> Poll<Option<Result<RecordBatch>>> {
+ // Return any remaining completed batches before final termination
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // HACK for the doc test in https://github.com/apache/datafusion/blob/main/datafusion/core/src/dataframe/mod.rs#L1265
+ // If this operator directly return `Poll::Ready(None)`
+ // for empty result, the final result will become an empty
+ // batch with empty schema, however the expected result
+ // should be with the expected schema for this operator
+ if !self.handled_empty_output {
+ let zero_count = Count::new();
+ if *self.join_metrics.baseline.output_rows() == zero_count {
+ let empty_batch = RecordBatch::new_empty(Arc::clone(&self.output_schema));
+ self.handled_empty_output = true;
+ return Poll::Ready(Some(Ok(empty_batch)));
}
+ }
- let res = Some(build_batch_from_indices(
- &self.schema,
- left_batch,
- right_batch,
- left_indices,
- &right_indices.slice(current_start, end - current_start),
- &self.column_indices,
- JoinSide::Left,
- )?);
+ Poll::Ready(None)
+ }
+
+ // ==== Core logic handling for each state ====
+
+ /// Returns bool to indicate should it continue probing
+ /// true -> continue in the same ProbeRight state
+ /// false -> It has done with the (buffered_left x cur_right_batch), go to
+ /// next state (ProbeRight)
+ fn process_probe_batch(&mut self) -> Result<bool> {
+ let left_data = Arc::clone(self.get_left_data()?);
+ let right_batch = self
+ .current_right_batch
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("Right batch should be available"))?
+ .clone();
- status.processed_count = end;
- return Ok(res);
+ // stop probing, the caller will go to the next state
+ if self.l_probe_idx >= left_data.batch().num_rows() {
+ return Ok(false);
}
- if current_start >= left_indices.len() {
- return Ok(None);
+ // ========
+ // Join (l_row x right_batch)
+ // and push the result into output_buffer
+ // ========
+
+ let l_idx = self.l_probe_idx;
+ let join_batch =
+ self.process_single_left_row_join(&left_data, &right_batch, l_idx)?;
+
+ if let Some(batch) = join_batch {
+ self.output_buffer.push_batch(batch)?;
}
- let end = min(
- current_start + self.intermediate_batch_size,
- left_indices.len(),
- );
+ // ==== Prepare for the next iteration ====
- let left_indices = &left_indices.slice(current_start, end - current_start);
- let right_indices = &right_indices.slice(current_start, end - current_start);
+ // Advance left cursor
+ self.l_probe_idx += 1;
- // Switch around the build side and probe side for `JoinType::RightMark`
- // because in a RightMark join, we want to mark rows on the right table
- // by looking for matches in the left.
- let res = if self.join_type == JoinType::RightMark {
- build_batch_from_indices(
- &self.schema,
+ // Return true to continue probing
+ Ok(true)
+ }
+
+ /// Process a single left row join with the current right batch.
+ /// Returns a RecordBatch containing the join results (None if empty)
+ fn process_single_left_row_join(
+ &mut self,
+ left_data: &JoinLeftData,
+ right_batch: &RecordBatch,
+ l_index: usize,
+ ) -> Result<Option<RecordBatch>> {
+ let right_row_count = right_batch.num_rows();
+ if right_row_count == 0 {
+ return Ok(None);
+ }
+
+ let cur_right_bitmap = if let Some(filter) = &self.join_filter {
+ apply_filter_to_row_join_batch(
+ left_data.batch(),
+ l_index,
right_batch,
- left_batch,
- left_indices,
- right_indices,
- &self.column_indices,
- JoinSide::Right,
- )
+ filter,
+ )?
} else {
- build_batch_from_indices(
- &self.schema,
- left_batch,
+ BooleanArray::from(vec![true; right_row_count])
+ };
+
+ self.update_matched_bitmap(l_index, &cur_right_bitmap)?;
+
+ // For the following join types: here we only have to set the left/right
+ // bitmap, and no need to output result
+ if matches!(
+ self.join_type,
+ JoinType::LeftAnti
+ | JoinType::LeftSemi
+ | JoinType::LeftMark
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ ) {
+ return Ok(None);
+ }
+
+ if cur_right_bitmap.true_count() == 0 {
+ // If none of the pairs has passed the join predicate/filter
+ Ok(None)
+ } else {
+ // Use the optimized approach similar to build_intermediate_batch_for_single_left_row
+ let join_batch = build_row_join_batch(
+ &self.output_schema,
+ left_data.batch(),
+ l_index,
right_batch,
- left_indices,
- right_indices,
+ Some(cur_right_bitmap),
&self.column_indices,
JoinSide::Left,
- )
- }?;
+ )?;
+ Ok(join_batch)
+ }
+ }
- status.processed_count = end;
+ /// Returns bool to indicate should it continue processing unmatched rows
+ /// true -> continue in the same EmitLeftUnmatched state
+ /// false -> next state (Done)
+ fn process_left_unmatched(&mut self) -> Result<bool> {
+ let left_data = self.get_left_data()?;
+ let left_batch = left_data.batch();
- Ok(Some(res))
- }
+ // Early return if join type can't have unmatched rows
+ if !need_produce_result_in_final(self.join_type) {
Review Comment:
wondering should we combine all 3 ifs into single boolean flag and return `Ok(false)` if flag is not set? 🤔
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
- // Produces 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
- *left_indices_cache = UInt64Array::from_iter_values(
- (0..output_row_count as u64).map(|i| i % left_row_count as u64),
- );
+ /// See the comments [`NestedLoopJoinExec`] for high-level design ideas.
+ ///
+ /// # Implementation
+ ///
+ /// This function is the entry point of NLJ operator's state machine
+ /// transitions. The rough state transition graph is as follow, for more
+ /// details see the comment in each state's matching arm.
+ ///
+ /// ============================
+ /// State transition graph:
+ /// ============================
+ ///
+ /// (start) --> BufferingLeft
+ /// ----------------------------
+ /// BufferingLeft → FetchingRight
+ ///
+ /// FetchingRight → ProbeRight (if right batch available)
+ /// FetchingRight → EmitLeftUnmatched (if right exhausted)
+ ///
+ /// ProbeRight → ProbeRight (next left row or after yielding output)
+ /// ProbeRight → EmitRightUnmatched (for special join types like right join)
+ /// ProbeRight → FetchingRight (done with the current right batch)
+ ///
+ /// EmitRightUnmatched → FetchingRight
+ ///
+ /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each
+ /// iteration)
+ /// EmitLeftUnmatched → Done (if finished)
+ /// ----------------------------
+ /// Done → (end)
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ loop {
+ match self.state {
+ // # NLJState transitions
+ // --> FetchingRight
+ // This state will prepare the left side batches, next state
+ // `FetchingRight` is responsible for preparing a single probe
+ // side batch, before start joining.
+ NLJState::BufferingLeft => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // inside `collect_left_input` (the rountine to buffer build
+ // -side batches), related metrics except build time will be
+ // updated.
+ // stop on drop
+ let build_metric = self.join_metrics.build_time.clone();
+ let _build_timer = build_metric.timer();
+
+ match self.handle_buffering_left(cx) {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => return poll,
+ }
+ }
- // Produces 0, 0, 0, 1, 1, 1, 2, 2, 2, ...
- *right_indices_cache = UInt32Array::from_iter_values(
- (0..output_row_count as u32).map(|i| i / left_row_count as u32),
- );
+ // # NLJState transitions:
+ // 1. --> ProbeRight
+ // Start processing the join for the newly fetched right
+ // batch.
+ // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit
+ // unmatched left side rows.
+ //
+ // After fetching a new batch from the right side, it will
+ // process all rows from the buffered left data:
+ // ```text
+ // for batch in right_side:
+ // for row in left_buffer:
+ // join(batch, row)
+ // ```
+ // Note: the implementation does this step incrementally,
+ // instead of materializing all intermediate Cartesian products
+ // at once in memory.
+ //
+ // So after the right side input is exhausted, the join phase
+ // for the current buffered left data is finished. We can go to
+ // the next `EmitLeftUnmatched` phase to check if there is any
+ // special handling (e.g., in cases like left join).
+ NLJState::FetchingRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_fetching_right(cx) {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => return poll,
+ }
+ }
+
+ // NLJState transitions:
+ // 1. --> ProbeRight(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> ProbeRight(2)
+ // After probing one right batch, and evaluating the
+ // join filter on (left-row x right-batch), it will advance
+ // to the next left row, then re-enter the current state and
+ // continue joining.
+ // 3. --> FetchRight
+ // After it has done with the current right batch (to join
+ // with all rows in the left buffer), it will go to
+ // FetchRight state to check what to do next.
+ NLJState::ProbeRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_probe_right() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been output by the join. In this state we have to
+ // output unmatched rows for current right batch (with null
+ // padding for left relation)
+ // Precondition: we have checked the join type so that it's
+ // possible to output right unmatched (e.g. it's right join)
+ NLJState::EmitRightUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_emit_right_unmatched() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // NLJState transitions:
+ // 1. --> EmitLeftUnmatched(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> EmitLeftUnmatched(2)
+ // After processing some unmatched rows, it will re-enter
+ // the same state, to check if there are any more final
+ // results to output.
+ // 3. --> Done
+ // It has processed all data, go to the final state and ready
+ // to exit.
+ //
+ // TODO: For memory-limited case, go back to `BufferingLeft`
+ // state again.
+ NLJState::EmitLeftUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_emit_left_unmatched() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
- (left_indices_cache.clone(), right_indices_cache.clone())
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+ // counting it in join timer due to there might be some
+ // final resout batches to output in this state
+
+ let poll = self.handle_done();
+ return self.join_metrics.baseline.record_poll(poll);
+ }
}
- };
+ }
+ }
+}
- if let Some(filter) = filter {
- apply_join_filter_to_indices(
- left_batch,
- right_batch,
- left_indices,
- right_indices,
- filter,
- JoinSide::Left,
- Some(max_intermediate_batch_size),
- )
- } else {
- Ok((left_indices, right_indices))
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
}
}
-impl NestedLoopJoinStream {
- fn poll_next_impl(
+impl NLJStream {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new(
+ schema: Arc<Schema>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ outer_table: SendableRecordBatchStream,
+ inner_table: OnceFut<JoinLeftData>,
+ column_indices: Vec<ColumnIndex>,
+ join_metrics: BuildProbeJoinMetrics,
+ cfg_batch_size: usize,
+ ) -> Self {
+ let should_track_unmatched_right = matches!(
+ join_type,
+ JoinType::Full
+ | JoinType::Right
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ );
+
+ Self {
+ output_schema: Arc::clone(&schema),
+ join_filter: filter,
+ join_type,
+ outer_table,
+ column_indices,
+ inner_table,
+ join_metrics,
+ buffered_left_data: None,
+ output_buffer: Box::new(BatchCoalescer::new(schema, cfg_batch_size)),
+ cfg_batch_size,
+ current_right_batch: None,
+ current_right_batch_matched: None,
+ state: NLJState::BufferingLeft,
+ l_probe_idx: 0,
+ l_emit_idx: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
+
+ // ==== State handler functions ====
+
+ /// Handle BufferingLeft state - prepare left side batches
+ fn handle_buffering_left(
&mut self,
cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- NestedLoopJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- NestedLoopJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- NestedLoopJoinStreamState::ProcessProbeBatch(_) => {
- let poll = handle_state!(self.process_probe_batch());
- self.join_metrics.baseline.record_poll(poll)
- }
- NestedLoopJoinStreamState::PrepareUnmatchedBuildRows => {
- handle_state!(self.prepare_unmatched_output_indices())
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ match self.inner_table.get_shared(cx) {
+ Poll::Ready(Ok(left_data)) => {
+ self.buffered_left_data = Some(left_data);
+ // TODO: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ // Continue to next state immediately
+ ControlFlow::Continue(())
+ }
+ Poll::Ready(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ Poll::Pending => ControlFlow::Break(Poll::Pending),
+ }
+ }
+
+ /// Handle FetchingRight state - fetch next right batch and prepare for processing
+ fn handle_fetching_right(
+ &mut self,
+ cx: &mut std::task::Context<'_>,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ match self.outer_table.poll_next_unpin(cx) {
+ Poll::Ready(result) => match result {
+ Some(Ok(right_batch)) => {
+ // Update metrics
+ self.join_metrics.input_rows.add(right_batch.num_rows());
+ self.join_metrics.input_batches.add(1);
+
+ let right_batch_size = right_batch.num_rows();
+
+ // Skip the empty batch
+ if right_batch_size == 0 {
+ return ControlFlow::Continue(());
+ }
+
+ self.current_right_batch = Some(right_batch);
+
+ // Prepare right bitmap
+ if self.should_track_unmatched_right {
+ let zeroed_buf = BooleanBuffer::new_unset(right_batch_size);
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(zeroed_buf, None));
+ }
+
+ self.l_probe_idx = 0;
+ self.state = NLJState::ProbeRight;
+ ControlFlow::Continue(())
}
- NestedLoopJoinStreamState::OutputUnmatchedBuildRows(_) => {
- let poll = handle_state!(self.build_unmatched_output());
- self.join_metrics.baseline.record_poll(poll)
+ Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ None => {
+ // Right stream exhausted
+ self.state = NLJState::EmitLeftUnmatched;
+ ControlFlow::Continue(())
}
- NestedLoopJoinStreamState::Completed => Poll::Ready(None),
- };
+ },
+ Poll::Pending => ControlFlow::Break(Poll::Pending),
}
}
- // This function's main job is to construct an output `RecordBatch` based on pre-calculated join indices.
- // It operates in a chunk-based manner, meaning it processes a portion of the results in each call,
- // making it suitable for streaming large datasets without high memory consumption.
- // This function behaves like an iterator. It returns `Ok(None)`
- // to signal that the result stream is exhausted and there is no more data.
- fn get_next_join_result(&mut self) -> Result<Option<RecordBatch>> {
- let status = self.join_result_status.as_mut().ok_or_else(|| {
- internal_datafusion_err!(
- "get_next_join_result called without initializing join_result_status"
- )
- })?;
+ /// Handle ProbeRight state - process current probe batch
+ fn handle_probe_right(&mut self) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return ControlFlow::Break(poll);
+ }
- let (left_indices, right_indices, current_start) = (
- &status.build_indices,
- &status.probe_indices,
- status.processed_count,
- );
+ // Process current probe state
+ match self.process_probe_batch() {
+ // State unchanged (ProbeRight)
+ // Continue probing until we have done joining the
+ // current right batch with all buffered left rows.
+ Ok(true) => ControlFlow::Continue(()),
+ // To next FetchRightState
+ // We have finished joining
+ // (cur_right_batch x buffered_left_batches)
+ Ok(false) => {
+ // Left exhausted, transition to FetchingRight
+ self.l_probe_idx = 0;
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ self.state = NLJState::EmitRightUnmatched;
+ } else {
+ self.current_right_batch = None;
+ self.state = NLJState::FetchingRight;
+ }
+ ControlFlow::Continue(())
+ }
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ }
+ }
- let left_batch = self
- .left_data
- .as_ref()
- .ok_or_else(|| internal_datafusion_err!("should have left_batch"))?
- .batch();
+ /// Handle EmitRightUnmatched state - emit unmatched right rows
+ fn handle_emit_right_unmatched(
+ &mut self,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return ControlFlow::Break(poll);
+ }
- let right_batch = match &self.state {
- NestedLoopJoinStreamState::ProcessProbeBatch(record_batch)
- | NestedLoopJoinStreamState::OutputUnmatchedBuildRows(record_batch) => {
- record_batch
+ debug_assert!(self.current_right_batch.is_some());
Review Comment:
meaningful message or delete
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
- // Produces 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
- *left_indices_cache = UInt64Array::from_iter_values(
- (0..output_row_count as u64).map(|i| i % left_row_count as u64),
- );
+ /// See the comments [`NestedLoopJoinExec`] for high-level design ideas.
+ ///
+ /// # Implementation
+ ///
+ /// This function is the entry point of NLJ operator's state machine
+ /// transitions. The rough state transition graph is as follow, for more
+ /// details see the comment in each state's matching arm.
+ ///
+ /// ============================
+ /// State transition graph:
+ /// ============================
+ ///
+ /// (start) --> BufferingLeft
+ /// ----------------------------
+ /// BufferingLeft → FetchingRight
+ ///
+ /// FetchingRight → ProbeRight (if right batch available)
+ /// FetchingRight → EmitLeftUnmatched (if right exhausted)
+ ///
+ /// ProbeRight → ProbeRight (next left row or after yielding output)
+ /// ProbeRight → EmitRightUnmatched (for special join types like right join)
+ /// ProbeRight → FetchingRight (done with the current right batch)
+ ///
+ /// EmitRightUnmatched → FetchingRight
+ ///
+ /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each
+ /// iteration)
+ /// EmitLeftUnmatched → Done (if finished)
+ /// ----------------------------
+ /// Done → (end)
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ loop {
+ match self.state {
+ // # NLJState transitions
+ // --> FetchingRight
+ // This state will prepare the left side batches, next state
+ // `FetchingRight` is responsible for preparing a single probe
+ // side batch, before start joining.
+ NLJState::BufferingLeft => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // inside `collect_left_input` (the rountine to buffer build
+ // -side batches), related metrics except build time will be
+ // updated.
+ // stop on drop
+ let build_metric = self.join_metrics.build_time.clone();
+ let _build_timer = build_metric.timer();
+
+ match self.handle_buffering_left(cx) {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => return poll,
+ }
+ }
- // Produces 0, 0, 0, 1, 1, 1, 2, 2, 2, ...
- *right_indices_cache = UInt32Array::from_iter_values(
- (0..output_row_count as u32).map(|i| i / left_row_count as u32),
- );
+ // # NLJState transitions:
+ // 1. --> ProbeRight
+ // Start processing the join for the newly fetched right
+ // batch.
+ // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit
+ // unmatched left side rows.
+ //
+ // After fetching a new batch from the right side, it will
+ // process all rows from the buffered left data:
+ // ```text
+ // for batch in right_side:
+ // for row in left_buffer:
+ // join(batch, row)
+ // ```
+ // Note: the implementation does this step incrementally,
+ // instead of materializing all intermediate Cartesian products
+ // at once in memory.
+ //
+ // So after the right side input is exhausted, the join phase
+ // for the current buffered left data is finished. We can go to
+ // the next `EmitLeftUnmatched` phase to check if there is any
+ // special handling (e.g., in cases like left join).
+ NLJState::FetchingRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_fetching_right(cx) {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => return poll,
+ }
+ }
+
+ // NLJState transitions:
+ // 1. --> ProbeRight(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> ProbeRight(2)
+ // After probing one right batch, and evaluating the
+ // join filter on (left-row x right-batch), it will advance
+ // to the next left row, then re-enter the current state and
+ // continue joining.
+ // 3. --> FetchRight
+ // After it has done with the current right batch (to join
+ // with all rows in the left buffer), it will go to
+ // FetchRight state to check what to do next.
+ NLJState::ProbeRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_probe_right() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been output by the join. In this state we have to
+ // output unmatched rows for current right batch (with null
+ // padding for left relation)
+ // Precondition: we have checked the join type so that it's
+ // possible to output right unmatched (e.g. it's right join)
+ NLJState::EmitRightUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_emit_right_unmatched() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // NLJState transitions:
+ // 1. --> EmitLeftUnmatched(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> EmitLeftUnmatched(2)
+ // After processing some unmatched rows, it will re-enter
+ // the same state, to check if there are any more final
+ // results to output.
+ // 3. --> Done
+ // It has processed all data, go to the final state and ready
+ // to exit.
+ //
+ // TODO: For memory-limited case, go back to `BufferingLeft`
+ // state again.
+ NLJState::EmitLeftUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+
+ match self.handle_emit_left_unmatched() {
+ ControlFlow::Continue(()) => continue,
+ ControlFlow::Break(poll) => {
+ return self.join_metrics.baseline.record_poll(poll)
+ }
+ }
+ }
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
- (left_indices_cache.clone(), right_indices_cache.clone())
+ // stop on drop
+ let join_metric = self.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+ // counting it in join timer due to there might be some
+ // final resout batches to output in this state
+
+ let poll = self.handle_done();
+ return self.join_metrics.baseline.record_poll(poll);
+ }
}
- };
+ }
+ }
+}
- if let Some(filter) = filter {
- apply_join_filter_to_indices(
- left_batch,
- right_batch,
- left_indices,
- right_indices,
- filter,
- JoinSide::Left,
- Some(max_intermediate_batch_size),
- )
- } else {
- Ok((left_indices, right_indices))
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
}
}
-impl NestedLoopJoinStream {
- fn poll_next_impl(
+impl NLJStream {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new(
+ schema: Arc<Schema>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ outer_table: SendableRecordBatchStream,
+ inner_table: OnceFut<JoinLeftData>,
+ column_indices: Vec<ColumnIndex>,
+ join_metrics: BuildProbeJoinMetrics,
+ cfg_batch_size: usize,
+ ) -> Self {
+ let should_track_unmatched_right = matches!(
+ join_type,
+ JoinType::Full
+ | JoinType::Right
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ );
+
+ Self {
+ output_schema: Arc::clone(&schema),
+ join_filter: filter,
+ join_type,
+ outer_table,
+ column_indices,
+ inner_table,
+ join_metrics,
+ buffered_left_data: None,
+ output_buffer: Box::new(BatchCoalescer::new(schema, cfg_batch_size)),
+ cfg_batch_size,
+ current_right_batch: None,
+ current_right_batch_matched: None,
+ state: NLJState::BufferingLeft,
+ l_probe_idx: 0,
+ l_emit_idx: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
+
+ // ==== State handler functions ====
+
+ /// Handle BufferingLeft state - prepare left side batches
+ fn handle_buffering_left(
&mut self,
cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- NestedLoopJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- NestedLoopJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- NestedLoopJoinStreamState::ProcessProbeBatch(_) => {
- let poll = handle_state!(self.process_probe_batch());
- self.join_metrics.baseline.record_poll(poll)
- }
- NestedLoopJoinStreamState::PrepareUnmatchedBuildRows => {
- handle_state!(self.prepare_unmatched_output_indices())
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ match self.inner_table.get_shared(cx) {
+ Poll::Ready(Ok(left_data)) => {
+ self.buffered_left_data = Some(left_data);
+ // TODO: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ // Continue to next state immediately
+ ControlFlow::Continue(())
+ }
+ Poll::Ready(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ Poll::Pending => ControlFlow::Break(Poll::Pending),
+ }
+ }
+
+ /// Handle FetchingRight state - fetch next right batch and prepare for processing
+ fn handle_fetching_right(
+ &mut self,
+ cx: &mut std::task::Context<'_>,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ match self.outer_table.poll_next_unpin(cx) {
+ Poll::Ready(result) => match result {
+ Some(Ok(right_batch)) => {
+ // Update metrics
+ self.join_metrics.input_rows.add(right_batch.num_rows());
+ self.join_metrics.input_batches.add(1);
+
+ let right_batch_size = right_batch.num_rows();
+
+ // Skip the empty batch
+ if right_batch_size == 0 {
+ return ControlFlow::Continue(());
+ }
+
+ self.current_right_batch = Some(right_batch);
+
+ // Prepare right bitmap
+ if self.should_track_unmatched_right {
+ let zeroed_buf = BooleanBuffer::new_unset(right_batch_size);
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(zeroed_buf, None));
+ }
+
+ self.l_probe_idx = 0;
+ self.state = NLJState::ProbeRight;
+ ControlFlow::Continue(())
}
- NestedLoopJoinStreamState::OutputUnmatchedBuildRows(_) => {
- let poll = handle_state!(self.build_unmatched_output());
- self.join_metrics.baseline.record_poll(poll)
+ Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ None => {
+ // Right stream exhausted
+ self.state = NLJState::EmitLeftUnmatched;
+ ControlFlow::Continue(())
}
- NestedLoopJoinStreamState::Completed => Poll::Ready(None),
- };
+ },
+ Poll::Pending => ControlFlow::Break(Poll::Pending),
}
}
- // This function's main job is to construct an output `RecordBatch` based on pre-calculated join indices.
- // It operates in a chunk-based manner, meaning it processes a portion of the results in each call,
- // making it suitable for streaming large datasets without high memory consumption.
- // This function behaves like an iterator. It returns `Ok(None)`
- // to signal that the result stream is exhausted and there is no more data.
- fn get_next_join_result(&mut self) -> Result<Option<RecordBatch>> {
- let status = self.join_result_status.as_mut().ok_or_else(|| {
- internal_datafusion_err!(
- "get_next_join_result called without initializing join_result_status"
- )
- })?;
+ /// Handle ProbeRight state - process current probe batch
+ fn handle_probe_right(&mut self) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return ControlFlow::Break(poll);
+ }
- let (left_indices, right_indices, current_start) = (
- &status.build_indices,
- &status.probe_indices,
- status.processed_count,
- );
+ // Process current probe state
+ match self.process_probe_batch() {
+ // State unchanged (ProbeRight)
+ // Continue probing until we have done joining the
+ // current right batch with all buffered left rows.
+ Ok(true) => ControlFlow::Continue(()),
+ // To next FetchRightState
+ // We have finished joining
+ // (cur_right_batch x buffered_left_batches)
+ Ok(false) => {
+ // Left exhausted, transition to FetchingRight
+ self.l_probe_idx = 0;
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ self.state = NLJState::EmitRightUnmatched;
+ } else {
+ self.current_right_batch = None;
+ self.state = NLJState::FetchingRight;
+ }
+ ControlFlow::Continue(())
+ }
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ }
+ }
- let left_batch = self
- .left_data
- .as_ref()
- .ok_or_else(|| internal_datafusion_err!("should have left_batch"))?
- .batch();
+ /// Handle EmitRightUnmatched state - emit unmatched right rows
+ fn handle_emit_right_unmatched(
+ &mut self,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return ControlFlow::Break(poll);
+ }
- let right_batch = match &self.state {
- NestedLoopJoinStreamState::ProcessProbeBatch(record_batch)
- | NestedLoopJoinStreamState::OutputUnmatchedBuildRows(record_batch) => {
- record_batch
+ debug_assert!(self.current_right_batch.is_some());
+ debug_assert!(self.current_right_batch_matched.is_some());
+
+ // Construct the result batch for unmatched right rows using a utility function
+ match self.process_right_unmatched() {
+ Ok(Some(batch)) => {
+ match self.output_buffer.push_batch(batch) {
+ Ok(()) => {
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
+ ControlFlow::Continue(())
+ }
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(
+ DataFusionError::ArrowError(Box::new(e), None),
+ )))),
+ }
}
- _ => {
- return internal_err!(
- "State should be ProcessProbeBatch or OutputUnmatchedBuildRows"
- )
+ Ok(None) => {
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
+ ControlFlow::Continue(())
}
- };
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
+ }
+ }
- if left_indices.is_empty() && right_indices.is_empty() && current_start == 0 {
- // To match the behavior of the previous implementation, return an empty RecordBatch.
- let res = RecordBatch::new_empty(Arc::clone(&self.schema));
- status.processed_count = 1;
- return Ok(Some(res));
+ /// Handle EmitLeftUnmatched state - emit unmatched left rows
+ fn handle_emit_left_unmatched(
+ &mut self,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return ControlFlow::Break(poll);
}
- if matches!(self.join_type, JoinType::RightSemi | JoinType::RightAnti) {
- // in this case left_indices.num_rows() == 0
- let end = min(
- current_start + self.intermediate_batch_size,
- right_indices.len(),
- );
+ // Process current unmatched state
+ match self.process_left_unmatched() {
+ // State unchanged (EmitLeftUnmatched)
+ // Continue processing until we have processed all unmatched rows
+ Ok(true) => ControlFlow::Continue(()),
+ // To Done state
+ // We have finished processing all unmatched rows
+ Ok(false) => match self.output_buffer.finish_buffered_batch() {
+ Ok(()) => {
+ self.state = NLJState::Done;
+ ControlFlow::Continue(())
+ }
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(
+ DataFusionError::ArrowError(Box::new(e), None),
Review Comment:
please use `arrow_err!`
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2262474885
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
Review Comment:
I have updated the comment to better explain 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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3182175534
There are no new changes included. The speedup reaches 5× simply because the NLJ micro-benchmark is extended with cases where the join predicate is very cheap to evaluate (see https://github.com/apache/datafusion/pull/16996#discussion_r2253750647), and such cases favor the new implementation more. I changed the title to make it more click-baity 🤣
I plan to keep it open until Sunday in case others have additional suggestions to polish the implementation.
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2262497081
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
Review Comment:
I actually want to enforce this naming style at the project scope.
For rare configuration options that has been propagated many layers to some utility function, it would be easier to figure out the purpose of such argument by prefixing `cfg_`, usually in the config page we can find a more detailed explanation.
We can open issue, and write a style guide for it, WDYT?
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3173870087
> > I really appreciate your detailed review. I have addressed them in [4c111c3](https://github.com/apache/datafusion/pull/16996/commits/4c111c38e6b05553d78bf9c46346f4a47cb7e710)
> > BTW, I take code understandability very seriously, and I’m happy to make any small changes that improve it. I encourage others to do similar reviews -- point out anything that doesn’t make sense and to share even small nitpicks.
> > > Thanks @2010YOUY01 I'll continue later, I'm out of my mental capacity :) Appreciate for the comments it was much easier to navigate.
> > > Good sign we got fuzz testing working with the new implementation. For the post filtering it would be probably possible to split filter stage evalulation for certain types of join. Inner, outer joins filter can be evaluated early whereas SEMI, ANTI on late stage like now.
> >
> >
> > I don't get this changing filter evaluation time idea, could you elaborate?
>
> As far as I understood the filter evaluation happens in post phase, after tuples joined by key and this makes a lot of sense especially for SEMI, ANTI joins where you have to track filter eval results for the same key across all batches. However for Inner join and outer joins it maybe beneficial to evaluate record batches early and avoid join calculation and save some ticks
>
> For example
>
> ```
> > with t as (select 1 a, 2 b), t1 as (select 2 a, 3 b) select * from t left join t1 on (t.a < t1.b and t1.b = 4);
> +---+---+------+------+
> | a | b | a | b |
> +---+---+------+------+
> | 1 | 2 | NULL | NULL |
> +---+---+------+------+
> ```
>
> So the right record batch you can evaluate filter `t1.b = 4` and figure out that indices in the batch doesn't fit to join condition even without having actual join
I see, but this step should be handled below the join operator by filter pushdown:
See the `FilterExec` in the left input in the below example
```
> explain SELECT *
FROM range(10000) AS t1
JOIN range(10000) AS t2
ON ((t1.value + t2.value) % 1000 = 0) AND (t1.value > 5000);
+---------------+------------------------------------------------------------+
| plan_type | plan |
+---------------+------------------------------------------------------------+
| physical_plan | ┌───────────────────────────┐ |
| | │ NestedLoopJoinExec ├──────────────┐ |
| | └─────────────┬─────────────┘ │ |
| | ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ |
| | │ CoalescePartitionsExec ││ RepartitionExec │ |
| | │ ││ -------------------- │ |
| | │ ││ partition_count(in->out): │ |
| | │ ││ 1 -> 14 │ |
| | │ ││ │ |
| | │ ││ partitioning_scheme: │ |
| | │ ││ RoundRobinBatch(14) │ |
| | └─────────────┬─────────────┘└─────────────┬─────────────┘ |
| | ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ |
| | │ CoalesceBatchesExec ││ LazyMemoryExec │ |
| | │ -------------------- ││ -------------------- │ |
| | │ target_batch_size: ││ batch_generators: │ |
| | │ 8192 ││ range: start=0, end=10000,│ |
| | │ ││ batch_size=8192 │ |
| | └─────────────┬─────────────┘└───────────────────────────┘ |
| | ┌─────────────┴─────────────┐ |
| | │ FilterExec │ |
| | │ -------------------- │ |
| | │ predicate: │ |
| | │ value > 5000 │ |
| | └─────────────┬─────────────┘ |
| | ┌─────────────┴─────────────┐ |
| | │ RepartitionExec │ |
| | │ -------------------- │ |
| | │ partition_count(in->out): │ |
| | │ 1 -> 14 │ |
| | │ │ |
| | │ partitioning_scheme: │ |
| | │ RoundRobinBatch(14) │ |
| | └─────────────┬─────────────┘ |
| | ┌─────────────┴─────────────┐ |
| | │ LazyMemoryExec │ |
| | │ -------------------- │ |
| | │ batch_generators: │ |
| | │ range: start=0, end=10000,│ |
| | │ batch_size=8192 │ |
| | └───────────────────────────┘ |
| | |
+---------------+------------------------------------------------------------+
```
Perhaps it's better to rename the variable name from `filter` to `join_predicate` to make it more clear.
I'm not sure if such filter pushdowns would be incorrect for SEMI/ANTI joins, I'll double check that.
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2267654762
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +691,1168 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NestedLoopJoinStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
Review Comment:
I don't agree with this. "Inner" and "outer" table, in particular for a nested loop join, have fairly strict definitions. We should not confuse the issue by referring to what is historically referred to as the outer table as the "inner" table in this implementation. It makes it harder for contributors to understand the implementation.
--
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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3189449241
> > @2010YOUY01 I think you can run extended tests on your forked repo from `Actions` providing a branch?
>
> Ah, yes. I forget pushing to private clone's main branch can also trigger it.
>
> I have started it: https://github.com/2010YOUY01/arrow-datafusion/actions/runs/16970623908
It is passed. I think its safe to merge now, thanks everyone@
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3140384776
Draft PR and early discussions: https://github.com/apache/datafusion/pull/16889
Thanks @UBarney @ding-young and @jonathanc-n for the help 🙏🏼
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "ding-young (via GitHub)" <gi...@apache.org>.
ding-young commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3141976153
### Memory Usage Benchmark
#### current
<details>
```
Query Time (ms) Peak RSS Peak Commit Major Page Faults
----------------------------------------------------------------
1 761.84 1042.7 MB 5.9 GB 0
2 973.85 1486.3 MB 8.3 GB 0
3 1197.72 2.2 GB 11.9 GB 0
4 3644.36 12.2 GB 20.8 GB 0
5 3729.40 3.9 GB 13.4 GB 0
6 18882.52 22.9 GB 24.5 GB 0
7 3128.09 3.9 GB 13.4 GB 0
8 19002.03 22.9 GB 24.6 GB 0
9 2968.10 10.3 GB 15.1 GB 0
10 6064.70 19.4 GB 28.3 GB 0
```
</details>
#### nlj-rewrite
<details>
```
Query Time (ms) Peak RSS Peak Commit Major Page Faults
----------------------------------------------------------------
1 218.46 78.6 MB 1024.1 MB 0
2 318.69 76.7 MB 1024.2 MB 0
3 352.38 77.0 MB 1024.2 MB 0
4 1384.05 85.1 MB 1024.2 MB 0
5 2217.99 82.4 MB 1024.2 MB 0
6 4331.11 77.0 MB 1024.1 MB 0
7 2295.73 79.1 MB 1024.1 MB 0
8 4372.36 78.6 MB 1024.1 MB 0
9 911.73 81.1 MB 1024.2 MB 0
10 1212.94 71.6 MB 1024.1 MB 0
```
</details>
(reproducer branch & comments to be added)
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3146452338
@UBarney I'm so sorry about that — I completely let that issue discussion slip through.
For important matters, feel free to ping me multiple times or bring it up again in a PR discussion.
I have to admit I often forget to respond to issue replies, but I usually follow up on PR discussions. (I think probably it's true for others as well)
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2269761709
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +691,1168 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NestedLoopJoinStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
Review Comment:
I believe Snowflake flips the convention and calls the left table "inner" which confuses the issue, regardless of physical join implementation. I think left/build and right/probe is consistent and fine. I just was concerned about inner and outer getting confused particularly for a NLJ. Thanks!
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3172093273
Btw Would be that possible to calculate the cost of the join like in https://www.youtube.com/watch?v=RcEW0P8iVTc
?
The video shows multiple implementations for NLJ and how to calculate the cost and describe pseudo code, it would be super useful for community and further improvements.
From what I understood, the left side scanned once, and entirely saved in memory, what about right scans?
Perhaps in future we can play with blocks of input left batches to prevent OOM
--
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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3416791653
@ianthetechie reported a performance regression in 50.0.0, and I tracked it down to this PR. See details / reproducer on:
- https://github.com/apache/datafusion/issues/18070#issuecomment-3416788644
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2248805092
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1048 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
+
+ /// Should we use a bitmap to track each incoming right batch's each row's
+ /// 'joined' status.
+ /// For example in right joins, we have to use a bit map to track matched
+ /// right side rows, and later enter a `EmitRightUnmatched` stage to emit
+ /// unmatched right rows.
+ should_track_unmatched_right: bool,
+
+ // ========================================================================
+ // STATE FLAGS/BUFFERS:
+ // Fields that hold intermediate data/flags during execution
+ // ========================================================================
+ /// State Tracking
+ state: NLJState,
+ /// Output buffer holds the join result to output. It will emit eagerly when
+ /// the threshold is reached.
+ output_buffer: Box<BatchCoalescer>,
+ /// See comments in [`NLJState::Done`] for its purpose
+ handled_empty_output: bool,
+
+ // Buffer(left) side
+ // -----------------
+ /// The current buffered left data to join
+ buffered_left_data: Option<Arc<JoinLeftData>>,
+ /// Index into the left buffered batch. Used in `ProbeRight` state
+ l_probe_idx: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ l_emit_idx: usize,
+ /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched`
+ /// state is over.
+ left_exhausted: bool,
+ /// If we can buffer all left data in one pass
+ /// TODO(now): this is for the (unimplemented) memory-limited execution
+ #[allow(dead_code)]
+ left_buffered_in_one_pass: bool,
+
+ // Probe(right) side
+ // -----------------
+ /// The current probe batch to process
+ current_right_batch: Option<RecordBatch>,
+ // For right join, keep track of matched rows in `current_right_batch`
+ // Constructed when fetching each new incoming right batch in `FetchingRight` state.
+ current_right_batch_matched: Option<BooleanArray>,
}
-/// Creates a Cartesian product of two input batches, preserving the order of the right batch,
-/// and applying a join filter if provided.
-///
-/// # Example
-/// Input:
-/// left = [0, 1], right = [0, 1, 2]
-///
-/// Output:
-/// left_indices = [0, 1, 0, 1, 0, 1], right_indices = [0, 0, 1, 1, 2, 2]
-///
-/// Input:
-/// left = [0, 1, 2], right = [0, 1, 2, 3], filter = left.a != right.a
-///
-/// Output:
-/// left_indices = [1, 2, 0, 2, 0, 1, 0, 1, 2], right_indices = [0, 0, 1, 1, 2, 2, 3, 3, 3]
-fn build_join_indices(
- left_batch: &RecordBatch,
- right_batch: &RecordBatch,
- filter: Option<&JoinFilter>,
- indices_cache: &mut (UInt64Array, UInt32Array),
- max_intermediate_batch_size: usize,
-) -> Result<(UInt64Array, UInt32Array)> {
- let left_row_count = left_batch.num_rows();
- let right_row_count = right_batch.num_rows();
- let output_row_count = left_row_count * right_row_count;
-
- // We always use the same indices before applying the filter, so we can cache them
- let (left_indices_cache, right_indices_cache) = indices_cache;
- let cached_output_row_count = left_indices_cache.len();
-
- let (left_indices, right_indices) =
- match output_row_count.cmp(&cached_output_row_count) {
- std::cmp::Ordering::Equal => {
- // Reuse the cached indices
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- std::cmp::Ordering::Less => {
- // Left_row_count never changes because it's the build side. The changes to the
- // right_row_count can be handled trivially by taking the first output_row_count
- // elements of the cache because of how the indices are generated.
- // (See the Ordering::Greater match arm)
- (
- left_indices_cache.slice(0, output_row_count),
- right_indices_cache.slice(0, output_row_count),
- )
- }
- std::cmp::Ordering::Greater => {
- // Rebuild the indices cache
-
- // Produces 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
- *left_indices_cache = UInt64Array::from_iter_values(
- (0..output_row_count as u64).map(|i| i % left_row_count as u64),
- );
-
- // Produces 0, 0, 0, 1, 1, 1, 2, 2, 2, ...
- *right_indices_cache = UInt32Array::from_iter_values(
- (0..output_row_count as u32).map(|i| i / left_row_count as u32),
- );
-
- (left_indices_cache.clone(), right_indices_cache.clone())
- }
- };
-
- if let Some(filter) = filter {
- apply_join_filter_to_indices(
- left_batch,
- right_batch,
- left_indices,
- right_indices,
- filter,
- JoinSide::Left,
- Some(max_intermediate_batch_size),
- )
- } else {
- Ok((left_indices, right_indices))
- }
-}
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
-impl NestedLoopJoinStream {
- fn poll_next_impl(
- &mut self,
+ /// See the comments [`NestedLoopJoinExec`] for high-level design ideas.
+ ///
+ /// # Implementation
+ ///
+ /// This function is the entry point of NLJ operator's state machine
+ /// transitions. The rough state transition graph is as follow, for more
+ /// details see the comment in each state's matching arm.
+ ///
+ /// ============================
+ /// State transition graph:
+ /// ============================
+ ///
+ /// (start) --> BufferingLeft
+ /// ----------------------------
+ /// BufferingLeft → FetchingRight
+ ///
+ /// FetchingRight → ProbeRight (if right batch available)
+ /// FetchingRight → EmitLeftUnmatched (if right exhausted)
+ ///
+ /// ProbeRight → ProbeRight (next left row or after yielding output)
+ /// ProbeRight → EmitRightUnmatched (for special join types like right join)
+ /// ProbeRight → FetchingRight (done with the current right batch)
+ ///
+ /// EmitRightUnmatched → FetchingRight
+ ///
+ /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each
+ /// iteration)
+ /// EmitLeftUnmatched → Done (if finished)
+ /// ----------------------------
+ /// Done → (end)
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
+ ) -> Poll<Option<Self::Item>> {
loop {
- return match self.state {
- NestedLoopJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
+ match self.state {
+ // # NLJState transitions
+ // --> FetchingRight
+ // This state will prepare the left side batches, next state
+ // `FetchingRight` is responsible for preparing a single probe
+ // side batch, before start joining.
+ NLJState::BufferingLeft => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ match ready!(self.inner_table.get_shared(cx)) {
+ Ok(left_data) => {
+ self.buffered_left_data = Some(left_data);
+ // TODO: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
+
+ // # NLJState transitions:
+ // 1. --> ProbeRight
+ // Start processing the join for the newly fetched right
+ // batch.
+ // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit
+ // unmatched left side rows.
+ //
+ // After fetching a new batch from the right side, it will
+ // process all rows from the buffered left data:
+ // ```text
+ // for batch in right_side:
+ // for row in left_buffer:
+ // join(batch, row)
+ // ```
+ // Note: the implementation does this step incrementally,
+ // instead of materializing all intermediate Cartesian products
+ // at once in memory.
+ //
+ // So after the right side input is exhausted, the join phase
+ // for the current buffered left data is finished. We can go to
+ // the next `EmitLeftUnmatched` phase to check if there is any
+ // special handling (e.g., in cases like left join).
+ NLJState::FetchingRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ match ready!(self.outer_table.poll_next_unpin(cx)) {
+ Some(Ok(right_batch)) => {
+ let right_batch_size = right_batch.num_rows();
+
+ // Skip the empty batch
+ if right_batch_size == 0 {
+ continue;
+ }
+
+ self.current_right_batch = Some(right_batch);
+
+ // Prepare right bitmap
+ if self.should_track_unmatched_right {
+ let zeroed_buf =
+ BooleanBuffer::new_unset(right_batch_size);
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(zeroed_buf, None));
+ }
+
+ self.l_probe_idx = 0;
+ self.state = NLJState::ProbeRight;
+ continue;
+ }
+ Some(Err(e)) => return Poll::Ready(Some(Err(e))),
+ None => {
+ // Right stream exhausted/as
+ self.state = NLJState::EmitLeftUnmatched;
+ continue;
+ }
+ }
}
- NestedLoopJoinStreamState::ProcessProbeBatch(_) => {
- let poll = handle_state!(self.process_probe_batch());
- self.join_metrics.baseline.record_poll(poll)
+
+ // NLJState transitions:
+ // 1. --> ProbeRight(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> ProbeRight(2)
+ // After probing one right batch, and evaluating the
+ // join filter on (left-row x right-batch), it will advance
+ // to the next left row, then re-enter the current state and
+ // continue joining.
+ // 3. --> FetchRight
+ // After it has done with the current right batch (to join
+ // with all rows in the left buffer), it will go to
+ // FetchRight state to check what to do next.
+ NLJState::ProbeRight => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // Process current probe state
+ match self.process_probe_batch() {
+ // State unchanged (ProbeRight)
+ // Continue probing until we have done joining the
+ // current right batch with all buffered left rows.
+ Ok(true) => continue,
+ // To next FetchRightState
+ // We have finished joining
+ // (cur_right_batch x buffered_left_batches)
+ Ok(false) => {
+ // Left exhausted, transition to FetchingRight
+ self.l_probe_idx = 0;
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ self.state = NLJState::EmitRightUnmatched;
+ } else {
+ self.current_right_batch = None;
+ self.state = NLJState::FetchingRight;
+ }
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::PrepareUnmatchedBuildRows => {
- handle_state!(self.prepare_unmatched_output_indices())
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been output by the join. In this state we have to
+ // output unmatched rows for current right batch (with null
+ // padding for left relation)
+ // Precondition: we have checked the join type so that it's
+ // possible to output right unmatched (e.g. it's right join)
+ NLJState::EmitRightUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ debug_assert!(self.current_right_batch.is_some());
+ debug_assert!(self.current_right_batch_matched.is_some());
+
+ // Construct the result batch for unmatched right rows using a utility function
+ if let Some(batch) = self.process_right_unmatched()? {
+ self.output_buffer.push_batch(batch)?;
+ }
+
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
}
- NestedLoopJoinStreamState::OutputUnmatchedBuildRows(_) => {
- let poll = handle_state!(self.build_unmatched_output());
- self.join_metrics.baseline.record_poll(poll)
+
+ // NLJState transitions:
+ // 1. --> EmitLeftUnmatched(1)
+ // If we have already buffered enough output to yield, it
+ // will first give back control to the parent state machine,
+ // then resume at the same place.
+ // 2. --> EmitLeftUnmatched(2)
+ // After processing some unmatched rows, it will re-enter
+ // the same state, to check if there are any more final
+ // results to output.
+ // 3. --> Done
+ // It has processed all data, go to the final state and ready
+ // to exit.
+ //
+ // TODO: For memory-limited case, go back to `BufferingLeft`
+ // state again.
+ NLJState::EmitLeftUnmatched => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // Process current unmatched state
+ match self.process_left_unmatched() {
+ // State unchanged (EmitLeftUnmatched)
+ // Continue processing until we have processed all unmatched rows
+ Ok(true) => continue,
+ // To Done state
+ // We have finished processing all unmatched rows
+ Ok(false) => {
+ self.output_buffer.finish_buffered_batch()?;
+ self.state = NLJState::Done;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
}
- NestedLoopJoinStreamState::Completed => Poll::Ready(None),
- };
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any remaining completed batches before final termination
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return poll;
+ }
+
+ // HACK for the doc test in https://github.com/apache/datafusion/blob/main/datafusion/core/src/dataframe/mod.rs#L1265
+ // If this operator directly return `Poll::Ready(None)`
+ // for empty result, the final result will become an empty
+ // batch with empty schema, however the expected result
+ // should be with the expected schema for this operator
+ if !self.handled_empty_output {
+ let zero_count = Count::new();
+ if *self.join_metrics.baseline.output_rows() == zero_count {
+ let empty_batch =
+ RecordBatch::new_empty(Arc::clone(&self.output_schema));
+ self.handled_empty_output = true;
+ return Poll::Ready(Some(Ok(empty_batch)));
+ }
+ }
+
+ return Poll::Ready(None);
+ }
+ }
}
}
+}
- // This function's main job is to construct an output `RecordBatch` based on pre-calculated join indices.
- // It operates in a chunk-based manner, meaning it processes a portion of the results in each call,
- // making it suitable for streaming large datasets without high memory consumption.
- // This function behaves like an iterator. It returns `Ok(None)`
- // to signal that the result stream is exhausted and there is no more data.
- fn get_next_join_result(&mut self) -> Result<Option<RecordBatch>> {
- let status = self.join_result_status.as_mut().ok_or_else(|| {
- internal_datafusion_err!(
- "get_next_join_result called without initializing join_result_status"
- )
- })?;
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
+ }
+}
- let (left_indices, right_indices, current_start) = (
- &status.build_indices,
- &status.probe_indices,
- status.processed_count,
+impl NLJStream {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new(
+ schema: Arc<Schema>,
+ filter: Option<JoinFilter>,
+ join_type: JoinType,
+ outer_table: SendableRecordBatchStream,
+ inner_table: OnceFut<JoinLeftData>,
+ column_indices: Vec<ColumnIndex>,
+ join_metrics: BuildProbeJoinMetrics,
+ cfg_batch_size: usize,
+ ) -> Self {
+ let should_track_unmatched_right = matches!(
+ join_type,
+ JoinType::Full
+ | JoinType::Right
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
);
- let left_batch = self
- .left_data
- .as_ref()
- .ok_or_else(|| internal_datafusion_err!("should have left_batch"))?
- .batch();
+ Self {
+ output_schema: Arc::clone(&schema),
+ join_filter: filter,
+ join_type,
+ outer_table,
+ column_indices,
+ inner_table,
+ join_metrics,
+ buffered_left_data: None,
+ output_buffer: Box::new(BatchCoalescer::new(schema, cfg_batch_size)),
+ cfg_batch_size,
+ current_right_batch: None,
+ current_right_batch_matched: None,
+ state: NLJState::BufferingLeft,
+ l_probe_idx: 0,
+ l_emit_idx: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
- let right_batch = match &self.state {
- NestedLoopJoinStreamState::ProcessProbeBatch(record_batch)
- | NestedLoopJoinStreamState::OutputUnmatchedBuildRows(record_batch) => {
- record_batch
- }
- _ => {
- return internal_err!(
- "State should be ProcessProbeBatch or OutputUnmatchedBuildRows"
- )
- }
- };
+ // ==== Core logic handling for each state ====
- if left_indices.is_empty() && right_indices.is_empty() && current_start == 0 {
- // To match the behavior of the previous implementation, return an empty RecordBatch.
- let res = RecordBatch::new_empty(Arc::clone(&self.schema));
- status.processed_count = 1;
- return Ok(Some(res));
- }
+ /// Returns bool to indicate should it continue probing
+ /// true -> continue in the same ProbeRight state
+ /// false -> It has done with the (buffered_left x cur_right_batch), go to
+ /// next state (ProbeRight)
+ fn process_probe_batch(&mut self) -> Result<bool> {
+ let left_data = Arc::clone(self.get_left_data()?);
+ let right_batch = self
+ .current_right_batch
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("Right batch should be available"))?
+ .clone();
- if matches!(self.join_type, JoinType::RightSemi | JoinType::RightAnti) {
- // in this case left_indices.num_rows() == 0
- let end = min(
- current_start + self.intermediate_batch_size,
- right_indices.len(),
- );
+ // stop probing, the caller will go to the next state
+ if self.l_probe_idx >= left_data.batch().num_rows() {
+ return Ok(false);
+ }
- if current_start >= end {
- return Ok(None);
- }
+ // ========
+ // Join (l_row x right_batch)
+ // and push the result into output_buffer
+ // ========
- let res = Some(build_batch_from_indices(
- &self.schema,
- left_batch,
- right_batch,
- left_indices,
- &right_indices.slice(current_start, end - current_start),
- &self.column_indices,
- JoinSide::Left,
- )?);
+ let l_idx = self.l_probe_idx;
+ let join_batch =
+ self.process_single_left_row_join(&left_data, &right_batch, l_idx)?;
- status.processed_count = end;
- return Ok(res);
+ if let Some(batch) = join_batch {
+ self.output_buffer.push_batch(batch)?;
}
- if current_start >= left_indices.len() {
+ // ==== Prepare for the next iteration ====
+
+ // Advance left cursor
+ self.l_probe_idx += 1;
+
+ // Return true to continue probing
+ Ok(true)
+ }
+
+ /// Process a single left row join with the current right batch.
+ /// Returns a RecordBatch containing the join results (None if empty)
+ fn process_single_left_row_join(
+ &mut self,
+ left_data: &JoinLeftData,
+ right_batch: &RecordBatch,
+ l_index: usize,
+ ) -> Result<Option<RecordBatch>> {
+ let right_row_count = right_batch.num_rows();
+ if right_row_count == 0 {
return Ok(None);
}
- let end = min(
- current_start + self.intermediate_batch_size,
- left_indices.len(),
- );
+ let cur_right_bitmap = if let Some(filter) = &self.join_filter {
+ apply_filter_to_row_join_batch(
+ left_data.batch(),
+ l_index,
+ right_batch,
+ filter,
+ )?
+ } else {
+ BooleanArray::from(vec![true; right_row_count])
+ };
- let left_indices = &left_indices.slice(current_start, end - current_start);
- let right_indices = &right_indices.slice(current_start, end - current_start);
+ self.update_matched_bitmap(l_index, &cur_right_bitmap)?;
- // Switch around the build side and probe side for `JoinType::RightMark`
- // because in a RightMark join, we want to mark rows on the right table
- // by looking for matches in the left.
- let res = if self.join_type == JoinType::RightMark {
- build_batch_from_indices(
- &self.schema,
- right_batch,
- left_batch,
- left_indices,
- right_indices,
- &self.column_indices,
- JoinSide::Right,
- )
+ // For the following join types: here we only have to set the left/right
+ // bitmap, and no need to output result
+ if matches!(
+ self.join_type,
+ JoinType::LeftAnti
+ | JoinType::LeftSemi
+ | JoinType::LeftMark
+ | JoinType::RightAnti
+ | JoinType::RightMark
+ | JoinType::RightSemi
+ ) {
+ return Ok(None);
+ }
+
+ if cur_right_bitmap.true_count() == 0 {
+ // If none of the pairs has passed the join predicate/filter
+ Ok(None)
} else {
- build_batch_from_indices(
- &self.schema,
- left_batch,
+ // Use the optimized approach similar to build_intermediate_batch_for_single_left_row
+ let join_batch = build_row_join_batch(
+ &self.output_schema,
+ left_data.batch(),
+ l_index,
right_batch,
- left_indices,
- right_indices,
+ Some(cur_right_bitmap),
&self.column_indices,
JoinSide::Left,
- )
- }?;
+ )?;
+ Ok(join_batch)
+ }
+ }
- status.processed_count = end;
+ /// Returns bool to indicate should it continue processing unmatched rows
+ /// true -> continue in the same EmitLeftUnmatched state
+ /// false -> next state (Done)
+ fn process_left_unmatched(&mut self) -> Result<bool> {
+ let left_data = self.get_left_data()?;
+ let left_batch = left_data.batch();
- Ok(Some(res))
- }
+ // Early return if join type can't have unmatched rows
+ if !need_produce_result_in_final(self.join_type) {
+ return Ok(false);
+ }
- fn collect_build_side(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- let build_timer = self.join_metrics.build_time.timer();
- // build hash table from left (build) side, if not yet done
- self.left_data = Some(ready!(self.inner_table.get_shared(cx))?);
- build_timer.done();
+ // Early return if another thread is already processing unmatched rows
+ if self.l_emit_idx == 0 && !left_data.report_probe_completed() {
+ return Ok(false);
+ }
- self.state = NestedLoopJoinStreamState::FetchProbeBatch;
+ // Stop processing unmatched rows, the caller will go to the next state
+ if self.l_emit_idx >= left_batch.num_rows() {
+ return Ok(false);
+ }
- Poll::Ready(Ok(StatefulStreamResult::Continue))
- }
+ // ========
+ // Process unmatched rows and push the result into output_buffer
+ // Each time, the number to process is up to batch size
+ // ========
+ let start_idx = self.l_emit_idx;
+ let end_idx =
+ std::cmp::min(start_idx + self.cfg_batch_size, left_batch.num_rows());
- /// Fetches next batch from probe-side
- ///
- /// If a non-empty batch has been fetched, updates state to
- /// `ProcessProbeBatchState`, otherwise updates state to `ExhaustedProbeSide`.
- fn fetch_probe_batch(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- match ready!(self.outer_table.poll_next_unpin(cx)) {
- None => {
- self.state = NestedLoopJoinStreamState::PrepareUnmatchedBuildRows;
- }
- Some(Ok(right_batch)) => {
- self.join_metrics.input_batches.add(1);
- self.join_metrics.input_rows.add(right_batch.num_rows());
+ if let Some(batch) =
+ self.process_left_unmatched_range(left_data, start_idx, end_idx)?
+ {
+ self.output_buffer.push_batch(batch)?;
+ }
- self.state = NestedLoopJoinStreamState::ProcessProbeBatch(right_batch);
- }
- Some(Err(err)) => return Poll::Ready(Err(err)),
- };
+ // ==== Prepare for the next iteration ====
+ self.l_emit_idx = end_idx;
- Poll::Ready(Ok(StatefulStreamResult::Continue))
+ // Return true to continue processing unmatched rows
+ Ok(true)
}
- /// Joins current probe batch with build-side data and produces batch with
- /// matched output, updates state to `FetchProbeBatch`.
- fn process_probe_batch(
- &mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let Some(left_data) = self.left_data.clone() else {
- return internal_err!(
- "Expected left_data to be Some in ProcessProbeBatch state"
+ /// Process unmatched rows from the left data within the specified range.
+ /// Returns a RecordBatch containing the unmatched rows (None if empty).
+ ///
+ /// # Arguments
+ /// * `left_data` - The left side data containing the batch and bitmap
+ /// * `start_idx` - Start index (inclusive) of the range to process
+ /// * `end_idx` - End index (exclusive) of the range to process
+ ///
+ /// # Safety
+ /// The caller is responsible for ensuring that `start_idx` and `end_idx` are
+ /// within valid bounds of the left batch. This function does not perform
+ /// bounds checking.
+ fn process_left_unmatched_range(
+ &self,
+ left_data: &JoinLeftData,
+ start_idx: usize,
+ end_idx: usize,
+ ) -> Result<Option<RecordBatch>> {
+ if start_idx == end_idx {
+ return Ok(None);
+ }
+
+ // Slice both left batch, and bitmap to range [start_idx, end_idx)
+ // The range is bit index (not byte)
+ let left_batch = left_data.batch();
+ let left_batch_sliced = left_batch.slice(start_idx, end_idx - start_idx);
+
+ // Can this be more efficient?
+ let mut bitmap_sliced = BooleanBufferBuilder::new(end_idx - start_idx);
+ bitmap_sliced.append_n(end_idx - start_idx, false);
+ let bitmap = left_data.bitmap().lock();
+ for i in start_idx..end_idx {
+ assert!(
+ i - start_idx < bitmap_sliced.capacity(),
+ "DBG: {start_idx}, {end_idx}"
);
- };
- let visited_left_side = left_data.bitmap();
- let batch = self.state.try_as_process_probe_batch()?;
+ bitmap_sliced.set_bit(i - start_idx, bitmap.get_bit(i));
+ }
+ let bitmap_sliced = BooleanArray::new(bitmap_sliced.finish(), None);
+
+ build_unmatched_batch(
+ Arc::clone(&self.output_schema),
+ &left_batch_sliced,
+ bitmap_sliced,
+ self.outer_table.schema(),
+ &self.column_indices,
+ self.join_type,
+ JoinSide::Left,
+ )
+ }
- let binding = self.join_metrics.join_time.clone();
- let _timer = binding.timer();
+ /// Process unmatched rows from the current right batch and reset the bitmap.
+ /// Returns a RecordBatch containing the unmatched right rows (None if empty).
+ fn process_right_unmatched(&mut self) -> Result<Option<RecordBatch>> {
+ // ==== Take current right batch and its bitmap ====
+ let right_batch_bitmap: BooleanArray =
+ std::mem::take(&mut self.current_right_batch_matched).ok_or_else(|| {
+ internal_datafusion_err!("right bitmap should be available")
+ })?;
+
+ let right_batch = self.current_right_batch.take();
+ let cur_right_batch = unwrap_or_internal_err!(right_batch);
+
+ let left_data = self.get_left_data()?;
+ let left_schema = left_data.batch().schema();
+
+ let res = build_unmatched_batch(
+ Arc::clone(&self.output_schema),
+ &cur_right_batch,
+ right_batch_bitmap,
+ left_schema,
+ &self.column_indices,
+ self.join_type,
+ JoinSide::Right,
+ );
- if self.join_result_status.is_none() {
- let (left_side_indices, right_side_indices) = join_left_and_right_batch(
- left_data.batch(),
- batch,
- self.join_type,
- self.filter.as_ref(),
- visited_left_side,
- &mut self.indices_cache,
- self.right_side_ordered,
- self.intermediate_batch_size,
- )?;
- self.join_result_status = Some(JoinResultProgress {
- build_indices: left_side_indices,
- probe_indices: right_side_indices,
- processed_count: 0,
- })
- }
+ // ==== Clean-up ====
+ self.current_right_batch_matched = None;
- let join_result = self.get_next_join_result()?;
+ res
+ }
- match join_result {
- Some(res) => {
- self.join_metrics.output_batches.add(1);
- Ok(StatefulStreamResult::Ready(Some(res)))
- }
- None => {
- self.state = NestedLoopJoinStreamState::FetchProbeBatch;
- self.join_result_status = None;
- Ok(StatefulStreamResult::Continue)
- }
- }
+ // ==== Utilities ====
+
+ /// Get the inner data of left data, errors if it's None
+ fn get_left_data(&self) -> Result<&Arc<JoinLeftData>> {
+ self.buffered_left_data
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("LeftData should be available"))
}
- fn build_unmatched_output(
- &mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let binding = self.join_metrics.join_time.clone();
- let _timer = binding.timer();
-
- let res = self.get_next_join_result()?;
- match res {
- Some(res) => {
- self.join_metrics.output_batches.add(1);
- Ok(StatefulStreamResult::Ready(Some(res)))
- }
- None => {
- self.state = NestedLoopJoinStreamState::Completed;
- Ok(StatefulStreamResult::Ready(None))
+ /// Flush the `output_buffer` if there are batches ready to output
+ /// None if no result batch ready.
+ fn maybe_flush_ready_batch(&mut self) -> Option<Poll<Option<Result<RecordBatch>>>> {
+ if self.output_buffer.has_completed_batch() {
+ if let Some(batch) = self.output_buffer.next_completed_batch() {
+ let poll = Poll::Ready(Some(Ok(batch)));
+ return Some(self.join_metrics.baseline.record_poll(poll));
}
}
+
+ None
}
- /// This function's primary purpose is to handle the final output stage required by specific join types after all right-side (probe) data has been exhausted.
- /// It is critically important for LEFT*/FULL joins, which must emit left-side (build) rows that found no match. For these cases, it identifies the unmatched rows and prepares the necessary state to output them.
- fn prepare_unmatched_output_indices(
+ /// After joining (l_index@left_buffer x current_right_batch), it will result
+ /// in a bitmap (the same length as current_right_batch) as the join match
+ /// result. Use this bitmap to update the global bitmap, for special join
+ /// types like full joins.
+ ///
+ /// Example:
+ /// After joining l_index=1 (1-indexed row in the left buffer), and the
+ /// current right batch with 3 elements, this function will be called with
+ /// arguments: l_index = 1, r_matched = [false, false, true]
+ /// - If the join type is FullJoin, the 1-index in the left bitmap will be
+ /// set to true, and also the right bitmap will be bitwise-ORed with the
+ /// input r_matched bitmap.
+ /// - For join types that don't require output unmatched rows, this
+ /// function can be a no-op. For inner joins, this function is a no-op; for left
+ /// joins, only the left bitmap may be updated.
+ fn update_matched_bitmap(
&mut self,
- ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let Some(left_data) = self.left_data.clone() else {
- return internal_err!(
- "Expected left_data to be Some in ExhaustedProbeSide state"
- );
- };
- let visited_left_side = left_data.bitmap();
- if need_produce_result_in_final(self.join_type) {
- // At this stage `visited_left_side` won't be updated, so it's
- // safe to report about probe completion.
- //
- // Setting `is_exhausted` / returning None will prevent from
- // multiple calls of `report_probe_completed()`
- if !left_data.report_probe_completed() {
- self.state = NestedLoopJoinStreamState::Completed;
- return Ok(StatefulStreamResult::Ready(None));
- };
+ l_index: usize,
+ r_matched: &BooleanArray,
+ ) -> Result<()> {
+ let left_data = self.get_left_data()?;
- // Only setting up timer, input is exhausted
- let _timer = self.join_metrics.join_time.timer();
- // use the global left bitmap to produce the left indices and right indices
- let (left_side, right_side) =
- get_final_indices_from_shared_bitmap(visited_left_side, self.join_type);
-
- self.join_result_status = Some(JoinResultProgress {
- build_indices: left_side,
- probe_indices: right_side,
- processed_count: 0,
- });
- self.state = NestedLoopJoinStreamState::OutputUnmatchedBuildRows(
- RecordBatch::new_empty(self.outer_table.schema()),
- );
+ // number of successfully joined pairs from (l_index x cur_right_batch)
+ let joined_len = r_matched.true_count();
- Ok(StatefulStreamResult::Continue)
- } else {
- // end of the join loop
- self.state = NestedLoopJoinStreamState::Completed;
- Ok(StatefulStreamResult::Ready(None))
+ // 1. Maybe update the left bitmap
+ if need_produce_result_in_final(self.join_type) && (joined_len > 0) {
+ let mut bitmap = left_data.bitmap().lock();
+ bitmap.set_bit(l_index, true);
+ }
+
+ // 2. Maybe updateh the right bitmap
+ if self.should_track_unmatched_right {
+ debug_assert!(self.current_right_batch_matched.is_some());
+ // after bit-wise or, it will be put back
+ let right_bitmap = std::mem::take(&mut self.current_right_batch_matched)
+ .ok_or_else(|| {
+ internal_datafusion_err!("right batch's bitmap should be present")
+ })?;
+ let (buf, nulls) = right_bitmap.into_parts();
+ debug_assert!(nulls.is_none());
+ let updated_right_bitmap = buf.bitor(r_matched.values());
+
+ self.current_right_batch_matched =
+ Some(BooleanArray::new(updated_right_bitmap, None));
}
+
+ Ok(())
}
}
-#[allow(clippy::too_many_arguments)]
-fn join_left_and_right_batch(
+// ==== Utilities ====
+
+/// Apply the join filter between:
+/// (l_index th row in left buffer) x (right batch)
+/// Returns a bitmap, with successfully joined indices set to true
+fn apply_filter_to_row_join_batch(
left_batch: &RecordBatch,
+ l_index: usize,
right_batch: &RecordBatch,
- join_type: JoinType,
- filter: Option<&JoinFilter>,
- visited_left_side: &SharedBitmapBuilder,
- indices_cache: &mut (UInt64Array, UInt32Array),
- right_side_ordered: bool,
- max_intermediate_batch_size: usize,
-) -> Result<(PrimitiveArray<UInt64Type>, PrimitiveArray<UInt32Type>)> {
- let (left_side, right_side) = build_join_indices(
- left_batch,
- right_batch,
- filter,
- indices_cache,
- max_intermediate_batch_size,
- )
- .map_err(|e| {
- exec_datafusion_err!(
- "Fail to build join indices in NestedLoopJoinExec, error: {e}"
- )
- })?;
-
- // set the left bitmap
- // and only full join need the left bitmap
- if need_produce_result_in_final(join_type) {
- let mut bitmap = visited_left_side.lock();
- left_side.values().iter().for_each(|x| {
- bitmap.set_bit(*x as usize, true);
- });
- }
- // adjust the two side indices base on the join type
- let (left_side, right_side) = adjust_indices_by_join_type(
- left_side,
- right_side,
- 0..right_batch.num_rows(),
- join_type,
- right_side_ordered,
- )?;
+ filter: &JoinFilter,
+) -> Result<BooleanArray> {
+ debug_assert!(left_batch.num_rows() != 0 && right_batch.num_rows() != 0);
+
+ let intermediate_batch = if filter.schema.fields().is_empty() {
+ // If filter is constant (e.g. literal `true`), empty batch can be used
+ // in the later filter step.
+ create_record_batch_with_empty_schema(
+ Arc::new((*filter.schema).clone()),
+ right_batch.num_rows(),
+ )?
+ } else {
+ build_row_join_batch(
+ &filter.schema,
+ left_batch,
+ l_index,
+ right_batch,
+ None,
+ &filter.column_indices,
+ JoinSide::Left,
+ )?
+ .ok_or_else(|| internal_datafusion_err!("This function assume input batch is not empty, so the intermediate batch can't be empty too"))?
+ };
+
+ let filter_result = filter
+ .expression()
+ .evaluate(&intermediate_batch)?
+ .into_array(intermediate_batch.num_rows())?;
- Ok((left_side, right_side))
+ Ok(as_boolean_array(&filter_result)?.clone())
}
-impl Stream for NestedLoopJoinStream {
- type Item = Result<RecordBatch>;
+/// This function performs the following steps:
+/// 1. Apply filter to probe-side batch
+/// 2. Broadcast the left row (build_side_batch\[build_side_index\]) to the
+/// filtered probe-side batch
+/// 3. Concat them together according to `col_indices`, and return the result
+/// (None if the result is empty)
+///
+/// Example:
+/// build_side_batch:
+/// a
+/// ----
+/// 1
+/// 2
+/// 3
+///
+/// # 0 index element in the build_side_batch (that is `1`) will be used
+/// build_side_index: 0
+///
+/// probe_side_batch:
+/// b
+/// ----
+/// 10
+/// 20
+/// 30
+/// 40
+///
+/// # After applying it, only index 1 and 3 elemnt in probe_side_batch will be
+/// # kept
+/// probe_side_filter:
+/// false
+/// true
+/// false
+/// true
+///
+///
+/// # Projections to the build/probe side batch, to construct the output batch
+/// col_indices:
+/// [(left, 0), (right, 0)]
+///
+/// build_side: left
+///
+/// ====
+/// Result batch:
+/// a b
+/// ----
+/// 1 20
+/// 1 40
+fn build_row_join_batch(
+ output_schema: &Schema,
+ build_side_batch: &RecordBatch,
+ build_side_index: usize,
+ probe_side_batch: &RecordBatch,
+ probe_side_filter: Option<BooleanArray>,
+ // See [`NLJStream`] struct's `column_indices` field for more detail
+ col_indices: &[ColumnIndex],
+ // If the build side is left or right, used to interpret the side information
+ // in `col_indices`
+ build_side: JoinSide,
+) -> Result<Option<RecordBatch>> {
+ debug_assert!(build_side != JoinSide::None);
+
+ // TODO(perf): since the output might be projection of right batch, this
+ // filtering step is more efficient to be done inside the column_index loop
+ let filtered_probe_batch = if let Some(filter) = probe_side_filter {
+ &filter_record_batch(probe_side_batch, &filter)?
+ } else {
+ probe_side_batch
+ };
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
+ if filtered_probe_batch.num_rows() == 0 {
+ return Ok(None);
+ }
+
+ let mut columns: Vec<Arc<dyn Array>> =
+ Vec::with_capacity(output_schema.fields().len());
+
+ for column_index in col_indices {
+ let array = if column_index.side == build_side {
+ // Broadcast the single build-side row to match the filtered
+ // probe-side batch length
+ let original_left_array = build_side_batch.column(column_index.index);
+ let scalar_value = ScalarValue::try_from_array(
+ original_left_array.as_ref(),
+ build_side_index,
+ )?;
+ scalar_value.to_array_of_size(filtered_probe_batch.num_rows())?
Review Comment:
I think it would be neat if we can avoid generating this array and directly apply the filter based on the scalar, rather than converting to an array of the same size.
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3173879649
> Btw Would be that possible to calculate the cost of the join like in https://www.youtube.com/watch?v=RcEW0P8iVTc ?
>
> The video shows multiple implementations for NLJ and how to calculate the cost and describe pseudo code, it would be super useful for community and further improvements.
>
> From what I understood, the left side scanned once, and entirely saved in memory, what about right scans? Perhaps in future we can play with blocks of input left batches to prevent OOM
Yes, that's exactly the idea for the future memory-limited NLJ implementation -- for each buffered left batches (under memory limit), do 1 round of right scan.
Though I don't think there are much tuning opportunities here, I think input scanning would be expensive if it's a parquet file, so the goal here is to minimize the number of right scans, and we should buffer as much left batches as possible.
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3172083106
> I really appreciate your detailed review. I have addressed them in [4c111c3](https://github.com/apache/datafusion/pull/16996/commits/4c111c38e6b05553d78bf9c46346f4a47cb7e710)
>
> BTW, I take code understandability very seriously, and I’m happy to make any small changes that improve it. I encourage others to do similar reviews -- point out anything that doesn’t make sense and to share even small nitpicks.
>
> > Thanks @2010YOUY01 I'll continue later, I'm out of my mental capacity :) Appreciate for the comments it was much easier to navigate.
> > Good sign we got fuzz testing working with the new implementation. For the post filtering it would be probably possible to split filter stage evalulation for certain types of join. Inner, outer joins filter can be evaluated early whereas SEMI, ANTI on late stage like now.
>
> I don't get this changing filter evaluation time idea, could you elaborate?
As far as I understood the filter evaluation happens in post phase, after tuples joined by key and this makes a lot of sense especially for SEMI, ANTI joins where you have to track filter eval results for the same key across all batches.
However for Inner join and outer joins it maybe beneficial to evaluate record batches early and avoid join calculation and save some ticks
For example
```
> with t as (select 1 a, 2 b), t1 as (select 2 a, 3 b) select * from t left join t1 on (t.a < t1.b and t1.b = 4);
+---+---+------+------+
| a | b | a | b |
+---+---+------+------+
| 1 | 2 | NULL | NULL |
+---+---+------+------+
```
So the right record batch you can evaluate filter `t1.b = 4` and figure out that indices in the batch doesn't fit to join condition even without having actual join
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2264956305
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +684,1180 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
+ // right/outer/probe-side table interchangeably.
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
/// join filter
- filter: Option<JoinFilter>,
+ pub(crate) join_filter: Option<JoinFilter>,
/// type of the join
- join_type: JoinType,
- /// the outer table data of the nested loop join
- outer_table: SendableRecordBatchStream,
- /// the inner table data of the nested loop join
- inner_table: OnceFut<JoinLeftData>,
- /// Information of index and left / right placement of columns
- column_indices: Vec<ColumnIndex>,
- // TODO: support null aware equal
- // null_equality: NullEquality,
+ pub(crate) join_type: JoinType,
+ /// the outer(right) table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner(left) table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Projection to construct the output schema from the left and right tables.
+ /// Example:
+ /// - output_schema: ['a', 'c']
+ /// - left_schema: ['a', 'b']
+ /// - right_schema: ['c']
+ ///
+ /// The column indices would be [(left, 0), (right, 0)] -- taking the left
+ /// 0th column and right 0th column can construct the output schema.
+ ///
+ /// Note there are other columns ('b' in the example) still kept after
+ /// projection pushdown; this is because they might be used to evaluate
+ /// the join filter (e.g., `JOIN ON (b+c)>0`).
+ pub(crate) column_indices: Vec<ColumnIndex>,
/// Join execution metrics
- join_metrics: BuildProbeJoinMetrics,
- /// Cache for join indices calculations
- indices_cache: (UInt64Array, UInt32Array),
- /// Whether the right side is ordered
- right_side_ordered: bool,
- /// Current state of the stream
- state: NestedLoopJoinStreamState,
- /// Result of the left data future
- left_data: Option<Arc<JoinLeftData>>,
-
- /// Tracks progress when building join result batches incrementally.
- join_result_status: Option<JoinResultProgress>,
-
- intermediate_batch_size: usize,
+ pub(crate) join_metrics: BuildProbeJoinMetrics,
+
+ /// `batch_size` from configuration
+ cfg_batch_size: usize,
Review Comment:
it might be good idea to standardize naming, although it is typically a painful process. IMO I'd prefer `batch_size` for now as people got used to it and it would be a confusion for community to understand if there is any diff between `batch_size` and `cfg_batch_size`
--
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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3186587861
Run extended tests
--
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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3189039109
> @2010YOUY01 I think you can run extended tests on your forked repo from `Actions` providing a branch?
Ah, yes. I forget pushing to private clone's main branch can also trigger it.
I have started it:
https://github.com/2010YOUY01/arrow-datafusion/actions/runs/16970623908
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "UBarney (via GitHub)" <gi...@apache.org>.
UBarney commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3146402155
Hi @2010YOUY01
Thanks for creating this amazing PR and for the detailed explanation on why we don't need to maintain the right_side order. This is a great optimization!
Just to provide some context for future collaboration, this approach (joining only one left row at a time) was actually the initial direction I considered. I was uncertain about the right_side_ordered requirement, so I raised this specific concern and @-mentioned you in [my comment](https://github.com/apache/datafusion/issues/16364#issuecomment-2975520489) seeking clarification.
Since I didn't get a response, I proceeded with the alternative implementation to keep things moving. A quick clarification on that point back then would have saved the effort on my [merged PR](https://github.com/apache/datafusion/pull/16443) and allowed us to get to this optimal solution directly.
No worries, the main thing is we've landed the best solution for the project. Just thought this feedback on the process could be helpful for our future interactions.
Great work on this!
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3167244769
I really appreciate your detailed review. I have addressed them in [4c111c3](https://github.com/apache/datafusion/pull/16996/commits/4c111c38e6b05553d78bf9c46346f4a47cb7e710)
BTW, I take code understandability very seriously, and I’m happy to make any small changes that improve it. I encourage others to do similar reviews -- point out anything that doesn’t make sense and to share even small nitpicks.
> Thanks @2010YOUY01 I'll continue later, I'm out of my mental capacity :) Appreciate for the comments it was much easier to navigate.
>
> Good sign we got fuzz testing working with the new implementation. For the post filtering it would be probably possible to split filter stage evalulation for certain types of join. Inner, outer joins filter can be evaluated early whereas SEMI, ANTI on late stage like now.
I don't get this changing filter evaluation time idea, could you elaborate?
--
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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3188943040
@2010YOUY01 I think you can run extended tests on your forked repo from `Actions` providing a branch?
--
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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead merged PR #16996:
URL: https://github.com/apache/datafusion/pull/16996
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2268515283
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -660,529 +691,1168 @@ async fn collect_left_input(
))
}
-/// This enumeration represents various states of the nested loop join algorithm.
-#[derive(Debug, Clone)]
-enum NestedLoopJoinStreamState {
- /// The initial state, indicating that build-side data not collected yet
- WaitBuildSide,
- /// Indicates that build-side has been collected, and stream is ready for
- /// fetching probe-side
- FetchProbeBatch,
- /// Indicates that a non-empty batch has been fetched from probe-side, and
- /// is ready to be processed
- ProcessProbeBatch(RecordBatch),
- /// Preparation phase: Gathers the indices of unmatched rows from the build-side.
- /// This state is entered for join types that emit unmatched build-side rows
- /// (e.g., LEFT and FULL joins) after the entire probe-side input has been consumed.
- PrepareUnmatchedBuildRows,
- /// Output unmatched build-side rows.
- /// The indices for rows to output has already been calculated in the previous
- /// `PrepareUnmatchedBuildRows` state. In this state the final batch will be materialized incrementally.
- // The inner `RecordBatch` is an empty dummy batch used to get right schema.
- OutputUnmatchedBuildRows(RecordBatch),
- /// Indicates that NestedLoopJoinStream execution is completed
- Completed,
-}
-
-impl NestedLoopJoinStreamState {
- /// Tries to extract a `ProcessProbeBatchState` from the
- /// `NestedLoopJoinStreamState` enum. Returns an error if state is not
- /// `ProcessProbeBatchState`.
- fn try_as_process_probe_batch(&mut self) -> Result<&RecordBatch> {
- match self {
- NestedLoopJoinStreamState::ProcessProbeBatch(state) => Ok(state),
- _ => internal_err!("Expected join stream in ProcessProbeBatch state"),
- }
- }
-}
-
-/// Tracks incremental output of join result batches.
-///
-/// Initialized with all matching pairs that satisfy the join predicate.
-/// Pairs are stored as indices in `build_indices` and `probe_indices`
-/// Each poll outputs a batch within the configured size limit and updates
-/// processed_count until all pairs are consumed.
-///
-/// Example: 5000 matches, batch size limit is 100
-/// - Poll 1: output batch[0..100], processed_count = 100
-/// - Poll 2: output batch[100..200], processed_count = 200
-/// - ...continues until processed_count = 5000
-struct JoinResultProgress {
- /// Row indices from build-side table (left table).
- build_indices: PrimitiveArray<UInt64Type>,
- /// Row indices from probe-side table (right table).
- probe_indices: PrimitiveArray<UInt32Type>,
- /// Number of index pairs already processed into output batches.
- /// We have completed join result for indices [0..processed_count).
- processed_count: usize,
+/// States for join processing. See `poll_next()` comment for more details about
+/// state transitions.
+#[derive(Debug, Clone, Copy)]
+enum NLJState {
+ BufferingLeft,
+ FetchingRight,
+ ProbeRight,
+ EmitRightUnmatched,
+ EmitLeftUnmatched,
+ Done,
}
-
-/// A stream that issues [RecordBatch]es as they arrive from the right of the join.
-struct NestedLoopJoinStream {
- /// Input schema
- schema: Arc<Schema>,
+pub(crate) struct NestedLoopJoinStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ //
+ // Note: The implementation uses the terms left/inner/build-side table and
Review Comment:
I see. What do you think about using only the terms left/build and right/probe in the implementation?
I think build and probe are unambiguous, and in DataFusion, left == build is a convention: it always swaps the smaller table to the left as the build side.
--
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] Rewrite Nested Loop Join executor for 3.5× speed and 1% memory usage [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#discussion_r2267606728
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -307,29 +318,9 @@ impl NestedLoopJoinExec {
))
}
- /// Returns a vector indicating whether the left and right inputs maintain their order.
- /// The first element corresponds to the left input, and the second to the right.
- ///
- /// The left (build-side) input's order may change, but the right (probe-side) input's
- /// order is maintained for INNER, RIGHT, RIGHT ANTI, and RIGHT SEMI joins.
- ///
- /// Maintaining the right input's order helps optimize the nodes down the pipeline
- /// (See [`ExecutionPlan::maintains_input_order`]).
- ///
- /// This is a separate method because it is also called when computing properties, before
- /// a [`NestedLoopJoinExec`] is created. It also takes [`JoinType`] as an argument, as
- /// opposed to `Self`, for the same reason.
- fn maintains_input_order(join_type: JoinType) -> Vec<bool> {
- vec![
- false,
- matches!(
- join_type,
- JoinType::Inner
- | JoinType::Right
- | JoinType::RightAnti
- | JoinType::RightSemi
- ),
- ]
+ /// This join implementation does not preserve the input order of either side.
Review Comment:
Not part of your diff, but spotted while reviewing, line 183:
```
/// the hash table creation.
inner_table: OnceAsync<JoinLeftData>,
```
The hash table comment seems like an old copy-paste error from hash_join.rs.
--
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] Rewrite Nested Loop Join executor for 5× speed and 1% memory usage [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16996:
URL: https://github.com/apache/datafusion/pull/16996#issuecomment-3186586752
> Here is an example of a query that passes in main and fails on this branch:
>
> ```sql
> CREATE TABLE tab0(col0 INTEGER, col1 INTEGER, col2 INTEGER);
> CREATE TABLE tab1(col0 INTEGER, col1 INTEGER, col2 INTEGER);
> CREATE TABLE tab2(col0 INTEGER, col1 INTEGER, col2 INTEGER);
> INSERT INTO tab0 VALUES(97,1,99);
> INSERT INTO tab0 VALUES(15,81,47);
> INSERT INTO tab0 VALUES(87,21,10);
> INSERT INTO tab1 VALUES(51,14,96);
> INSERT INTO tab1 VALUES(85,5,59);
> INSERT INTO tab1 VALUES(91,47,68);
> INSERT INTO tab2 VALUES(64,77,40);
> INSERT INTO tab2 VALUES(75,67,58);
> INSERT INTO tab2 VALUES(46,51,23);
>
> SELECT DISTINCT 32 AS col2 FROM tab0 AS cor0 LEFT OUTER JOIN tab2 AS cor1 ON ( NULL ) IS NULL;
>
> Arrow error: Invalid argument error: must either specify a row count or at least one column
> ```
>
> There are a bunch of other failures where the result count between main and this branch do not match. I verified a bunch of the queries do have NestedLoopJoinExec in the explain result.
Thanks for the reporting, I'll work on that today. I have ran the extended test on a early version, those failures must be introduced by my recent changes 🤔
--
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