## 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/union.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::iter::Peekable;
    19
    use std::sync::Arc;
    20
    21
    use super::EquivalenceProperties;
    22
    use crate::equivalence::class::AcrossPartitions;
    23
    use crate::{ConstExpr, PhysicalSortExpr};
    24
    25
    use arrow::datatypes::SchemaRef;
    26
    use datafusion_common::{internal_err, Result};
    27
    use datafusion_physical_expr_common::sort_expr::LexOrdering;
    28
    29
    /// Computes the union (in the sense of `UnionExec`) `EquivalenceProperties`
    30
    /// of `lhs` and `rhs` according to the schema of `lhs`.
    31
    ///
    32
    /// Rules: The `UnionExec` does not interleave its inputs, instead it passes
    33
    /// each input partition from the children as its own output.
    34
    ///
    35
    /// Since the output equivalence properties are properties that are true for
    36
    /// *all* output partitions, that is the same as being true for all *input*
    37
    /// partitions.
    38
    0
    fn calculate_union_binary(
    39
    0
        lhs: EquivalenceProperties,
    40
    0
        mut rhs: EquivalenceProperties,
    41
    0
    ) -> Result<EquivalenceProperties> {
    42
        // Harmonize the schema of the rhs with the schema of the lhs (which is the accumulator schema):
    43
    0
        if !rhs.schema.eq(&lhs.schema) {
    44
    0
            rhs = rhs.with_new_schema(Arc::clone(&lhs.schema))?;
    45
    0
        }
    46
    47
        // First, calculate valid constants for the union. An expression is constant
    48
        // at the output of the union if it is constant in both sides with matching values.
    49
    0
        let rhs_constants = rhs.constants();
    50
    0
        let constants = lhs
    51
    0
            .constants()
    52
    0
            .into_iter()
    53
    0
            .filter_map(|lhs_const| {
    54
                // Find matching constant expression in RHS
    55
    0
                rhs_constants
    56
    0
                    .iter()
    57
    0
                    .find(|rhs_const| rhs_const.expr.eq(&lhs_const.expr))
    58
    0
                    .map(|rhs_const| {
    59
    0
                        let mut const_expr = lhs_const.clone();
    60
                        // If both sides have matching constant values, preserve it.
    61
                        // Otherwise, set fall back to heterogeneous values.
    62
    0
                        if lhs_const.across_partitions != rhs_const.across_partitions {
    63
    0
                            const_expr.across_partitions = AcrossPartitions::Heterogeneous;
    64
    0
                        }
    65
    0
                        const_expr
    66
    0
                    })
    67
    0
            })
    68
    0
            .collect::<Vec<_>>();
    69
    70
        // Next, calculate valid orderings for the union by searching for prefixes
    71
        // in both sides.
    72
    0
        let mut orderings = UnionEquivalentOrderingBuilder::new();
    73
    0
        orderings.add_satisfied_orderings(&lhs, &rhs)?;
    74
    0
        orderings.add_satisfied_orderings(&rhs, &lhs)?;
    75
    0
        let orderings = orderings.build();
    76
    77
    0
        let mut eq_properties = EquivalenceProperties::new(lhs.schema);
    78
    0
        eq_properties.add_constants(constants)?;
    79
    0
        eq_properties.add_orderings(orderings);
    80
    0
        Ok(eq_properties)
    81
    0
    }
    82
    83
    /// Calculates the union (in the sense of `UnionExec`) `EquivalenceProperties`
    84
    /// of the given `EquivalenceProperties` in `eqps` according to the given
    85
    /// output `schema` (which need not be the same with those of `lhs` and `rhs`
    86
    /// as details such as nullability may be different).
    87
    0
    pub fn calculate_union(
    88
    0
        eqps: Vec<EquivalenceProperties>,
    89
    0
        schema: SchemaRef,
    90
    0
    ) -> Result<EquivalenceProperties> {
    91
        // TODO: In some cases, we should be able to preserve some equivalence
    92
        //       classes. Add support for such cases.
    93
    0
        let mut iter = eqps.into_iter();
    94
    0
        let Some(mut acc) = iter.next() else {
    95
    0
            return internal_err!(
    96
    0
                "Cannot calculate EquivalenceProperties for a union with no inputs"
    97
            );
    98
        };
    99
    100
        // Harmonize the schema of the init with the schema of the union:
    101
    0
        if !acc.schema.eq(&schema) {
    102
    0
            acc = acc.with_new_schema(schema)?;
    103
    0
        }
    104
        // Fold in the rest of the EquivalenceProperties:
    105
    0
        for props in iter {
    106
    0
            acc = calculate_union_binary(acc, props)?;
    107
        }
    108
    0
        Ok(acc)
    109
    0
    }
    110
    111
    #[derive(Debug)]
    112
    enum AddedOrdering {
    113
        /// The ordering was added to the in progress result
    114
        Yes,
    115
        /// The ordering was not added
    116
        No(LexOrdering),
    117
    }
    118
    119
    /// Builds valid output orderings of a `UnionExec`
    120
    #[derive(Debug)]
    121
    struct UnionEquivalentOrderingBuilder {
    122
        orderings: Vec<LexOrdering>,
    123
    }
    124
    125
    impl UnionEquivalentOrderingBuilder {
    126
    0
        fn new() -> Self {
    127
    0
            Self { orderings: vec![] }
    128
    0
        }
    129
    130
        /// Add all orderings from `source` that satisfy `properties`,
    131
        /// potentially augmented with the constants in `source`.
    132
        ///
    133
        /// Note: Any column that is known to be constant can be inserted into the
    134
        /// ordering without changing its meaning.
    135
        ///
    136
        /// For example:
    137
        /// * Orderings in `source` contains `[a ASC, c ASC]` and constants contains
    138
        ///   `b`,
    139
        /// * `properties` has the ordering `[a ASC, b ASC]`.
    140
        ///
    141
        /// Then this will add `[a ASC, b ASC]` to the `orderings` list (as `a` was
    142
        /// in the sort order and `b` was a constant).
    143
    0
        fn add_satisfied_orderings(
    144
    0
            &mut self,
    145
    0
            source: &EquivalenceProperties,
    146
    0
            properties: &EquivalenceProperties,
    147
    0
        ) -> Result<()> {
    148
    0
            let constants = source.constants();
    149
    0
            let properties_constants = properties.constants();
    150
    0
            for mut ordering in source.oeq_cache.normal_cls.clone() {
    151
                // Progressively shorten the ordering to search for a satisfied prefix:
    152
                loop {
    153
    0
                    ordering = match self.try_add_ordering(
    154
    0
                        ordering,
    155
    0
                        &constants,
    156
    0
                        properties,
    157
    0
                        &properties_constants,
    158
    0
                    )? {
    159
    0
                        AddedOrdering::Yes => break,
    160
    0
                        AddedOrdering::No(ordering) => {
    161
    0
                            let mut sort_exprs: Vec<_> = ordering.into();
    162
    0
                            sort_exprs.pop();
    163
    0
                            if let Some(ordering) = LexOrdering::new(sort_exprs) {
    164
    0
                                ordering
    165
                            } else {
    166
    0
                                break;
    167
                            }
    168
                        }
    169
                    }
    170
                }
    171
            }
    172
    0
            Ok(())
    173
    0
        }
    174
    175
        /// Adds `ordering`, potentially augmented with `constants`, if it satisfies
    176
        /// the given `properties`.
    177
        ///
    178
        /// # Returns
    179
        ///
    180
        /// An [`AddedOrdering::Yes`] instance if the ordering was added (either
    181
        /// directly or augmented), or was empty. An [`AddedOrdering::No`] instance
    182
        /// otherwise.
    183
    0
        fn try_add_ordering(
    184
    0
            &mut self,
    185
    0
            ordering: LexOrdering,
    186
    0
            constants: &[ConstExpr],
    187
    0
            properties: &EquivalenceProperties,
    188
    0
            properties_constants: &[ConstExpr],
    189
    0
        ) -> Result<AddedOrdering> {
    190
    0
            if properties.ordering_satisfy(ordering.clone())? {
    191
                // If the ordering satisfies the target properties, no need to
    192
                // augment it with constants.
    193
    0
                self.orderings.push(ordering);
    194
    0
                Ok(AddedOrdering::Yes)
    195
    0
            } else if self.try_find_augmented_ordering(
    196
    0
                &ordering,
    197
    0
                constants,
    198
    0
                properties,
    199
    0
                properties_constants,
    200
            ) {
    201
                // Augmented with constants to match the properties.
    202
    0
                Ok(AddedOrdering::Yes)
    203
            } else {
    204
    0
                Ok(AddedOrdering::No(ordering))
    205
            }
    206
    0
        }
    207
    208
        /// Attempts to add `constants` to `ordering` to satisfy the properties.
    209
        /// Returns `true` if augmentation took place, `false` otherwise.
    210
    0
        fn try_find_augmented_ordering(
    211
    0
            &mut self,
    212
    0
            ordering: &LexOrdering,
    213
    0
            constants: &[ConstExpr],
    214
    0
            properties: &EquivalenceProperties,
    215
    0
            properties_constants: &[ConstExpr],
    216
    0
        ) -> bool {
    217
    0
            let mut result = false;
    218
            // Can only augment if there are constants.
    219
    0
            if !constants.is_empty() {
    220
                // For each equivalent ordering in properties, try and augment
    221
                // `ordering` with the constants to match `existing_ordering`:
    222
    0
                for existing_ordering in properties.oeq_class.iter() {
    223
    0
                    if let Some(augmented_ordering) = Self::augment_ordering(
    224
    0
                        ordering,
    225
    0
                        constants,
    226
    0
                        existing_ordering,
    227
    0
                        properties_constants,
    228
    0
                    ) {
    229
    0
                        self.orderings.push(augmented_ordering);
    230
    0
                        result = true;
    231
    0
                    }
    232
                }
    233
    0
            }
    234
    0
            result
    235
    0
        }
    236
    237
        /// Attempts to augment the ordering with constants to match `existing_ordering`.
    238
        /// Returns `Some(ordering)` if an augmented ordering was found, `None` otherwise.
    239
    0
        fn augment_ordering(
    240
    0
            ordering: &LexOrdering,
    241
    0
            constants: &[ConstExpr],
    242
    0
            existing_ordering: &LexOrdering,
    243
    0
            existing_constants: &[ConstExpr],
    244
    0
        ) -> Option<LexOrdering> {
    245
    0
            let mut augmented_ordering = vec![];
    246
    0
            let mut sort_exprs = ordering.iter().peekable();
    247
    0
            let mut existing_sort_exprs = existing_ordering.iter().peekable();
    248
    249
            // Walk in parallel down the two orderings, trying to match them up:
    250
    0
            while sort_exprs.peek().is_some() || existing_sort_exprs.peek().is_some() {
    251
                // If the next expressions are equal, add the next match. Otherwise,
    252
                // try and match with a constant.
    253
    0
                if let Some(expr) =
    254
    0
                    advance_if_match(&mut sort_exprs, &mut existing_sort_exprs)
    255
    0
                {
    256
    0
                    augmented_ordering.push(expr);
    257
    0
                } else if let Some(expr) =
    258
    0
                    advance_if_matches_constant(&mut sort_exprs, existing_constants)
    259
    0
                {
    260
    0
                    augmented_ordering.push(expr);
    261
    0
                } else if let Some(expr) =
    262
    0
                    advance_if_matches_constant(&mut existing_sort_exprs, constants)
    263
    0
                {
    264
    0
                    augmented_ordering.push(expr);
    265
    0
                } else {
    266
                    // no match, can't continue the ordering, return what we have
    267
    0
                    break;
    268
                }
    269
            }
    270
    271
    0
            LexOrdering::new(augmented_ordering)
    272
    0
        }
    273
    274
    0
        fn build(self) -> Vec<LexOrdering> {
    275
    0
            self.orderings
    276
    0
        }
    277
    }
    278
    279
    /// Advances two iterators in parallel if the next expressions are equal.
    280
    /// Otherwise, the iterators are left unchanged and returns `None`.
    281
    0
    fn advance_if_match<'a>(
    282
    0
        iter1: &mut Peekable<impl Iterator<Item = &'a PhysicalSortExpr>>,
    283
    0
        iter2: &mut Peekable<impl Iterator<Item = &'a PhysicalSortExpr>>,
    284
    0
    ) -> Option<PhysicalSortExpr> {
    285
    0
        let (expr1, expr2) = (iter1.peek()?, iter2.peek()?);
    286
    0
        if expr1.eq(expr2) {
    287
    0
            iter1.next();
    288
    0
            iter2.next().cloned()
    289
        } else {
    290
    0
            None
    291
        }
    292
    0
    }
    293
    294
    /// Advances the iterator with a constant if the next expression matches one of
    295
    /// the constants. Otherwise, the iterator is left unchanged and returns `None`.
    296
    0
    fn advance_if_matches_constant<'a>(
    297
    0
        iter: &mut Peekable<impl Iterator<Item = &'a PhysicalSortExpr>>,
    298
    0
        constants: &[ConstExpr],
    299
    0
    ) -> Option<PhysicalSortExpr> {
    300
    0
        let expr = iter.peek()?;
    301
    0
        let const_expr = constants.iter().find(|c| expr.expr.eq(&c.expr))?;
    302
    0
        let found_expr = PhysicalSortExpr::new(Arc::clone(&const_expr.expr), expr.options);
    303
    0
        iter.next();
    304
    0
        Some(found_expr)
    305
    0
    }
    306
    307
    #[cfg(test)]
    308
    mod tests {
    309
        use super::*;
    310
        use crate::equivalence::tests::{create_test_schema, parse_sort_expr};
    311
        use crate::expressions::col;
    312
        use crate::PhysicalExpr;
    313
    314
        use arrow::datatypes::{DataType, Field, Schema};
    315
        use datafusion_common::ScalarValue;
    316
    317
        use itertools::Itertools;
    318
    319
        /// Checks whether `expr` is among in the `const_exprs`.
    320
    0
        fn const_exprs_contains(
    321
    0
            const_exprs: &[ConstExpr],
    322
    0
            expr: &Arc<dyn PhysicalExpr>,
    323
    0
        ) -> bool {
    324
    0
            const_exprs
    325
    0
                .iter()
    326
    0
                .any(|const_expr| const_expr.expr.eq(expr))
    327
    0
        }
    328
    329
        #[test]
    330
    0
        fn test_union_equivalence_properties_multi_children_1() -> Result<()> {
    331
    0
            let schema = create_test_schema().unwrap();
    332
    0
            let schema2 = append_fields(&schema, "1");
    333
    0
            let schema3 = append_fields(&schema, "2");
    334
    0
            UnionEquivalenceTest::new(&schema)
    335
                // Children 1
    336
    0
                .with_child_sort(vec![vec!["a", "b", "c"]], &schema)?
    337
                // Children 2
    338
    0
                .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2)?
    339
                // Children 3
    340
    0
                .with_child_sort(vec![vec!["a2", "b2"]], &schema3)?
    341
    0
                .with_expected_sort(vec![vec!["a", "b"]])?
    342
    0
                .run()
    343
    0
        }
    344
    345
        #[test]
    346
    0
        fn test_union_equivalence_properties_multi_children_2() -> Result<()> {
    347
    0
            let schema = create_test_schema().unwrap();
    348
    0
            let schema2 = append_fields(&schema, "1");
    349
    0
            let schema3 = append_fields(&schema, "2");
    350
    0
            UnionEquivalenceTest::new(&schema)
    351
                // Children 1
    352
    0
                .with_child_sort(vec![vec!["a", "b", "c"]], &schema)?
    353
                // Children 2
    354
    0
                .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2)?
    355
                // Children 3
    356
    0
                .with_child_sort(vec![vec!["a2", "b2", "c2"]], &schema3)?
    357
    0
                .with_expected_sort(vec![vec!["a", "b", "c"]])?
    358
    0
                .run()
    359
    0
        }
    360
    361
        #[test]
    362
    0
        fn test_union_equivalence_properties_multi_children_3() -> Result<()> {
    363
    0
            let schema = create_test_schema().unwrap();
    364
    0
            let schema2 = append_fields(&schema, "1");
    365
    0
            let schema3 = append_fields(&schema, "2");
    366
    0
            UnionEquivalenceTest::new(&schema)
    367
                // Children 1
    368
    0
                .with_child_sort(vec![vec!["a", "b"]], &schema)?
    369
                // Children 2
    370
    0
                .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2)?
    371
                // Children 3
    372
    0
                .with_child_sort(vec![vec!["a2", "b2", "c2"]], &schema3)?
    373
    0
                .with_expected_sort(vec![vec!["a", "b"]])?
    374
    0
                .run()
    375
    0
        }
    376
    377
        #[test]
    378
    0
        fn test_union_equivalence_properties_multi_children_4() -> Result<()> {
    379
    0
            let schema = create_test_schema().unwrap();
    380
    0
            let schema2 = append_fields(&schema, "1");
    381
    0
            let schema3 = append_fields(&schema, "2");
    382
    0
            UnionEquivalenceTest::new(&schema)
    383
                // Children 1
    384
    0
                .with_child_sort(vec![vec!["a", "b"]], &schema)?
    385
                // Children 2
    386
    0
                .with_child_sort(vec![vec!["a1", "b1"]], &schema2)?
    387
                // Children 3
    388
    0
                .with_child_sort(vec![vec!["b2", "c2"]], &schema3)?
    389
    0
                .with_expected_sort(vec![])?
    390
    0
                .run()
    391
    0
        }
    392
    393
        #[test]
    394
    0
        fn test_union_equivalence_properties_multi_children_5() -> Result<()> {
    395
    0
            let schema = create_test_schema().unwrap();
    396
    0
            let schema2 = append_fields(&schema, "1");
    397
    0
            UnionEquivalenceTest::new(&schema)
    398
                // Children 1
    399
    0
                .with_child_sort(vec![vec!["a", "b"], vec!["c"]], &schema)?
    400
                // Children 2
    401
    0
                .with_child_sort(vec![vec!["a1", "b1"], vec!["c1"]], &schema2)?
    402
    0
                .with_expected_sort(vec![vec!["a", "b"], vec!["c"]])?
    403
    0
                .run()
    404
    0
        }
    405
    406
        #[test]
    407
    0
        fn test_union_equivalence_properties_constants_common_constants() -> Result<()> {
    408
    0
            let schema = create_test_schema().unwrap();
    409
    0
            UnionEquivalenceTest::new(&schema)
    410
    0
                .with_child_sort_and_const_exprs(
    411
                    // First child: [a ASC], const [b, c]
    412
    0
                    vec![vec!["a"]],
    413
    0
                    vec!["b", "c"],
    414
    0
                    &schema,
    415
    0
                )?
    416
    0
                .with_child_sort_and_const_exprs(
    417
                    // Second child: [b ASC], const [a, c]
    418
    0
                    vec![vec!["b"]],
    419
    0
                    vec!["a", "c"],
    420
    0
                    &schema,
    421
    0
                )?
    422
    0
                .with_expected_sort_and_const_exprs(
    423
                    // Union expected orderings: [[a ASC], [b ASC]], const [c]
    424
    0
                    vec![vec!["a"], vec!["b"]],
    425
    0
                    vec!["c"],
    426
    0
                )?
    427
    0
                .run()
    428
    0
        }
    429
    430
        #[test]
    431
    0
        fn test_union_equivalence_properties_constants_prefix() -> Result<()> {
    432
    0
            let schema = create_test_schema().unwrap();
    433
    0
            UnionEquivalenceTest::new(&schema)
    434
    0
                .with_child_sort_and_const_exprs(
    435
                    // First child: [a ASC], const []
    436
    0
                    vec![vec!["a"]],
    437
    0
                    vec![],
    438
    0
                    &schema,
    439
    0
                )?
    440
    0
                .with_child_sort_and_const_exprs(
    441
                    // Second child: [a ASC, b ASC], const []
    442
    0
                    vec![vec!["a", "b"]],
    443
    0
                    vec![],
    444
    0
                    &schema,
    445
    0
                )?
    446
    0
                .with_expected_sort_and_const_exprs(
    447
                    // Union orderings: [a ASC], const []
    448
    0
                    vec![vec!["a"]],
    449
    0
                    vec![],
    450
    0
                )?
    451
    0
                .run()
    452
    0
        }
    453
    454
        #[test]
    455
    0
        fn test_union_equivalence_properties_constants_asc_desc_mismatch() -> Result<()> {
    456
    0
            let schema = create_test_schema().unwrap();
    457
    0
            UnionEquivalenceTest::new(&schema)
    458
    0
                .with_child_sort_and_const_exprs(
    459
                    // First child: [a ASC], const []
    460
    0
                    vec![vec!["a"]],
    461
    0
                    vec![],
    462
    0
                    &schema,
    463
    0
                )?
    464
    0
                .with_child_sort_and_const_exprs(
    465
                    // Second child orderings: [a DESC], const []
    466
    0
                    vec![vec!["a DESC"]],
    467
    0
                    vec![],
    468
    0
                    &schema,
    469
    0
                )?
    470
    0
                .with_expected_sort_and_const_exprs(
    471
                    // Union doesn't have any ordering or constant
    472
    0
                    vec![],
    473
    0
                    vec![],
    474
    0
                )?
    475
    0
                .run()
    476
    0
        }
    477
    478
        #[test]
    479
    0
        fn test_union_equivalence_properties_constants_different_schemas() -> Result<()> {
    480
    0
            let schema = create_test_schema().unwrap();
    481
    0
            let schema2 = append_fields(&schema, "1");
    482
    0
            UnionEquivalenceTest::new(&schema)
    483
    0
                .with_child_sort_and_const_exprs(
    484
                    // First child orderings: [a ASC], const []
    485
    0
                    vec![vec!["a"]],
    486
    0
                    vec![],
    487
    0
                    &schema,
    488
    0
                )?
    489
    0
                .with_child_sort_and_const_exprs(
    490
                    // Second child orderings: [a1 ASC, b1 ASC], const []
    491
    0
                    vec![vec!["a1", "b1"]],
    492
    0
                    vec![],
    493
    0
                    &schema2,
    494
    0
                )?
    495
    0
                .with_expected_sort_and_const_exprs(
    496
                    // Union orderings: [a ASC]
    497
                    //
    498
                    // Note that a, and a1 are at the same index for their
    499
                    // corresponding schemas.
    500
    0
                    vec![vec!["a"]],
    501
    0
                    vec![],
    502
    0
                )?
    503
    0
                .run()
    504
    0
        }
    505
    506
        #[test]
    507
    0
        fn test_union_equivalence_properties_constants_fill_gaps() -> Result<()> {
    508
    0
            let schema = create_test_schema().unwrap();
    509
    0
            UnionEquivalenceTest::new(&schema)
    510
    0
                .with_child_sort_and_const_exprs(
    511
                    // First child orderings: [a ASC, c ASC], const [b]
    512
    0
                    vec![vec!["a", "c"]],
    513
    0
                    vec!["b"],
    514
    0
                    &schema,
    515
    0
                )?
    516
    0
                .with_child_sort_and_const_exprs(
    517
                    // Second child orderings: [b ASC, c ASC], const [a]
    518
    0
                    vec![vec!["b", "c"]],
    519
    0
                    vec!["a"],
    520
    0
                    &schema,
    521
    0
                )?
    522
    0
                .with_expected_sort_and_const_exprs(
    523
                    // Union orderings: [
    524
                    //   [a ASC, b ASC, c ASC],
    525
                    //   [b ASC, a ASC, c ASC]
    526
                    // ], const []
    527
    0
                    vec![vec!["a", "b", "c"], vec!["b", "a", "c"]],
    528
    0
                    vec![],
    529
    0
                )?
    530
    0
                .run()
    531
    0
        }
    532
    533
        #[test]
    534
    0
        fn test_union_equivalence_properties_constants_no_fill_gaps() -> Result<()> {
    535
    0
            let schema = create_test_schema().unwrap();
    536
    0
            UnionEquivalenceTest::new(&schema)
    537
    0
                .with_child_sort_and_const_exprs(
    538
                    // First child orderings: [a ASC, c ASC], const [d] // some other constant
    539
    0
                    vec![vec!["a", "c"]],
    540
    0
                    vec!["d"],
    541
    0
                    &schema,
    542
    0
                )?
    543
    0
                .with_child_sort_and_const_exprs(
    544
                    // Second child orderings: [b ASC, c ASC], const [a]
    545
    0
                    vec![vec!["b", "c"]],
    546
    0
                    vec!["a"],
    547
    0
                    &schema,
    548
    0
                )?
    549
    0
                .with_expected_sort_and_const_exprs(
    550
                    // Union orderings: [[a]] (only a is constant)
    551
    0
                    vec![vec!["a"]],
    552
    0
                    vec![],
    553
    0
                )?
    554
    0
                .run()
    555
    0
        }
    556
    557
        #[test]
    558
    0
        fn test_union_equivalence_properties_constants_fill_some_gaps() -> Result<()> {
    559
    0
            let schema = create_test_schema().unwrap();
    560
    0
            UnionEquivalenceTest::new(&schema)
    561
    0
                .with_child_sort_and_const_exprs(
    562
                    // First child orderings: [c ASC], const [a, b] // some other constant
    563
    0
                    vec![vec!["c"]],
    564
    0
                    vec!["a", "b"],
    565
    0
                    &schema,
    566
    0
                )?
    567
    0
                .with_child_sort_and_const_exprs(
    568
                    // Second child orderings: [a DESC, b], const []
    569
    0
                    vec![vec!["a DESC", "b"]],
    570
    0
                    vec![],
    571
    0
                    &schema,
    572
    0
                )?
    573
    0
                .with_expected_sort_and_const_exprs(
    574
                    // Union orderings: [[a, b]] (can fill in the a/b with constants)
    575
    0
                    vec![vec!["a DESC", "b"]],
    576
    0
                    vec![],
    577
    0
                )?
    578
    0
                .run()
    579
    0
        }
    580
    581
        #[test]
    582
    0
        fn test_union_equivalence_properties_constants_fill_gaps_non_symmetric() -> Result<()>
    583
        {
    584
    0
            let schema = create_test_schema().unwrap();
    585
    0
            UnionEquivalenceTest::new(&schema)
    586
    0
                .with_child_sort_and_const_exprs(
    587
                    // First child orderings: [a ASC, c ASC], const [b]
    588
    0
                    vec![vec!["a", "c"]],
    589
    0
                    vec!["b"],
    590
    0
                    &schema,
    591
    0
                )?
    592
    0
                .with_child_sort_and_const_exprs(
    593
                    // Second child orderings: [b ASC, c ASC], const [a]
    594
    0
                    vec![vec!["b DESC", "c"]],
    595
    0
                    vec!["a"],
    596
    0
                    &schema,
    597
    0
                )?
    598
    0
                .with_expected_sort_and_const_exprs(
    599
                    // Union orderings: [
    600
                    //   [a ASC, b ASC, c ASC],
    601
                    //   [b ASC, a ASC, c ASC]
    602
                    // ], const []
    603
    0
                    vec![vec!["a", "b DESC", "c"], vec!["b DESC", "a", "c"]],
    604
    0
                    vec![],
    605
    0
                )?
    606
    0
                .run()
    607
    0
        }
    608
    609
        #[test]
    610
    0
        fn test_union_equivalence_properties_constants_gap_fill_symmetric() -> Result<()> {
    611
    0
            let schema = create_test_schema().unwrap();
    612
    0
            UnionEquivalenceTest::new(&schema)
    613
    0
                .with_child_sort_and_const_exprs(
    614
                    // First child: [a ASC, b ASC, d ASC], const [c]
    615
    0
                    vec![vec!["a", "b", "d"]],
    616
    0
                    vec!["c"],
    617
    0
                    &schema,
    618
    0
                )?
    619
    0
                .with_child_sort_and_const_exprs(
    620
                    // Second child: [a ASC, c ASC, d ASC], const [b]
    621
    0
                    vec![vec!["a", "c", "d"]],
    622
    0
                    vec!["b"],
    623
    0
                    &schema,
    624
    0
                )?
    625
    0
                .with_expected_sort_and_const_exprs(
    626
                    // Union orderings:
    627
                    // [a, b, c, d]
    628
                    // [a, c, b, d]
    629
    0
                    vec![vec!["a", "c", "b", "d"], vec!["a", "b", "c", "d"]],
    630
    0
                    vec![],
    631
    0
                )?
    632
    0
                .run()
    633
    0
        }
    634
    635
        #[test]
    636
    0
        fn test_union_equivalence_properties_constants_gap_fill_and_common() -> Result<()> {
    637
    0
            let schema = create_test_schema().unwrap();
    638
    0
            UnionEquivalenceTest::new(&schema)
    639
    0
                .with_child_sort_and_const_exprs(
    640
                    // First child: [a DESC, d ASC], const [b, c]
    641
    0
                    vec![vec!["a DESC", "d"]],
    642
    0
                    vec!["b", "c"],
    643
    0
                    &schema,
    644
    0
                )?
    645
    0
                .with_child_sort_and_const_exprs(
    646
                    // Second child: [a DESC, c ASC, d ASC], const [b]
    647
    0
                    vec![vec!["a DESC", "c", "d"]],
    648
    0
                    vec!["b"],
    649
    0
                    &schema,
    650
    0
                )?
    651
    0
                .with_expected_sort_and_const_exprs(
    652
                    // Union orderings:
    653
                    // [a DESC, c, d]  [b]
    654
    0
                    vec![vec!["a DESC", "c", "d"]],
    655
    0
                    vec!["b"],
    656
    0
                )?
    657
    0
                .run()
    658
    0
        }
    659
    660
        #[test]
    661
    0
        fn test_union_equivalence_properties_constants_middle_desc() -> Result<()> {
    662
    0
            let schema = create_test_schema().unwrap();
    663
    0
            UnionEquivalenceTest::new(&schema)
    664
    0
                .with_child_sort_and_const_exprs(
    665
                    // NB `b DESC` in the first child
    666
                    //
    667
                    // First child: [a ASC, b DESC, d ASC], const [c]
    668
    0
                    vec![vec!["a", "b DESC", "d"]],
    669
    0
                    vec!["c"],
    670
    0
                    &schema,
    671
    0
                )?
    672
    0
                .with_child_sort_and_const_exprs(
    673
                    // Second child: [a ASC, c ASC, d ASC], const [b]
    674
    0
                    vec![vec!["a", "c", "d"]],
    675
    0
                    vec!["b"],
    676
    0
                    &schema,
    677
    0
                )?
    678
    0
                .with_expected_sort_and_const_exprs(
    679
                    // Union orderings:
    680
                    // [a, b, d] (c constant)
    681
                    // [a, c, d] (b constant)
    682
    0
                    vec![vec!["a", "c", "b DESC", "d"], vec!["a", "b DESC", "c", "d"]],
    683
    0
                    vec![],
    684
    0
                )?
    685
    0
                .run()
    686
    0
        }
    687
    688
        // TODO tests with multiple constants
    689
    690
        #[derive(Debug)]
    691
        struct UnionEquivalenceTest {
    692
            /// The schema of the output of the Union
    693
            output_schema: SchemaRef,
    694
            /// The equivalence properties of each child to the union
    695
            child_properties: Vec<EquivalenceProperties>,
    696
            /// The expected output properties of the union. Must be set before
    697
            /// running `build`
    698
            expected_properties: Option<EquivalenceProperties>,
    699
        }
    700
    701
        impl UnionEquivalenceTest {
    702
    0
            fn new(output_schema: &SchemaRef) -> Self {
    703
    0
                Self {
    704
    0
                    output_schema: Arc::clone(output_schema),
    705
    0
                    child_properties: vec![],
    706
    0
                    expected_properties: None,
    707
    0
                }
    708
    0
            }
    709
    710
            /// Add a union input with the specified orderings
    711
            ///
    712
            /// See [`Self::make_props`] for the format of the strings in `orderings`
    713
    0
            fn with_child_sort(
    714
    0
                mut self,
    715
    0
                orderings: Vec<Vec<&str>>,
    716
    0
                schema: &SchemaRef,
    717
    0
            ) -> Result<Self> {
    718
    0
                let properties = self.make_props(orderings, vec![], schema)?;
    719
    0
                self.child_properties.push(properties);
    720
    0
                Ok(self)
    721
    0
            }
    722
    723
            /// Add a union input with the specified orderings and constant
    724
            /// equivalences
    725
            ///
    726
            /// See [`Self::make_props`] for the format of the strings in
    727
            /// `orderings` and `constants`
    728
    0
            fn with_child_sort_and_const_exprs(
    729
    0
                mut self,
    730
    0
                orderings: Vec<Vec<&str>>,
    731
    0
                constants: Vec<&str>,
    732
    0
                schema: &SchemaRef,
    733
    0
            ) -> Result<Self> {
    734
    0
                let properties = self.make_props(orderings, constants, schema)?;
    735
    0
                self.child_properties.push(properties);
    736
    0
                Ok(self)
    737
    0
            }
    738
    739
            /// Set the expected output sort order for the union of the children
    740
            ///
    741
            /// See [`Self::make_props`] for the format of the strings in `orderings`
    742
    0
            fn with_expected_sort(mut self, orderings: Vec<Vec<&str>>) -> Result<Self> {
    743
    0
                let properties = self.make_props(orderings, vec![], &self.output_schema)?;
    744
    0
                self.expected_properties = Some(properties);
    745
    0
                Ok(self)
    746
    0
            }
    747
    748
            /// Set the expected output sort order and constant expressions for the
    749
            /// union of the children
    750
            ///
    751
            /// See [`Self::make_props`] for the format of the strings in
    752
            /// `orderings` and `constants`.
    753
    0
            fn with_expected_sort_and_const_exprs(
    754
    0
                mut self,
    755
    0
                orderings: Vec<Vec<&str>>,
    756
    0
                constants: Vec<&str>,
    757
    0
            ) -> Result<Self> {
    758
    0
                let properties =
    759
    0
                    self.make_props(orderings, constants, &self.output_schema)?;
    760
    0
                self.expected_properties = Some(properties);
    761
    0
                Ok(self)
    762
    0
            }
    763
    764
            /// compute the union's output equivalence properties from the child
    765
            /// properties, and compare them to the expected properties
    766
    0
            fn run(self) -> Result<()> {
    767
                let Self {
    768
    0
                    output_schema,
    769
    0
                    child_properties,
    770
    0
                    expected_properties,
    771
    0
                } = self;
    772
    773
    0
                let expected_properties =
    774
    0
                    expected_properties.expect("expected_properties not set");
    775
    776
                // try all permutations of the children
    777
                // as the code treats lhs and rhs differently
    778
    0
                for child_properties in child_properties
    779
    0
                    .iter()
    780
    0
                    .cloned()
    781
    0
                    .permutations(child_properties.len())
    782
                {
    783
    0
                    println!("--- permutation ---");
    784
    0
                    for c in &child_properties {
    785
    0
                        println!("{c}");
    786
    0
                    }
    787
    0
                    let actual_properties =
    788
    0
                        calculate_union(child_properties, Arc::clone(&output_schema))
    789
    0
                            .expect("failed to calculate union equivalence properties");
    790
    0
                    Self::assert_eq_properties_same(
    791
    0
                        &actual_properties,
    792
    0
                        &expected_properties,
    793
    0
                        format!(
    794
    0
                            "expected: {expected_properties:?}\nactual:  {actual_properties:?}"
    795
                        ),
    796
                    );
    797
                }
    798
    0
                Ok(())
    799
    0
            }
    800
    801
    0
            fn assert_eq_properties_same(
    802
    0
                lhs: &EquivalenceProperties,
    803
    0
                rhs: &EquivalenceProperties,
    804
    0
                err_msg: String,
    805
    0
            ) {
    806
                // Check whether constants are same
    807
    0
                let lhs_constants = lhs.constants();
    808
    0
                let rhs_constants = rhs.constants();
    809
    0
                for rhs_constant in &rhs_constants {
    810
    0
                    assert!(
    811
    0
                        const_exprs_contains(&lhs_constants, &rhs_constant.expr),
    812
    0
                        "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
    813
                    );
    814
                }
    815
    0
                assert_eq!(
    816
    0
                    lhs_constants.len(),
    817
    0
                    rhs_constants.len(),
    818
    0
                    "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
    819
                );
    820
    821
                // Check whether orderings are same.
    822
    0
                let lhs_orderings = lhs.oeq_class();
    823
    0
                let rhs_orderings = rhs.oeq_class();
    824
    0
                for rhs_ordering in rhs_orderings.iter() {
    825
    0
                    assert!(
    826
    0
                        lhs_orderings.contains(rhs_ordering),
    827
    0
                        "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
    828
                    );
    829
                }
    830
    0
                assert_eq!(
    831
    0
                    lhs_orderings.len(),
    832
    0
                    rhs_orderings.len(),
    833
    0
                    "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
    834
                );
    835
    0
            }
    836
    837
            /// Make equivalence properties for the specified columns named in orderings and constants
    838
            ///
    839
            /// orderings: strings formatted like `"a"` or `"a DESC"`. See [`parse_sort_expr`]
    840
            /// constants: strings formatted like `"a"`.
    841
    0
            fn make_props(
    842
    0
                &self,
    843
    0
                orderings: Vec<Vec<&str>>,
    844
    0
                constants: Vec<&str>,
    845
    0
                schema: &SchemaRef,
    846
    0
            ) -> Result<EquivalenceProperties> {
    847
    0
                let orderings = orderings.iter().map(|ordering| {
    848
    0
                    ordering.iter().map(|name| parse_sort_expr(name, schema))
    849
    0
                });
    850
    851
    0
                let constants = constants
    852
    0
                    .iter()
    853
    0
                    .map(|col_name| ConstExpr::from(col(col_name, schema).unwrap()));
    854
    855
    0
                let mut props =
    856
    0
                    EquivalenceProperties::new_with_orderings(Arc::clone(schema), orderings);
    857
    0
                props.add_constants(constants)?;
    858
    0
                Ok(props)
    859
    0
            }
    860
        }
    861
    862
        #[test]
    863
    0
        fn test_union_constant_value_preservation() -> Result<()> {
    864
    0
            let schema = Arc::new(Schema::new(vec![
    865
    0
                Field::new("a", DataType::Int32, true),
    866
    0
                Field::new("b", DataType::Int32, true),
    867
            ]));
    868
    869
    0
            let col_a = col("a", &schema)?;
    870
    0
            let literal_10 = ScalarValue::Int32(Some(10));
    871
    872
            // Create first input with a=10
    873
    0
            let const_expr1 = ConstExpr::new(
    874
    0
                Arc::clone(&col_a),
    875
    0
                AcrossPartitions::Uniform(Some(literal_10.clone())),
    876
            );
    877
    0
            let mut input1 = EquivalenceProperties::new(Arc::clone(&schema));
    878
    0
            input1.add_constants(vec![const_expr1])?;
    879
    880
            // Create second input with a=10
    881
    0
            let const_expr2 = ConstExpr::new(
    882
    0
                Arc::clone(&col_a),
    883
    0
                AcrossPartitions::Uniform(Some(literal_10.clone())),
    884
            );
    885
    0
            let mut input2 = EquivalenceProperties::new(Arc::clone(&schema));
    886
    0
            input2.add_constants(vec![const_expr2])?;
    887
    888
            // Calculate union properties
    889
    0
            let union_props = calculate_union(vec![input1, input2], schema)?;
    890
    891
            // Verify column 'a' remains constant with value 10
    892
    0
            let const_a = &union_props.constants()[0];
    893
    0
            assert!(const_a.expr.eq(&col_a));
    894
    0
            assert_eq!(
    895
                const_a.across_partitions,
    896
    0
                AcrossPartitions::Uniform(Some(literal_10))
    897
            );
    898
    899
    0
            Ok(())
    900
    0
        }
    901
    902
        /// Return a new schema with the same types, but new field names
    903
        ///
    904
        /// The new field names are the old field names with `text` appended.
    905
        ///
    906
        /// For example, the schema "a", "b", "c" becomes "a1", "b1", "c1"
    907
        /// if `text` is "1".
    908
    0
        fn append_fields(schema: &SchemaRef, text: &str) -> SchemaRef {
    909
    0
            Arc::new(Schema::new(
    910
    0
                schema
    911
    0
                    .fields()
    912
    0
                    .iter()
    913
    0
                    .map(|field| {
    914
    0
                        Field::new(
    915
                            // Annotate name with `text`:
    916
    0
                            format!("{}{}", field.name(), text),
    917
    0
                            field.data_type().clone(),
    918
    0
                            field.is_nullable(),
    919
                        )
    920
    0
                    })
    921
    0
                    .collect::<Vec<_>>(),
    922
            ))
    923
    0
        }
    924
    925
        #[test]
    926
    0
        fn test_constants_share_values() -> Result<()> {
    927
    0
            let schema = Arc::new(Schema::new(vec![
    928
    0
                Field::new("const_1", DataType::Utf8, false),
    929
    0
                Field::new("const_2", DataType::Utf8, false),
    930
            ]));
    931
    932
    0
            let col_const_1 = col("const_1", &schema)?;
    933
    0
            let col_const_2 = col("const_2", &schema)?;
    934
    935
    0
            let literal_foo = ScalarValue::Utf8(Some("foo".to_owned()));
    936
    0
            let literal_bar = ScalarValue::Utf8(Some("bar".to_owned()));
    937
    938
    0
            let const_expr_1_foo = ConstExpr::new(
    939
    0
                Arc::clone(&col_const_1),
    940
    0
                AcrossPartitions::Uniform(Some(literal_foo.clone())),
    941
            );
    942
    0
            let const_expr_2_foo = ConstExpr::new(
    943
    0
                Arc::clone(&col_const_2),
    944
    0
                AcrossPartitions::Uniform(Some(literal_foo.clone())),
    945
            );
    946
    0
            let const_expr_2_bar = ConstExpr::new(
    947
    0
                Arc::clone(&col_const_2),
    948
    0
                AcrossPartitions::Uniform(Some(literal_bar.clone())),
    949
            );
    950
    951
    0
            let mut input1 = EquivalenceProperties::new(Arc::clone(&schema));
    952
    0
            let mut input2 = EquivalenceProperties::new(Arc::clone(&schema));
    953
    954
            // | Input | Const_1 | Const_2 |
    955
            // | ----- | ------- | ------- |
    956
            // |     1 | foo     | foo     |
    957
            // |     2 | foo     | bar     |
    958
    0
            input1.add_constants(vec![const_expr_1_foo.clone(), const_expr_2_foo.clone()])?;
    959
    0
            input2.add_constants(vec![const_expr_1_foo.clone(), const_expr_2_bar.clone()])?;
    960
    961
            // Calculate union properties
    962
    0
            let union_props = calculate_union(vec![input1, input2], schema)?;
    963
    964
            // This should result in:
    965
            //   const_1 = Uniform("foo")
    966
            //   const_2 = Heterogeneous
    967
    0
            assert_eq!(union_props.constants().len(), 2);
    968
    0
            let union_const_1 = &union_props.constants()[0];
    969
    0
            assert!(union_const_1.expr.eq(&col_const_1));
    970
    0
            assert_eq!(
    971
                union_const_1.across_partitions,
    972
    0
                AcrossPartitions::Uniform(Some(literal_foo)),
    973
            );
    974
    0
            let union_const_2 = &union_props.constants()[1];
    975
    0
            assert!(union_const_2.expr.eq(&col_const_2));
    976
    0
            assert_eq!(
    977
                union_const_2.across_partitions,
    978
                AcrossPartitions::Heterogeneous,
    979
            );
    980
    981
    0
            Ok(())
    982
    0
        }
    983
    }