Skip to main content

datafusion_physical_plan/joins/
symmetric_hash_join.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! This file implements the symmetric hash join algorithm with range-based
19//! data pruning to join two (potentially infinite) streams.
20//!
21//! A [`SymmetricHashJoinExec`] plan takes two children plan (with appropriate
22//! output ordering) and produces the join output according to the given join
23//! type and other options.
24//!
25//! This plan uses the [`OneSideHashJoiner`] object to facilitate join calculations
26//! for both its children.
27
28use std::fmt::{self, Debug};
29use std::mem::{size_of, size_of_val};
30use std::sync::Arc;
31use std::task::{Context, Poll};
32use std::vec;
33
34use crate::common::SharedMemoryReservation;
35use crate::execution_plan::{boundedness_from_children, emission_type_from_children};
36use crate::joins::stream_join_utils::{
37    PruningJoinHashMap, SortedFilterExpr, StreamJoinMetrics,
38    calculate_filter_expr_intervals, combine_two_batches,
39    convert_sort_expr_with_filter_schema, get_pruning_anti_indices,
40    get_pruning_semi_indices, prepare_sorted_exprs, record_visited_indices,
41};
42use crate::joins::utils::{
43    BatchSplitter, BatchTransformer, ColumnIndex, JoinFilter, JoinHashMapType, JoinOn,
44    JoinOnRef, NoopBatchTransformer, StatefulStreamResult, apply_join_filter_to_indices,
45    build_batch_from_indices, build_join_schema, check_join_is_valid, equal_rows_arr,
46    matchable_join_keys, symmetric_join_output_partitioning, update_hash,
47};
48use crate::projection::{
49    JoinData, ProjectionExec, try_pushdown_through_join_with_column_indices,
50};
51use crate::stream::EmptyRecordBatchStream;
52use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count};
53use crate::{
54    DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
55    InputDistributionRequirements, PlanProperties, RecordBatchStream,
56    SendableRecordBatchStream,
57    joins::StreamJoinPartitionMode,
58    metrics::{ExecutionPlanMetricsSet, MetricsSet},
59};
60
61use arrow::array::{
62    ArrowPrimitiveType, NativeAdapter, PrimitiveArray, PrimitiveBuilder, UInt32Array,
63    UInt64Array,
64};
65use arrow::compute::concat_batches;
66use arrow::datatypes::{ArrowNativeType, Schema, SchemaRef};
67use arrow::record_batch::RecordBatch;
68use datafusion_common::hash_utils::create_hashes;
69use datafusion_common::tree_node::TreeNodeRecursion;
70use datafusion_common::utils::bisect;
71use datafusion_common::{
72    HashSet, JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err,
73    plan_err,
74};
75use datafusion_execution::TaskContext;
76use datafusion_execution::memory_pool::MemoryConsumer;
77use datafusion_expr::interval_arithmetic::Interval;
78use datafusion_physical_expr::equivalence::join_equivalence_properties;
79use datafusion_physical_expr::intervals::cp_solver::ExprIntervalGraph;
80use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql};
81use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements};
82
83use datafusion_common::hash_utils::RandomState;
84use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays;
85use futures::{Stream, StreamExt, ready};
86
87const HASHMAP_SHRINK_SCALE_FACTOR: usize = 4;
88
89/// A symmetric hash join with range conditions is when both streams are hashed on the
90/// join key and the resulting hash tables are used to join the streams.
91/// The join is considered symmetric because the hash table is built on the join keys from both
92/// streams, and the matching of rows is based on the values of the join keys in both streams.
93/// This type of join is efficient in streaming context as it allows for fast lookups in the hash
94/// table, rather than having to scan through one or both of the streams to find matching rows, also it
95/// only considers the elements from the stream that fall within a certain sliding window (w/ range conditions),
96/// making it more efficient and less likely to store stale data. This enables operating on unbounded streaming
97/// data without any memory issues.
98///
99/// For each input stream, create a hash table.
100///   - For each new [RecordBatch] in build side, hash and insert into inputs hash table. Update offsets.
101///   - Test if input is equal to a predefined set of other inputs.
102///   - If so record the visited rows. If the matched row results must be produced (INNER, LEFT), output the [RecordBatch].
103///   - Try to prune other side (probe) with new [RecordBatch].
104///   - If the join type indicates that the unmatched rows results must be produced (LEFT, FULL etc.),
105///     output the [RecordBatch] when a pruning happens or at the end of the data.
106///
107///
108/// ``` text
109///                        +-------------------------+
110///                        |                         |
111///   left stream ---------|  Left OneSideHashJoiner |---+
112///                        |                         |   |
113///                        +-------------------------+   |
114///                                                      |
115///                                                      |--------- Joined output
116///                                                      |
117///                        +-------------------------+   |
118///                        |                         |   |
119///  right stream ---------| Right OneSideHashJoiner |---+
120///                        |                         |
121///                        +-------------------------+
122///
123/// Prune build side when the new RecordBatch comes to the probe side. We utilize interval arithmetic
124/// on JoinFilter's sorted PhysicalExprs to calculate the joinable range.
125///
126///
127///               PROBE SIDE          BUILD SIDE
128///                 BUFFER              BUFFER
129///             +-------------+     +------------+
130///             |             |     |            |    Unjoinable
131///             |             |     |            |    Range
132///             |             |     |            |
133///             |             |  |---------------------------------
134///             |             |  |  |            |
135///             |             |  |  |            |
136///             |             | /   |            |
137///             |             | |   |            |
138///             |             | |   |            |
139///             |             | |   |            |
140///             |             | |   |            |
141///             |             | |   |            |    Joinable
142///             |             |/    |            |    Range
143///             |             ||    |            |
144///             |+-----------+||    |            |
145///             || Record    ||     |            |
146///             || Batch     ||     |            |
147///             |+-----------+||    |            |
148///             +-------------+\    +------------+
149///                             |
150///                             \
151///                              |---------------------------------
152///
153///  This happens when range conditions are provided on sorted columns. E.g.
154///
155///        SELECT * FROM left_table, right_table
156///        ON
157///          left_key = right_key AND
158///          left_time > right_time - INTERVAL 12 MINUTES AND left_time < right_time + INTERVAL 2 HOUR
159///
160/// or
161///       SELECT * FROM left_table, right_table
162///        ON
163///          left_key = right_key AND
164///          left_sorted > right_sorted - 3 AND left_sorted < right_sorted + 10
165///
166/// For general purpose, in the second scenario, when the new data comes to probe side, the conditions can be used to
167/// determine a specific threshold for discarding rows from the inner buffer. For example, if the sort order the
168/// two columns ("left_sorted" and "right_sorted") are ascending (it can be different in another scenarios)
169/// and the join condition is "left_sorted > right_sorted - 3" and the latest value on the right input is 1234, meaning
170/// that the left side buffer must only keep rows where "leftTime > rightTime - 3 > 1234 - 3 > 1231" ,
171/// making the smallest value in 'left_sorted' 1231 and any rows below (since ascending)
172/// than that can be dropped from the inner buffer.
173/// ```
174#[derive(Debug, Clone)]
175pub struct SymmetricHashJoinExec {
176    /// Left side stream
177    pub(crate) left: Arc<dyn ExecutionPlan>,
178    /// Right side stream
179    pub(crate) right: Arc<dyn ExecutionPlan>,
180    /// Set of common columns used to join on
181    pub(crate) on: Vec<(PhysicalExprRef, PhysicalExprRef)>,
182    /// Filters applied when finding matching rows
183    pub(crate) filter: Option<JoinFilter>,
184    /// How the join is performed
185    pub(crate) join_type: JoinType,
186    /// Shares the `RandomState` for the hashing algorithm
187    random_state: RandomState,
188    /// Execution metrics
189    metrics: ExecutionPlanMetricsSet,
190    /// Information of index and left / right placement of columns
191    column_indices: Vec<ColumnIndex>,
192    /// Defines the null equality for the join.
193    pub(crate) null_equality: NullEquality,
194    /// Left side sort expression(s)
195    pub(crate) left_sort_exprs: Option<LexOrdering>,
196    /// Right side sort expression(s)
197    pub(crate) right_sort_exprs: Option<LexOrdering>,
198    /// Partition Mode
199    mode: StreamJoinPartitionMode,
200    /// Cache holding plan properties like equivalences, output partitioning etc.
201    cache: Arc<PlanProperties>,
202}
203
204impl SymmetricHashJoinExec {
205    /// Tries to create a new [SymmetricHashJoinExec].
206    /// # Error
207    /// This function errors when:
208    /// - It is not possible to join the left and right sides on keys `on`, or
209    /// - It fails to construct `SortedFilterExpr`s, or
210    /// - It fails to create the [ExprIntervalGraph].
211    #[expect(clippy::too_many_arguments)]
212    pub fn try_new(
213        left: Arc<dyn ExecutionPlan>,
214        right: Arc<dyn ExecutionPlan>,
215        on: JoinOn,
216        filter: Option<JoinFilter>,
217        join_type: &JoinType,
218        null_equality: NullEquality,
219        left_sort_exprs: Option<LexOrdering>,
220        right_sort_exprs: Option<LexOrdering>,
221        mode: StreamJoinPartitionMode,
222    ) -> Result<Self> {
223        let left_schema = left.schema();
224        let right_schema = right.schema();
225
226        // Error out if no "on" constraints are given:
227        if on.is_empty() {
228            return plan_err!(
229                "On constraints in SymmetricHashJoinExec should be non-empty"
230            );
231        }
232
233        // Check if the join is valid with the given on constraints:
234        check_join_is_valid(&left_schema, &right_schema, &on)?;
235
236        // Build the join schema from the left and right schemas:
237        let (schema, column_indices) =
238            build_join_schema(&left_schema, &right_schema, join_type);
239
240        // Initialize the random state for the join operation:
241        let random_state = RandomState::with_seed(0);
242        let schema = Arc::new(schema);
243        let cache = Self::compute_properties(&left, &right, schema, *join_type, &on)?;
244        Ok(SymmetricHashJoinExec {
245            left,
246            right,
247            on,
248            filter,
249            join_type: *join_type,
250            random_state,
251            metrics: ExecutionPlanMetricsSet::new(),
252            column_indices,
253            null_equality,
254            left_sort_exprs,
255            right_sort_exprs,
256            mode,
257            cache: Arc::new(cache),
258        })
259    }
260
261    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
262    fn compute_properties(
263        left: &Arc<dyn ExecutionPlan>,
264        right: &Arc<dyn ExecutionPlan>,
265        schema: SchemaRef,
266        join_type: JoinType,
267        join_on: JoinOnRef,
268    ) -> Result<PlanProperties> {
269        // Calculate equivalence properties:
270        let eq_properties = join_equivalence_properties(
271            left.equivalence_properties().clone(),
272            right.equivalence_properties().clone(),
273            &join_type,
274            schema,
275            &[false, false],
276            // Has alternating probe side
277            None,
278            join_on,
279        )?;
280
281        let output_partitioning =
282            symmetric_join_output_partitioning(left, right, &join_type)?;
283
284        Ok(PlanProperties::new(
285            eq_properties,
286            output_partitioning,
287            emission_type_from_children([left, right]),
288            boundedness_from_children([left, right]),
289        ))
290    }
291
292    /// left stream
293    pub fn left(&self) -> &Arc<dyn ExecutionPlan> {
294        &self.left
295    }
296
297    /// right stream
298    pub fn right(&self) -> &Arc<dyn ExecutionPlan> {
299        &self.right
300    }
301
302    /// Set of common columns used to join on
303    pub fn on(&self) -> &[(PhysicalExprRef, PhysicalExprRef)] {
304        &self.on
305    }
306
307    /// Filters applied before join output
308    pub fn filter(&self) -> Option<&JoinFilter> {
309        self.filter.as_ref()
310    }
311
312    /// How the join is performed
313    pub fn join_type(&self) -> &JoinType {
314        &self.join_type
315    }
316
317    /// Get null_equality
318    pub fn null_equality(&self) -> NullEquality {
319        self.null_equality
320    }
321
322    /// Get partition mode
323    pub fn partition_mode(&self) -> StreamJoinPartitionMode {
324        self.mode
325    }
326
327    /// Get left_sort_exprs
328    pub fn left_sort_exprs(&self) -> Option<&LexOrdering> {
329        self.left_sort_exprs.as_ref()
330    }
331
332    /// Get right_sort_exprs
333    pub fn right_sort_exprs(&self) -> Option<&LexOrdering> {
334        self.right_sort_exprs.as_ref()
335    }
336
337    /// Check if order information covers every column in the filter expression.
338    pub fn check_if_order_information_available(&self) -> Result<bool> {
339        if let Some(filter) = self.filter() {
340            let left = self.left();
341            if let Some(left_ordering) = left.output_ordering() {
342                let right = self.right();
343                if let Some(right_ordering) = right.output_ordering() {
344                    let left_convertible = convert_sort_expr_with_filter_schema(
345                        &JoinSide::Left,
346                        filter,
347                        &left.schema(),
348                        &left_ordering[0],
349                    )?
350                    .is_some();
351                    let right_convertible = convert_sort_expr_with_filter_schema(
352                        &JoinSide::Right,
353                        filter,
354                        &right.schema(),
355                        &right_ordering[0],
356                    )?
357                    .is_some();
358                    return Ok(left_convertible && right_convertible);
359                }
360            }
361        }
362        Ok(false)
363    }
364}
365
366impl DisplayAs for SymmetricHashJoinExec {
367    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
368        match t {
369            DisplayFormatType::Default | DisplayFormatType::Verbose => {
370                let display_filter = self.filter.as_ref().map_or_else(
371                    || "".to_string(),
372                    |f| format!(", filter={}", f.expression()),
373                );
374                let on = self
375                    .on
376                    .iter()
377                    .map(|(c1, c2)| format!("({c1}, {c2})"))
378                    .collect::<Vec<String>>()
379                    .join(", ");
380                write!(
381                    f,
382                    "SymmetricHashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}",
383                    self.mode, self.join_type, on, display_filter
384                )
385            }
386            DisplayFormatType::TreeRender => {
387                let on = self
388                    .on
389                    .iter()
390                    .map(|(c1, c2)| {
391                        format!("({} = {})", fmt_sql(c1.as_ref()), fmt_sql(c2.as_ref()))
392                    })
393                    .collect::<Vec<String>>()
394                    .join(", ");
395
396                writeln!(f, "mode={:?}", self.mode)?;
397                if *self.join_type() != JoinType::Inner {
398                    writeln!(f, "join_type={:?}", self.join_type)?;
399                }
400                writeln!(f, "on={on}")
401            }
402        }
403    }
404}
405
406impl ExecutionPlan for SymmetricHashJoinExec {
407    fn name(&self) -> &'static str {
408        "SymmetricHashJoinExec"
409    }
410
411    fn properties(&self) -> &Arc<PlanProperties> {
412        &self.cache
413    }
414
415    fn required_input_distribution(&self) -> Vec<Distribution> {
416        self.input_distribution_requirements().into_per_child()
417    }
418
419    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
420        match self.mode {
421            StreamJoinPartitionMode::Partitioned => {
422                let (left_expr, right_expr) = self
423                    .on
424                    .iter()
425                    .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _))
426                    .unzip();
427                InputDistributionRequirements::co_partitioned(vec![
428                    Distribution::KeyPartitioned(left_expr),
429                    Distribution::KeyPartitioned(right_expr),
430                ])
431            }
432            StreamJoinPartitionMode::SinglePartition => {
433                InputDistributionRequirements::new(vec![
434                    Distribution::SinglePartition,
435                    Distribution::SinglePartition,
436                ])
437            }
438        }
439    }
440
441    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
442        vec![
443            self.left_sort_exprs
444                .as_ref()
445                .map(|e| OrderingRequirements::from(e.clone())),
446            self.right_sort_exprs
447                .as_ref()
448                .map(|e| OrderingRequirements::from(e.clone())),
449        ]
450    }
451
452    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
453        vec![&self.left, &self.right]
454    }
455
456    fn apply_expressions(
457        &self,
458        f: &mut dyn FnMut(&Arc<dyn crate::PhysicalExpr>) -> Result<TreeNodeRecursion>,
459    ) -> Result<TreeNodeRecursion> {
460        let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]);
461        let filter = self.filter.iter().map(|filter| filter.expression());
462        crate::apply_expression_roots(join_keys.chain(filter), f)
463    }
464
465    fn replace_children(
466        self: Arc<Self>,
467        mut children: Vec<Arc<dyn ExecutionPlan>>,
468        options: ReplaceChildrenOptions,
469    ) -> Result<Arc<dyn ExecutionPlan>> {
470        validate_child_count!(self, children);
471        match options.children_properties {
472            ChildrenPropertiesMode::Keep => {
473                let left = children.swap_remove(0);
474                let right = children.swap_remove(0);
475                Ok(Arc::new(Self {
476                    left,
477                    right,
478                    metrics: ExecutionPlanMetricsSet::new(),
479                    ..Self::clone(&*self)
480                }))
481            }
482            ChildrenPropertiesMode::Recompute => {
483                Ok(Arc::new(SymmetricHashJoinExec::try_new(
484                    Arc::clone(&children[0]),
485                    Arc::clone(&children[1]),
486                    self.on.clone(),
487                    self.filter.clone(),
488                    &self.join_type,
489                    self.null_equality,
490                    self.left_sort_exprs.clone(),
491                    self.right_sort_exprs.clone(),
492                    self.mode,
493                )?))
494            }
495        }
496    }
497
498    fn with_new_children(
499        self: Arc<Self>,
500        children: Vec<Arc<dyn ExecutionPlan>>,
501    ) -> Result<Arc<dyn ExecutionPlan>> {
502        self.replace_children(
503            children,
504            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
505        )
506    }
507
508    fn with_new_children_and_same_properties(
509        self: Arc<Self>,
510        children: Vec<Arc<dyn ExecutionPlan>>,
511    ) -> Result<Arc<dyn ExecutionPlan>> {
512        self.replace_children(
513            children,
514            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
515        )
516    }
517
518    fn metrics(&self) -> Option<MetricsSet> {
519        Some(self.metrics.clone_inner())
520    }
521
522    fn execute(
523        &self,
524        partition: usize,
525        context: Arc<TaskContext>,
526    ) -> Result<SendableRecordBatchStream> {
527        let left_partitions = self.left.output_partitioning().partition_count();
528        let right_partitions = self.right.output_partitioning().partition_count();
529        assert_eq_or_internal_err!(
530            left_partitions,
531            right_partitions,
532            "Invalid SymmetricHashJoinExec, partition count mismatch {left_partitions}!={right_partitions},\
533                 consider using RepartitionExec"
534        );
535        // If `filter_state` and `filter` are both present, then calculate sorted
536        // filter expressions for both sides, and build an expression graph.
537        let (left_sorted_filter_expr, right_sorted_filter_expr, graph) = match (
538            self.left_sort_exprs(),
539            self.right_sort_exprs(),
540            &self.filter,
541        ) {
542            (Some(left_sort_exprs), Some(right_sort_exprs), Some(filter)) => {
543                let (left, right, graph) = prepare_sorted_exprs(
544                    filter,
545                    &self.left,
546                    &self.right,
547                    left_sort_exprs,
548                    right_sort_exprs,
549                )?;
550                (Some(left), Some(right), Some(graph))
551            }
552            // If `filter_state` or `filter` is not present, then return None
553            // for all three values:
554            _ => (None, None, None),
555        };
556
557        let (on_left, on_right) = self.on.iter().cloned().unzip();
558
559        let left_side_joiner =
560            OneSideHashJoiner::new(JoinSide::Left, on_left, self.left.schema());
561        let right_side_joiner =
562            OneSideHashJoiner::new(JoinSide::Right, on_right, self.right.schema());
563
564        let left_stream = self.left.execute(partition, Arc::clone(&context))?;
565
566        let right_stream = self.right.execute(partition, Arc::clone(&context))?;
567
568        let batch_size = context.session_config().batch_size();
569        let enforce_batch_size_in_joins =
570            context.session_config().enforce_batch_size_in_joins();
571
572        let reservation = Arc::new(
573            MemoryConsumer::new(format!("SymmetricHashJoinStream[{partition}]"))
574                .register(context.memory_pool()),
575        );
576        if let Some(g) = graph.as_ref() {
577            reservation.try_grow(g.size())?;
578        }
579
580        if enforce_batch_size_in_joins {
581            Ok(Box::pin(SymmetricHashJoinStream {
582                left_stream,
583                right_stream,
584                schema: self.schema(),
585                filter: self.filter.clone(),
586                join_type: self.join_type,
587                random_state: self.random_state.clone(),
588                left: left_side_joiner,
589                right: right_side_joiner,
590                column_indices: self.column_indices.clone(),
591                metrics: StreamJoinMetrics::new(partition, &self.metrics),
592                graph,
593                left_sorted_filter_expr,
594                right_sorted_filter_expr,
595                null_equality: self.null_equality,
596                state: SHJStreamState::PullRight,
597                reservation,
598                batch_transformer: BatchSplitter::new(batch_size),
599            }))
600        } else {
601            Ok(Box::pin(SymmetricHashJoinStream {
602                left_stream,
603                right_stream,
604                schema: self.schema(),
605                filter: self.filter.clone(),
606                join_type: self.join_type,
607                random_state: self.random_state.clone(),
608                left: left_side_joiner,
609                right: right_side_joiner,
610                column_indices: self.column_indices.clone(),
611                metrics: StreamJoinMetrics::new(partition, &self.metrics),
612                graph,
613                left_sorted_filter_expr,
614                right_sorted_filter_expr,
615                null_equality: self.null_equality,
616                state: SHJStreamState::PullRight,
617                reservation,
618                batch_transformer: NoopBatchTransformer::new(),
619            }))
620        }
621    }
622
623    /// Tries to swap the projection with its input [`SymmetricHashJoinExec`]. If it can be done,
624    /// it returns the new swapped version having the [`SymmetricHashJoinExec`] as the top plan.
625    /// Otherwise, it returns None.
626    fn try_swapping_with_projection(
627        &self,
628        projection: &ProjectionExec,
629    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
630        let schema = self.schema();
631        if let Some(JoinData {
632            projected_left_child,
633            projected_right_child,
634            join_filter,
635            join_on,
636        }) = try_pushdown_through_join_with_column_indices(
637            projection,
638            self.left(),
639            self.right(),
640            self.on(),
641            &schema,
642            self.filter(),
643            self.column_indices.as_slice(),
644        )? {
645            SymmetricHashJoinExec::try_new(
646                Arc::new(projected_left_child),
647                Arc::new(projected_right_child),
648                join_on,
649                join_filter,
650                self.join_type(),
651                self.null_equality(),
652                self.right().output_ordering().cloned(),
653                self.left().output_ordering().cloned(),
654                self.partition_mode(),
655            )
656            .map(|e| Some(Arc::new(e) as _))
657        } else {
658            Ok(None)
659        }
660    }
661
662    #[cfg(feature = "proto")]
663    fn try_to_proto(
664        &self,
665        ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
666    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
667        use datafusion_proto_models::protobuf;
668
669        let left = ctx.encode_child(self.left())?;
670        let right = ctx.encode_child(self.right())?;
671        let on = self
672            .on()
673            .iter()
674            .map(|(left, right)| {
675                Ok(protobuf::JoinOn {
676                    left: Some(ctx.encode_expr(left)?),
677                    right: Some(ctx.encode_expr(right)?),
678                })
679            })
680            .collect::<Result<Vec<_>>>()?;
681
682        let join_type = match self.join_type() {
683            JoinType::Inner => protobuf::JoinType::Inner,
684            JoinType::Left => protobuf::JoinType::Left,
685            JoinType::Right => protobuf::JoinType::Right,
686            JoinType::Full => protobuf::JoinType::Full,
687            JoinType::LeftSemi => protobuf::JoinType::Leftsemi,
688            JoinType::RightSemi => protobuf::JoinType::Rightsemi,
689            JoinType::LeftAnti => protobuf::JoinType::Leftanti,
690            JoinType::RightAnti => protobuf::JoinType::Rightanti,
691            JoinType::LeftMark => protobuf::JoinType::Leftmark,
692            JoinType::RightMark => protobuf::JoinType::Rightmark,
693        };
694        let null_equality = match self.null_equality() {
695            NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing,
696            NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull,
697        };
698        let partition_mode = match self.partition_mode() {
699            StreamJoinPartitionMode::SinglePartition => {
700                protobuf::StreamPartitionMode::SinglePartition
701            }
702            StreamJoinPartitionMode::Partitioned => {
703                protobuf::StreamPartitionMode::PartitionedExec
704            }
705        };
706        let filter = self
707            .filter()
708            .map(|filter| -> Result<protobuf::JoinFilter> {
709                let expression = ctx.encode_expr(filter.expression())?;
710                let column_indices = filter
711                    .column_indices()
712                    .iter()
713                    .map(|column_index| {
714                        let side = match column_index.side {
715                            JoinSide::Left => protobuf::JoinSide::LeftSide,
716                            JoinSide::Right => protobuf::JoinSide::RightSide,
717                            JoinSide::None => protobuf::JoinSide::None,
718                        };
719                        protobuf::ColumnIndex {
720                            index: column_index.index as u32,
721                            side: side.into(),
722                        }
723                    })
724                    .collect();
725                Ok(protobuf::JoinFilter {
726                    expression: Some(expression),
727                    column_indices,
728                    schema: Some(filter.schema().as_ref().try_into()?),
729                })
730            })
731            .transpose()?;
732        let expr_ctx = ctx.expr_ctx();
733        let left_sort_exprs =
734            datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto(
735                self.left_sort_exprs(),
736                &expr_ctx,
737            )?;
738        let right_sort_exprs =
739            datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto(
740                self.right_sort_exprs(),
741                &expr_ctx,
742            )?;
743
744        Ok(Some(protobuf::PhysicalPlanNode {
745            physical_plan_type: Some(
746                protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin(
747                    Box::new(protobuf::SymmetricHashJoinExecNode {
748                        left: Some(Box::new(left)),
749                        right: Some(Box::new(right)),
750                        on,
751                        join_type: join_type.into(),
752                        partition_mode: partition_mode.into(),
753                        null_equality: null_equality.into(),
754                        filter,
755                        left_sort_exprs,
756                        right_sort_exprs,
757                    }),
758                ),
759            ),
760        }))
761    }
762}
763
764#[cfg(feature = "proto")]
765impl SymmetricHashJoinExec {
766    /// Reconstruct a [`SymmetricHashJoinExec`] from its protobuf representation.
767    ///
768    /// The exact inverse of [`ExecutionPlan::try_to_proto`].
769    ///
770    /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto
771    pub fn try_from_proto(
772        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
773        ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
774    ) -> Result<Arc<dyn ExecutionPlan>> {
775        use datafusion_common::internal_datafusion_err;
776        use datafusion_proto_models::protobuf;
777
778        let sym_join = crate::expect_plan_variant!(
779            node,
780            protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin,
781            "SymmetricHashJoinExec",
782        );
783        let left = ctx.decode_required_child(
784            sym_join.left.as_deref(),
785            "SymmetricHashJoinExec",
786            "left",
787        )?;
788        let right = ctx.decode_required_child(
789            sym_join.right.as_deref(),
790            "SymmetricHashJoinExec",
791            "right",
792        )?;
793        let left_schema = left.schema();
794        let right_schema = right.schema();
795        let on = sym_join
796            .on
797            .iter()
798            .map(|columns| {
799                let left = ctx.decode_required_expr(
800                    columns.left.as_ref(),
801                    left_schema.as_ref(),
802                    "SymmetricHashJoinExec",
803                    "on.left",
804                )?;
805                let right = ctx.decode_required_expr(
806                    columns.right.as_ref(),
807                    right_schema.as_ref(),
808                    "SymmetricHashJoinExec",
809                    "on.right",
810                )?;
811                Ok((left, right))
812            })
813            .collect::<Result<JoinOn>>()?;
814
815        let join_type =
816            match protobuf::JoinType::try_from(sym_join.join_type).map_err(|_| {
817                internal_datafusion_err!(
818                    "SymmetricHashJoinExec: unknown JoinType {}",
819                    sym_join.join_type
820                )
821            })? {
822                protobuf::JoinType::Inner => JoinType::Inner,
823                protobuf::JoinType::Left => JoinType::Left,
824                protobuf::JoinType::Right => JoinType::Right,
825                protobuf::JoinType::Full => JoinType::Full,
826                protobuf::JoinType::Leftsemi => JoinType::LeftSemi,
827                protobuf::JoinType::Rightsemi => JoinType::RightSemi,
828                protobuf::JoinType::Leftanti => JoinType::LeftAnti,
829                protobuf::JoinType::Rightanti => JoinType::RightAnti,
830                protobuf::JoinType::Leftmark => JoinType::LeftMark,
831                protobuf::JoinType::Rightmark => JoinType::RightMark,
832            };
833        let null_equality = match protobuf::NullEquality::try_from(sym_join.null_equality)
834            .map_err(|_| {
835                internal_datafusion_err!(
836                    "SymmetricHashJoinExec: unknown NullEquality {}",
837                    sym_join.null_equality
838                )
839            })? {
840            protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing,
841            protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull,
842        };
843        let partition_mode =
844            match protobuf::StreamPartitionMode::try_from(sym_join.partition_mode)
845                .map_err(|_| {
846                    internal_datafusion_err!(
847                        "SymmetricHashJoinExec: unknown StreamPartitionMode {}",
848                        sym_join.partition_mode
849                    )
850                })? {
851                protobuf::StreamPartitionMode::SinglePartition => {
852                    StreamJoinPartitionMode::SinglePartition
853                }
854                protobuf::StreamPartitionMode::PartitionedExec => {
855                    StreamJoinPartitionMode::Partitioned
856                }
857            };
858        let filter = sym_join
859            .filter
860            .as_ref()
861            .map(|filter| -> Result<JoinFilter> {
862                let schema: Schema = filter
863                    .schema
864                    .as_ref()
865                    .ok_or_else(|| {
866                        internal_datafusion_err!(
867                            "SymmetricHashJoinExec: JoinFilter missing schema"
868                        )
869                    })?
870                    .try_into()?;
871                let expression = ctx.decode_required_expr(
872                    filter.expression.as_ref(),
873                    &schema,
874                    "SymmetricHashJoinExec",
875                    "filter.expression",
876                )?;
877                let column_indices = filter
878                    .column_indices
879                    .iter()
880                    .map(|column_index| {
881                        let side = protobuf::JoinSide::try_from(column_index.side)
882                            .map_err(|_| {
883                                internal_datafusion_err!(
884                                    "SymmetricHashJoinExec: unknown JoinSide {}",
885                                    column_index.side
886                                )
887                            })?;
888                        let side = match side {
889                            protobuf::JoinSide::LeftSide => JoinSide::Left,
890                            protobuf::JoinSide::RightSide => JoinSide::Right,
891                            protobuf::JoinSide::None => JoinSide::None,
892                        };
893                        Ok(ColumnIndex {
894                            index: column_index.index as usize,
895                            side,
896                        })
897                    })
898                    .collect::<Result<Vec<_>>>()?;
899                Ok(JoinFilter::new(
900                    expression,
901                    column_indices,
902                    Arc::new(schema),
903                ))
904            })
905            .transpose()?;
906        let left_sort_exprs =
907            datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto(
908                &sym_join.left_sort_exprs,
909                &ctx.expr_ctx(left_schema.as_ref()),
910            )?;
911        let right_sort_exprs =
912            datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto(
913                &sym_join.right_sort_exprs,
914                &ctx.expr_ctx(right_schema.as_ref()),
915            )?;
916
917        Self::try_new(
918            left,
919            right,
920            on,
921            filter,
922            &join_type,
923            null_equality,
924            left_sort_exprs,
925            right_sort_exprs,
926            partition_mode,
927        )
928        .map(|exec| Arc::new(exec) as _)
929    }
930}
931
932/// A stream that issues [RecordBatch]es as they arrive from the right  of the join.
933struct SymmetricHashJoinStream<T> {
934    /// Input streams
935    left_stream: SendableRecordBatchStream,
936    right_stream: SendableRecordBatchStream,
937    /// Input schema
938    schema: Arc<Schema>,
939    /// join filter
940    filter: Option<JoinFilter>,
941    /// type of the join
942    join_type: JoinType,
943    // left hash joiner
944    left: OneSideHashJoiner,
945    /// right hash joiner
946    right: OneSideHashJoiner,
947    /// Information of index and left / right placement of columns
948    column_indices: Vec<ColumnIndex>,
949    // Expression graph for range pruning.
950    graph: Option<ExprIntervalGraph>,
951    // Left globally sorted filter expr
952    left_sorted_filter_expr: Option<SortedFilterExpr>,
953    // Right globally sorted filter expr
954    right_sorted_filter_expr: Option<SortedFilterExpr>,
955    /// Random state used for hashing initialization
956    random_state: RandomState,
957    /// Defines the null equality for the join.
958    null_equality: NullEquality,
959    /// Metrics
960    metrics: StreamJoinMetrics,
961    /// Memory reservation
962    reservation: SharedMemoryReservation,
963    /// State machine for input execution
964    state: SHJStreamState,
965    /// Transforms the output batch before returning.
966    batch_transformer: T,
967}
968
969impl<T: BatchTransformer + Unpin + Send> RecordBatchStream
970    for SymmetricHashJoinStream<T>
971{
972    fn schema(&self) -> SchemaRef {
973        Arc::clone(&self.schema)
974    }
975}
976
977impl<T: BatchTransformer + Unpin + Send> Stream for SymmetricHashJoinStream<T> {
978    type Item = Result<RecordBatch>;
979
980    fn poll_next(
981        mut self: std::pin::Pin<&mut Self>,
982        cx: &mut Context<'_>,
983    ) -> Poll<Option<Self::Item>> {
984        self.poll_next_impl(cx)
985    }
986}
987
988/// Determine the pruning length for `buffer`.
989///
990/// This function evaluates the build side filter expression, converts the
991/// result into an array and determines the pruning length by performing a
992/// binary search on the array.
993///
994/// # Arguments
995///
996/// * `buffer`: The record batch to be pruned.
997/// * `build_side_filter_expr`: The filter expression on the build side used
998///   to determine the pruning length.
999///
1000/// # Returns
1001///
1002/// A [Result] object that contains the pruning length. The function will return
1003/// an error if
1004/// - there is an issue evaluating the build side filter expression;
1005/// - there is an issue converting the build side filter expression into an array
1006fn determine_prune_length(
1007    buffer: &RecordBatch,
1008    build_side_filter_expr: &SortedFilterExpr,
1009) -> Result<usize> {
1010    let origin_sorted_expr = build_side_filter_expr.origin_sorted_expr();
1011    let interval = build_side_filter_expr.interval();
1012    // Evaluate the build side filter expression and convert it into an array
1013    let batch_arr = origin_sorted_expr
1014        .expr
1015        .evaluate(buffer)?
1016        .into_array(buffer.num_rows())?;
1017
1018    // Get the lower or upper interval based on the sort direction
1019    let target = if origin_sorted_expr.options.descending {
1020        interval.upper().clone()
1021    } else {
1022        interval.lower().clone()
1023    };
1024
1025    // Perform binary search on the array to determine the length of the record batch to be pruned
1026    bisect::<true>(&[batch_arr], &[target], &[origin_sorted_expr.options])
1027}
1028
1029/// This method determines if the result of the join should be produced in the final step or not.
1030///
1031/// # Arguments
1032///
1033/// * `build_side` - Enum indicating the side of the join used as the build side.
1034/// * `join_type` - Enum indicating the type of join to be performed.
1035///
1036/// # Returns
1037///
1038/// A boolean indicating whether the result of the join should be produced in the final step or not.
1039/// The result will be true if the build side is JoinSide::Left and the join type is one of
1040/// JoinType::Left, JoinType::LeftAnti, JoinType::Full or JoinType::LeftSemi.
1041/// If the build side is JoinSide::Right, the result will be true if the join type
1042/// is one of JoinType::Right, JoinType::RightAnti, JoinType::Full, or JoinType::RightSemi.
1043fn need_to_produce_result_in_final(build_side: JoinSide, join_type: JoinType) -> bool {
1044    if build_side == JoinSide::Left {
1045        matches!(
1046            join_type,
1047            JoinType::Left
1048                | JoinType::LeftAnti
1049                | JoinType::Full
1050                | JoinType::LeftSemi
1051                | JoinType::LeftMark
1052        )
1053    } else {
1054        matches!(
1055            join_type,
1056            JoinType::Right
1057                | JoinType::RightAnti
1058                | JoinType::Full
1059                | JoinType::RightSemi
1060                | JoinType::RightMark
1061        )
1062    }
1063}
1064
1065/// Calculate indices by join type.
1066///
1067/// This method returns a tuple of two arrays: build and probe indices.
1068/// The length of both arrays will be the same.
1069///
1070/// # Arguments
1071///
1072/// * `build_side`: Join side which defines the build side.
1073/// * `prune_length`: Length of the prune data.
1074/// * `visited_rows`: Hash set of visited rows of the build side.
1075/// * `deleted_offset`: Deleted offset of the build side.
1076/// * `join_type`: The type of join to be performed.
1077///
1078/// # Returns
1079///
1080/// A tuple of two arrays of primitive types representing the build and probe indices.
1081fn calculate_indices_by_join_type<L: ArrowPrimitiveType, R: ArrowPrimitiveType>(
1082    build_side: JoinSide,
1083    prune_length: usize,
1084    visited_rows: &HashSet<usize>,
1085    deleted_offset: usize,
1086    join_type: JoinType,
1087) -> Result<(PrimitiveArray<L>, PrimitiveArray<R>)>
1088where
1089    NativeAdapter<L>: From<<L as ArrowPrimitiveType>::Native>,
1090{
1091    // Store the result in a tuple
1092    let result = match (build_side, join_type) {
1093        // For a mark join we “mark” each build‐side row with a dummy 0 in the probe‐side index
1094        // if it ever matched. For example, if
1095        //
1096        // prune_length = 5
1097        // deleted_offset = 0
1098        // visited_rows = {1, 3}
1099        //
1100        // then we produce:
1101        //
1102        // build_indices = [0, 1, 2, 3, 4]
1103        // probe_indices = [None, Some(0), None, Some(0), None]
1104        //
1105        // Example: for each build row i in [0..5):
1106        //   – We always output its own index i in `build_indices`
1107        //   – We output `Some(0)` in `probe_indices[i]` if row i was ever visited, else `None`
1108        (JoinSide::Left, JoinType::LeftMark) => {
1109            let build_indices = (0..prune_length)
1110                .map(L::Native::from_usize)
1111                .collect::<PrimitiveArray<L>>();
1112            let probe_indices = (0..prune_length)
1113                .map(|idx| {
1114                    // For mark join we output a dummy index 0 to indicate the row had a match
1115                    visited_rows
1116                        .contains(&(idx + deleted_offset))
1117                        .then_some(R::Native::from_usize(0).unwrap())
1118                })
1119                .collect();
1120            (build_indices, probe_indices)
1121        }
1122        (JoinSide::Right, JoinType::RightMark) => {
1123            let build_indices = (0..prune_length)
1124                .map(L::Native::from_usize)
1125                .collect::<PrimitiveArray<L>>();
1126            let probe_indices = (0..prune_length)
1127                .map(|idx| {
1128                    // For mark join we output a dummy index 0 to indicate the row had a match
1129                    visited_rows
1130                        .contains(&(idx + deleted_offset))
1131                        .then_some(R::Native::from_usize(0).unwrap())
1132                })
1133                .collect();
1134            (build_indices, probe_indices)
1135        }
1136        // In the case of `Left` or `Right` join, or `Full` join, get the anti indices
1137        (JoinSide::Left, JoinType::Left | JoinType::LeftAnti)
1138        | (JoinSide::Right, JoinType::Right | JoinType::RightAnti)
1139        | (_, JoinType::Full) => {
1140            let build_unmatched_indices =
1141                get_pruning_anti_indices(prune_length, deleted_offset, visited_rows);
1142            let mut builder =
1143                PrimitiveBuilder::<R>::with_capacity(build_unmatched_indices.len());
1144            builder.append_nulls(build_unmatched_indices.len());
1145            let probe_indices = builder.finish();
1146            (build_unmatched_indices, probe_indices)
1147        }
1148        // In the case of `LeftSemi` or `RightSemi` join, get the semi indices
1149        (JoinSide::Left, JoinType::LeftSemi) | (JoinSide::Right, JoinType::RightSemi) => {
1150            let build_unmatched_indices =
1151                get_pruning_semi_indices(prune_length, deleted_offset, visited_rows);
1152            let mut builder =
1153                PrimitiveBuilder::<R>::with_capacity(build_unmatched_indices.len());
1154            builder.append_nulls(build_unmatched_indices.len());
1155            let probe_indices = builder.finish();
1156            (build_unmatched_indices, probe_indices)
1157        }
1158        // The case of other join types is not considered
1159        _ => unreachable!(),
1160    };
1161    Ok(result)
1162}
1163
1164/// This function produces unmatched record results based on the build side,
1165/// join type and other parameters.
1166///
1167/// The method uses first `prune_length` rows from the build side input buffer
1168/// to produce results.
1169///
1170/// # Arguments
1171///
1172/// * `output_schema` - The schema of the final output record batch.
1173/// * `prune_length` - The length of the determined prune length.
1174/// * `probe_schema` - The schema of the probe [RecordBatch].
1175/// * `join_type` - The type of join to be performed.
1176/// * `column_indices` - Indices of columns that are being joined.
1177///
1178/// # Returns
1179///
1180/// * `Option<RecordBatch>` - The final output record batch if required, otherwise [None].
1181pub(crate) fn build_side_determined_results(
1182    build_hash_joiner: &OneSideHashJoiner,
1183    output_schema: &SchemaRef,
1184    prune_length: usize,
1185    probe_schema: SchemaRef,
1186    join_type: JoinType,
1187    column_indices: &[ColumnIndex],
1188) -> Result<Option<RecordBatch>> {
1189    // Check if we need to produce a result in the final output:
1190    if prune_length > 0
1191        && need_to_produce_result_in_final(build_hash_joiner.build_side, join_type)
1192    {
1193        // Calculate the indices for build and probe sides based on join type and build side:
1194        let (build_indices, probe_indices) = calculate_indices_by_join_type(
1195            build_hash_joiner.build_side,
1196            prune_length,
1197            &build_hash_joiner.visited_rows,
1198            build_hash_joiner.deleted_offset,
1199            join_type,
1200        )?;
1201
1202        // Create an empty probe record batch:
1203        let empty_probe_batch = RecordBatch::new_empty(probe_schema);
1204        // Build the final result from the indices of build and probe sides:
1205        build_batch_from_indices(
1206            output_schema.as_ref(),
1207            &build_hash_joiner.input_buffer,
1208            &empty_probe_batch,
1209            &build_indices,
1210            &probe_indices,
1211            column_indices,
1212            build_hash_joiner.build_side,
1213            join_type,
1214        )
1215        .map(|batch| (batch.num_rows() > 0).then_some(batch))
1216    } else {
1217        // If we don't need to produce a result, return None
1218        Ok(None)
1219    }
1220}
1221
1222/// This method performs a join between the build side input buffer and the probe side batch.
1223///
1224/// # Arguments
1225///
1226/// * `build_hash_joiner` - Build side hash joiner
1227/// * `probe_hash_joiner` - Probe side hash joiner
1228/// * `schema` - A reference to the schema of the output record batch.
1229/// * `join_type` - The type of join to be performed.
1230/// * `on_probe` - An array of columns on which the join will be performed. The columns are from the probe side of the join.
1231/// * `filter` - An optional filter on the join condition.
1232/// * `probe_batch` - The second record batch to be joined.
1233/// * `column_indices` - An array of columns to be selected for the result of the join.
1234/// * `random_state` - The random state for the join.
1235/// * `null_equality` - Indicates whether NULL values should be treated as equal when joining.
1236///
1237/// # Returns
1238///
1239/// A [Result] containing an optional record batch if the join type is not one of `LeftAnti`, `RightAnti`, `LeftSemi` or `RightSemi`.
1240/// If the join type is one of the above four, the function will return [None].
1241#[expect(clippy::too_many_arguments)]
1242pub(crate) fn join_with_probe_batch(
1243    build_hash_joiner: &mut OneSideHashJoiner,
1244    probe_hash_joiner: &mut OneSideHashJoiner,
1245    schema: &SchemaRef,
1246    join_type: JoinType,
1247    filter: Option<&JoinFilter>,
1248    probe_batch: &RecordBatch,
1249    column_indices: &[ColumnIndex],
1250    random_state: &RandomState,
1251    null_equality: NullEquality,
1252) -> Result<Option<RecordBatch>> {
1253    if build_hash_joiner.input_buffer.num_rows() == 0 || probe_batch.num_rows() == 0 {
1254        return Ok(None);
1255    }
1256    let (build_indices, probe_indices) = lookup_join_hashmap(
1257        &build_hash_joiner.hashmap,
1258        &build_hash_joiner.input_buffer,
1259        probe_batch,
1260        &build_hash_joiner.on,
1261        &probe_hash_joiner.on,
1262        random_state,
1263        null_equality,
1264        &mut build_hash_joiner.hashes_buffer,
1265        Some(build_hash_joiner.deleted_offset),
1266    )?;
1267
1268    let (build_indices, probe_indices) = if let Some(filter) = filter {
1269        apply_join_filter_to_indices(
1270            &build_hash_joiner.input_buffer,
1271            probe_batch,
1272            build_indices,
1273            probe_indices,
1274            filter,
1275            build_hash_joiner.build_side,
1276            None,
1277            join_type,
1278        )?
1279    } else {
1280        (build_indices, probe_indices)
1281    };
1282
1283    if need_to_produce_result_in_final(build_hash_joiner.build_side, join_type) {
1284        record_visited_indices(
1285            &mut build_hash_joiner.visited_rows,
1286            build_hash_joiner.deleted_offset,
1287            &build_indices,
1288        );
1289    }
1290    if need_to_produce_result_in_final(build_hash_joiner.build_side.negate(), join_type) {
1291        record_visited_indices(
1292            &mut probe_hash_joiner.visited_rows,
1293            probe_hash_joiner.offset,
1294            &probe_indices,
1295        );
1296    }
1297    if matches!(
1298        join_type,
1299        JoinType::LeftAnti
1300            | JoinType::RightAnti
1301            | JoinType::LeftSemi
1302            | JoinType::LeftMark
1303            | JoinType::RightSemi
1304            | JoinType::RightMark
1305    ) {
1306        Ok(None)
1307    } else {
1308        build_batch_from_indices(
1309            schema,
1310            &build_hash_joiner.input_buffer,
1311            probe_batch,
1312            &build_indices,
1313            &probe_indices,
1314            column_indices,
1315            build_hash_joiner.build_side,
1316            join_type,
1317        )
1318        .map(|batch| (batch.num_rows() > 0).then_some(batch))
1319    }
1320}
1321
1322/// This method performs lookups against JoinHashMap by hash values of join-key columns, and handles potential
1323/// hash collisions.
1324///
1325/// # Arguments
1326///
1327/// * `build_hashmap` - hashmap collected from build side data.
1328/// * `build_batch` - Build side record batch.
1329/// * `probe_batch` - Probe side record batch.
1330/// * `build_on` - An array of columns on which the join will be performed. The columns are from the build side of the join.
1331/// * `probe_on` - An array of columns on which the join will be performed. The columns are from the probe side of the join.
1332/// * `random_state` - The random state for the join.
1333/// * `null_equality` - Indicates whether NULL values should be treated as equal when joining.
1334/// * `hashes_buffer` - Buffer used for probe side keys hash calculation.
1335/// * `deleted_offset` - deleted offset for build side data.
1336///
1337/// # Returns
1338///
1339/// A [Result] containing a tuple with two equal length arrays, representing indices of rows from build and probe side,
1340/// matched by join key columns.
1341#[expect(clippy::too_many_arguments)]
1342fn lookup_join_hashmap(
1343    build_hashmap: &PruningJoinHashMap,
1344    build_batch: &RecordBatch,
1345    probe_batch: &RecordBatch,
1346    build_on: &[PhysicalExprRef],
1347    probe_on: &[PhysicalExprRef],
1348    random_state: &RandomState,
1349    null_equality: NullEquality,
1350    hashes_buffer: &mut Vec<u64>,
1351    deleted_offset: Option<usize>,
1352) -> Result<(UInt64Array, UInt32Array)> {
1353    let keys_values = evaluate_expressions_to_arrays(probe_on, probe_batch)?;
1354    let build_join_values = evaluate_expressions_to_arrays(build_on, build_batch)?;
1355
1356    hashes_buffer.clear();
1357    hashes_buffer.resize(probe_batch.num_rows(), 0);
1358    let hash_values = create_hashes(&keys_values, random_state, hashes_buffer)?;
1359
1360    // As SymmetricHashJoin uses LIFO JoinHashMap, the chained list algorithm
1361    // will return build indices for each probe row in a reverse order as such:
1362    // Build Indices: [5, 4, 3]
1363    // Probe Indices: [1, 1, 1]
1364    //
1365    // This affects the output sequence. Hypothetically, it's possible to preserve the lexicographic order on the build side.
1366    // Let's consider probe rows [0,1] as an example:
1367    //
1368    // When the probe iteration sequence is reversed, the following pairings can be derived:
1369    //
1370    // For probe row 1:
1371    //     (5, 1)
1372    //     (4, 1)
1373    //     (3, 1)
1374    //
1375    // For probe row 0:
1376    //     (5, 0)
1377    //     (4, 0)
1378    //     (3, 0)
1379    //
1380    // After reversing both sets of indices, we obtain reversed indices:
1381    //
1382    //     (3,0)
1383    //     (4,0)
1384    //     (5,0)
1385    //     (3,1)
1386    //     (4,1)
1387    //     (5,1)
1388    //
1389    // With this approach, the lexicographic order on both the probe side and the build side is preserved.
1390    //
1391    // Probe rows whose key contains a NULL cannot match any build row and are
1392    // skipped without a map lookup.
1393    let valid_keys = matchable_join_keys(&keys_values, null_equality);
1394    let (mut matched_probe, mut matched_build) = build_hashmap.get_matched_indices(
1395        Box::new(
1396            hash_values
1397                .iter()
1398                .enumerate()
1399                .filter(|(i, _)| {
1400                    valid_keys.as_ref().is_none_or(|valid| valid.is_valid(*i))
1401                })
1402                .rev(),
1403        ),
1404        deleted_offset,
1405    );
1406
1407    matched_probe.reverse();
1408    matched_build.reverse();
1409
1410    let build_indices: UInt64Array = matched_build.into();
1411    let probe_indices: UInt32Array = matched_probe.into();
1412
1413    let (build_indices, probe_indices) = equal_rows_arr(
1414        &build_indices,
1415        &probe_indices,
1416        &build_join_values,
1417        &keys_values,
1418        null_equality,
1419    )?;
1420
1421    Ok((build_indices, probe_indices))
1422}
1423
1424pub struct OneSideHashJoiner {
1425    /// Build side
1426    build_side: JoinSide,
1427    /// Input record batch buffer
1428    pub input_buffer: RecordBatch,
1429    /// Columns from the side
1430    pub(crate) on: Vec<PhysicalExprRef>,
1431    /// Hashmap
1432    pub(crate) hashmap: PruningJoinHashMap,
1433    /// Reuse the hashes buffer
1434    pub(crate) hashes_buffer: Vec<u64>,
1435    /// Matched rows
1436    pub(crate) visited_rows: HashSet<usize>,
1437    /// Offset
1438    pub(crate) offset: usize,
1439    /// Deleted offset
1440    pub(crate) deleted_offset: usize,
1441}
1442
1443impl OneSideHashJoiner {
1444    pub fn size(&self) -> usize {
1445        let mut size = 0;
1446        size += size_of_val(self);
1447        size += size_of_val(&self.build_side);
1448        size += self.input_buffer.get_array_memory_size();
1449        size += size_of_val(&self.on);
1450        size += self.hashmap.size();
1451        size += self.hashes_buffer.capacity() * size_of::<u64>();
1452        size += self.visited_rows.capacity() * size_of::<usize>();
1453        size += size_of_val(&self.offset);
1454        size += size_of_val(&self.deleted_offset);
1455        size
1456    }
1457    pub fn new(
1458        build_side: JoinSide,
1459        on: Vec<PhysicalExprRef>,
1460        schema: SchemaRef,
1461    ) -> Self {
1462        Self {
1463            build_side,
1464            input_buffer: RecordBatch::new_empty(schema),
1465            on,
1466            hashmap: PruningJoinHashMap::with_capacity(0),
1467            hashes_buffer: vec![],
1468            visited_rows: HashSet::new(),
1469            offset: 0,
1470            deleted_offset: 0,
1471        }
1472    }
1473
1474    /// Updates the internal state of the [OneSideHashJoiner] with the incoming batch.
1475    ///
1476    /// # Arguments
1477    ///
1478    /// * `batch` - The incoming [RecordBatch] to be merged with the internal input buffer
1479    /// * `random_state` - The random state used to hash values
1480    /// * `null_equality` - Null semantics to use
1481    ///
1482    /// # Returns
1483    ///
1484    /// Returns a [Result] encapsulating any intermediate errors.
1485    pub(crate) fn update_internal_state(
1486        &mut self,
1487        batch: &RecordBatch,
1488        random_state: &RandomState,
1489        null_equality: NullEquality,
1490    ) -> Result<()> {
1491        // Merge the incoming batch with the existing input buffer:
1492        self.input_buffer = concat_batches(&batch.schema(), [&self.input_buffer, batch])?;
1493        // Resize the hashes buffer to the number of rows in the incoming batch:
1494        self.hashes_buffer.resize(batch.num_rows(), 0);
1495        // Get allocation_info before adding the item
1496        // Update the hashmap with the join key values and hashes of the incoming batch:
1497        update_hash(
1498            &self.on,
1499            batch,
1500            &mut self.hashmap,
1501            self.offset,
1502            random_state,
1503            &mut self.hashes_buffer,
1504            self.deleted_offset,
1505            false,
1506            null_equality,
1507        )?;
1508        Ok(())
1509    }
1510
1511    /// Calculate prune length.
1512    ///
1513    /// # Arguments
1514    ///
1515    /// * `build_side_sorted_filter_expr` - Build side mutable sorted filter expression..
1516    /// * `probe_side_sorted_filter_expr` - Probe side mutable sorted filter expression.
1517    /// * `graph` - A mutable reference to the physical expression graph.
1518    ///
1519    /// # Returns
1520    ///
1521    /// A Result object that contains the pruning length.
1522    pub(crate) fn calculate_prune_length_with_probe_batch(
1523        &mut self,
1524        build_side_sorted_filter_expr: &mut SortedFilterExpr,
1525        probe_side_sorted_filter_expr: &mut SortedFilterExpr,
1526        graph: &mut ExprIntervalGraph,
1527    ) -> Result<usize> {
1528        // Return early if the input buffer is empty:
1529        if self.input_buffer.num_rows() == 0 {
1530            return Ok(0);
1531        }
1532        // Process the build and probe side sorted filter expressions if both are present:
1533        // Collect the sorted filter expressions into a vector of (node_index, interval) tuples:
1534        let mut filter_intervals = vec![];
1535        for expr in [
1536            &build_side_sorted_filter_expr,
1537            &probe_side_sorted_filter_expr,
1538        ] {
1539            filter_intervals.push((expr.node_index(), expr.interval().clone()))
1540        }
1541        // Update the physical expression graph using the join filter intervals:
1542        graph.update_ranges(&mut filter_intervals, Interval::TRUE)?;
1543        // Extract the new join filter interval for the build side:
1544        let calculated_build_side_interval = filter_intervals.remove(0).1;
1545        // If the intervals have not changed, return early without pruning:
1546        if calculated_build_side_interval.eq(build_side_sorted_filter_expr.interval()) {
1547            return Ok(0);
1548        }
1549        // Update the build side interval and determine the pruning length:
1550        build_side_sorted_filter_expr.set_interval(calculated_build_side_interval);
1551
1552        determine_prune_length(&self.input_buffer, build_side_sorted_filter_expr)
1553    }
1554
1555    pub(crate) fn prune_internal_state(&mut self, prune_length: usize) -> Result<()> {
1556        // Prune the hash values:
1557        self.hashmap.prune_hash_values(
1558            prune_length,
1559            self.deleted_offset as u64,
1560            HASHMAP_SHRINK_SCALE_FACTOR,
1561        );
1562        // Remove pruned rows from the visited rows set:
1563        for row in self.deleted_offset..(self.deleted_offset + prune_length) {
1564            self.visited_rows.remove(&row);
1565        }
1566        // Update the input buffer after pruning:
1567        self.input_buffer = self
1568            .input_buffer
1569            .slice(prune_length, self.input_buffer.num_rows() - prune_length);
1570        // Increment the deleted offset:
1571        self.deleted_offset += prune_length;
1572        Ok(())
1573    }
1574}
1575
1576/// `SymmetricHashJoinStream` manages incremental join operations between two
1577/// streams. Unlike traditional join approaches that need to scan one side of
1578/// the join fully before proceeding, `SymmetricHashJoinStream` facilitates
1579/// more dynamic join operations by working with streams as they emit data. This
1580/// approach allows for more efficient processing, particularly in scenarios
1581/// where waiting for complete data materialization is not feasible or optimal.
1582/// The trait provides a framework for handling various states of such a join
1583/// process, ensuring that join logic is efficiently executed as data becomes
1584/// available from either stream.
1585///
1586/// This implementation performs eager joins of data from two different asynchronous
1587/// streams, typically referred to as left and right streams. The implementation
1588/// provides a comprehensive set of methods to control and execute the join
1589/// process, leveraging the states defined in `SHJStreamState`. Methods are
1590/// primarily focused on asynchronously fetching data batches from each stream,
1591/// processing them, and managing transitions between various states of the join.
1592///
1593/// This implementations use a state machine approach to navigate different
1594/// stages of the join operation, handling data from both streams and determining
1595/// when the join completes.
1596///
1597/// State Transitions:
1598/// - From `PullLeft` to `PullRight` or `LeftExhausted`:
1599///   - In `fetch_next_from_left_stream`, when fetching a batch from the left stream:
1600///     - On success (`Some(Ok(batch))`), state transitions to `PullRight` for
1601///       processing the batch.
1602///     - On error (`Some(Err(e))`), the error is returned, and the state remains
1603///       unchanged.
1604///     - On no data (`None`), state changes to `LeftExhausted`, returning `Continue`
1605///       to proceed with the join process.
1606/// - From `PullRight` to `PullLeft` or `RightExhausted`:
1607///   - In `fetch_next_from_right_stream`, when fetching from the right stream:
1608///     - If a batch is available, state changes to `PullLeft` for processing.
1609///     - On error, the error is returned without changing the state.
1610///     - If right stream is exhausted (`None`), state transitions to `RightExhausted`,
1611///       with a `Continue` result.
1612/// - Handling `RightExhausted` and `LeftExhausted`:
1613///   - Methods `handle_right_stream_end` and `handle_left_stream_end` manage scenarios
1614///     when streams are exhausted:
1615///     - They attempt to continue processing with the other stream.
1616///     - If both streams are exhausted, state changes to `BothExhausted { final_result: false }`.
1617/// - Transition to `BothExhausted { final_result: true }`:
1618///   - Occurs in `prepare_for_final_results_after_exhaustion` when both streams are
1619///     exhausted, indicating completion of processing and availability of final results.
1620impl<T: BatchTransformer> SymmetricHashJoinStream<T> {
1621    /// Implements the main polling logic for the join stream.
1622    ///
1623    /// This method continuously checks the state of the join stream and
1624    /// acts accordingly by delegating the handling to appropriate sub-methods
1625    /// depending on the current state.
1626    ///
1627    /// # Arguments
1628    ///
1629    /// * `cx` - A context that facilitates cooperative non-blocking execution within a task.
1630    ///
1631    /// # Returns
1632    ///
1633    /// * `Poll<Option<Result<RecordBatch>>>` - A polled result, either a `RecordBatch` or None.
1634    fn poll_next_impl(
1635        &mut self,
1636        cx: &mut Context<'_>,
1637    ) -> Poll<Option<Result<RecordBatch>>> {
1638        loop {
1639            match self.batch_transformer.next() {
1640                None => {
1641                    let result = match self.state() {
1642                        SHJStreamState::PullRight => {
1643                            ready!(self.fetch_next_from_right_stream(cx))
1644                        }
1645                        SHJStreamState::PullLeft => {
1646                            ready!(self.fetch_next_from_left_stream(cx))
1647                        }
1648                        SHJStreamState::RightExhausted => {
1649                            ready!(self.handle_right_stream_end(cx))
1650                        }
1651                        SHJStreamState::LeftExhausted => {
1652                            ready!(self.handle_left_stream_end(cx))
1653                        }
1654                        SHJStreamState::BothExhausted {
1655                            final_result: false,
1656                        } => self.prepare_for_final_results_after_exhaustion(),
1657                        SHJStreamState::BothExhausted { final_result: true } => {
1658                            return Poll::Ready(None);
1659                        }
1660                    };
1661
1662                    match result? {
1663                        StatefulStreamResult::Ready(None) => {
1664                            return Poll::Ready(None);
1665                        }
1666                        StatefulStreamResult::Ready(Some(batch)) => {
1667                            self.batch_transformer.set_batch(batch);
1668                        }
1669                        _ => {}
1670                    }
1671                }
1672                Some((batch, _)) => {
1673                    return self
1674                        .metrics
1675                        .baseline_metrics
1676                        .record_poll(Poll::Ready(Some(Ok(batch))));
1677                }
1678            }
1679        }
1680    }
1681
1682    /// Release the right input pipeline's resources.
1683    fn cleanup_depleted_right_stream(&mut self) {
1684        let right_schema = self.right_stream.schema();
1685        self.right_stream = Box::pin(EmptyRecordBatchStream::new(right_schema));
1686    }
1687
1688    /// Release the left input pipeline's resources.
1689    fn cleanup_depleted_left_stream(&mut self) {
1690        let left_schema = self.left_stream.schema();
1691        self.left_stream = Box::pin(EmptyRecordBatchStream::new(left_schema));
1692    }
1693
1694    /// Asynchronously pulls the next batch from the right stream.
1695    ///
1696    /// This default implementation checks for the next value in the right stream.
1697    /// If a batch is found, the state is switched to `PullLeft`, and the batch handling
1698    /// is delegated to `process_batch_from_right`. If the stream ends, the state is set to `RightExhausted`.
1699    ///
1700    /// # Returns
1701    ///
1702    /// * `Result<StatefulStreamResult<Option<RecordBatch>>>` - The state result after pulling the batch.
1703    fn fetch_next_from_right_stream(
1704        &mut self,
1705        cx: &mut Context<'_>,
1706    ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
1707        match ready!(self.right_stream().poll_next_unpin(cx)) {
1708            Some(Ok(batch)) => {
1709                if batch.num_rows() == 0 {
1710                    return Poll::Ready(Ok(StatefulStreamResult::Continue));
1711                }
1712                self.set_state(SHJStreamState::PullLeft);
1713                Poll::Ready(self.process_batch_from_right(&batch))
1714            }
1715            Some(Err(e)) => Poll::Ready(Err(e)),
1716            None => {
1717                self.cleanup_depleted_right_stream();
1718                self.set_state(SHJStreamState::RightExhausted);
1719                Poll::Ready(Ok(StatefulStreamResult::Continue))
1720            }
1721        }
1722    }
1723
1724    /// Asynchronously pulls the next batch from the left stream.
1725    ///
1726    /// This default implementation checks for the next value in the left stream.
1727    /// If a batch is found, the state is switched to `PullRight`, and the batch handling
1728    /// is delegated to `process_batch_from_left`. If the stream ends, the state is set to `LeftExhausted`.
1729    ///
1730    /// # Returns
1731    ///
1732    /// * `Result<StatefulStreamResult<Option<RecordBatch>>>` - The state result after pulling the batch.
1733    fn fetch_next_from_left_stream(
1734        &mut self,
1735        cx: &mut Context<'_>,
1736    ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
1737        match ready!(self.left_stream().poll_next_unpin(cx)) {
1738            Some(Ok(batch)) => {
1739                if batch.num_rows() == 0 {
1740                    return Poll::Ready(Ok(StatefulStreamResult::Continue));
1741                }
1742                self.set_state(SHJStreamState::PullRight);
1743                Poll::Ready(self.process_batch_from_left(&batch))
1744            }
1745            Some(Err(e)) => Poll::Ready(Err(e)),
1746            None => {
1747                self.cleanup_depleted_left_stream();
1748                self.set_state(SHJStreamState::LeftExhausted);
1749                Poll::Ready(Ok(StatefulStreamResult::Continue))
1750            }
1751        }
1752    }
1753
1754    /// Asynchronously handles the scenario when the right stream is exhausted.
1755    ///
1756    /// In this default implementation, when the right stream is exhausted, it attempts
1757    /// to pull from the left stream. If a batch is found in the left stream, it delegates
1758    /// the handling to `process_batch_from_left`. If both streams are exhausted, the state is set
1759    /// to indicate both streams are exhausted without final results yet.
1760    ///
1761    /// # Returns
1762    ///
1763    /// * `Result<StatefulStreamResult<Option<RecordBatch>>>` - The state result after checking the exhaustion state.
1764    fn handle_right_stream_end(
1765        &mut self,
1766        cx: &mut Context<'_>,
1767    ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
1768        match ready!(self.left_stream().poll_next_unpin(cx)) {
1769            Some(Ok(batch)) => {
1770                if batch.num_rows() == 0 {
1771                    return Poll::Ready(Ok(StatefulStreamResult::Continue));
1772                }
1773                Poll::Ready(self.process_batch_after_right_end(&batch))
1774            }
1775            Some(Err(e)) => Poll::Ready(Err(e)),
1776            None => {
1777                self.cleanup_depleted_left_stream();
1778                self.set_state(SHJStreamState::BothExhausted {
1779                    final_result: false,
1780                });
1781                Poll::Ready(Ok(StatefulStreamResult::Continue))
1782            }
1783        }
1784    }
1785
1786    /// Asynchronously handles the scenario when the left stream is exhausted.
1787    ///
1788    /// When the left stream is exhausted, this default
1789    /// implementation tries to pull from the right stream and delegates the batch
1790    /// handling to `process_batch_after_left_end`. If both streams are exhausted, the state
1791    /// is updated to indicate so.
1792    ///
1793    /// # Returns
1794    ///
1795    /// * `Result<StatefulStreamResult<Option<RecordBatch>>>` - The state result after checking the exhaustion state.
1796    fn handle_left_stream_end(
1797        &mut self,
1798        cx: &mut Context<'_>,
1799    ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
1800        match ready!(self.right_stream().poll_next_unpin(cx)) {
1801            Some(Ok(batch)) => {
1802                if batch.num_rows() == 0 {
1803                    return Poll::Ready(Ok(StatefulStreamResult::Continue));
1804                }
1805                Poll::Ready(self.process_batch_after_left_end(&batch))
1806            }
1807            Some(Err(e)) => Poll::Ready(Err(e)),
1808            None => {
1809                self.cleanup_depleted_right_stream();
1810                self.set_state(SHJStreamState::BothExhausted {
1811                    final_result: false,
1812                });
1813                Poll::Ready(Ok(StatefulStreamResult::Continue))
1814            }
1815        }
1816    }
1817
1818    /// Handles the state when both streams are exhausted and final results are yet to be produced.
1819    ///
1820    /// This default implementation switches the state to indicate both streams are
1821    /// exhausted with final results and then invokes the handling for this specific
1822    /// scenario via `process_batches_before_finalization`.
1823    ///
1824    /// # Returns
1825    ///
1826    /// * `Result<StatefulStreamResult<Option<RecordBatch>>>` - The state result after both streams are exhausted.
1827    fn prepare_for_final_results_after_exhaustion(
1828        &mut self,
1829    ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1830        self.set_state(SHJStreamState::BothExhausted { final_result: true });
1831        self.process_batches_before_finalization()
1832    }
1833
1834    fn process_batch_from_right(
1835        &mut self,
1836        batch: &RecordBatch,
1837    ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1838        self.perform_join_for_given_side(batch, JoinSide::Right)
1839            .map(|maybe_batch| {
1840                if maybe_batch.is_some() {
1841                    StatefulStreamResult::Ready(maybe_batch)
1842                } else {
1843                    StatefulStreamResult::Continue
1844                }
1845            })
1846    }
1847
1848    fn process_batch_from_left(
1849        &mut self,
1850        batch: &RecordBatch,
1851    ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1852        self.perform_join_for_given_side(batch, JoinSide::Left)
1853            .map(|maybe_batch| {
1854                if maybe_batch.is_some() {
1855                    StatefulStreamResult::Ready(maybe_batch)
1856                } else {
1857                    StatefulStreamResult::Continue
1858                }
1859            })
1860    }
1861
1862    fn process_batch_after_left_end(
1863        &mut self,
1864        right_batch: &RecordBatch,
1865    ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1866        self.process_batch_from_right(right_batch)
1867    }
1868
1869    fn process_batch_after_right_end(
1870        &mut self,
1871        left_batch: &RecordBatch,
1872    ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1873        self.process_batch_from_left(left_batch)
1874    }
1875
1876    fn process_batches_before_finalization(
1877        &mut self,
1878    ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1879        // Get the left side results:
1880        let left_result = build_side_determined_results(
1881            &self.left,
1882            &self.schema,
1883            self.left.input_buffer.num_rows(),
1884            self.right.input_buffer.schema(),
1885            self.join_type,
1886            &self.column_indices,
1887        )?;
1888        // Get the right side results:
1889        let right_result = build_side_determined_results(
1890            &self.right,
1891            &self.schema,
1892            self.right.input_buffer.num_rows(),
1893            self.left.input_buffer.schema(),
1894            self.join_type,
1895            &self.column_indices,
1896        )?;
1897
1898        // Combine the left and right results:
1899        let result = combine_two_batches(&self.schema, left_result, right_result)?;
1900
1901        // Return the result:
1902        if result.is_some() {
1903            return Ok(StatefulStreamResult::Ready(result));
1904        }
1905        Ok(StatefulStreamResult::Continue)
1906    }
1907
1908    fn right_stream(&mut self) -> &mut SendableRecordBatchStream {
1909        &mut self.right_stream
1910    }
1911
1912    fn left_stream(&mut self) -> &mut SendableRecordBatchStream {
1913        &mut self.left_stream
1914    }
1915
1916    fn set_state(&mut self, state: SHJStreamState) {
1917        self.state = state;
1918    }
1919
1920    fn state(&mut self) -> SHJStreamState {
1921        self.state.clone()
1922    }
1923
1924    fn size(&self) -> usize {
1925        let mut size = 0;
1926        size += size_of_val(&self.schema);
1927        size += size_of_val(&self.filter);
1928        size += size_of_val(&self.join_type);
1929        size += self.left.size();
1930        size += self.right.size();
1931        size += size_of_val(&self.column_indices);
1932        size += self.graph.as_ref().map(|g| g.size()).unwrap_or(0);
1933        size += size_of_val(&self.left_sorted_filter_expr);
1934        size += size_of_val(&self.right_sorted_filter_expr);
1935        size += size_of_val(&self.random_state);
1936        size += size_of_val(&self.null_equality);
1937        size += size_of_val(&self.metrics);
1938        size
1939    }
1940
1941    /// Performs a join operation for the specified `probe_side` (either left or right).
1942    /// This function:
1943    /// 1. Determines which side is the probe and which is the build side.
1944    /// 2. Updates metrics based on the batch that was polled.
1945    /// 3. Executes the join with the given `probe_batch`.
1946    /// 4. Optionally computes anti-join results if all conditions are met.
1947    /// 5. Combines the results and returns a combined batch or `None` if no batch was produced.
1948    fn perform_join_for_given_side(
1949        &mut self,
1950        probe_batch: &RecordBatch,
1951        probe_side: JoinSide,
1952    ) -> Result<Option<RecordBatch>> {
1953        let (
1954            probe_hash_joiner,
1955            build_hash_joiner,
1956            probe_side_sorted_filter_expr,
1957            build_side_sorted_filter_expr,
1958            probe_side_metrics,
1959        ) = if probe_side.eq(&JoinSide::Left) {
1960            (
1961                &mut self.left,
1962                &mut self.right,
1963                &mut self.left_sorted_filter_expr,
1964                &mut self.right_sorted_filter_expr,
1965                &mut self.metrics.left,
1966            )
1967        } else {
1968            (
1969                &mut self.right,
1970                &mut self.left,
1971                &mut self.right_sorted_filter_expr,
1972                &mut self.left_sorted_filter_expr,
1973                &mut self.metrics.right,
1974            )
1975        };
1976        // Update the metrics for the stream that was polled:
1977        probe_side_metrics.input_batches.add(1);
1978        probe_side_metrics.input_rows.add(probe_batch.num_rows());
1979        // Update the internal state of the hash joiner for the build side:
1980        probe_hash_joiner.update_internal_state(
1981            probe_batch,
1982            &self.random_state,
1983            self.null_equality,
1984        )?;
1985        // Join the two sides:
1986        let equal_result = join_with_probe_batch(
1987            build_hash_joiner,
1988            probe_hash_joiner,
1989            &self.schema,
1990            self.join_type,
1991            self.filter.as_ref(),
1992            probe_batch,
1993            &self.column_indices,
1994            &self.random_state,
1995            self.null_equality,
1996        )?;
1997        // Increment the offset for the probe hash joiner:
1998        probe_hash_joiner.offset += probe_batch.num_rows();
1999
2000        let anti_result = if let (
2001            Some(build_side_sorted_filter_expr),
2002            Some(probe_side_sorted_filter_expr),
2003            Some(graph),
2004        ) = (
2005            build_side_sorted_filter_expr.as_mut(),
2006            probe_side_sorted_filter_expr.as_mut(),
2007            self.graph.as_mut(),
2008        ) {
2009            // Calculate filter intervals:
2010            calculate_filter_expr_intervals(
2011                &build_hash_joiner.input_buffer,
2012                build_side_sorted_filter_expr,
2013                probe_batch,
2014                probe_side_sorted_filter_expr,
2015            )?;
2016            let prune_length = build_hash_joiner
2017                .calculate_prune_length_with_probe_batch(
2018                    build_side_sorted_filter_expr,
2019                    probe_side_sorted_filter_expr,
2020                    graph,
2021                )?;
2022            let result = build_side_determined_results(
2023                build_hash_joiner,
2024                &self.schema,
2025                prune_length,
2026                probe_batch.schema(),
2027                self.join_type,
2028                &self.column_indices,
2029            )?;
2030            build_hash_joiner.prune_internal_state(prune_length)?;
2031            result
2032        } else {
2033            None
2034        };
2035
2036        // Combine results:
2037        let result = combine_two_batches(&self.schema, equal_result, anti_result)?;
2038        let capacity = self.size();
2039        self.metrics.stream_memory_usage.set(capacity);
2040        self.reservation.try_resize(capacity)?;
2041        Ok(result)
2042    }
2043}
2044
2045/// Represents the various states of an symmetric hash join stream operation.
2046///
2047/// This enum is used to track the current state of streaming during a join
2048/// operation. It provides indicators as to which side of the join needs to be
2049/// pulled next or if one (or both) sides have been exhausted. This allows
2050/// for efficient management of resources and optimal performance during the
2051/// join process.
2052#[derive(Clone, Debug)]
2053pub enum SHJStreamState {
2054    /// Indicates that the next step should pull from the right side of the join.
2055    PullRight,
2056
2057    /// Indicates that the next step should pull from the left side of the join.
2058    PullLeft,
2059
2060    /// State representing that the right side of the join has been fully processed.
2061    RightExhausted,
2062
2063    /// State representing that the left side of the join has been fully processed.
2064    LeftExhausted,
2065
2066    /// Represents a state where both sides of the join are exhausted.
2067    ///
2068    /// The `final_result` field indicates whether the join operation has
2069    /// produced a final result or not.
2070    BothExhausted { final_result: bool },
2071}
2072
2073#[cfg(test)]
2074mod tests {
2075    use std::collections::HashMap;
2076    use std::sync::{LazyLock, Mutex};
2077
2078    use super::*;
2079    use crate::joins::test_utils::{
2080        build_sides_record_batches, compare_batches, complicated_filter,
2081        create_memory_table, join_expr_tests_fixture_f64, join_expr_tests_fixture_i32,
2082        join_expr_tests_fixture_temporal, partitioned_hash_join_with_filter,
2083        partitioned_sym_join_with_filter, split_record_batches,
2084    };
2085
2086    use arrow::compute::SortOptions;
2087    use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit};
2088    use datafusion_common::ScalarValue;
2089    use datafusion_execution::config::SessionConfig;
2090    use datafusion_expr::Operator;
2091    use datafusion_physical_expr::expressions::{Column, binary, col, lit};
2092    use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
2093
2094    use rstest::*;
2095
2096    const TABLE_SIZE: i32 = 30;
2097
2098    type TableKey = (i32, i32, usize); // (cardinality.0, cardinality.1, batch_size)
2099    type TableValue = (Vec<RecordBatch>, Vec<RecordBatch>); // (left, right)
2100
2101    // Cache for storing tables
2102    static TABLE_CACHE: LazyLock<Mutex<HashMap<TableKey, TableValue>>> =
2103        LazyLock::new(|| Mutex::new(HashMap::new()));
2104
2105    fn get_or_create_table(
2106        cardinality: (i32, i32),
2107        batch_size: usize,
2108    ) -> Result<TableValue> {
2109        {
2110            let cache = TABLE_CACHE.lock().unwrap();
2111            if let Some(table) = cache.get(&(cardinality.0, cardinality.1, batch_size)) {
2112                return Ok(table.clone());
2113            }
2114        }
2115
2116        // If not, create the table
2117        let (left_batch, right_batch) =
2118            build_sides_record_batches(TABLE_SIZE, cardinality)?;
2119
2120        let (left_partition, right_partition) = (
2121            split_record_batches(&left_batch, batch_size)?,
2122            split_record_batches(&right_batch, batch_size)?,
2123        );
2124
2125        // Lock the cache again and store the table
2126        let mut cache = TABLE_CACHE.lock().unwrap();
2127
2128        // Store the table in the cache
2129        cache.insert(
2130            (cardinality.0, cardinality.1, batch_size),
2131            (left_partition.clone(), right_partition.clone()),
2132        );
2133
2134        Ok((left_partition, right_partition))
2135    }
2136
2137    pub async fn experiment(
2138        left: Arc<dyn ExecutionPlan>,
2139        right: Arc<dyn ExecutionPlan>,
2140        filter: Option<JoinFilter>,
2141        join_type: JoinType,
2142        on: JoinOn,
2143        task_ctx: Arc<TaskContext>,
2144    ) -> Result<()> {
2145        let first_batches = partitioned_sym_join_with_filter(
2146            Arc::clone(&left),
2147            Arc::clone(&right),
2148            on.clone(),
2149            filter.clone(),
2150            &join_type,
2151            NullEquality::NullEqualsNothing,
2152            Arc::clone(&task_ctx),
2153        )
2154        .await?;
2155        let second_batches = partitioned_hash_join_with_filter(
2156            left,
2157            right,
2158            on,
2159            filter,
2160            &join_type,
2161            NullEquality::NullEqualsNothing,
2162            task_ctx,
2163        )
2164        .await?;
2165        compare_batches(&first_batches, &second_batches);
2166        Ok(())
2167    }
2168
2169    #[rstest]
2170    #[tokio::test(flavor = "multi_thread")]
2171    async fn complex_join_all_one_ascending_numeric(
2172        #[values(
2173            JoinType::Inner,
2174            JoinType::Left,
2175            JoinType::Right,
2176            JoinType::RightSemi,
2177            JoinType::LeftSemi,
2178            JoinType::LeftAnti,
2179            JoinType::LeftMark,
2180            JoinType::RightAnti,
2181            JoinType::RightMark,
2182            JoinType::Full
2183        )]
2184        join_type: JoinType,
2185        #[values(
2186        (4, 5),
2187        (12, 17),
2188        )]
2189        cardinality: (i32, i32),
2190    ) -> Result<()> {
2191        // a + b > c + 10 AND a + b < c + 100
2192        let task_ctx = Arc::new(TaskContext::default());
2193
2194        let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2195
2196        let left_schema = &left_partition[0].schema();
2197        let right_schema = &right_partition[0].schema();
2198
2199        let left_sorted = [PhysicalSortExpr {
2200            expr: binary(
2201                col("la1", left_schema)?,
2202                Operator::Plus,
2203                col("la2", left_schema)?,
2204                left_schema,
2205            )?,
2206            options: SortOptions::default(),
2207        }]
2208        .into();
2209        let right_sorted = [PhysicalSortExpr {
2210            expr: col("ra1", right_schema)?,
2211            options: SortOptions::default(),
2212        }]
2213        .into();
2214        let (left, right) = create_memory_table(
2215            left_partition,
2216            right_partition,
2217            vec![left_sorted],
2218            vec![right_sorted],
2219        )?;
2220
2221        let on = vec![(
2222            binary(
2223                col("lc1", left_schema)?,
2224                Operator::Plus,
2225                lit(ScalarValue::Int32(Some(1))),
2226                left_schema,
2227            )?,
2228            Arc::new(Column::new_with_schema("rc1", right_schema)?) as _,
2229        )];
2230
2231        let intermediate_schema = Schema::new(vec![
2232            Field::new("0", DataType::Int32, true),
2233            Field::new("1", DataType::Int32, true),
2234            Field::new("2", DataType::Int32, true),
2235        ]);
2236        let filter_expr = complicated_filter(&intermediate_schema)?;
2237        let column_indices = vec![
2238            ColumnIndex {
2239                index: left_schema.index_of("la1")?,
2240                side: JoinSide::Left,
2241            },
2242            ColumnIndex {
2243                index: left_schema.index_of("la2")?,
2244                side: JoinSide::Left,
2245            },
2246            ColumnIndex {
2247                index: right_schema.index_of("ra1")?,
2248                side: JoinSide::Right,
2249            },
2250        ];
2251        let filter =
2252            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2253
2254        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2255        Ok(())
2256    }
2257
2258    #[rstest]
2259    #[tokio::test(flavor = "multi_thread")]
2260    async fn join_all_one_ascending_numeric(
2261        #[values(
2262            JoinType::Inner,
2263            JoinType::Left,
2264            JoinType::Right,
2265            JoinType::RightSemi,
2266            JoinType::LeftSemi,
2267            JoinType::LeftAnti,
2268            JoinType::LeftMark,
2269            JoinType::RightAnti,
2270            JoinType::RightMark,
2271            JoinType::Full
2272        )]
2273        join_type: JoinType,
2274        #[values(0, 1, 2, 3, 4, 5)] case_expr: usize,
2275    ) -> Result<()> {
2276        let task_ctx = Arc::new(TaskContext::default());
2277        let (left_partition, right_partition) = get_or_create_table((4, 5), 8)?;
2278
2279        let left_schema = &left_partition[0].schema();
2280        let right_schema = &right_partition[0].schema();
2281
2282        let left_sorted = [PhysicalSortExpr {
2283            expr: col("la1", left_schema)?,
2284            options: SortOptions::default(),
2285        }]
2286        .into();
2287        let right_sorted = [PhysicalSortExpr {
2288            expr: col("ra1", right_schema)?,
2289            options: SortOptions::default(),
2290        }]
2291        .into();
2292        let (left, right) = create_memory_table(
2293            left_partition,
2294            right_partition,
2295            vec![left_sorted],
2296            vec![right_sorted],
2297        )?;
2298
2299        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2300
2301        let intermediate_schema = Schema::new(vec![
2302            Field::new("left", DataType::Int32, true),
2303            Field::new("right", DataType::Int32, true),
2304        ]);
2305        let filter_expr = join_expr_tests_fixture_i32(
2306            case_expr,
2307            col("left", &intermediate_schema)?,
2308            col("right", &intermediate_schema)?,
2309        );
2310        let column_indices = vec![
2311            ColumnIndex {
2312                index: 0,
2313                side: JoinSide::Left,
2314            },
2315            ColumnIndex {
2316                index: 0,
2317                side: JoinSide::Right,
2318            },
2319        ];
2320        let filter =
2321            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2322
2323        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2324        Ok(())
2325    }
2326
2327    #[rstest]
2328    #[tokio::test(flavor = "multi_thread")]
2329    async fn join_without_sort_information(
2330        #[values(
2331            JoinType::Inner,
2332            JoinType::Left,
2333            JoinType::Right,
2334            JoinType::RightSemi,
2335            JoinType::LeftSemi,
2336            JoinType::LeftAnti,
2337            JoinType::LeftMark,
2338            JoinType::RightAnti,
2339            JoinType::RightMark,
2340            JoinType::Full
2341        )]
2342        join_type: JoinType,
2343        #[values(0, 1, 2, 3, 4, 5)] case_expr: usize,
2344    ) -> Result<()> {
2345        let task_ctx = Arc::new(TaskContext::default());
2346        let (left_partition, right_partition) = get_or_create_table((4, 5), 8)?;
2347
2348        let left_schema = &left_partition[0].schema();
2349        let right_schema = &right_partition[0].schema();
2350        let (left, right) =
2351            create_memory_table(left_partition, right_partition, vec![], vec![])?;
2352
2353        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2354
2355        let intermediate_schema = Schema::new(vec![
2356            Field::new("left", DataType::Int32, true),
2357            Field::new("right", DataType::Int32, true),
2358        ]);
2359        let filter_expr = join_expr_tests_fixture_i32(
2360            case_expr,
2361            col("left", &intermediate_schema)?,
2362            col("right", &intermediate_schema)?,
2363        );
2364        let column_indices = vec![
2365            ColumnIndex {
2366                index: 5,
2367                side: JoinSide::Left,
2368            },
2369            ColumnIndex {
2370                index: 5,
2371                side: JoinSide::Right,
2372            },
2373        ];
2374        let filter =
2375            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2376
2377        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2378        Ok(())
2379    }
2380
2381    #[rstest]
2382    #[tokio::test(flavor = "multi_thread")]
2383    async fn join_without_filter(
2384        #[values(
2385            JoinType::Inner,
2386            JoinType::Left,
2387            JoinType::Right,
2388            JoinType::RightSemi,
2389            JoinType::LeftSemi,
2390            JoinType::LeftAnti,
2391            JoinType::LeftMark,
2392            JoinType::RightAnti,
2393            JoinType::RightMark,
2394            JoinType::Full
2395        )]
2396        join_type: JoinType,
2397    ) -> Result<()> {
2398        let task_ctx = Arc::new(TaskContext::default());
2399        let (left_partition, right_partition) = get_or_create_table((11, 21), 8)?;
2400        let left_schema = &left_partition[0].schema();
2401        let right_schema = &right_partition[0].schema();
2402        let (left, right) =
2403            create_memory_table(left_partition, right_partition, vec![], vec![])?;
2404
2405        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2406        experiment(left, right, None, join_type, on, task_ctx).await?;
2407        Ok(())
2408    }
2409
2410    #[rstest]
2411    #[tokio::test(flavor = "multi_thread")]
2412    async fn join_all_one_descending_numeric_particular(
2413        #[values(
2414            JoinType::Inner,
2415            JoinType::Left,
2416            JoinType::Right,
2417            JoinType::RightSemi,
2418            JoinType::LeftSemi,
2419            JoinType::LeftAnti,
2420            JoinType::LeftMark,
2421            JoinType::RightAnti,
2422            JoinType::RightMark,
2423            JoinType::Full
2424        )]
2425        join_type: JoinType,
2426        #[values(0, 1, 2, 3, 4, 5)] case_expr: usize,
2427    ) -> Result<()> {
2428        let task_ctx = Arc::new(TaskContext::default());
2429        let (left_partition, right_partition) = get_or_create_table((11, 21), 8)?;
2430        let left_schema = &left_partition[0].schema();
2431        let right_schema = &right_partition[0].schema();
2432        let left_sorted = [PhysicalSortExpr {
2433            expr: col("la1_des", left_schema)?,
2434            options: SortOptions {
2435                descending: true,
2436                nulls_first: true,
2437            },
2438        }]
2439        .into();
2440        let right_sorted = [PhysicalSortExpr {
2441            expr: col("ra1_des", right_schema)?,
2442            options: SortOptions {
2443                descending: true,
2444                nulls_first: true,
2445            },
2446        }]
2447        .into();
2448        let (left, right) = create_memory_table(
2449            left_partition,
2450            right_partition,
2451            vec![left_sorted],
2452            vec![right_sorted],
2453        )?;
2454
2455        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2456
2457        let intermediate_schema = Schema::new(vec![
2458            Field::new("left", DataType::Int32, true),
2459            Field::new("right", DataType::Int32, true),
2460        ]);
2461        let filter_expr = join_expr_tests_fixture_i32(
2462            case_expr,
2463            col("left", &intermediate_schema)?,
2464            col("right", &intermediate_schema)?,
2465        );
2466        let column_indices = vec![
2467            ColumnIndex {
2468                index: 5,
2469                side: JoinSide::Left,
2470            },
2471            ColumnIndex {
2472                index: 5,
2473                side: JoinSide::Right,
2474            },
2475        ];
2476        let filter =
2477            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2478
2479        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2480        Ok(())
2481    }
2482
2483    #[tokio::test(flavor = "multi_thread")]
2484    async fn build_null_columns_first() -> Result<()> {
2485        let join_type = JoinType::Full;
2486        let case_expr = 1;
2487        let session_config = SessionConfig::new().with_repartition_joins(false);
2488        let task_ctx = TaskContext::default().with_session_config(session_config);
2489        let task_ctx = Arc::new(task_ctx);
2490        let (left_partition, right_partition) = get_or_create_table((10, 11), 8)?;
2491        let left_schema = &left_partition[0].schema();
2492        let right_schema = &right_partition[0].schema();
2493        let left_sorted = [PhysicalSortExpr {
2494            expr: col("l_asc_null_first", left_schema)?,
2495            options: SortOptions {
2496                descending: false,
2497                nulls_first: true,
2498            },
2499        }]
2500        .into();
2501        let right_sorted = [PhysicalSortExpr {
2502            expr: col("r_asc_null_first", right_schema)?,
2503            options: SortOptions {
2504                descending: false,
2505                nulls_first: true,
2506            },
2507        }]
2508        .into();
2509        let (left, right) = create_memory_table(
2510            left_partition,
2511            right_partition,
2512            vec![left_sorted],
2513            vec![right_sorted],
2514        )?;
2515
2516        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2517
2518        let intermediate_schema = Schema::new(vec![
2519            Field::new("left", DataType::Int32, true),
2520            Field::new("right", DataType::Int32, true),
2521        ]);
2522        let filter_expr = join_expr_tests_fixture_i32(
2523            case_expr,
2524            col("left", &intermediate_schema)?,
2525            col("right", &intermediate_schema)?,
2526        );
2527        let column_indices = vec![
2528            ColumnIndex {
2529                index: 6,
2530                side: JoinSide::Left,
2531            },
2532            ColumnIndex {
2533                index: 6,
2534                side: JoinSide::Right,
2535            },
2536        ];
2537        let filter =
2538            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2539        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2540        Ok(())
2541    }
2542
2543    #[tokio::test(flavor = "multi_thread")]
2544    async fn build_null_columns_last() -> Result<()> {
2545        let join_type = JoinType::Full;
2546        let case_expr = 1;
2547        let session_config = SessionConfig::new().with_repartition_joins(false);
2548        let task_ctx = TaskContext::default().with_session_config(session_config);
2549        let task_ctx = Arc::new(task_ctx);
2550        let (left_partition, right_partition) = get_or_create_table((10, 11), 8)?;
2551
2552        let left_schema = &left_partition[0].schema();
2553        let right_schema = &right_partition[0].schema();
2554        let left_sorted = [PhysicalSortExpr {
2555            expr: col("l_asc_null_last", left_schema)?,
2556            options: SortOptions {
2557                descending: false,
2558                nulls_first: false,
2559            },
2560        }]
2561        .into();
2562        let right_sorted = [PhysicalSortExpr {
2563            expr: col("r_asc_null_last", right_schema)?,
2564            options: SortOptions {
2565                descending: false,
2566                nulls_first: false,
2567            },
2568        }]
2569        .into();
2570        let (left, right) = create_memory_table(
2571            left_partition,
2572            right_partition,
2573            vec![left_sorted],
2574            vec![right_sorted],
2575        )?;
2576
2577        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2578
2579        let intermediate_schema = Schema::new(vec![
2580            Field::new("left", DataType::Int32, true),
2581            Field::new("right", DataType::Int32, true),
2582        ]);
2583        let filter_expr = join_expr_tests_fixture_i32(
2584            case_expr,
2585            col("left", &intermediate_schema)?,
2586            col("right", &intermediate_schema)?,
2587        );
2588        let column_indices = vec![
2589            ColumnIndex {
2590                index: 7,
2591                side: JoinSide::Left,
2592            },
2593            ColumnIndex {
2594                index: 7,
2595                side: JoinSide::Right,
2596            },
2597        ];
2598        let filter =
2599            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2600
2601        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2602        Ok(())
2603    }
2604
2605    #[tokio::test(flavor = "multi_thread")]
2606    async fn build_null_columns_first_descending() -> Result<()> {
2607        let join_type = JoinType::Full;
2608        let cardinality = (10, 11);
2609        let case_expr = 1;
2610        let session_config = SessionConfig::new().with_repartition_joins(false);
2611        let task_ctx = TaskContext::default().with_session_config(session_config);
2612        let task_ctx = Arc::new(task_ctx);
2613        let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2614
2615        let left_schema = &left_partition[0].schema();
2616        let right_schema = &right_partition[0].schema();
2617        let left_sorted = [PhysicalSortExpr {
2618            expr: col("l_desc_null_first", left_schema)?,
2619            options: SortOptions {
2620                descending: true,
2621                nulls_first: true,
2622            },
2623        }]
2624        .into();
2625        let right_sorted = [PhysicalSortExpr {
2626            expr: col("r_desc_null_first", right_schema)?,
2627            options: SortOptions {
2628                descending: true,
2629                nulls_first: true,
2630            },
2631        }]
2632        .into();
2633        let (left, right) = create_memory_table(
2634            left_partition,
2635            right_partition,
2636            vec![left_sorted],
2637            vec![right_sorted],
2638        )?;
2639
2640        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2641
2642        let intermediate_schema = Schema::new(vec![
2643            Field::new("left", DataType::Int32, true),
2644            Field::new("right", DataType::Int32, true),
2645        ]);
2646        let filter_expr = join_expr_tests_fixture_i32(
2647            case_expr,
2648            col("left", &intermediate_schema)?,
2649            col("right", &intermediate_schema)?,
2650        );
2651        let column_indices = vec![
2652            ColumnIndex {
2653                index: 8,
2654                side: JoinSide::Left,
2655            },
2656            ColumnIndex {
2657                index: 8,
2658                side: JoinSide::Right,
2659            },
2660        ];
2661        let filter =
2662            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2663
2664        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2665        Ok(())
2666    }
2667
2668    #[tokio::test(flavor = "multi_thread")]
2669    async fn complex_join_all_one_ascending_numeric_missing_stat() -> Result<()> {
2670        let cardinality = (3, 4);
2671        let join_type = JoinType::Full;
2672
2673        // a + b > c + 10 AND a + b < c + 100
2674        let session_config = SessionConfig::new().with_repartition_joins(false);
2675        let task_ctx = TaskContext::default().with_session_config(session_config);
2676        let task_ctx = Arc::new(task_ctx);
2677        let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2678
2679        let left_schema = &left_partition[0].schema();
2680        let right_schema = &right_partition[0].schema();
2681        let left_sorted = [PhysicalSortExpr {
2682            expr: col("la1", left_schema)?,
2683            options: SortOptions::default(),
2684        }]
2685        .into();
2686        let right_sorted = [PhysicalSortExpr {
2687            expr: col("ra1", right_schema)?,
2688            options: SortOptions::default(),
2689        }]
2690        .into();
2691        let (left, right) = create_memory_table(
2692            left_partition,
2693            right_partition,
2694            vec![left_sorted],
2695            vec![right_sorted],
2696        )?;
2697
2698        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2699
2700        let intermediate_schema = Schema::new(vec![
2701            Field::new("0", DataType::Int32, true),
2702            Field::new("1", DataType::Int32, true),
2703            Field::new("2", DataType::Int32, true),
2704        ]);
2705        let filter_expr = complicated_filter(&intermediate_schema)?;
2706        let column_indices = vec![
2707            ColumnIndex {
2708                index: 0,
2709                side: JoinSide::Left,
2710            },
2711            ColumnIndex {
2712                index: 4,
2713                side: JoinSide::Left,
2714            },
2715            ColumnIndex {
2716                index: 0,
2717                side: JoinSide::Right,
2718            },
2719        ];
2720        let filter =
2721            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2722
2723        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2724        Ok(())
2725    }
2726
2727    #[tokio::test(flavor = "multi_thread")]
2728    async fn complex_join_all_one_ascending_equivalence() -> Result<()> {
2729        let cardinality = (3, 4);
2730        let join_type = JoinType::Full;
2731
2732        // a + b > c + 10 AND a + b < c + 100
2733        let config = SessionConfig::new().with_repartition_joins(false);
2734        // let session_ctx = SessionContext::with_config(config);
2735        // let task_ctx = session_ctx.task_ctx();
2736        let task_ctx = Arc::new(TaskContext::default().with_session_config(config));
2737        let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2738        let left_schema = &left_partition[0].schema();
2739        let right_schema = &right_partition[0].schema();
2740        let left_sorted = vec![
2741            [PhysicalSortExpr {
2742                expr: col("la1", left_schema)?,
2743                options: SortOptions::default(),
2744            }]
2745            .into(),
2746            [PhysicalSortExpr {
2747                expr: col("la2", left_schema)?,
2748                options: SortOptions::default(),
2749            }]
2750            .into(),
2751        ];
2752
2753        let right_sorted = [PhysicalSortExpr {
2754            expr: col("ra1", right_schema)?,
2755            options: SortOptions::default(),
2756        }]
2757        .into();
2758
2759        let (left, right) = create_memory_table(
2760            left_partition,
2761            right_partition,
2762            left_sorted,
2763            vec![right_sorted],
2764        )?;
2765
2766        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2767
2768        let intermediate_schema = Schema::new(vec![
2769            Field::new("0", DataType::Int32, true),
2770            Field::new("1", DataType::Int32, true),
2771            Field::new("2", DataType::Int32, true),
2772        ]);
2773        let filter_expr = complicated_filter(&intermediate_schema)?;
2774        let column_indices = vec![
2775            ColumnIndex {
2776                index: 0,
2777                side: JoinSide::Left,
2778            },
2779            ColumnIndex {
2780                index: 4,
2781                side: JoinSide::Left,
2782            },
2783            ColumnIndex {
2784                index: 0,
2785                side: JoinSide::Right,
2786            },
2787        ];
2788        let filter =
2789            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2790
2791        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2792        Ok(())
2793    }
2794
2795    #[rstest]
2796    #[tokio::test(flavor = "multi_thread")]
2797    async fn testing_with_temporal_columns(
2798        #[values(
2799            JoinType::Inner,
2800            JoinType::Left,
2801            JoinType::Right,
2802            JoinType::RightSemi,
2803            JoinType::LeftSemi,
2804            JoinType::LeftAnti,
2805            JoinType::LeftMark,
2806            JoinType::RightAnti,
2807            JoinType::RightMark,
2808            JoinType::Full
2809        )]
2810        join_type: JoinType,
2811        #[values(
2812            (4, 5),
2813            (12, 17),
2814        )]
2815        cardinality: (i32, i32),
2816        #[values(0, 1, 2)] case_expr: usize,
2817    ) -> Result<()> {
2818        let session_config = SessionConfig::new().with_repartition_joins(false);
2819        let task_ctx = TaskContext::default().with_session_config(session_config);
2820        let task_ctx = Arc::new(task_ctx);
2821        let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2822
2823        let left_schema = &left_partition[0].schema();
2824        let right_schema = &right_partition[0].schema();
2825        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2826        let left_sorted = [PhysicalSortExpr {
2827            expr: col("lt1", left_schema)?,
2828            options: SortOptions {
2829                descending: false,
2830                nulls_first: true,
2831            },
2832        }]
2833        .into();
2834        let right_sorted = [PhysicalSortExpr {
2835            expr: col("rt1", right_schema)?,
2836            options: SortOptions {
2837                descending: false,
2838                nulls_first: true,
2839            },
2840        }]
2841        .into();
2842        let (left, right) = create_memory_table(
2843            left_partition,
2844            right_partition,
2845            vec![left_sorted],
2846            vec![right_sorted],
2847        )?;
2848        let intermediate_schema = Schema::new(vec![
2849            Field::new(
2850                "left",
2851                DataType::Timestamp(TimeUnit::Millisecond, None),
2852                false,
2853            ),
2854            Field::new(
2855                "right",
2856                DataType::Timestamp(TimeUnit::Millisecond, None),
2857                false,
2858            ),
2859        ]);
2860        let filter_expr = join_expr_tests_fixture_temporal(
2861            case_expr,
2862            col("left", &intermediate_schema)?,
2863            col("right", &intermediate_schema)?,
2864            &intermediate_schema,
2865        )?;
2866        let column_indices = vec![
2867            ColumnIndex {
2868                index: 3,
2869                side: JoinSide::Left,
2870            },
2871            ColumnIndex {
2872                index: 3,
2873                side: JoinSide::Right,
2874            },
2875        ];
2876        let filter =
2877            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2878        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2879        Ok(())
2880    }
2881
2882    #[rstest]
2883    #[tokio::test(flavor = "multi_thread")]
2884    async fn test_with_interval_columns(
2885        #[values(
2886            JoinType::Inner,
2887            JoinType::Left,
2888            JoinType::Right,
2889            JoinType::RightSemi,
2890            JoinType::LeftSemi,
2891            JoinType::LeftAnti,
2892            JoinType::LeftMark,
2893            JoinType::RightAnti,
2894            JoinType::RightMark,
2895            JoinType::Full
2896        )]
2897        join_type: JoinType,
2898        #[values(
2899            (4, 5),
2900            (12, 17),
2901        )]
2902        cardinality: (i32, i32),
2903    ) -> Result<()> {
2904        let session_config = SessionConfig::new().with_repartition_joins(false);
2905        let task_ctx = TaskContext::default().with_session_config(session_config);
2906        let task_ctx = Arc::new(task_ctx);
2907        let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2908
2909        let left_schema = &left_partition[0].schema();
2910        let right_schema = &right_partition[0].schema();
2911        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2912        let left_sorted = [PhysicalSortExpr {
2913            expr: col("li1", left_schema)?,
2914            options: SortOptions {
2915                descending: false,
2916                nulls_first: true,
2917            },
2918        }]
2919        .into();
2920        let right_sorted = [PhysicalSortExpr {
2921            expr: col("ri1", right_schema)?,
2922            options: SortOptions {
2923                descending: false,
2924                nulls_first: true,
2925            },
2926        }]
2927        .into();
2928        let (left, right) = create_memory_table(
2929            left_partition,
2930            right_partition,
2931            vec![left_sorted],
2932            vec![right_sorted],
2933        )?;
2934        let intermediate_schema = Schema::new(vec![
2935            Field::new("left", DataType::Interval(IntervalUnit::DayTime), false),
2936            Field::new("right", DataType::Interval(IntervalUnit::DayTime), false),
2937        ]);
2938        let filter_expr = join_expr_tests_fixture_temporal(
2939            0,
2940            col("left", &intermediate_schema)?,
2941            col("right", &intermediate_schema)?,
2942            &intermediate_schema,
2943        )?;
2944        let column_indices = vec![
2945            ColumnIndex {
2946                index: 9,
2947                side: JoinSide::Left,
2948            },
2949            ColumnIndex {
2950                index: 9,
2951                side: JoinSide::Right,
2952            },
2953        ];
2954        let filter =
2955            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2956        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2957
2958        Ok(())
2959    }
2960
2961    #[rstest]
2962    #[tokio::test(flavor = "multi_thread")]
2963    async fn testing_ascending_float_pruning(
2964        #[values(
2965            JoinType::Inner,
2966            JoinType::Left,
2967            JoinType::Right,
2968            JoinType::RightSemi,
2969            JoinType::LeftSemi,
2970            JoinType::LeftAnti,
2971            JoinType::LeftMark,
2972            JoinType::RightAnti,
2973            JoinType::RightMark,
2974            JoinType::Full
2975        )]
2976        join_type: JoinType,
2977        #[values(
2978            (4, 5),
2979            (12, 17),
2980        )]
2981        cardinality: (i32, i32),
2982        #[values(0, 1, 2, 3, 4, 5)] case_expr: usize,
2983    ) -> Result<()> {
2984        let session_config = SessionConfig::new().with_repartition_joins(false);
2985        let task_ctx = TaskContext::default().with_session_config(session_config);
2986        let task_ctx = Arc::new(task_ctx);
2987        let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2988
2989        let left_schema = &left_partition[0].schema();
2990        let right_schema = &right_partition[0].schema();
2991        let left_sorted = [PhysicalSortExpr {
2992            expr: col("l_float", left_schema)?,
2993            options: SortOptions::default(),
2994        }]
2995        .into();
2996        let right_sorted = [PhysicalSortExpr {
2997            expr: col("r_float", right_schema)?,
2998            options: SortOptions::default(),
2999        }]
3000        .into();
3001        let (left, right) = create_memory_table(
3002            left_partition,
3003            right_partition,
3004            vec![left_sorted],
3005            vec![right_sorted],
3006        )?;
3007
3008        let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
3009
3010        let intermediate_schema = Schema::new(vec![
3011            Field::new("left", DataType::Float64, true),
3012            Field::new("right", DataType::Float64, true),
3013        ]);
3014        let filter_expr = join_expr_tests_fixture_f64(
3015            case_expr,
3016            col("left", &intermediate_schema)?,
3017            col("right", &intermediate_schema)?,
3018        );
3019        let column_indices = vec![
3020            ColumnIndex {
3021                index: 10, // l_float
3022                side: JoinSide::Left,
3023            },
3024            ColumnIndex {
3025                index: 10, // r_float
3026                side: JoinSide::Right,
3027            },
3028        ];
3029        let filter =
3030            JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
3031
3032        experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
3033        Ok(())
3034    }
3035}