## 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/utils/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
    mod guarantee;
    19
    pub use guarantee::{Guarantee, LiteralGuarantee};
    20
    21
    use std::borrow::Borrow;
    22
    use std::sync::Arc;
    23
    24
    use crate::expressions::{BinaryExpr, Column};
    25
    use crate::tree_node::ExprContext;
    26
    use crate::PhysicalExpr;
    27
    use crate::PhysicalSortExpr;
    28
    29
    use arrow::datatypes::Schema;
    30
    use datafusion_common::tree_node::{
    31
        Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
    32
    };
    33
    use datafusion_common::{HashMap, HashSet, Result};
    34
    use datafusion_expr::Operator;
    35
    36
    use petgraph::graph::NodeIndex;
    37
    use petgraph::stable_graph::StableGraph;
    38
    39
    /// Assume the predicate is in the form of CNF, split the predicate to a Vec of PhysicalExprs.
    40
    ///
    41
    /// For example, split "a1 = a2 AND b1 <= b2 AND c1 != c2" into ["a1 = a2", "b1 <= b2", "c1 != c2"]
    42
    0
    pub fn split_conjunction(
    43
    0
        predicate: &Arc<dyn PhysicalExpr>,
    44
    0
    ) -> Vec<&Arc<dyn PhysicalExpr>> {
    45
    0
        split_impl(Operator::And, predicate, vec![])
    46
    0
    }
    47
    48
    /// Create a conjunction of the given predicates.
    49
    /// If the input is empty, return a literal true.
    50
    /// If the input contains a single predicate, return the predicate.
    51
    /// Otherwise, return a conjunction of the predicates (e.g. `a AND b AND c`).
    52
    0
    pub fn conjunction(
    53
    0
        predicates: impl IntoIterator<Item = Arc<dyn PhysicalExpr>>,
    54
    0
    ) -> Arc<dyn PhysicalExpr> {
    55
    0
        conjunction_opt(predicates).unwrap_or_else(|| crate::expressions::lit(true))
    56
    0
    }
    57
    58
    /// Create a conjunction of the given predicates.
    59
    /// If the input is empty or the return None.
    60
    /// If the input contains a single predicate, return Some(predicate).
    61
    /// Otherwise, return a Some(..) of a conjunction of the predicates (e.g. `Some(a AND b AND c)`).
    62
    0
    pub fn conjunction_opt(
    63
    0
        predicates: impl IntoIterator<Item = Arc<dyn PhysicalExpr>>,
    64
    0
    ) -> Option<Arc<dyn PhysicalExpr>> {
    65
    0
        predicates
    66
    0
            .into_iter()
    67
    0
            .fold(None, |acc, predicate| match acc {
    68
    0
                None => Some(predicate),
    69
    0
                Some(acc) => Some(Arc::new(BinaryExpr::new(acc, Operator::And, predicate))),
    70
    0
            })
    71
    0
    }
    72
    73
    /// Assume the predicate is in the form of DNF, split the predicate to a Vec of PhysicalExprs.
    74
    ///
    75
    /// For example, split "a1 = a2 OR b1 <= b2 OR c1 != c2" into ["a1 = a2", "b1 <= b2", "c1 != c2"]
    76
    0
    pub fn split_disjunction(
    77
    0
        predicate: &Arc<dyn PhysicalExpr>,
    78
    0
    ) -> Vec<&Arc<dyn PhysicalExpr>> {
    79
    0
        split_impl(Operator::Or, predicate, vec![])
    80
    0
    }
    81
    82
    0
    fn split_impl<'a>(
    83
    0
        operator: Operator,
    84
    0
        predicate: &'a Arc<dyn PhysicalExpr>,
    85
    0
        mut exprs: Vec<&'a Arc<dyn PhysicalExpr>>,
    86
    0
    ) -> Vec<&'a Arc<dyn PhysicalExpr>> {
    87
    0
        match predicate.as_any().downcast_ref::<BinaryExpr>() {
    88
    0
            Some(binary) if binary.op() == &operator => {
    89
    0
                let exprs = split_impl(operator, binary.left(), exprs);
    90
    0
                split_impl(operator, binary.right(), exprs)
    91
            }
    92
            Some(_) | None => {
    93
    0
                exprs.push(predicate);
    94
    0
                exprs
    95
            }
    96
        }
    97
    0
    }
    98
    99
    /// This function maps back requirement after ProjectionExec
    100
    /// to the Executor for its input.
    101
    // Specifically, `ProjectionExec` changes index of `Column`s in the schema of its input executor.
    102
    // This function changes requirement given according to ProjectionExec schema to the requirement
    103
    // according to schema of input executor to the ProjectionExec.
    104
    // For instance, Column{"a", 0} would turn to Column{"a", 1}. Please note that this function assumes that
    105
    // name of the Column is unique. If we have a requirement such that Column{"a", 0}, Column{"a", 1}.
    106
    // This function will produce incorrect result (It will only emit single Column as a result).
    107
    0
    pub fn map_columns_before_projection(
    108
    0
        parent_required: &[Arc<dyn PhysicalExpr>],
    109
    0
        proj_exprs: &[(Arc<dyn PhysicalExpr>, String)],
    110
    0
    ) -> Vec<Arc<dyn PhysicalExpr>> {
    111
    0
        if parent_required.is_empty() {
    112
            // No need to build mapping.
    113
    0
            return vec![];
    114
    0
        }
    115
    0
        let column_mapping = proj_exprs
    116
    0
            .iter()
    117
    0
            .filter_map(|(expr, name)| {
    118
    0
                expr.as_any()
    119
    0
                    .downcast_ref::<Column>()
    120
    0
                    .map(|column| (name.clone(), column.clone()))
    121
    0
            })
    122
    0
            .collect::<HashMap<_, _>>();
    123
    0
        parent_required
    124
    0
            .iter()
    125
    0
            .filter_map(|r| {
    126
    0
                r.as_any()
    127
    0
                    .downcast_ref::<Column>()
    128
    0
                    .and_then(|c| column_mapping.get(c.name()))
    129
    0
            })
    130
    0
            .map(|e| Arc::new(e.clone()) as _)
    131
    0
            .collect()
    132
    0
    }
    133
    134
    /// This function returns all `Arc<dyn PhysicalExpr>`s inside the given
    135
    /// `PhysicalSortExpr` sequence.
    136
    0
    pub fn convert_to_expr<T: Borrow<PhysicalSortExpr>>(
    137
    0
        sequence: impl IntoIterator<Item = T>,
    138
    0
    ) -> Vec<Arc<dyn PhysicalExpr>> {
    139
    0
        sequence
    140
    0
            .into_iter()
    141
    0
            .map(|elem| Arc::clone(&elem.borrow().expr))
    142
    0
            .collect()
    143
    0
    }
    144
    145
    /// This function finds the indices of `targets` within `items` using strict
    146
    /// equality.
    147
    0
    pub fn get_indices_of_exprs_strict<T: Borrow<Arc<dyn PhysicalExpr>>>(
    148
    0
        targets: impl IntoIterator<Item = T>,
    149
    0
        items: &[Arc<dyn PhysicalExpr>],
    150
    0
    ) -> Vec<usize> {
    151
    0
        targets
    152
    0
            .into_iter()
    153
    0
            .filter_map(|target| items.iter().position(|e| e.eq(target.borrow())))
    154
    0
            .collect()
    155
    0
    }
    156
    157
    pub type ExprTreeNode<T> = ExprContext<Option<T>>;
    158
    159
    /// This struct is used to convert a [`PhysicalExpr`] tree into a DAEG (i.e. an expression
    160
    /// DAG) by collecting identical expressions in one node. Caller specifies the node type
    161
    /// in the DAEG via the `constructor` argument, which constructs nodes in the DAEG from
    162
    /// the [`ExprTreeNode`] ancillary object.
    163
    struct PhysicalExprDAEGBuilder<'a, T, F: Fn(&ExprTreeNode<NodeIndex>) -> Result<T>> {
    164
        // The resulting DAEG (expression DAG).
    165
        graph: StableGraph<T, usize>,
    166
        // A vector of visited expression nodes and their corresponding node indices.
    167
        visited_plans: Vec<(Arc<dyn PhysicalExpr>, NodeIndex)>,
    168
        // A function to convert an input expression node to T.
    169
        constructor: &'a F,
    170
    }
    171
    172
    impl<T, F: Fn(&ExprTreeNode<NodeIndex>) -> Result<T>> PhysicalExprDAEGBuilder<'_, T, F> {
    173
        // This method mutates an expression node by transforming it to a physical expression
    174
        // and adding it to the graph. The method returns the mutated expression node.
    175
    0
        fn mutate(
    176
    0
            &mut self,
    177
    0
            mut node: ExprTreeNode<NodeIndex>,
    178
    0
        ) -> Result<Transformed<ExprTreeNode<NodeIndex>>> {
    179
            // Get the expression associated with the input expression node.
    180
    0
            let expr = &node.expr;
    181
    182
            // Check if the expression has already been visited.
    183
    0
            let node_idx = match self.visited_plans.iter().find(|(e, _)| expr.eq(e)) {
    184
                // If the expression has been visited, return the corresponding node index.
    185
    0
                Some((_, idx)) => *idx,
    186
                // If the expression has not been visited, add a new node to the graph and
    187
                // add edges to its child nodes. Add the visited expression to the vector
    188
                // of visited expressions and return the newly created node index.
    189
                None => {
    190
    0
                    let node_idx = self.graph.add_node((self.constructor)(&node)?);
    191
    0
                    for expr_node in node.children.iter() {
    192
    0
                        self.graph.add_edge(node_idx, expr_node.data.unwrap(), 0);
    193
    0
                    }
    194
    0
                    self.visited_plans.push((Arc::clone(expr), node_idx));
    195
    0
                    node_idx
    196
                }
    197
            };
    198
            // Set the data field of the input expression node to the corresponding node index.
    199
    0
            node.data = Some(node_idx);
    200
            // Return the mutated expression node.
    201
    0
            Ok(Transformed::yes(node))
    202
    0
        }
    203
    }
    204
    205
    // A function that builds a directed acyclic graph of physical expression trees.
    206
    0
    pub fn build_dag<T, F>(
    207
    0
        expr: Arc<dyn PhysicalExpr>,
    208
    0
        constructor: &F,
    209
    0
    ) -> Result<(NodeIndex, StableGraph<T, usize>)>
    210
    0
    where
    211
    0
        F: Fn(&ExprTreeNode<NodeIndex>) -> Result<T>,
    212
    {
    213
        // Create a new expression tree node from the input expression.
    214
    0
        let init = ExprTreeNode::new_default(expr);
    215
        // Create a new `PhysicalExprDAEGBuilder` instance.
    216
    0
        let mut builder = PhysicalExprDAEGBuilder {
    217
    0
            graph: StableGraph::<T, usize>::new(),
    218
    0
            visited_plans: Vec::<(Arc<dyn PhysicalExpr>, NodeIndex)>::new(),
    219
    0
            constructor,
    220
    0
        };
    221
        // Use the builder to transform the expression tree node into a DAG.
    222
    0
        let root = init.transform_up(|node| builder.mutate(node)).data()?;
    223
        // Return a tuple containing the root node index and the DAG.
    224
    0
        Ok((root.data.unwrap(), builder.graph))
    225
    0
    }
    226
    227
    /// Recursively extract referenced [`Column`]s within a [`PhysicalExpr`].
    228
    0
    pub fn collect_columns(expr: &Arc<dyn PhysicalExpr>) -> HashSet<Column> {
    229
    0
        let mut columns = HashSet::<Column>::new();
    230
    0
        expr.apply(|expr| {
    231
    0
            if let Some(column) = expr.as_any().downcast_ref::<Column>() {
    232
    0
                columns.get_or_insert_owned(column);
    233
    0
            }
    234
    0
            Ok(TreeNodeRecursion::Continue)
    235
    0
        })
    236
        // pre_visit always returns OK, so this will always too
    237
    0
        .expect("no way to return error during recursion");
    238
    0
        columns
    239
    0
    }
    240
    241
    /// Re-assign indices of [`Column`]s within the given [`PhysicalExpr`] according to
    242
    /// the provided [`Schema`].
    243
    ///
    244
    /// This can be useful when attempting to map an expression onto a different schema.
    245
    ///
    246
    /// # Errors
    247
    ///
    248
    /// This function will return an error if any column in the expression cannot be found
    249
    /// in the provided schema.
    250
    1
    pub fn reassign_expr_columns(
    251
    1
        expr: Arc<dyn PhysicalExpr>,
    252
    1
        schema: &Schema,
    253
    1
    ) -> Result<Arc<dyn PhysicalExpr>> {
    254
    3
        expr1.transform_down1(|expr| {
    255
    3
            if let Some(column1) = expr.as_any().downcast_ref::<Column>() {
    256
    1
                let index = schema.index_of(column.name())?0;
    257
    258
    1
                return Ok(Transformed::yes(Arc::new(Column::new(
    259
    1
                    column.name(),
    260
    1
                    index,
    261
    1
                ))));
    262
    2
            }
    263
    2
            Ok(Transformed::no(expr))
    264
    3
        })
    265
    1
        .data()
    266
    1
    }
    267
    268
    #[cfg(test)]
    269
    pub(crate) mod tests {
    270
        use std::any::Any;
    271
        use std::fmt::{Display, Formatter};
    272
    273
        use super::*;
    274
        use crate::expressions::{binary, cast, col, in_list, lit, Literal};
    275
    276
        use arrow::array::{ArrayRef, Float32Array, Float64Array};
    277
        use arrow::datatypes::{DataType, Field, Schema};
    278
        use datafusion_common::{exec_err, internal_datafusion_err, ScalarValue};
    279
        use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
    280
        use datafusion_expr::{
    281
            ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
    282
        };
    283
    284
        use petgraph::visit::Bfs;
    285
    286
        #[derive(Debug, PartialEq, Eq, Hash)]
    287
        pub struct TestScalarUDF {
    288
            pub(crate) signature: Signature,
    289
        }
    290
    291
        impl TestScalarUDF {
    292
    0
            pub fn new() -> Self {
    293
                use DataType::*;
    294
    0
                Self {
    295
    0
                    signature: Signature::uniform(
    296
    0
                        1,
    297
    0
                        vec![Float64, Float32],
    298
    0
                        Volatility::Immutable,
    299
    0
                    ),
    300
    0
                }
    301
    0
            }
    302
        }
    303
    304
        impl ScalarUDFImpl for TestScalarUDF {
    305
    0
            fn as_any(&self) -> &dyn Any {
    306
    0
                self
    307
    0
            }
    308
    0
            fn name(&self) -> &str {
    309
    0
                "test-scalar-udf"
    310
    0
            }
    311
    312
    0
            fn signature(&self) -> &Signature {
    313
    0
                &self.signature
    314
    0
            }
    315
    316
    0
            fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
    317
    0
                let arg_type = &arg_types[0];
    318
    319
    0
                match arg_type {
    320
    0
                    DataType::Float32 => Ok(DataType::Float32),
    321
    0
                    _ => Ok(DataType::Float64),
    322
                }
    323
    0
            }
    324
    325
    0
            fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
    326
    0
                Ok(input[0].sort_properties)
    327
    0
            }
    328
    329
    0
            fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    330
    0
                let args = ColumnarValue::values_to_arrays(&args.args)?;
    331
    332
    0
                let arr: ArrayRef = match args[0].data_type() {
    333
    0
                    DataType::Float64 => Arc::new({
    334
    0
                        let arg = &args[0]
    335
    0
                            .as_any()
    336
    0
                            .downcast_ref::<Float64Array>()
    337
    0
                            .ok_or_else(|| {
    338
    0
                                internal_datafusion_err!(
    339
    0
                                    "could not cast {} to {}",
    340
    0
                                    self.name(),
    341
    0
                                    std::any::type_name::<Float64Array>()
    342
                                )
    343
    0
                            })?;
    344
    345
    0
                        arg.iter()
    346
    0
                            .map(|a| a.map(f64::floor))
    347
    0
                            .collect::<Float64Array>()
    348
                    }),
    349
    0
                    DataType::Float32 => Arc::new({
    350
    0
                        let arg = &args[0]
    351
    0
                            .as_any()
    352
    0
                            .downcast_ref::<Float32Array>()
    353
    0
                            .ok_or_else(|| {
    354
    0
                                internal_datafusion_err!(
    355
    0
                                    "could not cast {} to {}",
    356
    0
                                    self.name(),
    357
    0
                                    std::any::type_name::<Float32Array>()
    358
                                )
    359
    0
                            })?;
    360
    361
    0
                        arg.iter()
    362
    0
                            .map(|a| a.map(f32::floor))
    363
    0
                            .collect::<Float32Array>()
    364
                    }),
    365
    0
                    other => {
    366
    0
                        return exec_err!(
    367
    0
                            "Unsupported data type {other:?} for function {}",
    368
    0
                            self.name()
    369
                        );
    370
                    }
    371
                };
    372
    0
                Ok(ColumnarValue::Array(arr))
    373
    0
            }
    374
        }
    375
    376
        #[derive(Clone)]
    377
        struct DummyProperty {
    378
            expr_type: String,
    379
        }
    380
    381
        /// This is a dummy node in the DAEG; it stores a reference to the actual
    382
        /// [PhysicalExpr] as well as a dummy property.
    383
        #[derive(Clone)]
    384
        struct PhysicalExprDummyNode {
    385
            pub expr: Arc<dyn PhysicalExpr>,
    386
            pub property: DummyProperty,
    387
        }
    388
    389
        impl Display for PhysicalExprDummyNode {
    390
    0
            fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
    391
    0
                write!(f, "{}", self.expr)
    392
    0
            }
    393
        }
    394
    395
    0
        fn make_dummy_node(node: &ExprTreeNode<NodeIndex>) -> Result<PhysicalExprDummyNode> {
    396
    0
            let expr = Arc::clone(&node.expr);
    397
    0
            let dummy_property = if expr.as_any().is::<BinaryExpr>() {
    398
    0
                "Binary"
    399
    0
            } else if expr.as_any().is::<Column>() {
    400
    0
                "Column"
    401
    0
            } else if expr.as_any().is::<Literal>() {
    402
    0
                "Literal"
    403
            } else {
    404
    0
                "Other"
    405
            }
    406
    0
            .to_owned();
    407
    0
            Ok(PhysicalExprDummyNode {
    408
    0
                expr,
    409
    0
                property: DummyProperty {
    410
    0
                    expr_type: dummy_property,
    411
    0
                },
    412
    0
            })
    413
    0
        }
    414
    415
        #[test]
    416
    0
        fn test_build_dag() -> Result<()> {
    417
    0
            let schema = Schema::new(vec![
    418
    0
                Field::new("0", DataType::Int32, true),
    419
    0
                Field::new("1", DataType::Int32, true),
    420
    0
                Field::new("2", DataType::Int32, true),
    421
            ]);
    422
    0
            let expr = binary(
    423
    0
                cast(
    424
    0
                    binary(
    425
    0
                        col("0", &schema)?,
    426
    0
                        Operator::Plus,
    427
    0
                        col("1", &schema)?,
    428
    0
                        &schema,
    429
    0
                    )?,
    430
    0
                    &schema,
    431
    0
                    DataType::Int64,
    432
    0
                )?,
    433
    0
                Operator::Gt,
    434
    0
                binary(
    435
    0
                    cast(col("2", &schema)?, &schema, DataType::Int64)?,
    436
    0
                    Operator::Plus,
    437
    0
                    lit(ScalarValue::Int64(Some(10))),
    438
    0
                    &schema,
    439
    0
                )?,
    440
    0
                &schema,
    441
    0
            )?;
    442
    0
            let mut vector_dummy_props = vec![];
    443
    0
            let (root, graph) = build_dag(expr, &make_dummy_node)?;
    444
    0
            let mut bfs = Bfs::new(&graph, root);
    445
    0
            while let Some(node_index) = bfs.next(&graph) {
    446
    0
                let node = &graph[node_index];
    447
    0
                vector_dummy_props.push(node.property.clone());
    448
    0
            }
    449
    450
    0
            assert_eq!(
    451
    0
                vector_dummy_props
    452
    0
                    .iter()
    453
    0
                    .filter(|property| property.expr_type == "Binary")
    454
    0
                    .count(),
    455
                3
    456
            );
    457
    0
            assert_eq!(
    458
    0
                vector_dummy_props
    459
    0
                    .iter()
    460
    0
                    .filter(|property| property.expr_type == "Column")
    461
    0
                    .count(),
    462
                3
    463
            );
    464
    0
            assert_eq!(
    465
    0
                vector_dummy_props
    466
    0
                    .iter()
    467
    0
                    .filter(|property| property.expr_type == "Literal")
    468
    0
                    .count(),
    469
                1
    470
            );
    471
    0
            assert_eq!(
    472
    0
                vector_dummy_props
    473
    0
                    .iter()
    474
    0
                    .filter(|property| property.expr_type == "Other")
    475
    0
                    .count(),
    476
                2
    477
            );
    478
    0
            Ok(())
    479
    0
        }
    480
    481
        #[test]
    482
    0
        fn test_convert_to_expr() -> Result<()> {
    483
    0
            let schema = Schema::new(vec![Field::new("a", DataType::UInt64, false)]);
    484
    0
            let sort_expr = vec![PhysicalSortExpr {
    485
    0
                expr: col("a", &schema)?,
    486
    0
                options: Default::default(),
    487
            }];
    488
    0
            assert!(convert_to_expr(&sort_expr)[0].eq(&sort_expr[0].expr));
    489
    0
            Ok(())
    490
    0
        }
    491
    492
        #[test]
    493
    0
        fn test_get_indices_of_exprs_strict() {
    494
    0
            let list1: Vec<Arc<dyn PhysicalExpr>> = vec![
    495
    0
                Arc::new(Column::new("a", 0)),
    496
    0
                Arc::new(Column::new("b", 1)),
    497
    0
                Arc::new(Column::new("c", 2)),
    498
    0
                Arc::new(Column::new("d", 3)),
    499
            ];
    500
    0
            let list2: Vec<Arc<dyn PhysicalExpr>> = vec![
    501
    0
                Arc::new(Column::new("b", 1)),
    502
    0
                Arc::new(Column::new("c", 2)),
    503
    0
                Arc::new(Column::new("a", 0)),
    504
            ];
    505
    0
            assert_eq!(get_indices_of_exprs_strict(&list1, &list2), vec![2, 0, 1]);
    506
    0
            assert_eq!(get_indices_of_exprs_strict(&list2, &list1), vec![1, 2, 0]);
    507
    0
        }
    508
    509
        #[test]
    510
    1
        fn test_reassign_expr_columns_in_list() {
    511
    1
            let int_field = Field::new("should_not_matter", DataType::Int64, true);
    512
    1
            let dict_field = Field::new(
    513
                "id",
    514
    1
                DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
    515
                true,
    516
            );
    517
    1
            let schema_small = Arc::new(Schema::new(vec![dict_field.clone()]));
    518
    1
            let schema_big = Arc::new(Schema::new(vec![int_field, dict_field]));
    519
    1
            let pred = in_list(
    520
    1
                Arc::new(Column::new_with_schema("id", &schema_big).unwrap()),
    521
    1
                vec![lit(ScalarValue::Dictionary(
    522
    1
                    Box::new(DataType::Int32),
    523
    1
                    Box::new(ScalarValue::from("2")),
    524
    1
                ))],
    525
    1
                &false,
    526
    1
                &schema_big,
    527
            )
    528
    1
            .unwrap();
    529
    530
    1
            let actual = reassign_expr_columns(pred, &schema_small).unwrap();
    531
    532
    1
            let expected = in_list(
    533
    1
                Arc::new(Column::new_with_schema("id", &schema_small).unwrap()),
    534
    1
                vec![lit(ScalarValue::Dictionary(
    535
    1
                    Box::new(DataType::Int32),
    536
    1
                    Box::new(ScalarValue::from("2")),
    537
    1
                ))],
    538
    1
                &false,
    539
    1
                &schema_small,
    540
            )
    541
    1
            .unwrap();
    542
    543
    1
            assert_eq!(actual.as_ref(), expected.as_ref());
    544
    1
        }
    545
    546
        #[test]
    547
    0
        fn test_collect_columns() -> Result<()> {
    548
    0
            let expr1 = Arc::new(Column::new("col1", 2)) as _;
    549
    0
            let mut expected = HashSet::new();
    550
    0
            expected.insert(Column::new("col1", 2));
    551
    0
            assert_eq!(collect_columns(&expr1), expected);
    552
    553
    0
            let expr2 = Arc::new(Column::new("col2", 5)) as _;
    554
    0
            let mut expected = HashSet::new();
    555
    0
            expected.insert(Column::new("col2", 5));
    556
    0
            assert_eq!(collect_columns(&expr2), expected);
    557
    558
    0
            let expr3 = Arc::new(BinaryExpr::new(expr1, Operator::Plus, expr2)) as _;
    559
    0
            let mut expected = HashSet::new();
    560
    0
            expected.insert(Column::new("col1", 2));
    561
    0
            expected.insert(Column::new("col2", 5));
    562
    0
            assert_eq!(collect_columns(&expr3), expected);
    563
    0
            Ok(())
    564
    0
        }
    565
    }