## Coverage Report

#### Created: 2024-10-13 08:39

    /Users/andrewlamb/Software/datafusion/datafusion/physical-plan/src/joins/symmetric_hash_join.rs

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