## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/physical-expr/src/equivalence/properties/joins.rs

    Line
    Count
    Source
    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
    use super::EquivalenceProperties;
    19
    use crate::{equivalence::OrderingEquivalenceClass, PhysicalExprRef};
    20
    21
    use arrow::datatypes::SchemaRef;
    22
    use datafusion_common::{JoinSide, JoinType, Result};
    23
    24
    /// Calculate ordering equivalence properties for the given join operation.
    25
    0
    pub fn join_equivalence_properties(
    26
    0
        left: EquivalenceProperties,
    27
    0
        right: EquivalenceProperties,
    28
    0
        join_type: &JoinType,
    29
    0
        join_schema: SchemaRef,
    30
    0
        maintains_input_order: &[bool],
    31
    0
        probe_side: Option<JoinSide>,
    32
    0
        on: &[(PhysicalExprRef, PhysicalExprRef)],
    33
    0
    ) -> Result<EquivalenceProperties> {
    34
    0
        let left_size = left.schema.fields.len();
    35
    0
        let mut result = EquivalenceProperties::new(join_schema);
    36
    0
        result.add_equivalence_group(left.eq_group().join(
    37
    0
            right.eq_group(),
    38
    0
            join_type,
    39
    0
            left_size,
    40
    0
            on,
    41
    0
        )?)?;
    42
    43
        let EquivalenceProperties {
    44
    0
            oeq_class: left_oeq_class,
    45
            ..
    46
    0
        } = left;
    47
        let EquivalenceProperties {
    48
    0
            oeq_class: mut right_oeq_class,
    49
            ..
    50
    0
        } = right;
    51
    0
        match maintains_input_order {
    52
    0
            [true, false] => {
    53
                // In this special case, right side ordering can be prefixed with
    54
                // the left side ordering.
    55
    0
                if matches!(join_type, JoinType::Inner | JoinType::Left)
    56
    0
                    && probe_side == Some(JoinSide::Left)
    57
                {
    58
    0
                    updated_right_ordering_equivalence_class(
    59
    0
                        &mut right_oeq_class,
    60
    0
                        join_type,
    61
    0
                        left_size,
    62
    0
                    )?;
    63
    64
                    // Right side ordering equivalence properties should be prepended
    65
                    // with those of the left side while constructing output ordering
    66
                    // equivalence properties since stream side is the left side.
    67
                    //
    68
                    // For example, if the right side ordering equivalences contain
    69
                    // `b ASC`, and the left side ordering equivalences contain `a ASC`,
    70
                    // then we should add `a ASC, b ASC` to the ordering equivalences
    71
                    // of the join output.
    72
    0
                    let out_oeq_class = left_oeq_class.join_suffix(&right_oeq_class);
    73
    0
                    result.add_orderings(out_oeq_class);
    74
    0
                } else {
    75
    0
                    result.add_orderings(left_oeq_class);
    76
    0
                }
    77
            }
    78
            [false, true] => {
    79
    0
                updated_right_ordering_equivalence_class(
    80
    0
                    &mut right_oeq_class,
    81
    0
                    join_type,
    82
    0
                    left_size,
    83
    0
                )?;
    84
                // In this special case, left side ordering can be prefixed with
    85
                // the right side ordering.
    86
    0
                if matches!(join_type, JoinType::Inner | JoinType::Right)
    87
    0
                    && probe_side == Some(JoinSide::Right)
    88
    0
                {
    89
    0
                    // Left side ordering equivalence properties should be prepended
    90
    0
                    // with those of the right side while constructing output ordering
    91
    0
                    // equivalence properties since stream side is the right side.
    92
    0
                    //
    93
    0
                    // For example, if the left side ordering equivalences contain
    94
    0
                    // `a ASC`, and the right side ordering equivalences contain `b ASC`,
    95
    0
                    // then we should add `b ASC, a ASC` to the ordering equivalences
    96
    0
                    // of the join output.
    97
    0
                    let out_oeq_class = right_oeq_class.join_suffix(&left_oeq_class);
    98
    0
                    result.add_orderings(out_oeq_class);
    99
    0
                } else {
    100
    0
                    result.add_orderings(right_oeq_class);
    101
    0
                }
    102
            }
    103
    0
            [false, false] => {}
    104
    0
            [true, true] => unreachable!("Cannot maintain ordering of both sides"),
    105
    0
            _ => unreachable!("Join operators can not have more than two children"),
    106
        }
    107
    0
        Ok(result)
    108
    0
    }
    109
    110
    /// In the context of a join, update the right side `OrderingEquivalenceClass`
    111
    /// so that they point to valid indices in the join output schema.
    112
    ///
    113
    /// To do so, we increment column indices by the size of the left table when
    114
    /// join schema consists of a combination of the left and right schemas. This
    115
    /// is the case for `Inner`, `Left`, `Full` and `Right` joins. For other cases,
    116
    /// indices do not change.
    117
    0
    pub fn updated_right_ordering_equivalence_class(
    118
    0
        right_oeq_class: &mut OrderingEquivalenceClass,
    119
    0
        join_type: &JoinType,
    120
    0
        left_size: usize,
    121
    0
    ) -> Result<()> {
    122
    0
        if matches!(
    123
    0
            join_type,
    124
            JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right
    125
        ) {
    126
    0
            right_oeq_class.add_offset(left_size as _)?;
    127
    0
        }
    128
    0
        Ok(())
    129
    0
    }
    130
    131
    #[cfg(test)]
    132
    mod tests {
    133
        use std::sync::Arc;
    134
    135
        use super::*;
    136
        use crate::equivalence::convert_to_orderings;
    137
        use crate::equivalence::tests::create_test_schema;
    138
        use crate::expressions::col;
    139
        use crate::physical_expr::add_offset_to_expr;
    140
    141
        use arrow::compute::SortOptions;
    142
        use arrow::datatypes::{DataType, Field, Fields, Schema};
    143
        use datafusion_common::Result;
    144
    145
        #[test]
    146
    0
        fn test_join_equivalence_properties() -> Result<()> {
    147
    0
            let schema = create_test_schema()?;
    148
    0
            let col_a = &col("a", &schema)?;
    149
    0
            let col_b = &col("b", &schema)?;
    150
    0
            let col_c = &col("c", &schema)?;
    151
    0
            let offset = schema.fields.len() as _;
    152
    0
            let col_a2 = &add_offset_to_expr(Arc::clone(col_a), offset)?;
    153
    0
            let col_b2 = &add_offset_to_expr(Arc::clone(col_b), offset)?;
    154
    0
            let option_asc = SortOptions {
    155
    0
                descending: false,
    156
    0
                nulls_first: false,
    157
    0
            };
    158
    0
            let test_cases = vec![
    159
                // ------- TEST CASE 1 --------
    160
                // [a ASC], [b ASC]
    161
    0
                (
    162
    0
                    // [a ASC], [b ASC]
    163
    0
                    vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]],
    164
    0
                    // [a ASC], [b ASC]
    165
    0
                    vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]],
    166
    0
                    // expected [a ASC, a2 ASC], [a ASC, b2 ASC], [b ASC, a2 ASC], [b ASC, b2 ASC]
    167
    0
                    vec![
    168
    0
                        vec![(col_a, option_asc), (col_a2, option_asc)],
    169
    0
                        vec![(col_a, option_asc), (col_b2, option_asc)],
    170
    0
                        vec![(col_b, option_asc), (col_a2, option_asc)],
    171
    0
                        vec![(col_b, option_asc), (col_b2, option_asc)],
    172
    0
                    ],
    173
    0
                ),
    174
                // ------- TEST CASE 2 --------
    175
                // [a ASC], [b ASC]
    176
    0
                (
    177
    0
                    // [a ASC], [b ASC], [c ASC]
    178
    0
                    vec![
    179
    0
                        vec![(col_a, option_asc)],
    180
    0
                        vec![(col_b, option_asc)],
    181
    0
                        vec![(col_c, option_asc)],
    182
    0
                    ],
    183
    0
                    // [a ASC], [b ASC]
    184
    0
                    vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]],
    185
    0
                    // expected [a ASC, a2 ASC], [a ASC, b2 ASC], [b ASC, a2 ASC], [b ASC, b2 ASC], [c ASC, a2 ASC], [c ASC, b2 ASC]
    186
    0
                    vec![
    187
    0
                        vec![(col_a, option_asc), (col_a2, option_asc)],
    188
    0
                        vec![(col_a, option_asc), (col_b2, option_asc)],
    189
    0
                        vec![(col_b, option_asc), (col_a2, option_asc)],
    190
    0
                        vec![(col_b, option_asc), (col_b2, option_asc)],
    191
    0
                        vec![(col_c, option_asc), (col_a2, option_asc)],
    192
    0
                        vec![(col_c, option_asc), (col_b2, option_asc)],
    193
    0
                    ],
    194
    0
                ),
    195
            ];
    196
    0
            for (left_orderings, right_orderings, expected) in test_cases {
    197
    0
                let mut left_eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
    198
    0
                let mut right_eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
    199
    0
                let left_orderings = convert_to_orderings(&left_orderings);
    200
    0
                let right_orderings = convert_to_orderings(&right_orderings);
    201
    0
                let expected = convert_to_orderings(&expected);
    202
    0
                left_eq_properties.add_orderings(left_orderings);
    203
    0
                right_eq_properties.add_orderings(right_orderings);
    204
    0
                let join_eq = join_equivalence_properties(
    205
    0
                    left_eq_properties,
    206
    0
                    right_eq_properties,
    207
    0
                    &JoinType::Inner,
    208
    0
                    Arc::new(Schema::empty()),
    209
    0
                    &[true, false],
    210
    0
                    Some(JoinSide::Left),
    211
    0
                    &[],
    212
    0
                )?;
    213
    0
                let err_msg =
    214
    0
                    format!("expected: {:?}, actual:{:?}", expected, &join_eq.oeq_class);
    215
    0
                assert_eq!(join_eq.oeq_class.len(), expected.len(), "{err_msg}");
    216
    0
                for ordering in join_eq.oeq_class {
    217
    0
                    assert!(
    218
    0
                        expected.contains(&ordering),
    219
    0
                        "{err_msg}, ordering: {ordering:?}"
    220
                    );
    221
                }
    222
            }
    223
    0
            Ok(())
    224
    0
        }
    225
    226
        #[test]
    227
    0
        fn test_get_updated_right_ordering_equivalence_properties() -> Result<()> {
    228
    0
            let join_type = JoinType::Inner;
    229
            // Join right child schema
    230
    0
            let child_fields: Fields = ["x", "y", "z", "w"]
    231
    0
                .into_iter()
    232
    0
                .map(|name| Field::new(name, DataType::Int32, true))
    233
    0
                .collect();
    234
    0
            let child_schema = Schema::new(child_fields);
    235
    0
            let col_x = &col("x", &child_schema)?;
    236
    0
            let col_y = &col("y", &child_schema)?;
    237
    0
            let col_z = &col("z", &child_schema)?;
    238
    0
            let col_w = &col("w", &child_schema)?;
    239
    0
            let option_asc = SortOptions {
    240
    0
                descending: false,
    241
    0
                nulls_first: false,
    242
    0
            };
    243
            // [x ASC, y ASC], [z ASC, w ASC]
    244
    0
            let orderings = vec![
    245
    0
                vec![(col_x, option_asc), (col_y, option_asc)],
    246
    0
                vec![(col_z, option_asc), (col_w, option_asc)],
    247
            ];
    248
    0
            let orderings = convert_to_orderings(&orderings);
    249
            // Right child ordering equivalences
    250
    0
            let mut right_oeq_class = OrderingEquivalenceClass::from(orderings);
    251
    252
    0
            let left_columns_len = 4;
    253
    254
    0
            let fields: Fields = ["a", "b", "c", "d", "x", "y", "z", "w"]
    255
    0
                .into_iter()
    256
    0
                .map(|name| Field::new(name, DataType::Int32, true))
    257
    0
                .collect();
    258
    259
            // Join Schema
    260
    0
            let schema = Schema::new(fields);
    261
    0
            let col_a = col("a", &schema)?;
    262
    0
            let col_d = col("d", &schema)?;
    263
    0
            let col_x = col("x", &schema)?;
    264
    0
            let col_y = col("y", &schema)?;
    265
    0
            let col_z = col("z", &schema)?;
    266
    0
            let col_w = col("w", &schema)?;
    267
    268
    0
            let mut join_eq_properties = EquivalenceProperties::new(Arc::new(schema));
    269
            // a=x and d=w
    270
    0
            join_eq_properties.add_equal_conditions(col_a, Arc::clone(&col_x))?;
    271
    0
            join_eq_properties.add_equal_conditions(col_d, Arc::clone(&col_w))?;
    272
    273
    0
            updated_right_ordering_equivalence_class(
    274
    0
                &mut right_oeq_class,
    275
    0
                &join_type,
    276
    0
                left_columns_len,
    277
    0
            )?;
    278
    0
            join_eq_properties.add_orderings(right_oeq_class);
    279
    0
            let result = join_eq_properties.oeq_class().clone();
    280
    281
            // [x ASC, y ASC], [z ASC, w ASC]
    282
    0
            let orderings = vec![
    283
    0
                vec![(col_x, option_asc), (col_y, option_asc)],
    284
    0
                vec![(col_z, option_asc), (col_w, option_asc)],
    285
            ];
    286
    0
            let orderings = convert_to_orderings(&orderings);
    287
    0
            let expected = OrderingEquivalenceClass::from(orderings);
    288
    289
    0
            assert_eq!(result, expected);
    290
    291
    0
            Ok(())
    292
    0
        }
    293
    }