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/24 12:31:45 UTC
[PR] WIP: Limit intermediate data size for NestedLoopJoin (up to 2X faster) [datafusion]
2010YOUY01 opened a new pull request, #16889:
URL: https://github.com/apache/datafusion/pull/16889
## 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.
-->
Now the NLJ execution logic is
1. Buffer all left input
2. Read one right batch at a time, and do a Cartesian product (all-left-batches x right-batch) to fetch the index combination
3. Incrementally do the join on that pre-computed indices
This approach includes one extra step for a very large Cartesian product calculation, it can be both memory-consuming and in-efficient. A better approach can be directly perform the join on a small intermediate.
## 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.
-->
### Summary
This PR introduces a somewhat unusual change: it rewrites the Nested Loop Join operator from scratch.
### Motivation
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.
### Why a Rewrite?
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.
### Current Status and Outlook
The rewrite is smoother than I expected: The current minimal implementation, which reuses existing utilities (despite some inefficiencies), is already up to **2× faster** and is likely significantly more memory-efficient (although memory usage has not yet been measured)
I also noticed there are several more optimizations opportunities to make it even faster. Now the implementation got regular tests passed, and there are several `join_fuzz` tests failing, after fixing them we can iterate on more optimizations.
### Benchmark Result
The NLJ micro-bench: https://github.com/apache/datafusion/pull/16819
```
┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ pr-16819 ┃ nlj-rewrite-tmp ┃ Change ┃
┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 219.41 ms │ 139.04 ms │ +1.58x faster │
│ QQuery 2 │ 265.81 ms │ 210.78 ms │ +1.26x faster │
│ QQuery 3 │ 363.44 ms │ 323.62 ms │ +1.12x faster │
│ QQuery 4 │ 995.84 ms │ 687.71 ms │ +1.45x faster │
│ QQuery 5 │ 640.90 ms │ 359.25 ms │ +1.78x faster │
│ QQuery 6 │ 5677.61 ms │ 2761.88 ms │ +2.06x faster │
│ QQuery 7 │ 666.71 ms │ 356.39 ms │ +1.87x faster │
│ QQuery 8 │ 5632.40 ms │ 2769.81 ms │ +2.03x faster │
│ QQuery 9 │ 673.28 ms │ 425.60 ms │ +1.58x faster │
│ QQuery 10 │ 2185.64 ms │ 1532.81 ms │ +1.43x faster │
└──────────────┴────────────┴─────────────────┴───────────────┘
```
### Implementation
The design/implementation doc can be found in the source.
### Next Steps
Given there are still many opportunities to make it even faster, I plan to include several major optimizations into the current PR to save some review bandwidth (It's always a contiguous low hundreds LoC for the core logic, even if we include more optimizations into it)
- [ ] TODO: doc about the major potential optimizations
- [ ] Fix bug that caused `join_fuzz` tests to fail
- [ ] Memory-limited case (maybe this can be split to a separate PR 🤔 , but I want to first do some POC to ensure it won't cause big structural change)
Would love to hear your thoughts on this @UBarney @jonathanc-n @ding-young , and looking forward to collaborating to make NLJ as fast as possible.
## 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
5. 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)?
-->
## 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.
-->
--
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#issuecomment-3121297298
Re: https://github.com/apache/datafusion/pull/16889#issuecomment-3121216474
I think the reason is: `datafusion-cli` won't buffer the final output, and now the NLJ bench will buffer all output (I'll fix 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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 3.2X faster) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#discussion_r2235985094
##########
datafusion/physical-plan/src/joins/nlj.rs:
##########
@@ -0,0 +1,803 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Implementation of the Nested Loop Join operator.
+//!
+//! For detailed information regarding the operator's state machine and execution flow,
+//! please refer to the documentation provided in the `poll_next()` method.
+
+use arrow::buffer::MutableBuffer;
+use arrow::compute::BatchCoalescer;
+use futures::{ready, StreamExt};
+use log::debug;
+use std::sync::Arc;
+use std::task::Poll;
+
+use crate::joins::nested_loop_join::JoinLeftData;
+use crate::joins::utils::{
+ apply_join_filter_to_indices, build_batch_from_indices_maybe_empty,
+ need_produce_result_in_final, BuildProbeJoinMetrics, ColumnIndex, JoinFilter,
+ OnceFut,
+};
+use crate::metrics::Count;
+use crate::{RecordBatchStream, SendableRecordBatchStream};
+
+use arrow::array::{
+ BooleanArray, BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array,
+ UInt64Builder,
+};
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{
+ internal_datafusion_err, unwrap_or_internal_err, DataFusionError, JoinSide, Result,
+};
+use datafusion_expr::JoinType;
+
+use futures::Stream;
+
+/// 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,
+}
+
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
+ /// join filter
+ pub(crate) join_filter: Option<JoinFilter>,
+ /// type of the join
+ pub(crate) join_type: JoinType,
+ /// the outer table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Information of index and left / right placement of columns
+ pub(crate) column_indices: Vec<ColumnIndex>,
+ /// Join execution metrics
+ 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_index: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ emit_cursor: u64,
+ /// 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<BooleanBufferBuilder>,
+}
+
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
+
+ /// # Design
+ ///
+ /// The high-level control flow for this operator is:
+ /// 1. Buffer all batches from the left side (unless memory limit is reached,
+ /// in which case see notes at 'Memory-limited Execution').
+ /// - Rationale: The right side scanning can be expensive (it might
+ /// include decoding Parquet files), so it tries to buffer more left
+ /// batches at once to minimize the scan passes.
+ /// 2. Read right side batch one at a time. For each iteration, it only
+ /// evaluates the join filter on (1-left-row x right-batch), and puts the
Review Comment:
Yes, we should mark it as a breaking change and add to the upgrade guide.
If anyone is relying heavily on this ordering property for optimization, we can keep the old impl (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] WIP: Limit intermediate data size for NestedLoopJoin (up to 2X faster) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#issuecomment-3121126110
> I'll need some time to read through the code, but reduced memory usage is also impressive 👍
>
> * v1 (old)
>
> ```
> Query Time (ms) Peak RSS Peak Commit Page Faults
> --------------------------------------------------------------
> 1 657.40 1112.0 MB 10.6 GB 0
> 2 833.66 1691.1 MB 10.6 GB 0
> 3 1688.13 3.4 GB 10.8 GB 0
> 4 4176.38 16.2 GB 20.1 GB 0
> 5 1629.22 23.2 GB 57.2 GB 0
> 6 17687.00 18.9 GB 26.1 GB 0
> 7 1245.64 23.3 GB 39.2 GB 0
> 8 19420.73 20.1 GB 31.3 GB 0
> 9 3467.28 10.4 GB 19.9 GB 0
> 10 5414.65 44.7 GB 46.0 GB 0
> ```
>
> * v2 (nlj-rewrite-tmp)
>
> ```
> Query Time (ms) Peak RSS Peak Commit Page Faults
> --------------------------------------------------------------
> 1 383.61 104.4 MB 5.0 GB 0
> 2 645.60 471.2 MB 8.0 GB 0
> 3 763.10 1738.5 MB 6.0 GB 0
> 4 1973.21 5.9 GB 9.6 GB 0
> 5 644.86 145.1 MB 4.0 GB 0
> 6 6043.58 149.3 MB 2.0 GB 0
> 7 543.70 154.5 MB 7.0 GB 0
> 8 6516.91 177.4 MB 4.0 GB 0
> 9 1571.97 140.4 MB 6.0 GB 0
> 10 3098.54 22.8 GB 26.9 GB 0
> ```
>
> I ran `cargo run --profile release-nonlto --bin mem_profile nlj` in https://github.com/ding-young/datafusion/tree/nlj-rewrite-tmp-review and adjusted v1/v2 with `static USE_NLJ_V2: bool = false;`
Thanks for the profiler 👍🏼 I'll review that soon.
I think Q4 and Q10 should not consume that much memory 😅
```sql
-- q10:
SELECT *
FROM range(30000) AS t1
FULL JOIN range(30000) AS t2
ON (t1.value + t2.value) % 10 <> 0;
```
It is supposed to only buffer (30k x 8Bytes), 20G RSS is not expected.
let me double check 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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "UBarney (via GitHub)" <gi...@apache.org>.
UBarney commented on code in PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#discussion_r2234007180
##########
datafusion/physical-plan/src/joins/nlj.rs:
##########
@@ -0,0 +1,803 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Implementation of the Nested Loop Join operator.
+//!
+//! For detailed information regarding the operator's state machine and execution flow,
+//! please refer to the documentation provided in the `poll_next()` method.
+
+use arrow::buffer::MutableBuffer;
+use arrow::compute::BatchCoalescer;
+use futures::{ready, StreamExt};
+use log::debug;
+use std::sync::Arc;
+use std::task::Poll;
+
+use crate::joins::nested_loop_join::JoinLeftData;
+use crate::joins::utils::{
+ apply_join_filter_to_indices, build_batch_from_indices_maybe_empty,
+ need_produce_result_in_final, BuildProbeJoinMetrics, ColumnIndex, JoinFilter,
+ OnceFut,
+};
+use crate::metrics::Count;
+use crate::{RecordBatchStream, SendableRecordBatchStream};
+
+use arrow::array::{
+ BooleanArray, BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array,
+ UInt64Builder,
+};
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{
+ internal_datafusion_err, unwrap_or_internal_err, DataFusionError, JoinSide, Result,
+};
+use datafusion_expr::JoinType;
+
+use futures::Stream;
+
+/// 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,
+}
+
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
+ /// join filter
+ pub(crate) join_filter: Option<JoinFilter>,
+ /// type of the join
+ pub(crate) join_type: JoinType,
+ /// the outer table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Information of index and left / right placement of columns
+ pub(crate) column_indices: Vec<ColumnIndex>,
+ /// Join execution metrics
+ 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_index: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ emit_cursor: u64,
+ /// 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<BooleanBufferBuilder>,
+}
+
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
+
+ /// # Design
+ ///
+ /// The high-level control flow for this operator is:
+ /// 1. Buffer all batches from the left side (unless memory limit is reached,
+ /// in which case see notes at 'Memory-limited Execution').
+ /// - Rationale: The right side scanning can be expensive (it might
+ /// include decoding Parquet files), so it tries to buffer more left
+ /// batches at once to minimize the scan passes.
+ /// 2. Read right side batch one at a time. For each iteration, it only
+ /// evaluates the join filter on (1-left-row x right-batch), and puts the
Review Comment:
> The new implementation can also preserve left order in some cases
Only when `right_table.num_rows() < batch_size ` ?
For join types like inner join, we are no longer preserving the order of the right table. I'm not sure if this is a breaking change. It's possible that some users rely on this behavior, especially those who create an ExecutionPlan directly, bypassing the DataFusion optimizer. (Are there any users who do 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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "UBarney (via GitHub)" <gi...@apache.org>.
UBarney commented on code in PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#discussion_r2232819590
##########
datafusion/physical-plan/src/joins/nlj.rs:
##########
@@ -0,0 +1,803 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Implementation of the Nested Loop Join operator.
+//!
+//! For detailed information regarding the operator's state machine and execution flow,
+//! please refer to the documentation provided in the `poll_next()` method.
+
+use arrow::buffer::MutableBuffer;
+use arrow::compute::BatchCoalescer;
+use futures::{ready, StreamExt};
+use log::debug;
+use std::sync::Arc;
+use std::task::Poll;
+
+use crate::joins::nested_loop_join::JoinLeftData;
+use crate::joins::utils::{
+ apply_join_filter_to_indices, build_batch_from_indices_maybe_empty,
+ need_produce_result_in_final, BuildProbeJoinMetrics, ColumnIndex, JoinFilter,
+ OnceFut,
+};
+use crate::metrics::Count;
+use crate::{RecordBatchStream, SendableRecordBatchStream};
+
+use arrow::array::{
+ BooleanArray, BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array,
+ UInt64Builder,
+};
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{
+ internal_datafusion_err, unwrap_or_internal_err, DataFusionError, JoinSide, Result,
+};
+use datafusion_expr::JoinType;
+
+use futures::Stream;
+
+/// 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,
+}
+
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
+ /// join filter
+ pub(crate) join_filter: Option<JoinFilter>,
+ /// type of the join
+ pub(crate) join_type: JoinType,
+ /// the outer table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Information of index and left / right placement of columns
+ pub(crate) column_indices: Vec<ColumnIndex>,
+ /// Join execution metrics
+ 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_index: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ emit_cursor: u64,
+ /// 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<BooleanBufferBuilder>,
+}
+
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
+
+ /// # Design
+ ///
+ /// The high-level control flow for this operator is:
+ /// 1. Buffer all batches from the left side (unless memory limit is reached,
+ /// in which case see notes at 'Memory-limited Execution').
+ /// - Rationale: The right side scanning can be expensive (it might
+ /// include decoding Parquet files), so it tries to buffer more left
+ /// batches at once to minimize the scan passes.
+ /// 2. Read right side batch one at a time. For each iteration, it only
+ /// evaluates the join filter on (1-left-row x right-batch), and puts the
Review Comment:
With this join method, it seems we can't preserve the order of the right table.
https://github.com/apache/datafusion/issues/16364#issuecomment-2975520489
##########
datafusion/physical-plan/src/joins/nlj.rs:
##########
@@ -0,0 +1,803 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Implementation of the Nested Loop Join operator.
+//!
+//! For detailed information regarding the operator's state machine and execution flow,
+//! please refer to the documentation provided in the `poll_next()` method.
+
+use arrow::buffer::MutableBuffer;
+use arrow::compute::BatchCoalescer;
+use futures::{ready, StreamExt};
+use log::debug;
+use std::sync::Arc;
+use std::task::Poll;
+
+use crate::joins::nested_loop_join::JoinLeftData;
+use crate::joins::utils::{
+ apply_join_filter_to_indices, build_batch_from_indices_maybe_empty,
+ need_produce_result_in_final, BuildProbeJoinMetrics, ColumnIndex, JoinFilter,
+ OnceFut,
+};
+use crate::metrics::Count;
+use crate::{RecordBatchStream, SendableRecordBatchStream};
+
+use arrow::array::{
+ BooleanArray, BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array,
+ UInt64Builder,
+};
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{
+ internal_datafusion_err, unwrap_or_internal_err, DataFusionError, JoinSide, Result,
+};
+use datafusion_expr::JoinType;
+
+use futures::Stream;
+
+/// 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,
+}
+
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
+ /// join filter
+ pub(crate) join_filter: Option<JoinFilter>,
+ /// type of the join
+ pub(crate) join_type: JoinType,
+ /// the outer table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Information of index and left / right placement of columns
+ pub(crate) column_indices: Vec<ColumnIndex>,
+ /// Join execution metrics
+ 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_index: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ emit_cursor: u64,
+ /// 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<BooleanBufferBuilder>,
+}
+
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
+
+ /// # Design
+ ///
+ /// The high-level control flow for this operator is:
+ /// 1. Buffer all batches from the left side (unless memory limit is reached,
+ /// in which case see notes at 'Memory-limited Execution').
+ /// - Rationale: The right side scanning can be expensive (it might
+ /// include decoding Parquet files), so it tries to buffer more left
+ /// batches at once to minimize the scan passes.
+ /// 2. Read right side batch one at a time. For each iteration, it only
+ /// evaluates the join filter on (1-left-row x right-batch), and puts the
+ /// result into the output buffer. Once the output buffer has reached
+ /// the threshold, output immediately.
+ /// - Rationale: Making the intermediate data smaller can 1) be more cache
+ /// friendly for processing to execute faster, and 2) use less memory.
+ ///
+ /// Note: Currently, both the filter-evaluation granularity and output
+ /// buffer size are `batch_size` from the configuration (default 8192).
+ /// We might try to tune it slightly for performance in the future.
+ ///
+ ///
+ ///
+ /// # Memory-limited Execution
+ ///
+ /// TODO.
+ /// The idea is each time buffer as much batches from the left side as
+ /// possible, then scan the right side once for all buffered left data.
+ /// Then buffer another left batches, scan right side again until finish.
+ ///
+ ///
+ ///
+ /// # 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.
+ ///
+ /// Draft 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);
+ match ready!(self.inner_table.get_shared(cx)) {
+ Ok(left_data) => {
+ self.buffered_left_data = Some(left_data);
+ // TOOD: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
+ }
+
+ // # 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 new_size = right_batch_size;
+ let zeroed_buf = MutableBuffer::from_len_zeroed(new_size);
+ self.current_right_batch_matched =
+ Some(BooleanBufferBuilder::new_from_buffer(
+ zeroed_buf, new_size,
+ ));
+ }
+
+ self.l_index = 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;
+ }
+ }
+ }
+
+ // 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 self.output_buffer.has_completed_batch() {
+ if let Some(batch) = self.output_buffer.next_completed_batch() {
+ let poll = Poll::Ready(Some(Ok(batch)));
+ return self.join_metrics.baseline.record_poll(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
+ // TODO(polish): use a flag for clarity
+ self.l_index = 0;
+ if self.current_right_batch_matched.is_some() {
+ // Don't reset current_right_batch, it'll be
+ // cleared inside `EmitRightUnmatched` state
+ self.state = NLJState::EmitRightUnmatched;
+ } else {
+ self.current_right_batch = None;
+ self.state = NLJState::FetchingRight;
+ }
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
+ }
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been outputed 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
+ let result_batch = self.process_right_unmatched()?;
+ self.output_buffer.push_batch(result_batch)?;
+
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
+ }
+
+ // 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 self.output_buffer.has_completed_batch() {
+ if let Some(batch) = self.output_buffer.next_completed_batch() {
+ let poll = Poll::Ready(Some(Ok(batch)));
+ return self.join_metrics.baseline.record_poll(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))),
+ }
+ }
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any remaining completed batches before final termination
+ 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 self.join_metrics.baseline.record_poll(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);
+ }
+ }
+ }
+ }
+}
+
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
+ }
+}
+
+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_index: 0,
+ emit_cursor: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
+
+ // ==== 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.buffered_left_data.as_ref().ok_or_else(|| {
+ internal_datafusion_err!("LeftData should be available")
+ })?);
+ let right_batch = self
+ .current_right_batch
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("Right batch should be available"))?
+ .clone();
+
+ // stop probing, the caller will go to the next state
+ if self.l_index >= left_data.batch().num_rows() {
+ return Ok(false);
+ }
+
+ // ========
+ // Join (l_row x right_batch)
+ // and push the result into output_buffer
+ // ========
+
+ let l_idx = self.l_index;
+ let join_batch =
+ self.process_single_left_row_join(&left_data, &right_batch, l_idx)?;
+
+ self.output_buffer.push_batch(join_batch)?;
+
+ // ==== Prepare for the next iteration ====
+
+ // Advance left cursor
+ self.l_index += 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 (may be empty).
+ fn process_single_left_row_join(
+ &mut self,
+ left_data: &JoinLeftData,
+ right_batch: &RecordBatch,
+ l_index: usize,
+ ) -> Result<RecordBatch> {
+ let right_row_count = right_batch.num_rows();
+ if right_row_count == 0 {
+ return Ok(RecordBatch::new_empty(Arc::clone(&self.output_schema)));
+ }
+
+ // Create indices for cross-join: current left row with all right rows
+ let left_indices = UInt64Array::from(vec![l_index as u64; right_row_count]);
+ let right_indices = UInt32Array::from_iter_values(0..right_row_count as u32);
+
+ // Apply join filter if present
+ let (joined_left_indices, joined_right_indices) =
+ if let Some(ref filter) = self.join_filter {
+ apply_join_filter_to_indices(
+ left_data.batch(),
+ right_batch,
+ left_indices,
+ right_indices,
+ filter,
+ JoinSide::Left,
+ None,
+ )?
+ } else {
+ (left_indices, right_indices)
+ };
+
+ // Update left row match bitmap for outer join support
+ if need_produce_result_in_final(self.join_type) && !joined_left_indices.is_empty()
+ {
+ let mut bitmap = left_data.bitmap().lock();
+ bitmap.set_bit(l_index, true);
+ }
+
+ // TODO(now-perf): better vectorize it
+ if let Some(bitmap) = self.current_right_batch_matched.as_mut() {
+ for i in joined_right_indices.iter() {
+ // After the initial join, indices must all be Some
+ bitmap.set_bit(i.unwrap() as usize, true);
+ // println!("Setting bit {i:?} to true");
+ }
+ }
+
+ // 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(RecordBatch::new_empty(Arc::clone(&self.output_schema)));
Review Comment:
maybe we can change return type of this function to `Result<Option<RecordBatch>>`
##########
datafusion/physical-plan/src/joins/nlj.rs:
##########
@@ -0,0 +1,803 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Implementation of the Nested Loop Join operator.
+//!
+//! For detailed information regarding the operator's state machine and execution flow,
+//! please refer to the documentation provided in the `poll_next()` method.
+
+use arrow::buffer::MutableBuffer;
+use arrow::compute::BatchCoalescer;
+use futures::{ready, StreamExt};
+use log::debug;
+use std::sync::Arc;
+use std::task::Poll;
+
+use crate::joins::nested_loop_join::JoinLeftData;
+use crate::joins::utils::{
+ apply_join_filter_to_indices, build_batch_from_indices_maybe_empty,
+ need_produce_result_in_final, BuildProbeJoinMetrics, ColumnIndex, JoinFilter,
+ OnceFut,
+};
+use crate::metrics::Count;
+use crate::{RecordBatchStream, SendableRecordBatchStream};
+
+use arrow::array::{
+ BooleanArray, BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array,
+ UInt64Builder,
+};
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{
+ internal_datafusion_err, unwrap_or_internal_err, DataFusionError, JoinSide, Result,
+};
+use datafusion_expr::JoinType;
+
+use futures::Stream;
+
+/// 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,
+}
+
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
+ /// join filter
+ pub(crate) join_filter: Option<JoinFilter>,
+ /// type of the join
+ pub(crate) join_type: JoinType,
+ /// the outer table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Information of index and left / right placement of columns
+ pub(crate) column_indices: Vec<ColumnIndex>,
+ /// Join execution metrics
+ 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_index: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ emit_cursor: u64,
+ /// 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<BooleanBufferBuilder>,
+}
+
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
+
+ /// # Design
+ ///
+ /// The high-level control flow for this operator is:
+ /// 1. Buffer all batches from the left side (unless memory limit is reached,
+ /// in which case see notes at 'Memory-limited Execution').
+ /// - Rationale: The right side scanning can be expensive (it might
+ /// include decoding Parquet files), so it tries to buffer more left
+ /// batches at once to minimize the scan passes.
+ /// 2. Read right side batch one at a time. For each iteration, it only
+ /// evaluates the join filter on (1-left-row x right-batch), and puts the
+ /// result into the output buffer. Once the output buffer has reached
+ /// the threshold, output immediately.
+ /// - Rationale: Making the intermediate data smaller can 1) be more cache
+ /// friendly for processing to execute faster, and 2) use less memory.
+ ///
+ /// Note: Currently, both the filter-evaluation granularity and output
+ /// buffer size are `batch_size` from the configuration (default 8192).
+ /// We might try to tune it slightly for performance in the future.
+ ///
+ ///
+ ///
+ /// # Memory-limited Execution
+ ///
+ /// TODO.
+ /// The idea is each time buffer as much batches from the left side as
+ /// possible, then scan the right side once for all buffered left data.
+ /// Then buffer another left batches, scan right side again until finish.
+ ///
+ ///
+ ///
+ /// # 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.
+ ///
+ /// Draft 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);
+ match ready!(self.inner_table.get_shared(cx)) {
+ Ok(left_data) => {
+ self.buffered_left_data = Some(left_data);
+ // TOOD: implement memory-limited case
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
+ }
+
+ // # 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 new_size = right_batch_size;
+ let zeroed_buf = MutableBuffer::from_len_zeroed(new_size);
+ self.current_right_batch_matched =
+ Some(BooleanBufferBuilder::new_from_buffer(
+ zeroed_buf, new_size,
+ ));
+ }
+
+ self.l_index = 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;
+ }
+ }
+ }
+
+ // 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 self.output_buffer.has_completed_batch() {
+ if let Some(batch) = self.output_buffer.next_completed_batch() {
+ let poll = Poll::Ready(Some(Ok(batch)));
+ return self.join_metrics.baseline.record_poll(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
+ // TODO(polish): use a flag for clarity
+ self.l_index = 0;
+ if self.current_right_batch_matched.is_some() {
+ // Don't reset current_right_batch, it'll be
+ // cleared inside `EmitRightUnmatched` state
+ self.state = NLJState::EmitRightUnmatched;
+ } else {
+ self.current_right_batch = None;
+ self.state = NLJState::FetchingRight;
+ }
+ continue;
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
+ }
+
+ // In the `current_right_batch_matched` bitmap, all trues mean
+ // it has been outputed 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
+ let result_batch = self.process_right_unmatched()?;
+ self.output_buffer.push_batch(result_batch)?;
+
+ // Processed all in one pass
+ // cleared inside `process_right_unmatched`
+ debug_assert!(self.current_right_batch.is_none());
+ self.state = NLJState::FetchingRight;
+ }
+
+ // 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 self.output_buffer.has_completed_batch() {
+ if let Some(batch) = self.output_buffer.next_completed_batch() {
+ let poll = Poll::Ready(Some(Ok(batch)));
+ return self.join_metrics.baseline.record_poll(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))),
+ }
+ }
+
+ // The final state and the exit point
+ NLJState::Done => {
+ debug!("[NLJState] Entering: {:?}", self.state);
+ // Return any remaining completed batches before final termination
+ 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 self.join_metrics.baseline.record_poll(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);
+ }
+ }
+ }
+ }
+}
+
+impl RecordBatchStream for NLJStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
+ }
+}
+
+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_index: 0,
+ emit_cursor: 0,
+ left_exhausted: false,
+ left_buffered_in_one_pass: true,
+ handled_empty_output: false,
+ should_track_unmatched_right,
+ }
+ }
+
+ // ==== 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.buffered_left_data.as_ref().ok_or_else(|| {
+ internal_datafusion_err!("LeftData should be available")
+ })?);
+ let right_batch = self
+ .current_right_batch
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("Right batch should be available"))?
+ .clone();
+
+ // stop probing, the caller will go to the next state
+ if self.l_index >= left_data.batch().num_rows() {
+ return Ok(false);
+ }
+
+ // ========
+ // Join (l_row x right_batch)
+ // and push the result into output_buffer
+ // ========
+
+ let l_idx = self.l_index;
+ let join_batch =
+ self.process_single_left_row_join(&left_data, &right_batch, l_idx)?;
+
+ self.output_buffer.push_batch(join_batch)?;
+
+ // ==== Prepare for the next iteration ====
+
+ // Advance left cursor
+ self.l_index += 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 (may be empty).
+ fn process_single_left_row_join(
+ &mut self,
+ left_data: &JoinLeftData,
+ right_batch: &RecordBatch,
+ l_index: usize,
+ ) -> Result<RecordBatch> {
+ let right_row_count = right_batch.num_rows();
+ if right_row_count == 0 {
+ return Ok(RecordBatch::new_empty(Arc::clone(&self.output_schema)));
+ }
+
+ // Create indices for cross-join: current left row with all right rows
+ let left_indices = UInt64Array::from(vec![l_index as u64; right_row_count]);
+ let right_indices = UInt32Array::from_iter_values(0..right_row_count as u32);
+
+ // Apply join filter if present
+ let (joined_left_indices, joined_right_indices) =
+ if let Some(ref filter) = self.join_filter {
+ apply_join_filter_to_indices(
+ left_data.batch(),
+ right_batch,
+ left_indices,
+ right_indices,
+ filter,
+ JoinSide::Left,
+ None,
+ )?
+ } else {
+ (left_indices, right_indices)
+ };
+
+ // Update left row match bitmap for outer join support
+ if need_produce_result_in_final(self.join_type) && !joined_left_indices.is_empty()
+ {
+ let mut bitmap = left_data.bitmap().lock();
+ bitmap.set_bit(l_index, true);
+ }
+
+ // TODO(now-perf): better vectorize it
+ if let Some(bitmap) = self.current_right_batch_matched.as_mut() {
+ for i in joined_right_indices.iter() {
+ // After the initial join, indices must all be Some
+ bitmap.set_bit(i.unwrap() as usize, true);
+ // println!("Setting bit {i:?} to true");
+ }
+ }
+
+ // 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(RecordBatch::new_empty(Arc::clone(&self.output_schema)));
+ }
+
+ // TODO(now): check if we're missed something inside the original
+ // logic inside `adjust_indices_by_join_types`
+
+ // Build output batch from matching indices
+ if !joined_left_indices.is_empty() {
+ let join_batch = build_batch_from_indices_maybe_empty(
Review Comment:
since we check `joined_left_indices.is_empty`, I think we can use `build_batch_from_indices` directly
--
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] WIP: Limit intermediate data size for NestedLoopJoin (up to 2X faster) [datafusion]
Posted by "ding-young (via GitHub)" <gi...@apache.org>.
ding-young commented on PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#issuecomment-3117009445
I'll need some time to read through the code, but for now I’ll run it and share the memory usage results :smile:
--
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "UBarney (via GitHub)" <gi...@apache.org>.
UBarney commented on code in PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#discussion_r2234007180
##########
datafusion/physical-plan/src/joins/nlj.rs:
##########
@@ -0,0 +1,803 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Implementation of the Nested Loop Join operator.
+//!
+//! For detailed information regarding the operator's state machine and execution flow,
+//! please refer to the documentation provided in the `poll_next()` method.
+
+use arrow::buffer::MutableBuffer;
+use arrow::compute::BatchCoalescer;
+use futures::{ready, StreamExt};
+use log::debug;
+use std::sync::Arc;
+use std::task::Poll;
+
+use crate::joins::nested_loop_join::JoinLeftData;
+use crate::joins::utils::{
+ apply_join_filter_to_indices, build_batch_from_indices_maybe_empty,
+ need_produce_result_in_final, BuildProbeJoinMetrics, ColumnIndex, JoinFilter,
+ OnceFut,
+};
+use crate::metrics::Count;
+use crate::{RecordBatchStream, SendableRecordBatchStream};
+
+use arrow::array::{
+ BooleanArray, BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array,
+ UInt64Builder,
+};
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{
+ internal_datafusion_err, unwrap_or_internal_err, DataFusionError, JoinSide, Result,
+};
+use datafusion_expr::JoinType;
+
+use futures::Stream;
+
+/// 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,
+}
+
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
+ /// join filter
+ pub(crate) join_filter: Option<JoinFilter>,
+ /// type of the join
+ pub(crate) join_type: JoinType,
+ /// the outer table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Information of index and left / right placement of columns
+ pub(crate) column_indices: Vec<ColumnIndex>,
+ /// Join execution metrics
+ 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_index: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ emit_cursor: u64,
+ /// 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<BooleanBufferBuilder>,
+}
+
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
+
+ /// # Design
+ ///
+ /// The high-level control flow for this operator is:
+ /// 1. Buffer all batches from the left side (unless memory limit is reached,
+ /// in which case see notes at 'Memory-limited Execution').
+ /// - Rationale: The right side scanning can be expensive (it might
+ /// include decoding Parquet files), so it tries to buffer more left
+ /// batches at once to minimize the scan passes.
+ /// 2. Read right side batch one at a time. For each iteration, it only
+ /// evaluates the join filter on (1-left-row x right-batch), and puts the
Review Comment:
> The new implementation can also preserve left order in some cases
Only when `right_table.num_rows() < batch_size ` ?
For join types like inner join, we are no longer preserving the order of the right table. I'm **not** sure if this is a breaking change. It's possible that some users rely on this behavior, especially those who create an ExecutionPlan directly, bypassing the DataFusion optimizer. (Are there any users who do 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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 3.2X faster) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#issuecomment-3126813407
I tried to eliminate redundant batch transformations, the speedup goes from 2X -> 3X!
Through the flamegraph, I can see that over 50% of the time is now spent evaluating the join filter (e.g. `(c1 + c2) % 10 = 1`). It seems there may not be any more low-hanging fruits for further speed improvements.
After a few more chores, this should be ready for review:
- [ ] figure out the output ordering properties and update test
- [ ] Metrics accounting
- [ ] Polish(more comments and clean-up)
--
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 3.2X faster) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 closed pull request #16889: WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 3.2X faster)
URL: https://github.com/apache/datafusion/pull/16889
--
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "UBarney (via GitHub)" <gi...@apache.org>.
UBarney commented on PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#issuecomment-3124214924
> > I didn't expect this PR to have such a big performance improvement. Like #16443, I still don't understand why there is a performance improvement.
>
> I’m aware of two key differences:
>
> * Fewer redundant steps for `indices <--> batches`
> * Always keeping the right batch in cache (the original implementation performs a left-chunk × right-row iteration)
>
> However I was just hoping to cleanup the codebase a bit, I also didn’t expect it to be an easy 2X.
I can't use perf to analyze cache misses in my Hyper-V VM.
<details>
```
sudo perf stat -e cycles,instructions,cache-references,cache-misses ./target/release/1_left_row_join_right_batch -c ' SELECT * FROM range(10000) AS t1
JOIN range(200000) AS t2
ON (t1.value + t2.value) % 1000 = 0;'erf list cache -M
Performance counter stats for './target/release/1_left_row_join_right_batch -c SELECT * FROM range(10000) AS t1
JOIN range(200000) AS t2
ON (t1.value + t2.value) % 1000 = 0;':
<not supported> cycles
<not supported> instructions
<not supported> cache-references
<not supported> cache-misses
0.662652693 seconds time elapsed
11.216489000 seconds user
0.082434000 seconds sys
```
</details>
However, using the `time` command, I discovered that the previous version had a significantly higher number of `Minor (reclaiming a frame) page faults` (3,207,160 vs 5,133) and much greater system time (19.36s vs 0.06s).
<details>
```
/usr/bin/time -v ./target/release/join_limit_join_batch_size -c ' SELECT * FROM range(10000) AS t1
JOIN range(200000) AS t2
ON (t1.value + t2.value) % 1000 = 0;'
Command being timed: "./target/release/join_limit_join_batch_size -c SELECT * FROM range(10000) AS t1
JOIN range(200000) AS t2
ON (t1.value + t2.value) % 1000 = 0;"
User time (seconds): 27.68
System time (seconds): 19.36
Percent of CPU this job got: 2058%
Elapsed (wall clock) time (h:mm:ss or m:ss): 0:02.28
Average shared text size (kbytes): 0
Average unshared data size (kbytes): 0
Average stack size (kbytes): 0
Average total size (kbytes): 0
Maximum resident set size (kbytes): 23530920
Average resident set size (kbytes): 0
Major (requiring I/O) page faults: 0
Minor (reclaiming a frame) page faults: 3207160
Voluntary context switches: 629
Involuntary context switches: 1771
Swaps: 0
File system inputs: 0
File system outputs: 0
Socket messages sent: 0
Socket messages received: 0
Signals delivered: 0
Page size (bytes): 4096
Exit status: 0
/usr/bin/time -v ./target/release/1_left_row_join_right_batch -c ' SELECT * FROM range(10000) AS t1
JOIN range(200000) AS t2
ON (t1.value + t2.value) % 1000 = 0;'
Command being timed: "./target/release/1_left_row_join_right_batch -c SELECT * FROM range(10000) AS t1
JOIN range(200000) AS t2
ON (t1.value + t2.value) % 1000 = 0;"
User time (seconds): 11.50
System time (seconds): 0.06
Percent of CPU this job got: 1896%
Elapsed (wall clock) time (h:mm:ss or m:ss): 0:00.61
Average shared text size (kbytes): 0
Average unshared data size (kbytes): 0
Average stack size (kbytes): 0
Average total size (kbytes): 0
Maximum resident set size (kbytes): 135744
Average resident set size (kbytes): 0
Major (requiring I/O) page faults: 0
Minor (reclaiming a frame) page faults: 5133
Voluntary context switches: 461
Involuntary context switches: 574
Swaps: 0
File system inputs: 0
File system outputs: 0
Socket messages sent: 0
Socket messages received: 0
Signals delivered: 0
Page size (bytes): 4096
Exit status: 0
```
</details>
My speculation is that the previous version _**suffered from memory management overload**_ due to need alloc large memory, as `perf` also indicated that the single kernel function `clear_page_erms`( a kernel function that efficiently zeroes out a page of memory using a fast CPU instruction.) was the top CPU consumer.
```
sudo perf report --no-children
Samples: 215K of event 'cpu-clock:ppp', Event count (approx.): 53759750000
Overhead Command Shared Object Symbol
+ 25.15% tokio-runtime-w [kernel.kallsyms] [k] clear_page_erms
+ 6.90% tokio-runtime-w join_limit_join_batch_size [.] 0x0000000000df5b44
+ 6.31% swapper [kernel.kallsyms] [k] pv_native_safe_halt
+ 4.37% tokio-runtime-w join_limit_join_batch_size [.] 0x0000000002bd63e5
+ 4.32% tokio-runtime-w join_limit_join_batch_size [.] 0x0000000002bd63d5
+ 3.81% tokio-runtime-w [kernel.kallsyms] [k] _raw_spin_unlock_irqrestore
+ 3.10% tokio-runtime-w join_limit_join_batch_size [.] 0x0000000002bd61e7
+ 2.99% tokio-runtime-w join_limit_join_batch_size [.] 0x0000000000e088a4
+ 2.84% tokio-runtime-w join_limit_join_batch_size [.] 0x0000000002bd61bc
```
--
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#issuecomment-3121293562
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "UBarney (via GitHub)" <gi...@apache.org>.
UBarney commented on code in PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#discussion_r2232819590
##########
datafusion/physical-plan/src/joins/nlj.rs:
##########
@@ -0,0 +1,803 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Implementation of the Nested Loop Join operator.
+//!
+//! For detailed information regarding the operator's state machine and execution flow,
+//! please refer to the documentation provided in the `poll_next()` method.
+
+use arrow::buffer::MutableBuffer;
+use arrow::compute::BatchCoalescer;
+use futures::{ready, StreamExt};
+use log::debug;
+use std::sync::Arc;
+use std::task::Poll;
+
+use crate::joins::nested_loop_join::JoinLeftData;
+use crate::joins::utils::{
+ apply_join_filter_to_indices, build_batch_from_indices_maybe_empty,
+ need_produce_result_in_final, BuildProbeJoinMetrics, ColumnIndex, JoinFilter,
+ OnceFut,
+};
+use crate::metrics::Count;
+use crate::{RecordBatchStream, SendableRecordBatchStream};
+
+use arrow::array::{
+ BooleanArray, BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array,
+ UInt64Builder,
+};
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{
+ internal_datafusion_err, unwrap_or_internal_err, DataFusionError, JoinSide, Result,
+};
+use datafusion_expr::JoinType;
+
+use futures::Stream;
+
+/// 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,
+}
+
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
+ /// join filter
+ pub(crate) join_filter: Option<JoinFilter>,
+ /// type of the join
+ pub(crate) join_type: JoinType,
+ /// the outer table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Information of index and left / right placement of columns
+ pub(crate) column_indices: Vec<ColumnIndex>,
+ /// Join execution metrics
+ 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_index: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ emit_cursor: u64,
+ /// 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<BooleanBufferBuilder>,
+}
+
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
+
+ /// # Design
+ ///
+ /// The high-level control flow for this operator is:
+ /// 1. Buffer all batches from the left side (unless memory limit is reached,
+ /// in which case see notes at 'Memory-limited Execution').
+ /// - Rationale: The right side scanning can be expensive (it might
+ /// include decoding Parquet files), so it tries to buffer more left
+ /// batches at once to minimize the scan passes.
+ /// 2. Read right side batch one at a time. For each iteration, it only
+ /// evaluates the join filter on (1-left-row x right-batch), and puts the
Review Comment:
With this join method, it seems we can't preserve the order of the right table.
https://github.com/apache/datafusion/issues/16364#issuecomment-2975520489
--
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "UBarney (via GitHub)" <gi...@apache.org>.
UBarney commented on PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#issuecomment-3121216474
> > I'll need some time to read through the code, but reduced memory usage is also impressive 👍
> >
> > * v1 (old)
> >
> > ```
> > Query Time (ms) Peak RSS Peak Commit Page Faults
> > --------------------------------------------------------------
> > 1 657.40 1112.0 MB 10.6 GB 0
> > 2 833.66 1691.1 MB 10.6 GB 0
> > 3 1688.13 3.4 GB 10.8 GB 0
> > 4 4176.38 16.2 GB 20.1 GB 0
> > 5 1629.22 23.2 GB 57.2 GB 0
> > 6 17687.00 18.9 GB 26.1 GB 0
> > 7 1245.64 23.3 GB 39.2 GB 0
> > 8 19420.73 20.1 GB 31.3 GB 0
> > 9 3467.28 10.4 GB 19.9 GB 0
> > 10 5414.65 44.7 GB 46.0 GB 0
> > ```
> >
> >
> >
> >
> >
> >
> >
> >
> >
> >
> >
> >
> > * v2 (nlj-rewrite-tmp)
> >
> > ```
> > Query Time (ms) Peak RSS Peak Commit Page Faults
> > --------------------------------------------------------------
> > 1 383.61 104.4 MB 5.0 GB 0
> > 2 645.60 471.2 MB 8.0 GB 0
> > 3 763.10 1738.5 MB 6.0 GB 0
> > 4 1973.21 5.9 GB 9.6 GB 0
> > 5 644.86 145.1 MB 4.0 GB 0
> > 6 6043.58 149.3 MB 2.0 GB 0
> > 7 543.70 154.5 MB 7.0 GB 0
> > 8 6516.91 177.4 MB 4.0 GB 0
> > 9 1571.97 140.4 MB 6.0 GB 0
> > 10 3098.54 22.8 GB 26.9 GB 0
> > ```
> >
> >
> >
> >
> >
> >
> >
> >
> >
> >
> >
> > I ran `cargo run --profile release-nonlto --bin mem_profile nlj` in https://github.com/ding-young/datafusion/tree/nlj-rewrite-tmp-review and adjusted v1/v2 with `static USE_NLJ_V2: bool = false;`
>
> Thanks for the profiler 👍🏼 I'll review that soon.
>
> I think Q4 and Q10 should not consume that much memory 😅
>
> ```sql
> -- q10:
> SELECT *
> FROM range(30000) AS t1
> FULL JOIN range(30000) AS t2
> ON (t1.value + t2.value) % 10 <> 0;
> ```
>
> It is supposed to only buffer (30k x 8Bytes), 20G RSS is not expected. let me double check the implementation...
@ding-young For q10, the maxRSS I got using `/usr/bin/time -v` is lower than yours (16712272kb vs 44.7 GB).
The code for `./target/release/join_limit_join_batch_size` should be identical to the "v1 (old)" code.
```
√ devhomeinsp ~/c/datafusion > /usr/bin/time -v ./target/release/join_limit_join_batch_size --maxrows 1 -c ' SELECT *
FROM range(30000) AS t1
FULL JOIN range(30000) AS t2
ON (t1.value + t2.value) % 10 <> 0;'
DataFusion CLI v48.0.0
+-------+-------+
| value | value |
+-------+-------+
| 0 | 24576 |
| . |
| . |
| . |
+-------+-------+
810000000 row(s) fetched. (First 1 displayed. Use --maxrows to adjust)
Elapsed 4.202 seconds.
Command being timed: "./target/release/join_limit_join_batch_size --maxrows 1 -c SELECT *
FROM range(30000) AS t1
FULL JOIN range(30000) AS t2
ON (t1.value + t2.value) % 10 <> 0;"
User time (seconds): 9.26
System time (seconds): 4.74
Percent of CPU this job got: 322%
Elapsed (wall clock) time (h:mm:ss or m:ss): 0:04.34
Average shared text size (kbytes): 0
Average unshared data size (kbytes): 0
Average stack size (kbytes): 0
Average total size (kbytes): 0
Maximum resident set size (kbytes): 16712272
Average resident set size (kbytes): 0
Major (requiring I/O) page faults: 0
Minor (reclaiming a frame) page faults: 3028619
Voluntary context switches: 72488
Involuntary context switches: 28
Swaps: 0
File system inputs: 0
File system outputs: 0
Socket messages sent: 0
Socket messages received: 0
Signals delivered: 0
Page size (bytes): 4096
Exit status: 0
```
--
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#discussion_r2233663117
##########
datafusion/physical-plan/src/joins/nlj.rs:
##########
@@ -0,0 +1,803 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Implementation of the Nested Loop Join operator.
+//!
+//! For detailed information regarding the operator's state machine and execution flow,
+//! please refer to the documentation provided in the `poll_next()` method.
+
+use arrow::buffer::MutableBuffer;
+use arrow::compute::BatchCoalescer;
+use futures::{ready, StreamExt};
+use log::debug;
+use std::sync::Arc;
+use std::task::Poll;
+
+use crate::joins::nested_loop_join::JoinLeftData;
+use crate::joins::utils::{
+ apply_join_filter_to_indices, build_batch_from_indices_maybe_empty,
+ need_produce_result_in_final, BuildProbeJoinMetrics, ColumnIndex, JoinFilter,
+ OnceFut,
+};
+use crate::metrics::Count;
+use crate::{RecordBatchStream, SendableRecordBatchStream};
+
+use arrow::array::{
+ BooleanArray, BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array,
+ UInt64Builder,
+};
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{
+ internal_datafusion_err, unwrap_or_internal_err, DataFusionError, JoinSide, Result,
+};
+use datafusion_expr::JoinType;
+
+use futures::Stream;
+
+/// 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,
+}
+
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
+ /// join filter
+ pub(crate) join_filter: Option<JoinFilter>,
+ /// type of the join
+ pub(crate) join_type: JoinType,
+ /// the outer table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Information of index and left / right placement of columns
+ pub(crate) column_indices: Vec<ColumnIndex>,
+ /// Join execution metrics
+ 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_index: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ emit_cursor: u64,
+ /// 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<BooleanBufferBuilder>,
+}
+
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
+
+ /// # Design
+ ///
+ /// The high-level control flow for this operator is:
+ /// 1. Buffer all batches from the left side (unless memory limit is reached,
+ /// in which case see notes at 'Memory-limited Execution').
+ /// - Rationale: The right side scanning can be expensive (it might
+ /// include decoding Parquet files), so it tries to buffer more left
+ /// batches at once to minimize the scan passes.
+ /// 2. Read right side batch one at a time. For each iteration, it only
+ /// evaluates the join filter on (1-left-row x right-batch), and puts the
Review Comment:
I think it's okay to relax that property. The new implementation can also preserve left order in some cases.
--
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#issuecomment-3123897115
> I didn't expect this PR to have such a big performance improvement. Like #16443, I still don't understand why there is a performance improvement.
I’m aware of two key differences:
- Fewer redundant steps for `indices <--> batches`
- Always keeping the right batch in cache (the original implementation performs a left-chunk × right-row iteration)
However I was just hoping to cleanup the codebase a bit, I also didn’t expect it to be an easy 2X.
--
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 2X faster) [datafusion]
Posted by "UBarney (via GitHub)" <gi...@apache.org>.
UBarney commented on code in PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#discussion_r2232819590
##########
datafusion/physical-plan/src/joins/nlj.rs:
##########
@@ -0,0 +1,803 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Implementation of the Nested Loop Join operator.
+//!
+//! For detailed information regarding the operator's state machine and execution flow,
+//! please refer to the documentation provided in the `poll_next()` method.
+
+use arrow::buffer::MutableBuffer;
+use arrow::compute::BatchCoalescer;
+use futures::{ready, StreamExt};
+use log::debug;
+use std::sync::Arc;
+use std::task::Poll;
+
+use crate::joins::nested_loop_join::JoinLeftData;
+use crate::joins::utils::{
+ apply_join_filter_to_indices, build_batch_from_indices_maybe_empty,
+ need_produce_result_in_final, BuildProbeJoinMetrics, ColumnIndex, JoinFilter,
+ OnceFut,
+};
+use crate::metrics::Count;
+use crate::{RecordBatchStream, SendableRecordBatchStream};
+
+use arrow::array::{
+ BooleanArray, BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array,
+ UInt64Builder,
+};
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{
+ internal_datafusion_err, unwrap_or_internal_err, DataFusionError, JoinSide, Result,
+};
+use datafusion_expr::JoinType;
+
+use futures::Stream;
+
+/// 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,
+}
+
+pub(crate) struct NLJStream {
+ // ========================================================================
+ // PROPERTIES:
+ // Operator's properties that remain constant
+ // ========================================================================
+ /// Output schema
+ pub(crate) output_schema: Arc<Schema>,
+ /// join filter
+ pub(crate) join_filter: Option<JoinFilter>,
+ /// type of the join
+ pub(crate) join_type: JoinType,
+ /// the outer table data of the nested loop join
+ pub(crate) outer_table: SendableRecordBatchStream,
+ /// the inner table data of the nested loop join
+ pub(crate) inner_table: OnceFut<JoinLeftData>,
+ /// Information of index and left / right placement of columns
+ pub(crate) column_indices: Vec<ColumnIndex>,
+ /// Join execution metrics
+ 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_index: usize,
+ /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state
+ emit_cursor: u64,
+ /// 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<BooleanBufferBuilder>,
+}
+
+impl Stream for NLJStream {
+ type Item = Result<RecordBatch>;
+
+ /// # Design
+ ///
+ /// The high-level control flow for this operator is:
+ /// 1. Buffer all batches from the left side (unless memory limit is reached,
+ /// in which case see notes at 'Memory-limited Execution').
+ /// - Rationale: The right side scanning can be expensive (it might
+ /// include decoding Parquet files), so it tries to buffer more left
+ /// batches at once to minimize the scan passes.
+ /// 2. Read right side batch one at a time. For each iteration, it only
+ /// evaluates the join filter on (1-left-row x right-batch), and puts the
Review Comment:
With this join method, it seems we can't preserve the order of the right table.
https://github.com/apache/datafusion/issues/16364#issuecomment-2975520489
https://github.com/apache/datafusion/blob/6fef6aff7b0991b5fb003f959215e7a52dd69ed1/datafusion/physical-plan/src/joins/nested_loop_join.rs#L772-L787
--
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] WIP: Rewrite NestedLoopJoin to limit intermediate size (up to 3.2X faster) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16889:
URL: https://github.com/apache/datafusion/pull/16889#issuecomment-3140413077
Move to the final PR https://github.com/apache/datafusion/pull/16996, so closing this one.
--
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