## 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/expr/src/expr_rewriter/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
    //! Expression rewriter
    19
    20
    use std::collections::HashMap;
    21
    use std::collections::HashSet;
    22
    use std::fmt::Debug;
    23
    use std::sync::Arc;
    24
    25
    use crate::expr::{Alias, Sort, Unnest};
    26
    use crate::logical_plan::Projection;
    27
    use crate::{Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder};
    28
    29
    use datafusion_common::config::ConfigOptions;
    30
    use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
    31
    use datafusion_common::TableReference;
    32
    use datafusion_common::{Column, DFSchema, Result};
    33
    34
    mod guarantees;
    35
    pub use guarantees::rewrite_with_guarantees;
    36
    pub use guarantees::rewrite_with_guarantees_map;
    37
    pub use guarantees::GuaranteeRewriter;
    38
    mod order_by;
    39
    40
    pub use order_by::rewrite_sort_cols_by_aggs;
    41
    42
    /// Trait for rewriting [`Expr`]s into function calls.
    43
    ///
    44
    /// This trait is used with `FunctionRegistry::register_function_rewrite` to
    45
    /// to evaluating `Expr`s using functions that may not be built in to DataFusion
    46
    ///
    47
    /// For example, concatenating arrays `a || b` is represented as
    48
    /// `Operator::ArrowAt`, but can be implemented by calling a function
    49
    /// `array_concat` from the `functions-nested` crate.
    50
    // This is not used in datafusion internally, but it is still helpful for downstream project so don't remove it.
    51
    pub trait FunctionRewrite: Debug {
    52
        /// Return a human readable name for this rewrite
    53
        fn name(&self) -> &str;
    54
    55
        /// Potentially rewrite `expr` to some other expression
    56
        ///
    57
        /// Note that recursion is handled by the caller -- this method should only
    58
        /// handle `expr`, not recurse to its children.
    59
        fn rewrite(
    60
            &self,
    61
            expr: Expr,
    62
            schema: &DFSchema,
    63
            config: &ConfigOptions,
    64
        ) -> Result<Transformed<Expr>>;
    65
    }
    66
    67
    /// Recursively call `LogicalPlanBuilder::normalize` on all [`Column`] expressions
    68
    /// in the `expr` expression tree.
    69
    0
    pub fn normalize_col(expr: Expr, plan: &LogicalPlan) -> Result<Expr> {
    70
    0
        expr.transform(|expr| {
    71
            Ok({
    72
    0
                if let Expr::Column(c) = expr {
    73
    0
                    let col = LogicalPlanBuilder::normalize(plan, c)?;
    74
    0
                    Transformed::yes(Expr::Column(col))
    75
                } else {
    76
    0
                    Transformed::no(expr)
    77
                }
    78
            })
    79
    0
        })
    80
    0
        .data()
    81
    0
    }
    82
    83
    /// See [`Column::normalize_with_schemas_and_ambiguity_check`] for usage
    84
    0
    pub fn normalize_col_with_schemas_and_ambiguity_check(
    85
    0
        expr: Expr,
    86
    0
        schemas: &[&[&DFSchema]],
    87
    0
        using_columns: &[HashSet<Column>],
    88
    0
    ) -> Result<Expr> {
    89
        // Normalize column inside Unnest
    90
    0
        if let Expr::Unnest(Unnest { expr }) = expr {
    91
    0
            let e = normalize_col_with_schemas_and_ambiguity_check(
    92
    0
                expr.as_ref().clone(),
    93
    0
                schemas,
    94
    0
                using_columns,
    95
    0
            )?;
    96
    0
            return Ok(Expr::Unnest(Unnest { expr: Box::new(e) }));
    97
    0
        }
    98
    99
    0
        expr.transform(|expr| {
    100
            Ok({
    101
    0
                if let Expr::Column(c) = expr {
    102
    0
                    let col =
    103
    0
                        c.normalize_with_schemas_and_ambiguity_check(schemas, using_columns)?;
    104
    0
                    Transformed::yes(Expr::Column(col))
    105
                } else {
    106
    0
                    Transformed::no(expr)
    107
                }
    108
            })
    109
    0
        })
    110
    0
        .data()
    111
    0
    }
    112
    113
    /// Recursively normalize all [`Column`] expressions in a list of expression trees
    114
    0
    pub fn normalize_cols(
    115
    0
        exprs: impl IntoIterator<Item = impl Into<Expr>>,
    116
    0
        plan: &LogicalPlan,
    117
    0
    ) -> Result<Vec<Expr>> {
    118
    0
        exprs
    119
    0
            .into_iter()
    120
    0
            .map(|e| normalize_col(e.into(), plan))
    121
    0
            .collect()
    122
    0
    }
    123
    124
    0
    pub fn normalize_sorts(
    125
    0
        sorts: impl IntoIterator<Item = impl Into<Sort>>,
    126
    0
        plan: &LogicalPlan,
    127
    0
    ) -> Result<Vec<Sort>> {
    128
    0
        sorts
    129
    0
            .into_iter()
    130
    0
            .map(|e| {
    131
    0
                let sort = e.into();
    132
    0
                normalize_col(sort.expr, plan)
    133
    0
                    .map(|expr| Sort::new(expr, sort.asc, sort.nulls_first))
    134
    0
            })
    135
    0
            .collect()
    136
    0
    }
    137
    138
    /// Recursively replace all [`Column`] expressions in a given expression tree with
    139
    /// `Column` expressions provided by the hash map argument.
    140
    0
    pub fn replace_col(expr: Expr, replace_map: &HashMap<&Column, &Column>) -> Result<Expr> {
    141
    0
        expr.transform(|expr| {
    142
            Ok({
    143
    0
                if let Expr::Column(c) = &expr {
    144
    0
                    match replace_map.get(c) {
    145
    0
                        Some(new_c) => Transformed::yes(Expr::Column((*new_c).to_owned())),
    146
    0
                        None => Transformed::no(expr),
    147
                    }
    148
                } else {
    149
    0
                    Transformed::no(expr)
    150
                }
    151
            })
    152
    0
        })
    153
    0
        .data()
    154
    0
    }
    155
    156
    /// Recursively 'unnormalize' (remove all qualifiers) from an
    157
    /// expression tree.
    158
    ///
    159
    /// For example, if there were expressions like `foo.bar` this would
    160
    /// rewrite it to just `bar`.
    161
    0
    pub fn unnormalize_col(expr: Expr) -> Expr {
    162
    0
        expr.transform(|expr| {
    163
            Ok({
    164
    0
                if let Expr::Column(c) = expr {
    165
    0
                    let col = Column::new_unqualified(c.name);
    166
    0
                    Transformed::yes(Expr::Column(col))
    167
                } else {
    168
    0
                    Transformed::no(expr)
    169
                }
    170
            })
    171
    0
        })
    172
    0
        .data()
    173
    0
        .expect("Unnormalize is infallible")
    174
    0
    }
    175
    176
    /// Create a Column from the Scalar Expr
    177
    0
    pub fn create_col_from_scalar_expr(
    178
    0
        scalar_expr: &Expr,
    179
    0
        subqry_alias: String,
    180
    0
    ) -> Result<Column> {
    181
    0
        match scalar_expr {
    182
    0
            Expr::Alias(Alias { name, .. }) => Ok(Column::new(
    183
    0
                Some::<TableReference>(subqry_alias.into()),
    184
    0
                name,
    185
    0
            )),
    186
    0
            Expr::Column(col) => Ok(col.with_relation(subqry_alias.into())),
    187
            _ => {
    188
    0
                let scalar_column = scalar_expr.schema_name().to_string();
    189
    0
                Ok(Column::new(
    190
    0
                    Some::<TableReference>(subqry_alias.into()),
    191
    0
                    scalar_column,
    192
    0
                ))
    193
            }
    194
        }
    195
    0
    }
    196
    197
    /// Recursively un-normalize all [`Column`] expressions in a list of expression trees
    198
    #[inline]
    199
    0
    pub fn unnormalize_cols(exprs: impl IntoIterator<Item = Expr>) -> Vec<Expr> {
    200
    0
        exprs.into_iter().map(unnormalize_col).collect()
    201
    0
    }
    202
    203
    /// Recursively remove all the ['OuterReferenceColumn'] and return the inside Column
    204
    /// in the expression tree.
    205
    0
    pub fn strip_outer_reference(expr: Expr) -> Expr {
    206
    0
        expr.transform(|expr| {
    207
            Ok({
    208
    0
                if let Expr::OuterReferenceColumn(_, col) = expr {
    209
    0
                    Transformed::yes(Expr::Column(col))
    210
                } else {
    211
    0
                    Transformed::no(expr)
    212
                }
    213
            })
    214
    0
        })
    215
    0
        .data()
    216
    0
        .expect("strip_outer_reference is infallible")
    217
    0
    }
    218
    219
    /// Returns plan with expressions coerced to types compatible with
    220
    /// schema types
    221
    0
    pub fn coerce_plan_expr_for_schema(
    222
    0
        plan: LogicalPlan,
    223
    0
        schema: &DFSchema,
    224
    0
    ) -> Result<LogicalPlan> {
    225
    0
        match plan {
    226
            // special case Projection to avoid adding multiple projections
    227
    0
            LogicalPlan::Projection(Projection { expr, input, .. }) => {
    228
    0
                let new_exprs = coerce_exprs_for_schema(expr, input.schema(), schema)?;
    229
    0
                let projection = Projection::try_new(new_exprs, input)?;
    230
    0
                Ok(LogicalPlan::Projection(projection))
    231
            }
    232
            _ => {
    233
    0
                let exprs: Vec<Expr> = plan.schema().iter().map(Expr::from).collect();
    234
    0
                let new_exprs = coerce_exprs_for_schema(exprs, plan.schema(), schema)?;
    235
    0
                let add_project = new_exprs.iter().any(|expr| expr.try_as_col().is_none());
    236
    0
                if add_project {
    237
    0
                    let projection = Projection::try_new(new_exprs, Arc::new(plan))?;
    238
    0
                    Ok(LogicalPlan::Projection(projection))
    239
                } else {
    240
    0
                    Ok(plan)
    241
                }
    242
            }
    243
        }
    244
    0
    }
    245
    246
    0
    fn coerce_exprs_for_schema(
    247
    0
        exprs: Vec<Expr>,
    248
    0
        src_schema: &DFSchema,
    249
    0
        dst_schema: &DFSchema,
    250
    0
    ) -> Result<Vec<Expr>> {
    251
    0
        exprs
    252
    0
            .into_iter()
    253
    0
            .enumerate()
    254
    0
            .map(|(idx, expr)| {
    255
    0
                let new_type = dst_schema.field(idx).data_type();
    256
    0
                if new_type != &expr.get_type(src_schema)? {
    257
    0
                    match expr {
    258
    0
                        Expr::Alias(Alias { expr, name, .. }) => {
    259
    0
                            Ok(expr.cast_to(new_type, src_schema)?.alias(name))
    260
                        }
    261
                        #[expect(deprecated)]
    262
    0
                        Expr::Wildcard { .. } => Ok(expr),
    263
    0
                        _ => expr.cast_to(new_type, src_schema),
    264
                    }
    265
                } else {
    266
    0
                    Ok(expr)
    267
                }
    268
    0
            })
    269
    0
            .collect::<Result<_>>()
    270
    0
    }
    271
    272
    /// Recursively un-alias an expressions
    273
    #[inline]
    274
    0
    pub fn unalias(expr: Expr) -> Expr {
    275
    0
        match expr {
    276
    0
            Expr::Alias(Alias { expr, .. }) => unalias(*expr),
    277
    0
            _ => expr,
    278
        }
    279
    0
    }
    280
    281
    /// Handles ensuring the name of rewritten expressions is not changed.
    282
    ///
    283
    /// This is important when optimizing plans to ensure the output
    284
    /// schema of plan nodes don't change after optimization.
    285
    /// For example, if an expression `1 + 2` is rewritten to `3`, the name of the
    286
    /// expression should be preserved: `3 as "1 + 2"`
    287
    ///
    288
    /// See <https://github.com/apache/datafusion/issues/3555> for details
    289
    pub struct NamePreserver {
    290
        use_alias: bool,
    291
    }
    292
    293
    /// If the qualified name of an expression is remembered, it will be preserved
    294
    /// when rewriting the expression
    295
    #[derive(Debug)]
    296
    pub enum SavedName {
    297
        /// Saved qualified name to be preserved
    298
        Saved {
    299
            relation: Option<TableReference>,
    300
            name: String,
    301
        },
    302
        /// Name is not preserved
    303
        None,
    304
    }
    305
    306
    impl NamePreserver {
    307
        /// Create a new NamePreserver for rewriting the `expr` that is part of the specified plan
    308
    0
        pub fn new(plan: &LogicalPlan) -> Self {
    309
            Self {
    310
                // The expressions of these plans do not contribute to their output schema,
    311
                // so there is no need to preserve expression names to prevent a schema change.
    312
    0
                use_alias: !matches!(
    313
    0
                    plan,
    314
                    LogicalPlan::Filter(_)
    315
                        | LogicalPlan::Join(_)
    316
                        | LogicalPlan::TableScan(_)
    317
                        | LogicalPlan::Limit(_)
    318
                        | LogicalPlan::Statement(_)
    319
                ),
    320
            }
    321
    0
        }
    322
    323
        /// Create a new NamePreserver for rewriting the `expr`s in `Projection`
    324
        ///
    325
        /// This will use aliases
    326
    0
        pub fn new_for_projection() -> Self {
    327
    0
            Self { use_alias: true }
    328
    0
        }
    329
    330
    0
        pub fn save(&self, expr: &Expr) -> SavedName {
    331
    0
            if self.use_alias {
    332
    0
                let (relation, name) = expr.qualified_name();
    333
    0
                SavedName::Saved { relation, name }
    334
            } else {
    335
    0
                SavedName::None
    336
            }
    337
    0
        }
    338
    }
    339
    340
    impl SavedName {
    341
        /// Ensures the qualified name of the rewritten expression is preserved
    342
    0
        pub fn restore(self, expr: Expr) -> Expr {
    343
    0
            match self {
    344
    0
                SavedName::Saved { relation, name } => {
    345
    0
                    let (new_relation, new_name) = expr.qualified_name();
    346
    0
                    if new_relation != relation || new_name != name {
    347
    0
                        expr.alias_qualified(relation, name)
    348
                    } else {
    349
    0
                        expr
    350
                    }
    351
                }
    352
    0
                SavedName::None => expr,
    353
            }
    354
    0
        }
    355
    }
    356
    357
    #[cfg(test)]
    358
    mod test {
    359
        use std::ops::Add;
    360
    361
        use super::*;
    362
        use crate::literal::lit_with_metadata;
    363
        use crate::{col, lit, Cast};
    364
        use arrow::datatypes::{DataType, Field, Schema};
    365
        use datafusion_common::tree_node::TreeNodeRewriter;
    366
        use datafusion_common::ScalarValue;
    367
    368
        #[derive(Default)]
    369
        struct RecordingRewriter {
    370
            v: Vec<String>,
    371
        }
    372
    373
        impl TreeNodeRewriter for RecordingRewriter {
    374
            type Node = Expr;
    375
    376
            fn f_down(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
    377
                self.v.push(format!("Previsited {expr}"));
    378
                Ok(Transformed::no(expr))
    379
            }
    380
    381
            fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
    382
                self.v.push(format!("Mutated {expr}"));
    383
                Ok(Transformed::no(expr))
    384
            }
    385
        }
    386
    387
        #[test]
    388
        fn rewriter_rewrite() {
    389
            // rewrites all "foo" string literals to "bar"
    390
            let transformer = |expr: Expr| -> Result<Transformed<Expr>> {
    391
                match expr {
    392
                    Expr::Literal(ScalarValue::Utf8(Some(utf8_val)), metadata) => {
    393
                        let utf8_val = if utf8_val == "foo" {
    394
                            "bar".to_string()
    395
                        } else {
    396
                            utf8_val
    397
                        };
    398
                        Ok(Transformed::yes(lit_with_metadata(utf8_val, metadata)))
    399
                    }
    400
                    // otherwise, return None
    401
                    _ => Ok(Transformed::no(expr)),
    402
                }
    403
            };
    404
    405
            // rewrites "foo" --> "bar"
    406
            let rewritten = col("state")
    407
                .eq(lit("foo"))
    408
                .transform(transformer)
    409
                .data()
    410
                .unwrap();
    411
            assert_eq!(rewritten, col("state").eq(lit("bar")));
    412
    413
            // doesn't rewrite
    414
            let rewritten = col("state")
    415
                .eq(lit("baz"))
    416
                .transform(transformer)
    417
                .data()
    418
                .unwrap();
    419
            assert_eq!(rewritten, col("state").eq(lit("baz")));
    420
        }
    421
    422
        #[test]
    423
        fn normalize_cols() {
    424
            let expr = col("a") + col("b") + col("c");
    425
    426
            // Schemas with some matching and some non matching cols
    427
            let schema_a = make_schema_with_empty_metadata(
    428
                vec![Some("tableA".into()), Some("tableA".into())],
    429
                vec!["a", "aa"],
    430
            );
    431
            let schema_c = make_schema_with_empty_metadata(
    432
                vec![Some("tableC".into()), Some("tableC".into())],
    433
                vec!["cc", "c"],
    434
            );
    435
            let schema_b =
    436
                make_schema_with_empty_metadata(vec![Some("tableB".into())], vec!["b"]);
    437
            // non matching
    438
            let schema_f = make_schema_with_empty_metadata(
    439
                vec![Some("tableC".into()), Some("tableC".into())],
    440
                vec!["f", "ff"],
    441
            );
    442
            let schemas = [schema_c, schema_f, schema_b, schema_a];
    443
            let schemas = schemas.iter().collect::<Vec<_>>();
    444
    445
            let normalized_expr =
    446
                normalize_col_with_schemas_and_ambiguity_check(expr, &[&schemas], &[])
    447
                    .unwrap();
    448
            assert_eq!(
    449
                normalized_expr,
    450
                col("tableA.a") + col("tableB.b") + col("tableC.c")
    451
            );
    452
        }
    453
    454
        #[test]
    455
        fn normalize_cols_non_exist() {
    456
            // test normalizing columns when the name doesn't exist
    457
            let expr = col("a") + col("b");
    458
            let schema_a =
    459
                make_schema_with_empty_metadata(vec![Some("\"tableA\"".into())], vec!["a"]);
    460
            let schemas = [schema_a];
    461
            let schemas = schemas.iter().collect::<Vec<_>>();
    462
    463
            let error =
    464
                normalize_col_with_schemas_and_ambiguity_check(expr, &[&schemas], &[])
    465
                    .unwrap_err()
    466
                    .strip_backtrace();
    467
            let expected = "Schema error: No field named b. \
    468
                Valid fields are \"tableA\".a.";
    469
            assert_eq!(error, expected);
    470
        }
    471
    472
        #[test]
    473
        fn unnormalize_cols() {
    474
            let expr = col("tableA.a") + col("tableB.b");
    475
            let unnormalized_expr = unnormalize_col(expr);
    476
            assert_eq!(unnormalized_expr, col("a") + col("b"));
    477
        }
    478
    479
        fn make_schema_with_empty_metadata(
    480
            qualifiers: Vec<Option<TableReference>>,
    481
            fields: Vec<&str>,
    482
        ) -> DFSchema {
    483
            let fields = fields
    484
                .iter()
    485
                .map(|f| Arc::new(Field::new((*f).to_string(), DataType::Int8, false)))
    486
                .collect::<Vec<_>>();
    487
            let schema = Arc::new(Schema::new(fields));
    488
            DFSchema::from_field_specific_qualified_schema(qualifiers, &schema).unwrap()
    489
        }
    490
    491
        #[test]
    492
        fn rewriter_visit() {
    493
            let mut rewriter = RecordingRewriter::default();
    494
            col("state").eq(lit("CO")).rewrite(&mut rewriter).unwrap();
    495
    496
            assert_eq!(
    497
                rewriter.v,
    498
                vec![
    499
                    "Previsited state = Utf8(\"CO\")",
    500
                    "Previsited state",
    501
                    "Mutated state",
    502
                    "Previsited Utf8(\"CO\")",
    503
                    "Mutated Utf8(\"CO\")",
    504
                    "Mutated state = Utf8(\"CO\")"
    505
                ]
    506
            )
    507
        }
    508
    509
        #[test]
    510
        fn test_rewrite_preserving_name() {
    511
            test_rewrite(col("a"), col("a"));
    512
    513
            test_rewrite(col("a"), col("b"));
    514
    515
            // cast data types
    516
            test_rewrite(
    517
                col("a"),
    518
                Expr::Cast(Cast::new(Box::new(col("a")), DataType::Int32)),
    519
            );
    520
    521
            // change literal type from i32 to i64
    522
            test_rewrite(col("a").add(lit(1i32)), col("a").add(lit(1i64)));
    523
    524
            // test preserve qualifier
    525
            test_rewrite(
    526
                Expr::Column(Column::new(Some("test"), "a")),
    527
                Expr::Column(Column::new_unqualified("test.a")),
    528
            );
    529
            test_rewrite(
    530
                Expr::Column(Column::new_unqualified("test.a")),
    531
                Expr::Column(Column::new(Some("test"), "a")),
    532
            );
    533
        }
    534
    535
        /// rewrites `expr_from` to `rewrite_to` while preserving the original qualified name
    536
        /// by using the `NamePreserver`
    537
        fn test_rewrite(expr_from: Expr, rewrite_to: Expr) {
    538
            struct TestRewriter {
    539
                rewrite_to: Expr,
    540
            }
    541
    542
            impl TreeNodeRewriter for TestRewriter {
    543
                type Node = Expr;
    544
    545
                fn f_up(&mut self, _: Expr) -> Result<Transformed<Expr>> {
    546
                    Ok(Transformed::yes(self.rewrite_to.clone()))
    547
                }
    548
            }
    549
    550
            let mut rewriter = TestRewriter {
    551
                rewrite_to: rewrite_to.clone(),
    552
            };
    553
            let saved_name = NamePreserver { use_alias: true }.save(&expr_from);
    554
            let new_expr = expr_from.clone().rewrite(&mut rewriter).unwrap().data;
    555
            let new_expr = saved_name.restore(new_expr);
    556
    557
            let original_name = expr_from.qualified_name();
    558
            let new_name = new_expr.qualified_name();
    559
            assert_eq!(
    560
                original_name, new_name,
    561
                "mismatch rewriting expr_from: {expr_from} to {rewrite_to}"
    562
            )
    563
        }
    564
    }