## 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/mod.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 std::borrow::Borrow;
    19
    use std::sync::Arc;
    20
    21
    use crate::PhysicalExpr;
    22
    23
    use arrow::compute::SortOptions;
    24
    use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
    25
    26
    mod class;
    27
    mod ordering;
    28
    mod properties;
    29
    30
    pub use class::{AcrossPartitions, ConstExpr, EquivalenceClass, EquivalenceGroup};
    31
    pub use ordering::OrderingEquivalenceClass;
    32
    // Re-export for backwards compatibility, we recommend importing from
    33
    // datafusion_physical_expr::projection instead
    34
    pub use crate::projection::{project_ordering, project_orderings, ProjectionMapping};
    35
    pub use properties::{
    36
        calculate_union, join_equivalence_properties, EquivalenceProperties,
    37
    };
    38
    39
    // Convert each tuple to a `PhysicalSortExpr` and construct a vector.
    40
    0
    pub fn convert_to_sort_exprs<T: Borrow<Arc<dyn PhysicalExpr>>>(
    41
    0
        args: &[(T, SortOptions)],
    42
    0
    ) -> Vec<PhysicalSortExpr> {
    43
    0
        args.iter()
    44
    0
            .map(|(expr, options)| PhysicalSortExpr::new(Arc::clone(expr.borrow()), *options))
    45
    0
            .collect()
    46
    0
    }
    47
    48
    // Convert each vector of tuples to a `LexOrdering`.
    49
    0
    pub fn convert_to_orderings<T: Borrow<Arc<dyn PhysicalExpr>>>(
    50
    0
        args: &[Vec<(T, SortOptions)>],
    51
    0
    ) -> Vec<LexOrdering> {
    52
    0
        args.iter()
    53
    0
            .filter_map(|sort_exprs| LexOrdering::new(convert_to_sort_exprs(sort_exprs)))
    54
    0
            .collect()
    55
    0
    }
    56
    57
    #[cfg(test)]
    58
    mod tests {
    59
        use super::*;
    60
        use crate::expressions::{col, Column};
    61
        use crate::{LexRequirement, PhysicalSortExpr};
    62
    63
        use arrow::compute::SortOptions;
    64
        use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
    65
        use datafusion_common::Result;
    66
        use datafusion_physical_expr_common::sort_expr::PhysicalSortRequirement;
    67
    68
        /// Converts a string to a physical sort expression
    69
        ///
    70
        /// # Example
    71
        /// * `"a"` -> (`"a"`, `SortOptions::default()`)
    72
        /// * `"a ASC"` -> (`"a"`, `SortOptions { descending: false, nulls_first: false }`)
    73
    0
        pub fn parse_sort_expr(name: &str, schema: &SchemaRef) -> PhysicalSortExpr {
    74
    0
            let mut parts = name.split_whitespace();
    75
    0
            let name = parts.next().expect("empty sort expression");
    76
    0
            let mut sort_expr = PhysicalSortExpr::new(
    77
    0
                col(name, schema).expect("invalid column name"),
    78
    0
                SortOptions::default(),
    79
            );
    80
    81
    0
            if let Some(options) = parts.next() {
    82
    0
                sort_expr = match options {
    83
    0
                    "ASC" => sort_expr.asc(),
    84
    0
                    "DESC" => sort_expr.desc(),
    85
    0
                    _ => panic!(
    86
    0
                        "unknown sort options. Expected 'ASC' or 'DESC', got {options}"
    87
                    ),
    88
                }
    89
    0
            }
    90
    91
    0
            assert!(
    92
    0
                parts.next().is_none(),
    93
    0
                "unexpected tokens in column name. Expected 'name' / 'name ASC' / 'name DESC' but got  '{name}'"
    94
            );
    95
    96
    0
            sort_expr
    97
    0
        }
    98
    99
        // Generate a schema which consists of 8 columns (a, b, c, d, e, f, g, h)
    100
    0
        pub fn create_test_schema() -> Result<SchemaRef> {
    101
    0
            let a = Field::new("a", DataType::Int32, true);
    102
    0
            let b = Field::new("b", DataType::Int32, true);
    103
    0
            let c = Field::new("c", DataType::Int32, true);
    104
    0
            let d = Field::new("d", DataType::Int32, true);
    105
    0
            let e = Field::new("e", DataType::Int32, true);
    106
    0
            let f = Field::new("f", DataType::Int32, true);
    107
    0
            let g = Field::new("g", DataType::Int32, true);
    108
    0
            let h = Field::new("h", DataType::Int32, true);
    109
    0
            let schema = Arc::new(Schema::new(vec![a, b, c, d, e, f, g, h]));
    110
    111
    0
            Ok(schema)
    112
    0
        }
    113
    114
        /// Construct a schema with following properties
    115
        /// Schema satisfies following orderings:
    116
        /// [a ASC], [d ASC, b ASC], [e DESC, f ASC, g ASC]
    117
        /// and
    118
        /// Column [a=c] (e.g they are aliases).
    119
    0
        pub fn create_test_params() -> Result<(SchemaRef, EquivalenceProperties)> {
    120
    0
            let test_schema = create_test_schema()?;
    121
    0
            let col_a = col("a", &test_schema)?;
    122
    0
            let col_b = col("b", &test_schema)?;
    123
    0
            let col_c = col("c", &test_schema)?;
    124
    0
            let col_d = col("d", &test_schema)?;
    125
    0
            let col_e = col("e", &test_schema)?;
    126
    0
            let col_f = col("f", &test_schema)?;
    127
    0
            let col_g = col("g", &test_schema)?;
    128
    0
            let mut eq_properties = EquivalenceProperties::new(Arc::clone(&test_schema));
    129
    0
            eq_properties.add_equal_conditions(Arc::clone(&col_a), Arc::clone(&col_c))?;
    130
    131
    0
            let option_asc = SortOptions {
    132
    0
                descending: false,
    133
    0
                nulls_first: false,
    134
    0
            };
    135
    0
            let option_desc = SortOptions {
    136
    0
                descending: true,
    137
    0
                nulls_first: true,
    138
    0
            };
    139
    0
            let orderings = vec![
    140
                // [a ASC]
    141
    0
                vec![(col_a, option_asc)],
    142
                // [d ASC, b ASC]
    143
    0
                vec![(col_d, option_asc), (col_b, option_asc)],
    144
                // [e DESC, f ASC, g ASC]
    145
    0
                vec![
    146
    0
                    (col_e, option_desc),
    147
    0
                    (col_f, option_asc),
    148
    0
                    (col_g, option_asc),
    149
                ],
    150
            ];
    151
    0
            let orderings = convert_to_orderings(&orderings);
    152
    0
            eq_properties.add_orderings(orderings);
    153
    0
            Ok((test_schema, eq_properties))
    154
    0
        }
    155
    156
        // Convert each tuple to a `PhysicalSortRequirement` and construct a
    157
        // a `LexRequirement` from them.
    158
    0
        pub fn convert_to_sort_reqs(
    159
    0
            args: &[(&Arc<dyn PhysicalExpr>, Option<SortOptions>)],
    160
    0
        ) -> LexRequirement {
    161
    0
            let exprs = args.iter().map(|(expr, options)| {
    162
    0
                PhysicalSortRequirement::new(Arc::clone(*expr), *options)
    163
    0
            });
    164
    0
            LexRequirement::new(exprs).unwrap()
    165
    0
        }
    166
    167
        #[test]
    168
    0
        fn add_equal_conditions_test() -> Result<()> {
    169
    0
            let schema = Arc::new(Schema::new(vec![
    170
    0
                Field::new("a", DataType::Int64, true),
    171
    0
                Field::new("b", DataType::Int64, true),
    172
    0
                Field::new("c", DataType::Int64, true),
    173
    0
                Field::new("x", DataType::Int64, true),
    174
    0
                Field::new("y", DataType::Int64, true),
    175
            ]));
    176
    177
    0
            let mut eq_properties = EquivalenceProperties::new(schema);
    178
    0
            let col_a = Arc::new(Column::new("a", 0)) as _;
    179
    0
            let col_b = Arc::new(Column::new("b", 1)) as _;
    180
    0
            let col_c = Arc::new(Column::new("c", 2)) as _;
    181
    0
            let col_x = Arc::new(Column::new("x", 3)) as _;
    182
    0
            let col_y = Arc::new(Column::new("y", 4)) as _;
    183
    184
            // a and b are aliases
    185
    0
            eq_properties.add_equal_conditions(Arc::clone(&col_a), Arc::clone(&col_b))?;
    186
    0
            assert_eq!(eq_properties.eq_group().len(), 1);
    187
    188
            // This new entry is redundant, size shouldn't increase
    189
    0
            eq_properties.add_equal_conditions(Arc::clone(&col_b), Arc::clone(&col_a))?;
    190
    0
            assert_eq!(eq_properties.eq_group().len(), 1);
    191
    0
            let eq_groups = eq_properties.eq_group().iter().next().unwrap();
    192
    0
            assert_eq!(eq_groups.len(), 2);
    193
    0
            assert!(eq_groups.contains(&col_a));
    194
    0
            assert!(eq_groups.contains(&col_b));
    195
    196
            // b and c are aliases. Existing equivalence class should expand,
    197
            // however there shouldn't be any new equivalence class
    198
    0
            eq_properties.add_equal_conditions(Arc::clone(&col_b), Arc::clone(&col_c))?;
    199
    0
            assert_eq!(eq_properties.eq_group().len(), 1);
    200
    0
            let eq_groups = eq_properties.eq_group().iter().next().unwrap();
    201
    0
            assert_eq!(eq_groups.len(), 3);
    202
    0
            assert!(eq_groups.contains(&col_a));
    203
    0
            assert!(eq_groups.contains(&col_b));
    204
    0
            assert!(eq_groups.contains(&col_c));
    205
    206
            // This is a new set of equality. Hence equivalent class count should be 2.
    207
    0
            eq_properties.add_equal_conditions(Arc::clone(&col_x), Arc::clone(&col_y))?;
    208
    0
            assert_eq!(eq_properties.eq_group().len(), 2);
    209
    210
            // This equality bridges distinct equality sets.
    211
            // Hence equivalent class count should decrease from 2 to 1.
    212
    0
            eq_properties.add_equal_conditions(Arc::clone(&col_x), Arc::clone(&col_a))?;
    213
    0
            assert_eq!(eq_properties.eq_group().len(), 1);
    214
    0
            let eq_groups = eq_properties.eq_group().iter().next().unwrap();
    215
    0
            assert_eq!(eq_groups.len(), 5);
    216
    0
            assert!(eq_groups.contains(&col_a));
    217
    0
            assert!(eq_groups.contains(&col_b));
    218
    0
            assert!(eq_groups.contains(&col_c));
    219
    0
            assert!(eq_groups.contains(&col_x));
    220
    0
            assert!(eq_groups.contains(&col_y));
    221
    222
    0
            Ok(())
    223
    0
        }
    224
    }