2010YOUY01 commented on code in PR #24681: URL: https://github.com/apache/datafusion/pull/24681#discussion_r3888265615
########## datafusion/physical-plan/src/joins/nested_loop_join/semi_anti_mark_stream.rs: ########## @@ -0,0 +1,674 @@ +// 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. + +//! Nested loop join stream specifically for semi, anti, and mark joins +//! Instantiated by [`NestedLoopJoinExec`](crate::joins::nested_loop_join::NestedLoopJoinExec) +//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, `RightAnti`, +//! `LeftMark`, or `RightMark`. +//! +//! # Algorithm +//! +//! For each buffered left chunk: +//! ```text +//! for right_batch in right_side: +//! for left_row in left_chunk: // bitmap-only probing +//! update left/right match bitmaps +//! emit/accumulate right-side SAM result +//! report_probe_completed() once +//! emit left-side SAM result (emitter partition only) +//! ``` +//! +//! In memory-limited mode, left chunks are loaded incrementally and right-side +//! match bitmaps are accumulated globally; after all chunks, the right input +//! is replayed from spill for final right-side emission. + +use std::future::poll_fn; +use std::sync::Arc; + +use super::shared::{ + JoinLeftData, LeftBufferBatchDecision, NestedLoopJoinMetrics, SpillState, + apply_filter_to_row_join_batch, buffer_left_batch_in_chunk, + build_global_right_result_batch, build_unmatched_batch, finalize_buffered_left_chunk, + initiate_spill_fallback, probe_sam_left_range, update_sam_matched_bitmaps, +}; +use crate::SendableRecordBatchStream; +use crate::joins::utils::{ColumnIndex, JoinFilter, OnceFut}; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; + +use arrow::array::BooleanArray; +use arrow::buffer::BooleanBuffer; +use arrow::compute::BatchCoalescer; +use arrow::datatypes::{Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::instant::Instant; +use datafusion_common::{DataFusionError, JoinSide, Result, internal_datafusion_err}; +use datafusion_execution::{TryEmitter, async_try_stream}; +use datafusion_expr::JoinType; +use futures::StreamExt; +use log::debug; + +/// Nested loop join stream for Semi/Anti/Mark joins. +/// +/// Evaluates the join predicate for every relevant left/right combination but +/// does not emit `(left, right)` pairs. Instead it accumulates a Boolean value +/// for each row on the output side to check for any match. +pub(super) struct SemiAntiMarkNestedLoopJoinStream { + /// Output schema after applying the join projection. + output_schema: Arc<Schema>, + /// Optional non-equality join predicate. + join_filter: Option<JoinFilter>, + /// Semi, anti, or mark join type handled by this stream. + join_type: JoinType, + /// Side whose rows are produced by the join. + join_side: JoinSide, + /// Current probe-side input. Replaced by each replay pass after spilling. + right_data: Option<SendableRecordBatchStream>, + /// Shared future that collects the build side for the standard path. + left_data: OnceFut<JoinLeftData>, + /// Projection used to construct output columns from the input sides. + column_indices: Vec<ColumnIndex>, + /// Join, spill, and selectivity metrics. + metrics: NestedLoopJoinMetrics, + /// Target output batch size and probe range size. + batch_size: usize, + /// Coalesces result batches before yielding them to the consumer. + output_buffer: Box<BatchCoalescer>, + /// Disabled, pending, or active memory-limited spill execution. + spill_state: SpillState, + /// Start of the current join-time interval; `None` while paused. + join_time_start: Option<Instant>, + /// Number of right-side passes opened for buffered left chunks. + right_pass_count: usize, + /// Whether this stream has emitted at least one output row. + emitted_rows: bool, +} + +impl SemiAntiMarkNestedLoopJoinStream { + /// Create the SAM stream and wrap its generator with baseline observation. + #[expect(clippy::too_many_arguments)] + pub(super) fn try_new( + schema: SchemaRef, + filter: Option<JoinFilter>, + join_type: JoinType, + right_data: SendableRecordBatchStream, + left_data: OnceFut<JoinLeftData>, + column_indices: Vec<ColumnIndex>, + metrics: NestedLoopJoinMetrics, + batch_size: usize, + spill_state: SpillState, + ) -> Result<SendableRecordBatchStream> { + debug_assert!( + matches!( + join_type, + JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::LeftMark + | JoinType::RightMark + ), + "SemiAntiMarkNestedLoopJoinStream does not handle {join_type:?}" + ); + + let join_side = match join_type { + JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { + JoinSide::Left + } + _ => JoinSide::Right, + }; + + let baseline_metrics = metrics.join_metrics.baseline.clone(); + + let mut state = Self { + output_schema: Arc::clone(&schema), + join_filter: filter, + join_type, + join_side, + right_data: Some(right_data), + column_indices, + left_data, + metrics, + output_buffer: Box::new(BatchCoalescer::new(Arc::clone(&schema), batch_size)), + batch_size, + spill_state, + join_time_start: None, + right_pass_count: 0, + emitted_rows: false, + }; + + let stream = async_try_stream(|mut emitter| async move { + state.start_join_time(); + let result = state.join(&mut emitter).await; Review Comment: Is it possible to use the explicit states here like https://github.com/apache/datafusion/blob/61bf6b96cc07d6a0518014dec7fe86fbef76576e/datafusion/physical-plan/src/aggregates/hash_stream.rs#L175-L202 I think the rule to follow is: if we have to draw a state transition graph to understand a operator, then implement it with the explicit state representation. - For CrossJoin, its logic is simple, so explicit states are not necessary - For NLJ, we do have to draw state machine graph to understand it, so the `enum` state can make it easier to understand and maintain. The emitter pattern doesn't seem to be mutual exclusive with explicit state representation, and its purpose is make the control flow implementation less verbose. -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
