## 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/conditional_expressions.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
    //! Conditional expressions
    19
    use crate::expr::Case;
    20
    use crate::{expr_schema::ExprSchemable, Expr};
    21
    use arrow::datatypes::DataType;
    22
    use datafusion_common::{plan_err, DFSchema, HashSet, Result};
    23
    use itertools::Itertools as _;
    24
    25
    /// Helper struct for building [Expr::Case]
    26
    #[derive(Debug, Clone)]
    27
    pub struct CaseBuilder {
    28
        expr: Option<Box<Expr>>,
    29
        when_expr: Vec<Expr>,
    30
        then_expr: Vec<Expr>,
    31
        else_expr: Option<Box<Expr>>,
    32
    }
    33
    34
    impl CaseBuilder {
    35
    0
        pub fn new(
    36
    0
            expr: Option<Box<Expr>>,
    37
    0
            when_expr: Vec<Expr>,
    38
    0
            then_expr: Vec<Expr>,
    39
    0
            else_expr: Option<Box<Expr>>,
    40
    0
        ) -> Self {
    41
    0
            Self {
    42
    0
                expr,
    43
    0
                when_expr,
    44
    0
                then_expr,
    45
    0
                else_expr,
    46
    0
            }
    47
    0
        }
    48
    0
        pub fn when(&mut self, when: Expr, then: Expr) -> CaseBuilder {
    49
    0
            self.when_expr.push(when);
    50
    0
            self.then_expr.push(then);
    51
    0
            CaseBuilder {
    52
    0
                expr: self.expr.clone(),
    53
    0
                when_expr: self.when_expr.clone(),
    54
    0
                then_expr: self.then_expr.clone(),
    55
    0
                else_expr: self.else_expr.clone(),
    56
    0
            }
    57
    0
        }
    58
    0
        pub fn otherwise(&mut self, else_expr: Expr) -> Result<Expr> {
    59
    0
            self.else_expr = Some(Box::new(else_expr));
    60
    0
            self.build()
    61
    0
        }
    62
    63
    0
        pub fn end(&self) -> Result<Expr> {
    64
    0
            self.build()
    65
    0
        }
    66
    67
    0
        fn build(&self) -> Result<Expr> {
    68
            // Collect all "then" expressions
    69
    0
            let mut then_expr = self.then_expr.clone();
    70
    0
            if let Some(e) = &self.else_expr {
    71
    0
                then_expr.push(e.as_ref().to_owned());
    72
    0
            }
    73
    74
    0
            let then_types: Vec<DataType> = then_expr
    75
    0
                .iter()
    76
    0
                .map(|e| match e {
    77
    0
                    Expr::Literal(_, _) => e.get_type(&DFSchema::empty()),
    78
    0
                    _ => Ok(DataType::Null),
    79
    0
                })
    80
    0
                .collect::<Result<Vec<_>>>()?;
    81
    82
    0
            if then_types.contains(&DataType::Null) {
    83
    0
                // Cannot verify types until execution type
    84
    0
            } else {
    85
    0
                let unique_types: HashSet<&DataType> = then_types.iter().collect();
    86
    0
                if unique_types.is_empty() {
    87
    0
                    return plan_err!("CASE expression 'then' values had no data types");
    88
    0
                } else if unique_types.len() != 1 {
    89
    0
                    return plan_err!(
    90
    0
                        "CASE expression 'then' values had multiple data types: {}",
    91
    0
                        unique_types.iter().join(", ")
    92
                    );
    93
    0
                }
    94
            }
    95
    96
    0
            Ok(Expr::Case(Case::new(
    97
    0
                self.expr.clone(),
    98
    0
                self.when_expr
    99
    0
                    .iter()
    100
    0
                    .zip(self.then_expr.iter())
    101
    0
                    .map(|(w, t)| (Box::new(w.clone()), Box::new(t.clone())))
    102
    0
                    .collect(),
    103
    0
                self.else_expr.clone(),
    104
            )))
    105
    0
        }
    106
    }
    107
    108
    #[cfg(test)]
    109
    mod tests {
    110
        use super::*;
    111
        use crate::{col, lit, when};
    112
    113
        #[test]
    114
        fn case_when_same_literal_then_types() -> Result<()> {
    115
            let _ = when(col("state").eq(lit("CO")), lit(303))
    116
                .when(col("state").eq(lit("NY")), lit(212))
    117
                .end()?;
    118
            Ok(())
    119
        }
    120
    121
        #[test]
    122
        fn case_when_different_literal_then_types() {
    123
            let maybe_expr = when(col("state").eq(lit("CO")), lit(303))
    124
                .when(col("state").eq(lit("NY")), lit("212"))
    125
                .end();
    126
            assert!(maybe_expr.is_err());
    127
        }
    128
    }