## 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/expressions/case.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 super::{Column, Literal};
    19
    use crate::expressions::case::ResultState::{Complete, Empty, Partial};
    20
    use crate::expressions::{lit, try_cast};
    21
    use crate::PhysicalExpr;
    22
    use arrow::array::*;
    23
    use arrow::compute::kernels::zip::zip;
    24
    use arrow::compute::{
    25
        is_not_null, not, nullif, prep_null_mask_filter, FilterBuilder, FilterPredicate,
    26
        SlicesIterator,
    27
    };
    28
    use arrow::datatypes::{DataType, Schema, UInt32Type, UnionMode};
    29
    use arrow::error::ArrowError;
    30
    use datafusion_common::cast::as_boolean_array;
    31
    use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
    32
    use datafusion_common::{
    33
        assert_or_internal_err, exec_err, internal_datafusion_err, internal_err,
    34
        DataFusionError, HashMap, HashSet, Result, ScalarValue,
    35
    };
    36
    use datafusion_expr::ColumnarValue;
    37
    use std::borrow::Cow;
    38
    use std::hash::Hash;
    39
    use std::{any::Any, sync::Arc};
    40
    41
    use datafusion_physical_expr_common::datum::compare_with_eq;
    42
    use itertools::Itertools;
    43
    use std::fmt::{Debug, Formatter};
    44
    45
    type WhenThen = (Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>);
    46
    47
    #[derive(Debug, Hash, PartialEq, Eq)]
    48
    enum EvalMethod {
    49
        /// CASE WHEN condition THEN result
    50
        ///      [WHEN ...]
    51
        ///      [ELSE result]
    52
        /// END
    53
        NoExpression(ProjectedCaseBody),
    54
        /// CASE expression
    55
        ///     WHEN value THEN result
    56
        ///     [WHEN ...]
    57
        ///     [ELSE result]
    58
        /// END
    59
        WithExpression(ProjectedCaseBody),
    60
        /// This is a specialization for a specific use case where we can take a fast path
    61
        /// for expressions that are infallible and can be cheaply computed for the entire
    62
        /// record batch rather than just for the rows where the predicate is true.
    63
        ///
    64
        /// CASE WHEN condition THEN column [ELSE NULL] END
    65
        InfallibleExprOrNull,
    66
        /// This is a specialization for a specific use case where we can take a fast path
    67
        /// if there is just one when/then pair and both the `then` and `else` expressions
    68
        /// are literal values
    69
        /// CASE WHEN condition THEN literal ELSE literal END
    70
        ScalarOrScalar,
    71
        /// This is a specialization for a specific use case where we can take a fast path
    72
        /// if there is just one when/then pair and both the `then` and `else` are expressions
    73
        ///
    74
        /// CASE WHEN condition THEN expression ELSE expression END
    75
        ExpressionOrExpression(ProjectedCaseBody),
    76
    }
    77
    78
    /// The body of a CASE expression which consists of an optional base expression, the "when/then"
    79
    /// branches and an optional "else" branch.
    80
    #[derive(Debug, Hash, PartialEq, Eq)]
    81
    struct CaseBody {
    82
        /// Optional base expression that can be compared to literal values in the "when" expressions
    83
        expr: Option<Arc<dyn PhysicalExpr>>,
    84
        /// One or more when/then expressions
    85
        when_then_expr: Vec<WhenThen>,
    86
        /// Optional "else" expression
    87
        else_expr: Option<Arc<dyn PhysicalExpr>>,
    88
    }
    89
    90
    impl CaseBody {
    91
        /// Derives a [ProjectedCaseBody] from this [CaseBody].
    92
    0
        fn project(&self) -> Result<ProjectedCaseBody> {
    93
            // Determine the set of columns that are used in all the expressions of the case body.
    94
    0
            let mut used_column_indices = HashSet::<usize>::new();
    95
    0
            let mut collect_column_indices = |expr: &Arc<dyn PhysicalExpr>| {
    96
    0
                expr.apply(|expr| {
    97
    0
                    if let Some(column) = expr.as_any().downcast_ref::<Column>() {
    98
    0
                        used_column_indices.insert(column.index());
    99
    0
                    }
    100
    0
                    Ok(TreeNodeRecursion::Continue)
    101
    0
                })
    102
    0
                .expect("Closure cannot fail");
    103
    0
            };
    104
    105
    0
            if let Some(e) = &self.expr {
    106
    0
                collect_column_indices(e);
    107
    0
            }
    108
    0
            self.when_then_expr.iter().for_each(|(w, t)| {
    109
    0
                collect_column_indices(w);
    110
    0
                collect_column_indices(t);
    111
    0
            });
    112
    0
            if let Some(e) = &self.else_expr {
    113
    0
                collect_column_indices(e);
    114
    0
            }
    115
    116
            // Construct a mapping from the original column index to the projected column index.
    117
    0
            let column_index_map = used_column_indices
    118
    0
                .iter()
    119
    0
                .enumerate()
    120
    0
                .map(|(projected, original)| (*original, projected))
    121
    0
                .collect::<HashMap<usize, usize>>();
    122
    123
            // Construct the projected body by rewriting each expression from the original body
    124
            // using the column index mapping.
    125
    0
            let project = |expr: &Arc<dyn PhysicalExpr>| -> Result<Arc<dyn PhysicalExpr>> {
    126
    0
                Arc::clone(expr)
    127
    0
                    .transform_down(|e| {
    128
    0
                        if let Some(column) = e.as_any().downcast_ref::<Column>() {
    129
    0
                            let original = column.index();
    130
    0
                            let projected = *column_index_map.get(&original).unwrap();
    131
    0
                            if projected != original {
    132
    0
                                return Ok(Transformed::yes(Arc::new(Column::new(
    133
    0
                                    column.name(),
    134
    0
                                    projected,
    135
    0
                                ))));
    136
    0
                            }
    137
    0
                        }
    138
    0
                        Ok(Transformed::no(e))
    139
    0
                    })
    140
    0
                    .map(|t| t.data)
    141
    0
            };
    142
    143
    0
            let projected_body = CaseBody {
    144
    0
                expr: self.expr.as_ref().map(project).transpose()?,
    145
    0
                when_then_expr: self
    146
    0
                    .when_then_expr
    147
    0
                    .iter()
    148
    0
                    .map(|(e, t)| Ok((project(e)?, project(t)?)))
    149
    0
                    .collect::<Result<Vec<_>>>()?,
    150
    0
                else_expr: self.else_expr.as_ref().map(project).transpose()?,
    151
            };
    152
    153
            // Construct the projection vector
    154
    0
            let projection = column_index_map
    155
    0
                .iter()
    156
    0
                .sorted_by_key(|(_, v)| **v)
    157
    0
                .map(|(k, _)| *k)
    158
    0
                .collect::<Vec<_>>();
    159
    160
    0
            Ok(ProjectedCaseBody {
    161
    0
                projection,
    162
    0
                body: projected_body,
    163
    0
            })
    164
    0
        }
    165
    }
    166
    167
    /// A derived case body that can be used to evaluate a case expression after projecting
    168
    /// record batches using a projection vector.
    169
    ///
    170
    /// This is used to avoid filtering columns that are not used in the
    171
    /// input `RecordBatch` when progressively evaluating a `CASE` expression's
    172
    /// remainder batches. Filtering these columns is wasteful since for a record
    173
    /// batch of `n` rows, filtering requires at worst a copy of `n - 1` values
    174
    /// per array. If these filtered values will never be accessed, the time spent
    175
    /// producing them is better avoided.
    176
    ///
    177
    /// For example, if we are evaluating the following case expression that
    178
    /// only references columns B and D:
    179
    ///
    180
    /// ```sql
    181
    /// SELECT CASE WHEN B > 10 THEN D ELSE NULL END FROM (VALUES (...)) T(A, B, C, D)
    182
    /// ```
    183
    ///
    184
    /// Of the 4 input columns `[A, B, C, D]`, the `CASE` expression only access `B` and `D`.
    185
    /// Filtering `A` and `C` would be unnecessary and wasteful.
    186
    ///
    187
    /// If we only retain columns `B` and `D` using `RecordBatch::project` and the projection vector
    188
    /// `[1, 3]`, the indices of these two columns will change to `[0, 1]`. To evaluate the
    189
    /// case expression, it will need to be rewritten from `CASE WHEN B@1 > 10 THEN D@3 ELSE NULL END`
    190
    /// to `CASE WHEN B@0 > 10 THEN D@1 ELSE NULL END`.
    191
    ///
    192
    /// The projection vector and the rewritten expression (which only differs from the original in
    193
    /// column reference indices) are held in a `ProjectedCaseBody`.
    194
    #[derive(Debug, Hash, PartialEq, Eq)]
    195
    struct ProjectedCaseBody {
    196
        projection: Vec<usize>,
    197
        body: CaseBody,
    198
    }
    199
    200
    /// The CASE expression is similar to a series of nested if/else and there are two forms that
    201
    /// can be used. The first form consists of a series of boolean "when" expressions with
    202
    /// corresponding "then" expressions, and an optional "else" expression.
    203
    ///
    204
    /// CASE WHEN condition THEN result
    205
    ///      [WHEN ...]
    206
    ///      [ELSE result]
    207
    /// END
    208
    ///
    209
    /// The second form uses a base expression and then a series of "when" clauses that match on a
    210
    /// literal value.
    211
    ///
    212
    /// CASE expression
    213
    ///     WHEN value THEN result
    214
    ///     [WHEN ...]
    215
    ///     [ELSE result]
    216
    /// END
    217
    #[derive(Debug, Hash, PartialEq, Eq)]
    218
    pub struct CaseExpr {
    219
        /// The case expression body
    220
        body: CaseBody,
    221
        /// Evaluation method to use
    222
        eval_method: EvalMethod,
    223
    }
    224
    225
    impl std::fmt::Display for CaseExpr {
    226
    0
        fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
    227
    0
            write!(f, "CASE ")?;
    228
    0
            if let Some(e) = &self.body.expr {
    229
    0
                write!(f, "{e} ")?;
    230
    0
            }
    231
    0
            for (w, t) in &self.body.when_then_expr {
    232
    0
                write!(f, "WHEN {w} THEN {t} ")?;
    233
            }
    234
    0
            if let Some(e) = &self.body.else_expr {
    235
    0
                write!(f, "ELSE {e} ")?;
    236
    0
            }
    237
    0
            write!(f, "END")
    238
    0
        }
    239
    }
    240
    241
    /// This is a specialization for a specific use case where we can take a fast path
    242
    /// for expressions that are infallible and can be cheaply computed for the entire
    243
    /// record batch rather than just for the rows where the predicate is true. For now,
    244
    /// this is limited to use with Column expressions but could potentially be used for other
    245
    /// expressions in the future
    246
    0
    fn is_cheap_and_infallible(expr: &Arc<dyn PhysicalExpr>) -> bool {
    247
    0
        expr.as_any().is::<Column>()
    248
    0
    }
    249
    250
    /// Creates a [FilterPredicate] from a boolean array.
    251
    0
    fn create_filter(predicate: &BooleanArray, optimize: bool) -> FilterPredicate {
    252
    0
        let mut filter_builder = FilterBuilder::new(predicate);
    253
    0
        if optimize {
    254
    0
            // Always optimize the filter since we use them multiple times.
    255
    0
            filter_builder = filter_builder.optimize();
    256
    0
        }
    257
    0
        filter_builder.build()
    258
    0
    }
    259
    260
    0
    fn multiple_arrays(data_type: &DataType) -> bool {
    261
    0
        match data_type {
    262
    0
            DataType::Struct(fields) => {
    263
    0
                fields.len() > 1
    264
    0
                    || fields.len() == 1 && multiple_arrays(fields[0].data_type())
    265
            }
    266
    0
            DataType::Union(fields, UnionMode::Sparse) => !fields.is_empty(),
    267
    0
            _ => false,
    268
        }
    269
    0
    }
    270
    271
    // This should be removed when https://github.com/apache/arrow-rs/pull/8693
    272
    // is merged and becomes available.
    273
    0
    fn filter_record_batch(
    274
    0
        record_batch: &RecordBatch,
    275
    0
        filter: &FilterPredicate,
    276
    0
    ) -> std::result::Result<RecordBatch, ArrowError> {
    277
    0
        let filtered_columns = record_batch
    278
    0
            .columns()
    279
    0
            .iter()
    280
    0
            .map(|a| filter_array(a, filter))
    281
    0
            .collect::<std::result::Result<Vec<_>, _>>()?;
    282
        // SAFETY: since we start from a valid RecordBatch, there's no need to revalidate the schema
    283
        // since the set of columns has not changed.
    284
        // The input column arrays all had the same length (since they're coming from a valid RecordBatch)
    285
        // and the filtering them with the same filter will produces a new set of arrays with identical
    286
        // lengths.
    287
        unsafe {
    288
    0
            Ok(RecordBatch::new_unchecked(
    289
    0
                record_batch.schema(),
    290
    0
                filtered_columns,
    291
    0
                filter.count(),
    292
    0
            ))
    293
        }
    294
    0
    }
    295
    296
    // This function exists purely to be able to use the same call style
    297
    // for `filter_record_batch` and `filter_array` at the point of use.
    298
    // When https://github.com/apache/arrow-rs/pull/8693 is available, replace
    299
    // both with method calls on `FilterPredicate`.
    300
    #[inline(always)]
    301
    0
    fn filter_array(
    302
    0
        array: &dyn Array,
    303
    0
        filter: &FilterPredicate,
    304
    0
    ) -> std::result::Result<ArrayRef, ArrowError> {
    305
    0
        filter.filter(array)
    306
    0
    }
    307
    308
    0
    fn merge(
    309
    0
        mask: &BooleanArray,
    310
    0
        truthy: ColumnarValue,
    311
    0
        falsy: ColumnarValue,
    312
    0
    ) -> std::result::Result<ArrayRef, ArrowError> {
    313
    0
        let (truthy, truthy_is_scalar) = match truthy {
    314
    0
            ColumnarValue::Array(a) => (a, false),
    315
    0
            ColumnarValue::Scalar(s) => (s.to_array()?, true),
    316
        };
    317
    0
        let (falsy, falsy_is_scalar) = match falsy {
    318
    0
            ColumnarValue::Array(a) => (a, false),
    319
    0
            ColumnarValue::Scalar(s) => (s.to_array()?, true),
    320
        };
    321
    322
    0
        if truthy_is_scalar && falsy_is_scalar {
    323
    0
            return zip(mask, &Scalar::new(truthy), &Scalar::new(falsy));
    324
    0
        }
    325
    326
    0
        let falsy = falsy.to_data();
    327
    0
        let truthy = truthy.to_data();
    328
    329
    0
        let mut mutable = MutableArrayData::new(vec![&truthy, &falsy], false, truthy.len());
    330
    331
        // the SlicesIterator slices only the true values. So the gaps left by this iterator we need to
    332
        // fill with falsy values
    333
    334
        // keep track of how much is filled
    335
    0
        let mut filled = 0;
    336
    0
        let mut falsy_offset = 0;
    337
    0
        let mut truthy_offset = 0;
    338
    339
    0
        SlicesIterator::new(mask).for_each(|(start, end)| {
    340
            // the gap needs to be filled with falsy values
    341
    0
            if start > filled {
    342
    0
                if falsy_is_scalar {
    343
    0
                    for _ in filled..start {
    344
    0
                        // Copy the first item from the 'falsy' array into the output buffer.
    345
    0
                        mutable.extend(1, 0, 1);
    346
    0
                    }
    347
    0
                } else {
    348
    0
                    let falsy_length = start - filled;
    349
    0
                    let falsy_end = falsy_offset + falsy_length;
    350
    0
                    mutable.extend(1, falsy_offset, falsy_end);
    351
    0
                    falsy_offset = falsy_end;
    352
    0
                }
    353
    0
            }
    354
            // fill with truthy values
    355
    0
            if truthy_is_scalar {
    356
    0
                for _ in start..end {
    357
    0
                    // Copy the first item from the 'truthy' array into the output buffer.
    358
    0
                    mutable.extend(0, 0, 1);
    359
    0
                }
    360
    0
            } else {
    361
    0
                let truthy_length = end - start;
    362
    0
                let truthy_end = truthy_offset + truthy_length;
    363
    0
                mutable.extend(0, truthy_offset, truthy_end);
    364
    0
                truthy_offset = truthy_end;
    365
    0
            }
    366
    0
            filled = end;
    367
    0
        });
    368
        // the remaining part is falsy
    369
    0
        if filled < mask.len() {
    370
    0
            if falsy_is_scalar {
    371
    0
                for _ in filled..mask.len() {
    372
    0
                    // Copy the first item from the 'falsy' array into the output buffer.
    373
    0
                    mutable.extend(1, 0, 1);
    374
    0
                }
    375
    0
            } else {
    376
    0
                let falsy_length = mask.len() - filled;
    377
    0
                let falsy_end = falsy_offset + falsy_length;
    378
    0
                mutable.extend(1, falsy_offset, falsy_end);
    379
    0
            }
    380
    0
        }
    381
    382
    0
        let data = mutable.freeze();
    383
    0
        Ok(make_array(data))
    384
    0
    }
    385
    386
    /// Merges elements by index from a list of [`ArrayData`], creating a new [`ColumnarValue`] from
    387
    /// those values.
    388
    ///
    389
    /// Each element in `indices` is the index of an array in `values`. The `indices` array is processed
    390
    /// sequentially. The first occurrence of index value `n` will be mapped to the first
    391
    /// value of the array at index `n`. The second occurrence to the second value, and so on.
    392
    /// An index value where `PartialResultIndex::is_none` is `true` is used to indicate null values.
    393
    ///
    394
    /// # Implementation notes
    395
    ///
    396
    /// This algorithm is similar in nature to both `zip` and `interleave`, but there are some important
    397
    /// differences.
    398
    ///
    399
    /// In contrast to `zip`, this function supports multiple input arrays. Instead of a boolean
    400
    /// selection vector, an index array is to take values from the input arrays, and a special marker
    401
    /// value is used to indicate null values.
    402
    ///
    403
    /// In contrast to `interleave`, this function does not use pairs of indices. The values in
    404
    /// `indices` serve the same purpose as the first value in the pairs passed to `interleave`.
    405
    /// The index in the array is implicit and is derived from the number of times a particular array
    406
    /// index occurs.
    407
    /// The more constrained indexing mechanism used by this algorithm makes it easier to copy values
    408
    /// in contiguous slices. In the example below, the two subsequent elements from array `2` can be
    409
    /// copied in a single operation from the source array instead of copying them one by one.
    410
    /// Long spans of null values are also especially cheap because they do not need to be represented
    411
    /// in an input array.
    412
    ///
    413
    /// # Safety
    414
    ///
    415
    /// This function does not check that the number of occurrences of any particular array index matches
    416
    /// the length of the corresponding input array. If an array contains more values than required, the
    417
    /// spurious values will be ignored. If an array contains fewer values than necessary, this function
    418
    /// will panic.
    419
    ///
    420
    /// # Example
    421
    ///
    422
    /// ```text
    423
    /// ┌───────────┐  ┌─────────┐                             ┌─────────┐
    424
    /// │┌─────────┐│  │   None  │                             │   NULL  │
    425
    /// ││    A    ││  ├─────────┤                             ├─────────┤
    426
    /// │└─────────┘│  │    1    │                             │    B    │
    427
    /// │┌─────────┐│  ├─────────┤                             ├─────────┤
    428
    /// ││    B    ││  │    0    │    merge(values, indices)   │    A    │
    429
    /// │└─────────┘│  ├─────────┤  ─────────────────────────▶ ├─────────┤
    430
    /// │┌─────────┐│  │   None  │                             │   NULL  │
    431
    /// ││    C    ││  ├─────────┤                             ├─────────┤
    432
    /// │├─────────┤│  │    2    │                             │    C    │
    433
    /// ││    D    ││  ├─────────┤                             ├─────────┤
    434
    /// │└─────────┘│  │    2    │                             │    D    │
    435
    /// └───────────┘  └─────────┘                             └─────────┘
    436
    ///    values        indices                                  result
    437
    /// ```
    438
    0
    fn merge_n(values: &[ArrayData], indices: &[PartialResultIndex]) -> Result<ArrayRef> {
    439
        #[cfg(debug_assertions)]
    440
    0
        for ix in indices {
    441
    0
            if let Some(index) = ix.index() {
    442
    0
                assert!(
    443
    0
                    index < values.len(),
    444
    0
                    "Index out of bounds: {} >= {}",
    445
                    index,
    446
    0
                    values.len()
    447
                );
    448
    0
            }
    449
        }
    450
    451
    0
        let data_refs = values.iter().collect();
    452
    0
        let mut mutable = MutableArrayData::new(data_refs, true, indices.len());
    453
    454
        // This loop extends the mutable array by taking slices from the partial results.
    455
        //
    456
        // take_offsets keeps track of how many values have been taken from each array.
    457
    0
        let mut take_offsets = vec![0; values.len() + 1];
    458
    0
        let mut start_row_ix = 0;
    459
        loop {
    460
    0
            let array_ix = indices[start_row_ix];
    461
    462
            // Determine the length of the slice to take.
    463
    0
            let mut end_row_ix = start_row_ix + 1;
    464
    0
            while end_row_ix < indices.len() && indices[end_row_ix] == array_ix {
    465
    0
                end_row_ix += 1;
    466
    0
            }
    467
    0
            let slice_length = end_row_ix - start_row_ix;
    468
    469
            // Extend mutable with either nulls or with values from the array.
    470
    0
            match array_ix.index() {
    471
    0
                None => mutable.extend_nulls(slice_length),
    472
    0
                Some(index) => {
    473
    0
                    let start_offset = take_offsets[index];
    474
    0
                    let end_offset = start_offset + slice_length;
    475
    0
                    mutable.extend(index, start_offset, end_offset);
    476
    0
                    take_offsets[index] = end_offset;
    477
    0
                }
    478
            }
    479
    480
    0
            if end_row_ix == indices.len() {
    481
    0
                break;
    482
    0
            } else {
    483
    0
                // Set the start_row_ix for the next slice.
    484
    0
                start_row_ix = end_row_ix;
    485
    0
            }
    486
        }
    487
    488
    0
        Ok(make_array(mutable.freeze()))
    489
    0
    }
    490
    491
    /// An index into the partial results array that's more compact than `usize`.
    492
    ///
    493
    /// `u32::MAX` is reserved as a special 'none' value. This is used instead of
    494
    /// `Option` to keep the array of indices as compact as possible.
    495
    #[derive(Copy, Clone, PartialEq, Eq)]
    496
    struct PartialResultIndex {
    497
        index: u32,
    498
    }
    499
    500
    const NONE_VALUE: u32 = u32::MAX;
    501
    502
    impl PartialResultIndex {
    503
        /// Returns the 'none' placeholder value.
    504
    0
        fn none() -> Self {
    505
    0
            Self { index: NONE_VALUE }
    506
    0
        }
    507
    508
    0
        fn zero() -> Self {
    509
    0
            Self { index: 0 }
    510
    0
        }
    511
    512
        /// Creates a new partial result index.
    513
        ///
    514
        /// If the provided value is greater than or equal to `u32::MAX`
    515
        /// an error will be returned.
    516
    0
        fn try_new(index: usize) -> Result<Self> {
    517
    0
            let Ok(index) = u32::try_from(index) else {
    518
    0
                return internal_err!("Partial result index exceeds limit");
    519
            };
    520
    521
    0
            assert_or_internal_err!(
    522
    0
                index != NONE_VALUE,
    523
    0
                "Partial result index exceeds limit"
    524
            );
    525
    526
    0
            Ok(Self { index })
    527
    0
        }
    528
    529
        /// Determines if this index is the 'none' placeholder value or not.
    530
    0
        fn is_none(&self) -> bool {
    531
    0
            self.index == NONE_VALUE
    532
    0
        }
    533
    534
        /// Returns `Some(index)` if this value is not the 'none' placeholder, `None` otherwise.
    535
    0
        fn index(&self) -> Option<usize> {
    536
    0
            if self.is_none() {
    537
    0
                None
    538
            } else {
    539
    0
                Some(self.index as usize)
    540
            }
    541
    0
        }
    542
    }
    543
    544
    impl Debug for PartialResultIndex {
    545
    0
        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
    546
    0
            if self.is_none() {
    547
    0
                write!(f, "null")
    548
            } else {
    549
    0
                write!(f, "{}", self.index)
    550
            }
    551
    0
        }
    552
    }
    553
    554
    enum ResultState {
    555
        /// The final result is an array containing only null values.
    556
        Empty,
    557
        /// The final result needs to be computed by merging the data in `arrays`.
    558
        Partial {
    559
            // A `Vec` of partial results that should be merged.
    560
            // `partial_result_indices` contains indexes into this vec.
    561
            arrays: Vec<ArrayData>,
    562
            // Indicates per result row from which array in `partial_results` a value should be taken.
    563
            indices: Vec<PartialResultIndex>,
    564
        },
    565
        /// A single branch matched all input rows. When creating the final result, no further merging
    566
        /// of partial results is necessary.
    567
        Complete(ColumnarValue),
    568
    }
    569
    570
    /// A builder for constructing result arrays for CASE expressions.
    571
    ///
    572
    /// Rather than building a monolithic array containing all results, it maintains a set of
    573
    /// partial result arrays and a mapping that indicates for each row which partial array
    574
    /// contains the result value for that row.
    575
    ///
    576
    /// On finish(), the builder will merge all partial results into a single array if necessary.
    577
    /// If all rows evaluated to the same array, that array can be returned directly without
    578
    /// any merging overhead.
    579
    struct ResultBuilder {
    580
        data_type: DataType,
    581
        /// The number of rows in the final result.
    582
        row_count: usize,
    583
        state: ResultState,
    584
    }
    585
    586
    impl ResultBuilder {
    587
        /// Creates a new ResultBuilder that will produce arrays of the given data type.
    588
        ///
    589
        /// The `row_count` parameter indicates the number of rows in the final result.
    590
    0
        fn new(data_type: &DataType, row_count: usize) -> Self {
    591
    0
            Self {
    592
    0
                data_type: data_type.clone(),
    593
    0
                row_count,
    594
    0
                state: Empty,
    595
    0
            }
    596
    0
        }
    597
    598
        /// Adds a result for one branch of the case expression.
    599
        ///
    600
        /// `row_indices` should be a [UInt32Array] containing [RecordBatch] relative row indices
    601
        /// for which `value` contains result values.
    602
        ///
    603
        /// If `value` is a scalar, the scalar value will be used as the value for each row in `row_indices`.
    604
        ///
    605
        /// If `value` is an array, the values from the array and the indices from `row_indices` will be
    606
        /// processed pairwise. The lengths of `value` and `row_indices` must match.
    607
        ///
    608
        /// The diagram below shows a situation where a when expression matched rows 1 and 4 of the
    609
        /// record batch. The then expression produced the value array `[A, D]`.
    610
        /// After adding this result, the result array will have been added to `partial arrays` and
    611
        /// `partial indices` will have been updated at indexes `1` and `4`.
    612
        ///
    613
        /// ```text
    614
        ///  ┌─────────┐     ┌─────────┐┌───────────┐                            ┌─────────┐┌───────────┐
    615
        ///  │    C    │     │ 0: None ││┌ 0 ──────┐│                            │ 0: None ││┌ 0 ──────┐│
    616
        ///  ├─────────┤     ├─────────┤││    A    ││                            ├─────────┤││    A    ││
    617
        ///  │    D    │     │ 1: None ││└─────────┘│                            │ 1:  2   ││└─────────┘│
    618
        ///  └─────────┘     ├─────────┤│┌ 1 ──────┐│   add_branch_result(       ├─────────┤│┌ 1 ──────┐│
    619
        ///   matching       │ 2:  0   │││    B    ││     row indices,           │ 2:  0   │││    B    ││
    620
        /// 'then' values    ├─────────┤│└─────────┘│     value                  ├─────────┤│└─────────┘│
    621
        ///                  │ 3: None ││           │   )                        │ 3: None ││┌ 2 ──────┐│
    622
        ///  ┌─────────┐     ├─────────┤│           │ ─────────────────────────▶ ├─────────┤││    C    ││
    623
        ///  │    1    │     │ 4: None ││           │                            │ 4:  2   ││├─────────┤│
    624
        ///  ├─────────┤     ├─────────┤│           │                            ├─────────┤││    D    ││
    625
        ///  │    4    │     │ 5:  1   ││           │                            │ 5:  1   ││└─────────┘│
    626
        ///  └─────────┘     └─────────┘└───────────┘                            └─────────┘└───────────┘
    627
        /// row indices        partial     partial                                 partial     partial
    628
        ///                    indices     arrays                                  indices     arrays
    629
        /// ```
    630
    0
        fn add_branch_result(
    631
    0
            &mut self,
    632
    0
            row_indices: &ArrayRef,
    633
    0
            value: ColumnarValue,
    634
    0
        ) -> Result<()> {
    635
    0
            match value {
    636
    0
                ColumnarValue::Array(a) => {
    637
    0
                    if a.len() != row_indices.len() {
    638
    0
                        internal_err!("Array length must match row indices length")
    639
    0
                    } else if row_indices.len() == self.row_count {
    640
    0
                        self.set_complete_result(ColumnarValue::Array(a))
    641
                    } else {
    642
    0
                        self.add_partial_result(row_indices, a.to_data())
    643
                    }
    644
                }
    645
    0
                ColumnarValue::Scalar(s) => {
    646
    0
                    if row_indices.len() == self.row_count {
    647
    0
                        self.set_complete_result(ColumnarValue::Scalar(s))
    648
                    } else {
    649
    0
                        self.add_partial_result(
    650
    0
                            row_indices,
    651
    0
                            s.to_array_of_size(row_indices.len())?.to_data(),
    652
                        )
    653
                    }
    654
                }
    655
            }
    656
    0
        }
    657
    658
        /// Adds a partial result array.
    659
        ///
    660
        /// This method adds the given array data as a partial result and updates the index mapping
    661
        /// to indicate that the specified rows should take their values from this array.
    662
        /// The partial results will be merged into a single array when finish() is called.
    663
    0
        fn add_partial_result(
    664
    0
            &mut self,
    665
    0
            row_indices: &ArrayRef,
    666
    0
            row_values: ArrayData,
    667
    0
        ) -> Result<()> {
    668
    0
            assert_or_internal_err!(
    669
    0
                row_indices.null_count() == 0,
    670
    0
                "Row indices must not contain nulls"
    671
            );
    672
    673
    0
            match &mut self.state {
    674
                Empty => {
    675
    0
                    let array_index = PartialResultIndex::zero();
    676
    0
                    let mut indices = vec![PartialResultIndex::none(); self.row_count];
    677
    0
                    for row_ix in row_indices.as_primitive::<UInt32Type>().values().iter() {
    678
    0
                        indices[*row_ix as usize] = array_index;
    679
    0
                    }
    680
    681
    0
                    self.state = Partial {
    682
    0
                        arrays: vec![row_values],
    683
    0
                        indices,
    684
    0
                    };
    685
    686
    0
                    Ok(())
    687
                }
    688
    0
                Partial { arrays, indices } => {
    689
    0
                    let array_index = PartialResultIndex::try_new(arrays.len())?;
    690
    691
    0
                    arrays.push(row_values);
    692
    693
    0
                    for row_ix in row_indices.as_primitive::<UInt32Type>().values().iter() {
    694
                        // This is check is only active for debug config because the callers of this method,
    695
                        // `case_when_with_expr` and `case_when_no_expr`, already ensure that
    696
                        // they only calculate a value for each row at most once.
    697
                        #[cfg(debug_assertions)]
    698
    0
                        assert_or_internal_err!(
    699
    0
                            indices[*row_ix as usize].is_none(),
    700
    0
                            "Duplicate value for row {}",
    701
                            *row_ix
    702
                        );
    703
    704
    0
                        indices[*row_ix as usize] = array_index;
    705
                    }
    706
    0
                    Ok(())
    707
                }
    708
    0
                Complete(_) => internal_err!(
    709
    0
                    "Cannot add a partial result when complete result is already set"
    710
                ),
    711
            }
    712
    0
        }
    713
    714
        /// Sets a result that applies to all rows.
    715
        ///
    716
        /// This is an optimization for cases where all rows evaluate to the same result.
    717
        /// When a complete result is set, the builder will return it directly from finish()
    718
        /// without any merging overhead.
    719
    0
        fn set_complete_result(&mut self, value: ColumnarValue) -> Result<()> {
    720
    0
            match &self.state {
    721
                Empty => {
    722
    0
                    self.state = Complete(value);
    723
    0
                    Ok(())
    724
                }
    725
                Partial { .. } => {
    726
    0
                    internal_err!(
    727
    0
                        "Cannot set a complete result when there are already partial results"
    728
                    )
    729
                }
    730
    0
                Complete(_) => internal_err!("Complete result already set"),
    731
            }
    732
    0
        }
    733
    734
        /// Finishes building the result and returns the final array.
    735
    0
        fn finish(self) -> Result<ColumnarValue> {
    736
    0
            match self.state {
    737
                Empty => {
    738
                    // No complete result and no partial results.
    739
                    // This can happen for case expressions with no else branch where no rows
    740
                    // matched.
    741
    0
                    Ok(ColumnarValue::Scalar(ScalarValue::try_new_null(
    742
    0
                        &self.data_type,
    743
    0
                    )?))
    744
                }
    745
    0
                Partial { arrays, indices } => {
    746
                    // Merge partial results into a single array.
    747
    0
                    Ok(ColumnarValue::Array(merge_n(&arrays, &indices)?))
    748
                }
    749
    0
                Complete(v) => {
    750
                    // If we have a complete result, we can just return it.
    751
    0
                    Ok(v)
    752
                }
    753
            }
    754
    0
        }
    755
    }
    756
    757
    impl CaseExpr {
    758
        /// Create a new CASE WHEN expression
    759
    0
        pub fn try_new(
    760
    0
            expr: Option<Arc<dyn PhysicalExpr>>,
    761
    0
            when_then_expr: Vec<WhenThen>,
    762
    0
            else_expr: Option<Arc<dyn PhysicalExpr>>,
    763
    0
        ) -> Result<Self> {
    764
            // normalize null literals to None in the else_expr (this already happens
    765
            // during SQL planning, but not necessarily for other use cases)
    766
    0
            let else_expr = match &else_expr {
    767
    0
                Some(e) => match e.as_any().downcast_ref::<Literal>() {
    768
    0
                    Some(lit) if lit.value().is_null() => None,
    769
    0
                    _ => else_expr,
    770
                },
    771
    0
                _ => else_expr,
    772
            };
    773
    774
    0
            if when_then_expr.is_empty() {
    775
    0
                return exec_err!("There must be at least one WHEN clause");
    776
    0
            }
    777
    778
    0
            let body = CaseBody {
    779
    0
                expr,
    780
    0
                when_then_expr,
    781
    0
                else_expr,
    782
    0
            };
    783
    784
    0
            let eval_method = if body.expr.is_some() {
    785
    0
                EvalMethod::WithExpression(body.project()?)
    786
    0
            } else if body.when_then_expr.len() == 1
    787
    0
                && is_cheap_and_infallible(&(body.when_then_expr[0].1))
    788
    0
                && body.else_expr.is_none()
    789
            {
    790
    0
                EvalMethod::InfallibleExprOrNull
    791
    0
            } else if body.when_then_expr.len() == 1
    792
    0
                && body.when_then_expr[0].1.as_any().is::<Literal>()
    793
    0
                && body.else_expr.is_some()
    794
    0
                && body.else_expr.as_ref().unwrap().as_any().is::<Literal>()
    795
            {
    796
    0
                EvalMethod::ScalarOrScalar
    797
    0
            } else if body.when_then_expr.len() == 1 && body.else_expr.is_some() {
    798
    0
                EvalMethod::ExpressionOrExpression(body.project()?)
    799
            } else {
    800
    0
                EvalMethod::NoExpression(body.project()?)
    801
            };
    802
    803
    0
            Ok(Self { body, eval_method })
    804
    0
        }
    805
    806
        /// Optional base expression that can be compared to literal values in the "when" expressions
    807
    0
        pub fn expr(&self) -> Option<&Arc<dyn PhysicalExpr>> {
    808
    0
            self.body.expr.as_ref()
    809
    0
        }
    810
    811
        /// One or more when/then expressions
    812
    0
        pub fn when_then_expr(&self) -> &[WhenThen] {
    813
    0
            &self.body.when_then_expr
    814
    0
        }
    815
    816
        /// Optional "else" expression
    817
    0
        pub fn else_expr(&self) -> Option<&Arc<dyn PhysicalExpr>> {
    818
    0
            self.body.else_expr.as_ref()
    819
    0
        }
    820
    }
    821
    822
    impl CaseBody {
    823
    0
        fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
    824
            // since all then results have the same data type, we can choose any one as the
    825
            // return data type except for the null.
    826
    0
            let mut data_type = DataType::Null;
    827
    0
            for i in 0..self.when_then_expr.len() {
    828
    0
                data_type = self.when_then_expr[i].1.data_type(input_schema)?;
    829
    0
                if !data_type.equals_datatype(&DataType::Null) {
    830
    0
                    break;
    831
    0
                }
    832
            }
    833
            // if all then results are null, we use data type of else expr instead if possible.
    834
    0
            if data_type.equals_datatype(&DataType::Null) {
    835
    0
                if let Some(e) = &self.else_expr {
    836
    0
                    data_type = e.data_type(input_schema)?;
    837
    0
                }
    838
    0
            }
    839
    840
    0
            Ok(data_type)
    841
    0
        }
    842
    843
        /// See [CaseExpr::case_when_with_expr].
    844
    0
        fn case_when_with_expr(
    845
    0
            &self,
    846
    0
            batch: &RecordBatch,
    847
    0
            return_type: &DataType,
    848
    0
        ) -> Result<ColumnarValue> {
    849
    0
            let mut result_builder = ResultBuilder::new(return_type, batch.num_rows());
    850
    851
            // `remainder_rows` contains the indices of the rows that need to be evaluated
    852
    0
            let mut remainder_rows: ArrayRef =
    853
    0
                Arc::new(UInt32Array::from_iter_values(0..batch.num_rows() as u32));
    854
            // `remainder_batch` contains the rows themselves that need to be evaluated
    855
    0
            let mut remainder_batch = Cow::Borrowed(batch);
    856
    857
            // evaluate the base expression
    858
    0
            let mut base_values = self
    859
    0
                .expr
    860
    0
                .as_ref()
    861
    0
                .unwrap()
    862
    0
                .evaluate(batch)?
    863
    0
                .into_array(batch.num_rows())?;
    864
    865
            // Fill in a result value already for rows where the base expression value is null
    866
            // Since each when expression is tested against the base expression using the equality
    867
            // operator, null base values can never match any when expression. `x = NULL` is falsy,
    868
            // for all possible values of `x`.
    869
    0
            if base_values.null_count() > 0 {
    870
                // Use `is_not_null` since this is a cheap clone of the null buffer from 'base_value'.
    871
                // We already checked there are nulls, so we can be sure a new buffer will not be
    872
                // created.
    873
    0
                let base_not_nulls = is_not_null(base_values.as_ref())?;
    874
    0
                let base_all_null = base_values.null_count() == remainder_batch.num_rows();
    875
    876
                // If there is an else expression, use that as the default value for the null rows
    877
                // Otherwise the default `null` value from the result builder will be used.
    878
    0
                if let Some(e) = &self.else_expr {
    879
    0
                    let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?;
    880
    881
    0
                    if base_all_null {
    882
                        // All base values were null, so no need to filter
    883
    0
                        let nulls_value = expr.evaluate(&remainder_batch)?;
    884
    0
                        result_builder.add_branch_result(&remainder_rows, nulls_value)?;
    885
                    } else {
    886
                        // Filter out the null rows and evaluate the else expression for those
    887
    0
                        let nulls_filter = create_filter(&not(&base_not_nulls)?, true);
    888
    0
                        let nulls_batch =
    889
    0
                            filter_record_batch(&remainder_batch, &nulls_filter)?;
    890
    0
                        let nulls_rows = filter_array(&remainder_rows, &nulls_filter)?;
    891
    0
                        let nulls_value = expr.evaluate(&nulls_batch)?;
    892
    0
                        result_builder.add_branch_result(&nulls_rows, nulls_value)?;
    893
                    }
    894
    0
                }
    895
    896
                // All base values are null, so we can return early
    897
    0
                if base_all_null {
    898
    0
                    return result_builder.finish();
    899
    0
                }
    900
    901
                // Remove the null rows from the remainder batch
    902
    0
                let not_null_filter = create_filter(&base_not_nulls, true);
    903
    0
                remainder_batch =
    904
    0
                    Cow::Owned(filter_record_batch(&remainder_batch, &not_null_filter)?);
    905
    0
                remainder_rows = filter_array(&remainder_rows, &not_null_filter)?;
    906
    0
                base_values = filter_array(&base_values, &not_null_filter)?;
    907
    0
            }
    908
    909
            // The types of case and when expressions will be coerced to match.
    910
            // We only need to check if the base_value is nested.
    911
    0
            let base_value_is_nested = base_values.data_type().is_nested();
    912
    913
    0
            for i in 0..self.when_then_expr.len() {
    914
                // Evaluate the 'when' predicate for the remainder batch
    915
                // This results in a boolean array with the same length as the remaining number of rows
    916
    0
                let when_expr = &self.when_then_expr[i].0;
    917
    0
                let when_value = match when_expr.evaluate(&remainder_batch)? {
    918
    0
                    ColumnarValue::Array(a) => {
    919
    0
                        compare_with_eq(&a, &base_values, base_value_is_nested)
    920
                    }
    921
    0
                    ColumnarValue::Scalar(s) => {
    922
    0
                        compare_with_eq(&s.to_scalar()?, &base_values, base_value_is_nested)
    923
                    }
    924
    0
                }?;
    925
    926
                // `true_count` ignores `true` values where the validity bit is not set, so there's
    927
                // no need to call `prep_null_mask_filter`.
    928
    0
                let when_true_count = when_value.true_count();
    929
    930
                // If the 'when' predicate did not match any rows, continue to the next branch immediately
    931
    0
                if when_true_count == 0 {
    932
    0
                    continue;
    933
    0
                }
    934
    935
                // If the 'when' predicate matched all remaining rows, there is no need to filter
    936
    0
                if when_true_count == remainder_batch.num_rows() {
    937
    0
                    let then_expression = &self.when_then_expr[i].1;
    938
    0
                    let then_value = then_expression.evaluate(&remainder_batch)?;
    939
    0
                    result_builder.add_branch_result(&remainder_rows, then_value)?;
    940
    0
                    return result_builder.finish();
    941
    0
                }
    942
    943
                // Filter the remainder batch based on the 'when' value
    944
                // This results in a batch containing only the rows that need to be evaluated
    945
                // for the current branch
    946
                // Still no need to call `prep_null_mask_filter` since `create_filter` will already do
    947
                // this unconditionally.
    948
    0
                let then_filter = create_filter(&when_value, true);
    949
    0
                let then_batch = filter_record_batch(&remainder_batch, &then_filter)?;
    950
    0
                let then_rows = filter_array(&remainder_rows, &then_filter)?;
    951
    952
    0
                let then_expression = &self.when_then_expr[i].1;
    953
    0
                let then_value = then_expression.evaluate(&then_batch)?;
    954
    0
                result_builder.add_branch_result(&then_rows, then_value)?;
    955
    956
                // If this is the last 'when' branch and there is no 'else' expression, there's no
    957
                // point in calculating the remaining rows.
    958
    0
                if self.else_expr.is_none() && i == self.when_then_expr.len() - 1 {
    959
    0
                    return result_builder.finish();
    960
    0
                }
    961
    962
                // Prepare the next when branch (or the else branch)
    963
    0
                let next_selection = match when_value.null_count() {
    964
    0
                    0 => not(&when_value),
    965
                    _ => {
    966
                        // `prep_null_mask_filter` is required to ensure the not operation treats nulls
    967
                        // as false
    968
    0
                        not(&prep_null_mask_filter(&when_value))
    969
                    }
    970
    0
                }?;
    971
    0
                let next_filter = create_filter(&next_selection, true);
    972
    0
                remainder_batch =
    973
    0
                    Cow::Owned(filter_record_batch(&remainder_batch, &next_filter)?);
    974
    0
                remainder_rows = filter_array(&remainder_rows, &next_filter)?;
    975
    0
                base_values = filter_array(&base_values, &next_filter)?;
    976
            }
    977
    978
            // If we reached this point, some rows were left unmatched.
    979
            // Check if those need to be evaluated using the 'else' expression.
    980
    0
            if let Some(e) = &self.else_expr {
    981
                // keep `else_expr`'s data type and return type consistent
    982
    0
                let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?;
    983
    0
                let else_value = expr.evaluate(&remainder_batch)?;
    984
    0
                result_builder.add_branch_result(&remainder_rows, else_value)?;
    985
    0
            }
    986
    987
    0
            result_builder.finish()
    988
    0
        }
    989
    990
        /// See [CaseExpr::case_when_no_expr].
    991
    0
        fn case_when_no_expr(
    992
    0
            &self,
    993
    0
            batch: &RecordBatch,
    994
    0
            return_type: &DataType,
    995
    0
        ) -> Result<ColumnarValue> {
    996
    0
            let mut result_builder = ResultBuilder::new(return_type, batch.num_rows());
    997
    998
            // `remainder_rows` contains the indices of the rows that need to be evaluated
    999
    0
            let mut remainder_rows: ArrayRef =
    1000
    0
                Arc::new(UInt32Array::from_iter(0..batch.num_rows() as u32));
    1001
            // `remainder_batch` contains the rows themselves that need to be evaluated
    1002
    0
            let mut remainder_batch = Cow::Borrowed(batch);
    1003
    1004
    0
            for i in 0..self.when_then_expr.len() {
    1005
                // Evaluate the 'when' predicate for the remainder batch
    1006
                // This results in a boolean array with the same length as the remaining number of rows
    1007
    0
                let when_predicate = &self.when_then_expr[i].0;
    1008
    0
                let when_value = when_predicate
    1009
    0
                    .evaluate(&remainder_batch)?
    1010
    0
                    .into_array(remainder_batch.num_rows())?;
    1011
    0
                let when_value = as_boolean_array(&when_value).map_err(|_| {
    1012
    0
                    internal_datafusion_err!("WHEN expression did not return a BooleanArray")
    1013
    0
                })?;
    1014
    1015
                // `true_count` ignores `true` values where the validity bit is not set, so there's
    1016
                // no need to call `prep_null_mask_filter`.
    1017
    0
                let when_true_count = when_value.true_count();
    1018
    1019
                // If the 'when' predicate did not match any rows, continue to the next branch immediately
    1020
    0
                if when_true_count == 0 {
    1021
    0
                    continue;
    1022
    0
                }
    1023
    1024
                // If the 'when' predicate matched all remaining rows, there is no need to filter
    1025
    0
                if when_true_count == remainder_batch.num_rows() {
    1026
    0
                    let then_expression = &self.when_then_expr[i].1;
    1027
    0
                    let then_value = then_expression.evaluate(&remainder_batch)?;
    1028
    0
                    result_builder.add_branch_result(&remainder_rows, then_value)?;
    1029
    0
                    return result_builder.finish();
    1030
    0
                }
    1031
    1032
                // Filter the remainder batch based on the 'when' value
    1033
                // This results in a batch containing only the rows that need to be evaluated
    1034
                // for the current branch
    1035
                // Still no need to call `prep_null_mask_filter` since `create_filter` will already do
    1036
                // this unconditionally.
    1037
    0
                let then_filter = create_filter(when_value, true);
    1038
    0
                let then_batch = filter_record_batch(&remainder_batch, &then_filter)?;
    1039
    0
                let then_rows = filter_array(&remainder_rows, &then_filter)?;
    1040
    1041
    0
                let then_expression = &self.when_then_expr[i].1;
    1042
    0
                let then_value = then_expression.evaluate(&then_batch)?;
    1043
    0
                result_builder.add_branch_result(&then_rows, then_value)?;
    1044
    1045
                // If this is the last 'when' branch and there is no 'else' expression, there's no
    1046
                // point in calculating the remaining rows.
    1047
    0
                if self.else_expr.is_none() && i == self.when_then_expr.len() - 1 {
    1048
    0
                    return result_builder.finish();
    1049
    0
                }
    1050
    1051
                // Prepare the next when branch (or the else branch)
    1052
    0
                let next_selection = match when_value.null_count() {
    1053
    0
                    0 => not(when_value),
    1054
                    _ => {
    1055
                        // `prep_null_mask_filter` is required to ensure the not operation treats nulls
    1056
                        // as false
    1057
    0
                        not(&prep_null_mask_filter(when_value))
    1058
                    }
    1059
    0
                }?;
    1060
    0
                let next_filter = create_filter(&next_selection, true);
    1061
    0
                remainder_batch =
    1062
    0
                    Cow::Owned(filter_record_batch(&remainder_batch, &next_filter)?);
    1063
    0
                remainder_rows = filter_array(&remainder_rows, &next_filter)?;
    1064
            }
    1065
    1066
            // If we reached this point, some rows were left unmatched.
    1067
            // Check if those need to be evaluated using the 'else' expression.
    1068
    0
            if let Some(e) = &self.else_expr {
    1069
                // keep `else_expr`'s data type and return type consistent
    1070
    0
                let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?;
    1071
    0
                let else_value = expr.evaluate(&remainder_batch)?;
    1072
    0
                result_builder.add_branch_result(&remainder_rows, else_value)?;
    1073
    0
            }
    1074
    1075
    0
            result_builder.finish()
    1076
    0
        }
    1077
    1078
        /// See [CaseExpr::expr_or_expr].
    1079
    0
        fn expr_or_expr(
    1080
    0
            &self,
    1081
    0
            batch: &RecordBatch,
    1082
    0
            when_value: &BooleanArray,
    1083
    0
        ) -> Result<ColumnarValue> {
    1084
    0
            let when_value = match when_value.null_count() {
    1085
    0
                0 => Cow::Borrowed(when_value),
    1086
                _ => {
    1087
                    // `prep_null_mask_filter` is required to ensure null is treated as false
    1088
    0
                    Cow::Owned(prep_null_mask_filter(when_value))
    1089
                }
    1090
            };
    1091
    1092
    0
            let optimize_filter = batch.num_columns() > 1
    1093
    0
                || (batch.num_columns() == 1 && multiple_arrays(batch.column(0).data_type()));
    1094
    1095
    0
            let when_filter = create_filter(&when_value, optimize_filter);
    1096
    0
            let then_batch = filter_record_batch(batch, &when_filter)?;
    1097
    0
            let then_value = self.when_then_expr[0].1.evaluate(&then_batch)?;
    1098
    1099
    0
            let else_selection = not(&when_value)?;
    1100
    0
            let else_filter = create_filter(&else_selection, optimize_filter);
    1101
    0
            let else_batch = filter_record_batch(batch, &else_filter)?;
    1102
    1103
            // keep `else_expr`'s data type and return type consistent
    1104
    0
            let e = self.else_expr.as_ref().unwrap();
    1105
    0
            let return_type = self.data_type(&batch.schema())?;
    1106
    0
            let else_expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())
    1107
    0
                .unwrap_or_else(|_| Arc::clone(e));
    1108
    1109
    0
            let else_value = else_expr.evaluate(&else_batch)?;
    1110
    1111
    0
            Ok(ColumnarValue::Array(merge(
    1112
    0
                &when_value,
    1113
    0
                then_value,
    1114
    0
                else_value,
    1115
    0
            )?))
    1116
    0
        }
    1117
    }
    1118
    1119
    impl CaseExpr {
    1120
        /// This function evaluates the form of CASE that matches an expression to fixed values.
    1121
        ///
    1122
        /// CASE expression
    1123
        ///     WHEN value THEN result
    1124
        ///     [WHEN ...]
    1125
        ///     [ELSE result]
    1126
        /// END
    1127
    0
        fn case_when_with_expr(
    1128
    0
            &self,
    1129
    0
            batch: &RecordBatch,
    1130
    0
            projected: &ProjectedCaseBody,
    1131
    0
        ) -> Result<ColumnarValue> {
    1132
    0
            let return_type = self.data_type(&batch.schema())?;
    1133
    0
            if projected.projection.len() < batch.num_columns() {
    1134
    0
                let projected_batch = batch.project(&projected.projection)?;
    1135
    0
                projected
    1136
    0
                    .body
    1137
    0
                    .case_when_with_expr(&projected_batch, &return_type)
    1138
            } else {
    1139
    0
                self.body.case_when_with_expr(batch, &return_type)
    1140
            }
    1141
    0
        }
    1142
    1143
        /// This function evaluates the form of CASE where each WHEN expression is a boolean
    1144
        /// expression.
    1145
        ///
    1146
        /// CASE WHEN condition THEN result
    1147
        ///      [WHEN ...]
    1148
        ///      [ELSE result]
    1149
        /// END
    1150
    0
        fn case_when_no_expr(
    1151
    0
            &self,
    1152
    0
            batch: &RecordBatch,
    1153
    0
            projected: &ProjectedCaseBody,
    1154
    0
        ) -> Result<ColumnarValue> {
    1155
    0
            let return_type = self.data_type(&batch.schema())?;
    1156
    0
            if projected.projection.len() < batch.num_columns() {
    1157
    0
                let projected_batch = batch.project(&projected.projection)?;
    1158
    0
                projected
    1159
    0
                    .body
    1160
    0
                    .case_when_no_expr(&projected_batch, &return_type)
    1161
            } else {
    1162
    0
                self.body.case_when_no_expr(batch, &return_type)
    1163
            }
    1164
    0
        }
    1165
    1166
        /// This function evaluates the specialized case of:
    1167
        ///
    1168
        /// CASE WHEN condition THEN column
    1169
        ///      [ELSE NULL]
    1170
        /// END
    1171
        ///
    1172
        /// Note that this function is only safe to use for "then" expressions
    1173
        /// that are infallible because the expression will be evaluated for all
    1174
        /// rows in the input batch.
    1175
    0
        fn case_column_or_null(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
    1176
    0
            let when_expr = &self.body.when_then_expr[0].0;
    1177
    0
            let then_expr = &self.body.when_then_expr[0].1;
    1178
    1179
    0
            match when_expr.evaluate(batch)? {
    1180
                // WHEN true --> column
    1181
                ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) => {
    1182
    0
                    then_expr.evaluate(batch)
    1183
                }
    1184
                // WHEN [false | null] --> NULL
    1185
                ColumnarValue::Scalar(_) => {
    1186
                    // return scalar NULL value
    1187
    0
                    ScalarValue::try_from(self.data_type(&batch.schema())?)
    1188
    0
                        .map(ColumnarValue::Scalar)
    1189
                }
    1190
                // WHEN column --> column
    1191
    0
                ColumnarValue::Array(bit_mask) => {
    1192
    0
                    let bit_mask = bit_mask
    1193
    0
                        .as_any()
    1194
    0
                        .downcast_ref::<BooleanArray>()
    1195
    0
                        .expect("predicate should evaluate to a boolean array");
    1196
                    // invert the bitmask
    1197
    0
                    let bit_mask = match bit_mask.null_count() {
    1198
    0
                        0 => not(bit_mask)?,
    1199
    0
                        _ => not(&prep_null_mask_filter(bit_mask))?,
    1200
                    };
    1201
    0
                    match then_expr.evaluate(batch)? {
    1202
    0
                        ColumnarValue::Array(array) => {
    1203
    0
                            Ok(ColumnarValue::Array(nullif(&array, &bit_mask)?))
    1204
                        }
    1205
                        ColumnarValue::Scalar(_) => {
    1206
    0
                            internal_err!("expression did not evaluate to an array")
    1207
                        }
    1208
                    }
    1209
                }
    1210
            }
    1211
    0
        }
    1212
    1213
    0
        fn scalar_or_scalar(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
    1214
    0
            let return_type = self.data_type(&batch.schema())?;
    1215
    1216
            // evaluate when expression
    1217
    0
            let when_value = self.body.when_then_expr[0].0.evaluate(batch)?;
    1218
    0
            let when_value = when_value.into_array(batch.num_rows())?;
    1219
    0
            let when_value = as_boolean_array(&when_value).map_err(|_| {
    1220
    0
                internal_datafusion_err!("WHEN expression did not return a BooleanArray")
    1221
    0
            })?;
    1222
    1223
            // Treat 'NULL' as false value
    1224
    0
            let when_value = match when_value.null_count() {
    1225
    0
                0 => Cow::Borrowed(when_value),
    1226
    0
                _ => Cow::Owned(prep_null_mask_filter(when_value)),
    1227
            };
    1228
    1229
            // evaluate then_value
    1230
    0
            let then_value = self.body.when_then_expr[0].1.evaluate(batch)?;
    1231
    0
            let then_value = Scalar::new(then_value.into_array(1)?);
    1232
    1233
    0
            let Some(e) = &self.body.else_expr else {
    1234
    0
                return internal_err!("expression did not evaluate to an array");
    1235
            };
    1236
            // keep `else_expr`'s data type and return type consistent
    1237
    0
            let expr = try_cast(Arc::clone(e), &batch.schema(), return_type)?;
    1238
    0
            let else_ = Scalar::new(expr.evaluate(batch)?.into_array(1)?);
    1239
    0
            Ok(ColumnarValue::Array(zip(&when_value, &then_value, &else_)?))
    1240
    0
        }
    1241
    1242
    0
        fn expr_or_expr(
    1243
    0
            &self,
    1244
    0
            batch: &RecordBatch,
    1245
    0
            projected: &ProjectedCaseBody,
    1246
    0
        ) -> Result<ColumnarValue> {
    1247
            // evaluate when condition on batch
    1248
    0
            let when_value = self.body.when_then_expr[0].0.evaluate(batch)?;
    1249
            // `num_rows == 1` is intentional to avoid expanding scalars.
    1250
            // If the `when_value` is effectively a scalar, the 'all true' and 'all false' checks
    1251
            // below will avoid incorrectly using the scalar as a merge/zip mask.
    1252
    0
            let when_value = when_value.into_array(1)?;
    1253
    0
            let when_value = as_boolean_array(&when_value).map_err(|e| {
    1254
    0
                DataFusionError::Context(
    1255
    0
                    "WHEN expression did not return a BooleanArray".to_string(),
    1256
    0
                    Box::new(e),
    1257
    0
                )
    1258
    0
            })?;
    1259
    1260
    0
            let true_count = when_value.true_count();
    1261
    0
            if true_count == when_value.len() {
    1262
                // All input rows are true, just call the 'then' expression
    1263
    0
                self.body.when_then_expr[0].1.evaluate(batch)
    1264
    0
            } else if true_count == 0 {
    1265
                // All input rows are false/null, just call the 'else' expression
    1266
    0
                self.body.else_expr.as_ref().unwrap().evaluate(batch)
    1267
    0
            } else if projected.projection.len() < batch.num_columns() {
    1268
                // The case expressions do not use all the columns of the input batch.
    1269
                // Project first to reduce time spent filtering.
    1270
    0
                let projected_batch = batch.project(&projected.projection)?;
    1271
    0
                projected.body.expr_or_expr(&projected_batch, when_value)
    1272
            } else {
    1273
                // All columns are used in the case expressions, so there is no need to project.
    1274
    0
                self.body.expr_or_expr(batch, when_value)
    1275
            }
    1276
    0
        }
    1277
    }
    1278
    1279
    impl PhysicalExpr for CaseExpr {
    1280
        /// Return a reference to Any that can be used for down-casting
    1281
    0
        fn as_any(&self) -> &dyn Any {
    1282
    0
            self
    1283
    0
        }
    1284
    1285
    0
        fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
    1286
    0
            self.body.data_type(input_schema)
    1287
    0
        }
    1288
    1289
    0
        fn nullable(&self, input_schema: &Schema) -> Result<bool> {
    1290
    0
            let nullable_then = self
    1291
    0
                .body
    1292
    0
                .when_then_expr
    1293
    0
                .iter()
    1294
    0
                .filter_map(|(w, t)| {
    1295
    0
                    let is_nullable = match t.nullable(input_schema) {
    1296
                        // Pass on error determining nullability verbatim
    1297
    0
                        Err(e) => return Some(Err(e)),
    1298
    0
                        Ok(n) => n,
    1299
                    };
    1300
    1301
                    // Branches with a then expression that is not nullable do not impact the
    1302
                    // nullability of the case expression.
    1303
    0
                    if !is_nullable {
    1304
    0
                        return None;
    1305
    0
                    }
    1306
    1307
                    // For case-with-expression assume all 'then' expressions are reachable
    1308
    0
                    if self.body.expr.is_some() {
    1309
    0
                        return Some(Ok(()));
    1310
    0
                    }
    1311
    1312
                    // For branches with a nullable 'then' expression, try to determine
    1313
                    // if the 'then' expression is ever reachable in the situation where
    1314
                    // it would evaluate to null.
    1315
    1316
                    // Replace the `then` expression with `NULL` in the `when` expression
    1317
    0
                    let with_null = match replace_with_null(w, t.as_ref(), input_schema) {
    1318
    0
                        Err(e) => return Some(Err(e)),
    1319
    0
                        Ok(e) => e,
    1320
                    };
    1321
    1322
                    // Try to const evaluate the modified `when` expression.
    1323
    0
                    let predicate_result = match evaluate_predicate(&with_null) {
    1324
    0
                        Err(e) => return Some(Err(e)),
    1325
    0
                        Ok(b) => b,
    1326
                    };
    1327
    1328
    0
                    match predicate_result {
    1329
                        // Evaluation was inconclusive or true, so the 'then' expression is reachable
    1330
    0
                        None | Some(true) => Some(Ok(())),
    1331
                        // Evaluation proves the branch will never be taken.
    1332
                        // The most common pattern for this is `WHEN x IS NOT NULL THEN x`.
    1333
    0
                        Some(false) => None,
    1334
                    }
    1335
    0
                })
    1336
    0
                .next();
    1337
    1338
    0
            if let Some(nullable_then) = nullable_then {
    1339
                // There is at least one reachable nullable 'then' expression, so the case
    1340
                // expression itself is nullable.
    1341
                // Use `Result::map` to propagate the error from `nullable_then` if there is one.
    1342
    0
                nullable_then.map(|_| true)
    1343
    0
            } else if let Some(e) = &self.body.else_expr {
    1344
                // There are no reachable nullable 'then' expressions, so all we still need to
    1345
                // check is the 'else' expression's nullability.
    1346
    0
                e.nullable(input_schema)
    1347
            } else {
    1348
                // CASE produces NULL if there is no `else` expr
    1349
                // (aka when none of the `when_then_exprs` match)
    1350
    0
                Ok(true)
    1351
            }
    1352
    0
        }
    1353
    1354
    0
        fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
    1355
    0
            match &self.eval_method {
    1356
    0
                EvalMethod::WithExpression(p) => {
    1357
                    // this use case evaluates "expr" and then compares the values with the "when"
    1358
                    // values
    1359
    0
                    self.case_when_with_expr(batch, p)
    1360
                }
    1361
    0
                EvalMethod::NoExpression(p) => {
    1362
                    // The "when" conditions all evaluate to boolean in this use case and can be
    1363
                    // arbitrary expressions
    1364
    0
                    self.case_when_no_expr(batch, p)
    1365
                }
    1366
                EvalMethod::InfallibleExprOrNull => {
    1367
                    // Specialization for CASE WHEN expr THEN column [ELSE NULL] END
    1368
    0
                    self.case_column_or_null(batch)
    1369
                }
    1370
    0
                EvalMethod::ScalarOrScalar => self.scalar_or_scalar(batch),
    1371
    0
                EvalMethod::ExpressionOrExpression(p) => self.expr_or_expr(batch, p),
    1372
            }
    1373
    0
        }
    1374
    1375
    0
        fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
    1376
    0
            let mut children = vec![];
    1377
    0
            if let Some(expr) = &self.body.expr {
    1378
    0
                children.push(expr)
    1379
    0
            }
    1380
    0
            self.body.when_then_expr.iter().for_each(|(cond, value)| {
    1381
    0
                children.push(cond);
    1382
    0
                children.push(value);
    1383
    0
            });
    1384
    1385
    0
            if let Some(else_expr) = &self.body.else_expr {
    1386
    0
                children.push(else_expr)
    1387
    0
            }
    1388
    0
            children
    1389
    0
        }
    1390
    1391
        // For physical CaseExpr, we do not allow modifying children size
    1392
    0
        fn with_new_children(
    1393
    0
            self: Arc<Self>,
    1394
    0
            children: Vec<Arc<dyn PhysicalExpr>>,
    1395
    0
        ) -> Result<Arc<dyn PhysicalExpr>> {
    1396
    0
            if children.len() != self.children().len() {
    1397
    0
                internal_err!("CaseExpr: Wrong number of children")
    1398
            } else {
    1399
    0
                let (expr, when_then_expr, else_expr) =
    1400
    0
                    match (self.expr().is_some(), self.body.else_expr.is_some()) {
    1401
    0
                        (true, true) => (
    1402
    0
                            Some(&children[0]),
    1403
    0
                            &children[1..children.len() - 1],
    1404
    0
                            Some(&children[children.len() - 1]),
    1405
    0
                        ),
    1406
                        (true, false) => {
    1407
    0
                            (Some(&children[0]), &children[1..children.len()], None)
    1408
                        }
    1409
    0
                        (false, true) => (
    1410
    0
                            None,
    1411
    0
                            &children[0..children.len() - 1],
    1412
    0
                            Some(&children[children.len() - 1]),
    1413
    0
                        ),
    1414
    0
                        (false, false) => (None, &children[0..children.len()], None),
    1415
                    };
    1416
    0
                Ok(Arc::new(CaseExpr::try_new(
    1417
    0
                    expr.cloned(),
    1418
    0
                    when_then_expr.iter().cloned().tuples().collect(),
    1419
    0
                    else_expr.cloned(),
    1420
    0
                )?))
    1421
            }
    1422
    0
        }
    1423
    1424
    0
        fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
    1425
    0
            write!(f, "CASE ")?;
    1426
    0
            if let Some(e) = &self.body.expr {
    1427
    0
                e.fmt_sql(f)?;
    1428
    0
                write!(f, " ")?;
    1429
    0
            }
    1430
    1431
    0
            for (w, t) in &self.body.when_then_expr {
    1432
    0
                write!(f, "WHEN ")?;
    1433
    0
                w.fmt_sql(f)?;
    1434
    0
                write!(f, " THEN ")?;
    1435
    0
                t.fmt_sql(f)?;
    1436
    0
                write!(f, " ")?;
    1437
            }
    1438
    1439
    0
            if let Some(e) = &self.body.else_expr {
    1440
    0
                write!(f, "ELSE ")?;
    1441
    0
                e.fmt_sql(f)?;
    1442
    0
                write!(f, " ")?;
    1443
    0
            }
    1444
    0
            write!(f, "END")
    1445
    0
        }
    1446
    }
    1447
    1448
    /// Attempts to const evaluate the given `predicate`.
    1449
    /// Returns:
    1450
    /// - `Some(true)` if the predicate evaluates to a truthy value.
    1451
    /// - `Some(false)` if the predicate evaluates to a falsy value.
    1452
    /// - `None` if the predicate could not be evaluated.
    1453
    0
    fn evaluate_predicate(predicate: &Arc<dyn PhysicalExpr>) -> Result<Option<bool>> {
    1454
        // Create a dummy record with no columns and one row
    1455
    0
        let batch = RecordBatch::try_new_with_options(
    1456
    0
            Arc::new(Schema::empty()),
    1457
    0
            vec![],
    1458
    0
            &RecordBatchOptions::new().with_row_count(Some(1)),
    1459
    0
        )?;
    1460
    1461
        // Evaluate the predicate and interpret the result as a boolean
    1462
    0
        let result = match predicate.evaluate(&batch) {
    1463
            // An error during evaluation means we couldn't const evaluate the predicate, so return `None`
    1464
    0
            Err(_) => None,
    1465
    0
            Ok(ColumnarValue::Array(array)) => Some(
    1466
    0
                ScalarValue::try_from_array(array.as_ref(), 0)?
    1467
    0
                    .cast_to(&DataType::Boolean)?,
    1468
            ),
    1469
    0
            Ok(ColumnarValue::Scalar(scalar)) => Some(scalar.cast_to(&DataType::Boolean)?),
    1470
        };
    1471
    0
        Ok(result.map(|v| matches!(v, ScalarValue::Boolean(Some(true)))))
    1472
    0
    }
    1473
    1474
    0
    fn replace_with_null(
    1475
    0
        expr: &Arc<dyn PhysicalExpr>,
    1476
    0
        expr_to_replace: &dyn PhysicalExpr,
    1477
    0
        input_schema: &Schema,
    1478
    0
    ) -> Result<Arc<dyn PhysicalExpr>, DataFusionError> {
    1479
    0
        let with_null = Arc::clone(expr)
    1480
    0
            .transform_down(|e| {
    1481
    0
                if e.as_ref().dyn_eq(expr_to_replace) {
    1482
    0
                    let data_type = e.data_type(input_schema)?;
    1483
    0
                    let null_literal = lit(ScalarValue::try_new_null(&data_type)?);
    1484
    0
                    Ok(Transformed::yes(null_literal))
    1485
                } else {
    1486
    0
                    Ok(Transformed::no(e))
    1487
                }
    1488
    0
            })?
    1489
            .data;
    1490
    0
        Ok(with_null)
    1491
    0
    }
    1492
    1493
    /// Create a CASE expression
    1494
    0
    pub fn case(
    1495
    0
        expr: Option<Arc<dyn PhysicalExpr>>,
    1496
    0
        when_thens: Vec<WhenThen>,
    1497
    0
        else_expr: Option<Arc<dyn PhysicalExpr>>,
    1498
    0
    ) -> Result<Arc<dyn PhysicalExpr>> {
    1499
    0
        Ok(Arc::new(CaseExpr::try_new(expr, when_thens, else_expr)?))
    1500
    0
    }
    1501
    1502
    #[cfg(test)]
    1503
    mod tests {
    1504
        use super::*;
    1505
    1506
        use crate::expressions;
    1507
        use crate::expressions::{binary, cast, col, is_not_null, lit, BinaryExpr};
    1508
        use arrow::buffer::Buffer;
    1509
        use arrow::datatypes::DataType::Float64;
    1510
        use arrow::datatypes::Field;
    1511
        use datafusion_common::cast::{as_float64_array, as_int32_array};
    1512
        use datafusion_common::plan_err;
    1513
        use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
    1514
        use datafusion_expr::type_coercion::binary::comparison_coercion;
    1515
        use datafusion_expr_common::operator::Operator;
    1516
        use datafusion_physical_expr_common::physical_expr::fmt_sql;
    1517
    1518
        #[test]
    1519
    0
        fn case_with_expr() -> Result<()> {
    1520
    0
            let batch = case_test_batch()?;
    1521
    0
            let schema = batch.schema();
    1522
    1523
            // CASE a WHEN 'foo' THEN 123 WHEN 'bar' THEN 456 END
    1524
    0
            let when1 = lit("foo");
    1525
    0
            let then1 = lit(123i32);
    1526
    0
            let when2 = lit("bar");
    1527
    0
            let then2 = lit(456i32);
    1528
    1529
    0
            let expr = generate_case_when_with_type_coercion(
    1530
    0
                Some(col("a", &schema)?),
    1531
    0
                vec![(when1, then1), (when2, then2)],
    1532
    0
                None,
    1533
    0
                schema.as_ref(),
    1534
    0
            )?;
    1535
    0
            let result = expr
    1536
    0
                .evaluate(&batch)?
    1537
    0
                .into_array(batch.num_rows())
    1538
    0
                .expect("Failed to convert to array");
    1539
    0
            let result = as_int32_array(&result)?;
    1540
    1541
    0
            let expected = &Int32Array::from(vec![Some(123), None, None, Some(456)]);
    1542
    1543
    0
            assert_eq!(expected, result);
    1544
    1545
    0
            Ok(())
    1546
    0
        }
    1547
    1548
        #[test]
    1549
    0
        fn case_with_expr_else() -> Result<()> {
    1550
    0
            let batch = case_test_batch()?;
    1551
    0
            let schema = batch.schema();
    1552
    1553
            // CASE a WHEN 'foo' THEN 123 WHEN 'bar' THEN 456 ELSE 999 END
    1554
    0
            let when1 = lit("foo");
    1555
    0
            let then1 = lit(123i32);
    1556
    0
            let when2 = lit("bar");
    1557
    0
            let then2 = lit(456i32);
    1558
    0
            let else_value = lit(999i32);
    1559
    1560
    0
            let expr = generate_case_when_with_type_coercion(
    1561
    0
                Some(col("a", &schema)?),
    1562
    0
                vec![(when1, then1), (when2, then2)],
    1563
    0
                Some(else_value),
    1564
    0
                schema.as_ref(),
    1565
    0
            )?;
    1566
    0
            let result = expr
    1567
    0
                .evaluate(&batch)?
    1568
    0
                .into_array(batch.num_rows())
    1569
    0
                .expect("Failed to convert to array");
    1570
    0
            let result = as_int32_array(&result)?;
    1571
    1572
    0
            let expected =
    1573
    0
                &Int32Array::from(vec![Some(123), Some(999), Some(999), Some(456)]);
    1574
    1575
    0
            assert_eq!(expected, result);
    1576
    1577
    0
            Ok(())
    1578
    0
        }
    1579
    1580
        #[test]
    1581
    0
        fn case_with_expr_divide_by_zero() -> Result<()> {
    1582
    0
            let batch = case_test_batch1()?;
    1583
    0
            let schema = batch.schema();
    1584
    1585
            // CASE a when 0 THEN float64(null) ELSE 25.0 / cast(a, float64)  END
    1586
    0
            let when1 = lit(0i32);
    1587
    0
            let then1 = lit(ScalarValue::Float64(None));
    1588
    0
            let else_value = binary(
    1589
    0
                lit(25.0f64),
    1590
    0
                Operator::Divide,
    1591
    0
                cast(col("a", &schema)?, &batch.schema(), Float64)?,
    1592
    0
                &batch.schema(),
    1593
    0
            )?;
    1594
    1595
    0
            let expr = generate_case_when_with_type_coercion(
    1596
    0
                Some(col("a", &schema)?),
    1597
    0
                vec![(when1, then1)],
    1598
    0
                Some(else_value),
    1599
    0
                schema.as_ref(),
    1600
    0
            )?;
    1601
    0
            let result = expr
    1602
    0
                .evaluate(&batch)?
    1603
    0
                .into_array(batch.num_rows())
    1604
    0
                .expect("Failed to convert to array");
    1605
    0
            let result =
    1606
    0
                as_float64_array(&result).expect("failed to downcast to Float64Array");
    1607
    1608
    0
            let expected = &Float64Array::from(vec![Some(25.0), None, None, Some(5.0)]);
    1609
    1610
    0
            assert_eq!(expected, result);
    1611
    1612
    0
            Ok(())
    1613
    0
        }
    1614
    1615
        #[test]
    1616
    0
        fn case_without_expr() -> Result<()> {
    1617
    0
            let batch = case_test_batch()?;
    1618
    0
            let schema = batch.schema();
    1619
    1620
            // CASE WHEN a = 'foo' THEN 123 WHEN a = 'bar' THEN 456 END
    1621
    0
            let when1 = binary(
    1622
    0
                col("a", &schema)?,
    1623
    0
                Operator::Eq,
    1624
    0
                lit("foo"),
    1625
    0
                &batch.schema(),
    1626
    0
            )?;
    1627
    0
            let then1 = lit(123i32);
    1628
    0
            let when2 = binary(
    1629
    0
                col("a", &schema)?,
    1630
    0
                Operator::Eq,
    1631
    0
                lit("bar"),
    1632
    0
                &batch.schema(),
    1633
    0
            )?;
    1634
    0
            let then2 = lit(456i32);
    1635
    1636
    0
            let expr = generate_case_when_with_type_coercion(
    1637
    0
                None,
    1638
    0
                vec![(when1, then1), (when2, then2)],
    1639
    0
                None,
    1640
    0
                schema.as_ref(),
    1641
    0
            )?;
    1642
    0
            let result = expr
    1643
    0
                .evaluate(&batch)?
    1644
    0
                .into_array(batch.num_rows())
    1645
    0
                .expect("Failed to convert to array");
    1646
    0
            let result = as_int32_array(&result)?;
    1647
    1648
    0
            let expected = &Int32Array::from(vec![Some(123), None, None, Some(456)]);
    1649
    1650
    0
            assert_eq!(expected, result);
    1651
    1652
    0
            Ok(())
    1653
    0
        }
    1654
    1655
        #[test]
    1656
    0
        fn case_with_expr_when_null() -> Result<()> {
    1657
    0
            let batch = case_test_batch()?;
    1658
    0
            let schema = batch.schema();
    1659
    1660
            // CASE a WHEN NULL THEN 0 WHEN a THEN 123 ELSE 999 END
    1661
    0
            let when1 = lit(ScalarValue::Utf8(None));
    1662
    0
            let then1 = lit(0i32);
    1663
    0
            let when2 = col("a", &schema)?;
    1664
    0
            let then2 = lit(123i32);
    1665
    0
            let else_value = lit(999i32);
    1666
    1667
    0
            let expr = generate_case_when_with_type_coercion(
    1668
    0
                Some(col("a", &schema)?),
    1669
    0
                vec![(when1, then1), (when2, then2)],
    1670
    0
                Some(else_value),
    1671
    0
                schema.as_ref(),
    1672
    0
            )?;
    1673
    0
            let result = expr
    1674
    0
                .evaluate(&batch)?
    1675
    0
                .into_array(batch.num_rows())
    1676
    0
                .expect("Failed to convert to array");
    1677
    0
            let result = as_int32_array(&result)?;
    1678
    1679
    0
            let expected =
    1680
    0
                &Int32Array::from(vec![Some(123), Some(123), Some(999), Some(123)]);
    1681
    1682
    0
            assert_eq!(expected, result);
    1683
    1684
    0
            Ok(())
    1685
    0
        }
    1686
    1687
        #[test]
    1688
    0
        fn case_without_expr_divide_by_zero() -> Result<()> {
    1689
    0
            let batch = case_test_batch1()?;
    1690
    0
            let schema = batch.schema();
    1691
    1692
            // CASE WHEN a > 0 THEN 25.0 / cast(a, float64) ELSE float64(null) END
    1693
    0
            let when1 = binary(col("a", &schema)?, Operator::Gt, lit(0i32), &batch.schema())?;
    1694
    0
            let then1 = binary(
    1695
    0
                lit(25.0f64),
    1696
    0
                Operator::Divide,
    1697
    0
                cast(col("a", &schema)?, &batch.schema(), Float64)?,
    1698
    0
                &batch.schema(),
    1699
    0
            )?;
    1700
    0
            let x = lit(ScalarValue::Float64(None));
    1701
    1702
    0
            let expr = generate_case_when_with_type_coercion(
    1703
    0
                None,
    1704
    0
                vec![(when1, then1)],
    1705
    0
                Some(x),
    1706
    0
                schema.as_ref(),
    1707
    0
            )?;
    1708
    0
            let result = expr
    1709
    0
                .evaluate(&batch)?
    1710
    0
                .into_array(batch.num_rows())
    1711
    0
                .expect("Failed to convert to array");
    1712
    0
            let result =
    1713
    0
                as_float64_array(&result).expect("failed to downcast to Float64Array");
    1714
    1715
    0
            let expected = &Float64Array::from(vec![Some(25.0), None, None, Some(5.0)]);
    1716
    1717
    0
            assert_eq!(expected, result);
    1718
    1719
    0
            Ok(())
    1720
    0
        }
    1721
    1722
    0
        fn case_test_batch1() -> Result<RecordBatch> {
    1723
    0
            let schema = Schema::new(vec![
    1724
    0
                Field::new("a", DataType::Int32, true),
    1725
    0
                Field::new("b", DataType::Int32, true),
    1726
    0
                Field::new("c", DataType::Int32, true),
    1727
            ]);
    1728
    0
            let a = Int32Array::from(vec![Some(1), Some(0), None, Some(5)]);
    1729
    0
            let b = Int32Array::from(vec![Some(3), None, Some(14), Some(7)]);
    1730
    0
            let c = Int32Array::from(vec![Some(0), Some(-3), Some(777), None]);
    1731
    0
            let batch = RecordBatch::try_new(
    1732
    0
                Arc::new(schema),
    1733
    0
                vec![Arc::new(a), Arc::new(b), Arc::new(c)],
    1734
    0
            )?;
    1735
    0
            Ok(batch)
    1736
    0
        }
    1737
    1738
        #[test]
    1739
    0
        fn case_without_expr_else() -> Result<()> {
    1740
    0
            let batch = case_test_batch()?;
    1741
    0
            let schema = batch.schema();
    1742
    1743
            // CASE WHEN a = 'foo' THEN 123 WHEN a = 'bar' THEN 456 ELSE 999 END
    1744
    0
            let when1 = binary(
    1745
    0
                col("a", &schema)?,
    1746
    0
                Operator::Eq,
    1747
    0
                lit("foo"),
    1748
    0
                &batch.schema(),
    1749
    0
            )?;
    1750
    0
            let then1 = lit(123i32);
    1751
    0
            let when2 = binary(
    1752
    0
                col("a", &schema)?,
    1753
    0
                Operator::Eq,
    1754
    0
                lit("bar"),
    1755
    0
                &batch.schema(),
    1756
    0
            )?;
    1757
    0
            let then2 = lit(456i32);
    1758
    0
            let else_value = lit(999i32);
    1759
    1760
    0
            let expr = generate_case_when_with_type_coercion(
    1761
    0
                None,
    1762
    0
                vec![(when1, then1), (when2, then2)],
    1763
    0
                Some(else_value),
    1764
    0
                schema.as_ref(),
    1765
    0
            )?;
    1766
    0
            let result = expr
    1767
    0
                .evaluate(&batch)?
    1768
    0
                .into_array(batch.num_rows())
    1769
    0
                .expect("Failed to convert to array");
    1770
    0
            let result = as_int32_array(&result)?;
    1771
    1772
    0
            let expected =
    1773
    0
                &Int32Array::from(vec![Some(123), Some(999), Some(999), Some(456)]);
    1774
    1775
    0
            assert_eq!(expected, result);
    1776
    1777
    0
            Ok(())
    1778
    0
        }
    1779
    1780
        #[test]
    1781
    0
        fn case_with_type_cast() -> Result<()> {
    1782
    0
            let batch = case_test_batch()?;
    1783
    0
            let schema = batch.schema();
    1784
    1785
            // CASE WHEN a = 'foo' THEN 123.3 ELSE 999 END
    1786
    0
            let when = binary(
    1787
    0
                col("a", &schema)?,
    1788
    0
                Operator::Eq,
    1789
    0
                lit("foo"),
    1790
    0
                &batch.schema(),
    1791
    0
            )?;
    1792
    0
            let then = lit(123.3f64);
    1793
    0
            let else_value = lit(999i32);
    1794
    1795
    0
            let expr = generate_case_when_with_type_coercion(
    1796
    0
                None,
    1797
    0
                vec![(when, then)],
    1798
    0
                Some(else_value),
    1799
    0
                schema.as_ref(),
    1800
    0
            )?;
    1801
    0
            let result = expr
    1802
    0
                .evaluate(&batch)?
    1803
    0
                .into_array(batch.num_rows())
    1804
    0
                .expect("Failed to convert to array");
    1805
    0
            let result =
    1806
    0
                as_float64_array(&result).expect("failed to downcast to Float64Array");
    1807
    1808
    0
            let expected =
    1809
    0
                &Float64Array::from(vec![Some(123.3), Some(999.0), Some(999.0), Some(999.0)]);
    1810
    1811
    0
            assert_eq!(expected, result);
    1812
    1813
    0
            Ok(())
    1814
    0
        }
    1815
    1816
        #[test]
    1817
    0
        fn case_with_matches_and_nulls() -> Result<()> {
    1818
    0
            let batch = case_test_batch_nulls()?;
    1819
    0
            let schema = batch.schema();
    1820
    1821
            // SELECT CASE WHEN load4 = 1.77 THEN load4 END
    1822
    0
            let when = binary(
    1823
    0
                col("load4", &schema)?,
    1824
    0
                Operator::Eq,
    1825
    0
                lit(1.77f64),
    1826
    0
                &batch.schema(),
    1827
    0
            )?;
    1828
    0
            let then = col("load4", &schema)?;
    1829
    1830
    0
            let expr = generate_case_when_with_type_coercion(
    1831
    0
                None,
    1832
    0
                vec![(when, then)],
    1833
    0
                None,
    1834
    0
                schema.as_ref(),
    1835
    0
            )?;
    1836
    0
            let result = expr
    1837
    0
                .evaluate(&batch)?
    1838
    0
                .into_array(batch.num_rows())
    1839
    0
                .expect("Failed to convert to array");
    1840
    0
            let result =
    1841
    0
                as_float64_array(&result).expect("failed to downcast to Float64Array");
    1842
    1843
    0
            let expected =
    1844
    0
                &Float64Array::from(vec![Some(1.77), None, None, None, None, Some(1.77)]);
    1845
    1846
    0
            assert_eq!(expected, result);
    1847
    1848
    0
            Ok(())
    1849
    0
        }
    1850
    1851
        #[test]
    1852
    0
        fn case_with_scalar_predicate() -> Result<()> {
    1853
    0
            let batch = case_test_batch_nulls()?;
    1854
    0
            let schema = batch.schema();
    1855
    1856
            // SELECT CASE WHEN TRUE THEN load4 END
    1857
    0
            let when = lit(true);
    1858
    0
            let then = col("load4", &schema)?;
    1859
    0
            let expr = generate_case_when_with_type_coercion(
    1860
    0
                None,
    1861
    0
                vec![(when, then)],
    1862
    0
                None,
    1863
    0
                schema.as_ref(),
    1864
    0
            )?;
    1865
    1866
            // many rows
    1867
    0
            let result = expr
    1868
    0
                .evaluate(&batch)?
    1869
    0
                .into_array(batch.num_rows())
    1870
    0
                .expect("Failed to convert to array");
    1871
    0
            let result =
    1872
    0
                as_float64_array(&result).expect("failed to downcast to Float64Array");
    1873
    0
            let expected = &Float64Array::from(vec![
    1874
    0
                Some(1.77),
    1875
    0
                None,
    1876
    0
                None,
    1877
    0
                Some(1.78),
    1878
    0
                None,
    1879
    0
                Some(1.77),
    1880
    0
            ]);
    1881
    0
            assert_eq!(expected, result);
    1882
    1883
            // one row
    1884
    0
            let expected = Float64Array::from(vec![Some(1.1)]);
    1885
    0
            let batch =
    1886
    0
                RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(expected.clone())])?;
    1887
    0
            let result = expr
    1888
    0
                .evaluate(&batch)?
    1889
    0
                .into_array(batch.num_rows())
    1890
    0
                .expect("Failed to convert to array");
    1891
    0
            let result =
    1892
    0
                as_float64_array(&result).expect("failed to downcast to Float64Array");
    1893
    0
            assert_eq!(&expected, result);
    1894
    1895
    0
            Ok(())
    1896
    0
        }
    1897
    1898
        #[test]
    1899
    0
        fn case_expr_matches_and_nulls() -> Result<()> {
    1900
    0
            let batch = case_test_batch_nulls()?;
    1901
    0
            let schema = batch.schema();
    1902
    1903
            // SELECT CASE load4 WHEN 1.77 THEN load4 END
    1904
    0
            let expr = col("load4", &schema)?;
    1905
    0
            let when = lit(1.77f64);
    1906
    0
            let then = col("load4", &schema)?;
    1907
    1908
    0
            let expr = generate_case_when_with_type_coercion(
    1909
    0
                Some(expr),
    1910
    0
                vec![(when, then)],
    1911
    0
                None,
    1912
    0
                schema.as_ref(),
    1913
    0
            )?;
    1914
    0
            let result = expr
    1915
    0
                .evaluate(&batch)?
    1916
    0
                .into_array(batch.num_rows())
    1917
    0
                .expect("Failed to convert to array");
    1918
    0
            let result =
    1919
    0
                as_float64_array(&result).expect("failed to downcast to Float64Array");
    1920
    1921
    0
            let expected =
    1922
    0
                &Float64Array::from(vec![Some(1.77), None, None, None, None, Some(1.77)]);
    1923
    1924
    0
            assert_eq!(expected, result);
    1925
    1926
    0
            Ok(())
    1927
    0
        }
    1928
    1929
        #[test]
    1930
    0
        fn test_when_null_and_some_cond_else_null() -> Result<()> {
    1931
    0
            let batch = case_test_batch()?;
    1932
    0
            let schema = batch.schema();
    1933
    1934
    0
            let when = binary(
    1935
    0
                Arc::new(Literal::new(ScalarValue::Boolean(None))),
    1936
    0
                Operator::And,
    1937
    0
                binary(col("a", &schema)?, Operator::Eq, lit("foo"), &schema)?,
    1938
    0
                &schema,
    1939
    0
            )?;
    1940
    0
            let then = col("a", &schema)?;
    1941
    1942
            // SELECT CASE WHEN (NULL AND a = 'foo') THEN a ELSE NULL END
    1943
    0
            let expr = Arc::new(CaseExpr::try_new(None, vec![(when, then)], None)?);
    1944
    0
            let result = expr
    1945
    0
                .evaluate(&batch)?
    1946
    0
                .into_array(batch.num_rows())
    1947
    0
                .expect("Failed to convert to array");
    1948
    0
            let result = as_string_array(&result);
    1949
    1950
            // all result values should be null
    1951
    0
            assert_eq!(result.logical_null_count(), batch.num_rows());
    1952
    0
            Ok(())
    1953
    0
        }
    1954
    1955
    0
        fn case_test_batch() -> Result<RecordBatch> {
    1956
    0
            let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
    1957
    0
            let a = StringArray::from(vec![Some("foo"), Some("baz"), None, Some("bar")]);
    1958
    0
            let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?;
    1959
    0
            Ok(batch)
    1960
    0
        }
    1961
    1962
        // Construct an array that has several NULL values whose
    1963
        // underlying buffer actually matches the where expr predicate
    1964
    0
        fn case_test_batch_nulls() -> Result<RecordBatch> {
    1965
    0
            let load4: Float64Array = vec![
    1966
    0
                Some(1.77), // 1.77
    1967
    0
                Some(1.77), // null <-- same value, but will be set to null
    1968
    0
                Some(1.77), // null <-- same value, but will be set to null
    1969
    0
                Some(1.78), // 1.78
    1970
    0
                None,       // null
    1971
    0
                Some(1.77), // 1.77
    1972
            ]
    1973
    0
            .into_iter()
    1974
    0
            .collect();
    1975
    1976
    0
            let null_buffer = Buffer::from([0b00101001u8]);
    1977
    0
            let load4 = load4
    1978
    0
                .into_data()
    1979
    0
                .into_builder()
    1980
    0
                .null_bit_buffer(Some(null_buffer))
    1981
    0
                .build()
    1982
    0
                .unwrap();
    1983
    0
            let load4: Float64Array = load4.into();
    1984
    1985
    0
            let batch =
    1986
    0
                RecordBatch::try_from_iter(vec![("load4", Arc::new(load4) as ArrayRef)])?;
    1987
    0
            Ok(batch)
    1988
    0
        }
    1989
    1990
        #[test]
    1991
    0
        fn case_test_incompatible() -> Result<()> {
    1992
            // 1 then is int64
    1993
            // 2 then is boolean
    1994
    0
            let batch = case_test_batch()?;
    1995
    0
            let schema = batch.schema();
    1996
    1997
            // CASE WHEN a = 'foo' THEN 123 WHEN a = 'bar' THEN true END
    1998
    0
            let when1 = binary(
    1999
    0
                col("a", &schema)?,
    2000
    0
                Operator::Eq,
    2001
    0
                lit("foo"),
    2002
    0
                &batch.schema(),
    2003
    0
            )?;
    2004
    0
            let then1 = lit(123i32);
    2005
    0
            let when2 = binary(
    2006
    0
                col("a", &schema)?,
    2007
    0
                Operator::Eq,
    2008
    0
                lit("bar"),
    2009
    0
                &batch.schema(),
    2010
    0
            )?;
    2011
    0
            let then2 = lit(true);
    2012
    2013
    0
            let expr = generate_case_when_with_type_coercion(
    2014
    0
                None,
    2015
    0
                vec![(when1, then1), (when2, then2)],
    2016
    0
                None,
    2017
    0
                schema.as_ref(),
    2018
            );
    2019
    0
            assert!(expr.is_err());
    2020
    2021
            // then 1 is int32
    2022
            // then 2 is int64
    2023
            // else is float
    2024
            // CASE WHEN a = 'foo' THEN 123 WHEN a = 'bar' THEN 456 ELSE 1.23 END
    2025
    0
            let when1 = binary(
    2026
    0
                col("a", &schema)?,
    2027
    0
                Operator::Eq,
    2028
    0
                lit("foo"),
    2029
    0
                &batch.schema(),
    2030
    0
            )?;
    2031
    0
            let then1 = lit(123i32);
    2032
    0
            let when2 = binary(
    2033
    0
                col("a", &schema)?,
    2034
    0
                Operator::Eq,
    2035
    0
                lit("bar"),
    2036
    0
                &batch.schema(),
    2037
    0
            )?;
    2038
    0
            let then2 = lit(456i64);
    2039
    0
            let else_expr = lit(1.23f64);
    2040
    2041
    0
            let expr = generate_case_when_with_type_coercion(
    2042
    0
                None,
    2043
    0
                vec![(when1, then1), (when2, then2)],
    2044
    0
                Some(else_expr),
    2045
    0
                schema.as_ref(),
    2046
            );
    2047
    0
            assert!(expr.is_ok());
    2048
    0
            let result_type = expr.unwrap().data_type(schema.as_ref())?;
    2049
    0
            assert_eq!(Float64, result_type);
    2050
    0
            Ok(())
    2051
    0
        }
    2052
    2053
        #[test]
    2054
    0
        fn case_eq() -> Result<()> {
    2055
    0
            let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
    2056
    2057
    0
            let when1 = lit("foo");
    2058
    0
            let then1 = lit(123i32);
    2059
    0
            let when2 = lit("bar");
    2060
    0
            let then2 = lit(456i32);
    2061
    0
            let else_value = lit(999i32);
    2062
    2063
    0
            let expr1 = generate_case_when_with_type_coercion(
    2064
    0
                Some(col("a", &schema)?),
    2065
    0
                vec![
    2066
    0
                    (Arc::clone(&when1), Arc::clone(&then1)),
    2067
    0
                    (Arc::clone(&when2), Arc::clone(&then2)),
    2068
                ],
    2069
    0
                Some(Arc::clone(&else_value)),
    2070
    0
                &schema,
    2071
    0
            )?;
    2072
    2073
    0
            let expr2 = generate_case_when_with_type_coercion(
    2074
    0
                Some(col("a", &schema)?),
    2075
    0
                vec![
    2076
    0
                    (Arc::clone(&when1), Arc::clone(&then1)),
    2077
    0
                    (Arc::clone(&when2), Arc::clone(&then2)),
    2078
                ],
    2079
    0
                Some(Arc::clone(&else_value)),
    2080
    0
                &schema,
    2081
    0
            )?;
    2082
    2083
    0
            let expr3 = generate_case_when_with_type_coercion(
    2084
    0
                Some(col("a", &schema)?),
    2085
    0
                vec![(Arc::clone(&when1), Arc::clone(&then1)), (when2, then2)],
    2086
    0
                None,
    2087
    0
                &schema,
    2088
    0
            )?;
    2089
    2090
    0
            let expr4 = generate_case_when_with_type_coercion(
    2091
    0
                Some(col("a", &schema)?),
    2092
    0
                vec![(when1, then1)],
    2093
    0
                Some(else_value),
    2094
    0
                &schema,
    2095
    0
            )?;
    2096
    2097
    0
            assert!(expr1.eq(&expr2));
    2098
    0
            assert!(expr2.eq(&expr1));
    2099
    2100
    0
            assert!(expr2.ne(&expr3));
    2101
    0
            assert!(expr3.ne(&expr2));
    2102
    2103
    0
            assert!(expr1.ne(&expr4));
    2104
    0
            assert!(expr4.ne(&expr1));
    2105
    2106
    0
            Ok(())
    2107
    0
        }
    2108
    2109
        #[test]
    2110
    0
        fn case_transform() -> Result<()> {
    2111
    0
            let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
    2112
    2113
    0
            let when1 = lit("foo");
    2114
    0
            let then1 = lit(123i32);
    2115
    0
            let when2 = lit("bar");
    2116
    0
            let then2 = lit(456i32);
    2117
    0
            let else_value = lit(999i32);
    2118
    2119
    0
            let expr = generate_case_when_with_type_coercion(
    2120
    0
                Some(col("a", &schema)?),
    2121
    0
                vec![
    2122
    0
                    (Arc::clone(&when1), Arc::clone(&then1)),
    2123
    0
                    (Arc::clone(&when2), Arc::clone(&then2)),
    2124
                ],
    2125
    0
                Some(Arc::clone(&else_value)),
    2126
    0
                &schema,
    2127
    0
            )?;
    2128
    2129
    0
            let expr2 = Arc::clone(&expr)
    2130
    0
                .transform(|e| {
    2131
    0
                    let transformed = match e.as_any().downcast_ref::<Literal>() {
    2132
    0
                        Some(lit_value) => match lit_value.value() {
    2133
    0
                            ScalarValue::Utf8(Some(str_value)) => {
    2134
    0
                                Some(lit(str_value.to_uppercase()))
    2135
                            }
    2136
    0
                            _ => None,
    2137
                        },
    2138
    0
                        _ => None,
    2139
                    };
    2140
    0
                    Ok(if let Some(transformed) = transformed {
    2141
    0
                        Transformed::yes(transformed)
    2142
                    } else {
    2143
    0
                        Transformed::no(e)
    2144
                    })
    2145
    0
                })
    2146
    0
                .data()
    2147
    0
                .unwrap();
    2148
    2149
    0
            let expr3 = Arc::clone(&expr)
    2150
    0
                .transform_down(|e| {
    2151
    0
                    let transformed = match e.as_any().downcast_ref::<Literal>() {
    2152
    0
                        Some(lit_value) => match lit_value.value() {
    2153
    0
                            ScalarValue::Utf8(Some(str_value)) => {
    2154
    0
                                Some(lit(str_value.to_uppercase()))
    2155
                            }
    2156
    0
                            _ => None,
    2157
                        },
    2158
    0
                        _ => None,
    2159
                    };
    2160
    0
                    Ok(if let Some(transformed) = transformed {
    2161
    0
                        Transformed::yes(transformed)
    2162
                    } else {
    2163
    0
                        Transformed::no(e)
    2164
                    })
    2165
    0
                })
    2166
    0
                .data()
    2167
    0
                .unwrap();
    2168
    2169
    0
            assert!(expr.ne(&expr2));
    2170
    0
            assert!(expr2.eq(&expr3));
    2171
    2172
    0
            Ok(())
    2173
    0
        }
    2174
    2175
        #[test]
    2176
    0
        fn test_column_or_null_specialization() -> Result<()> {
    2177
            // create input data
    2178
    0
            let mut c1 = Int32Builder::new();
    2179
    0
            let mut c2 = StringBuilder::new();
    2180
    0
            for i in 0..1000 {
    2181
    0
                c1.append_value(i);
    2182
    0
                if i % 7 == 0 {
    2183
    0
                    c2.append_null();
    2184
    0
                } else {
    2185
    0
                    c2.append_value(format!("string {i}"));
    2186
    0
                }
    2187
            }
    2188
    0
            let c1 = Arc::new(c1.finish());
    2189
    0
            let c2 = Arc::new(c2.finish());
    2190
    0
            let schema = Schema::new(vec![
    2191
    0
                Field::new("c1", DataType::Int32, true),
    2192
    0
                Field::new("c2", DataType::Utf8, true),
    2193
            ]);
    2194
    0
            let batch = RecordBatch::try_new(Arc::new(schema), vec![c1, c2]).unwrap();
    2195
    2196
            // CaseWhenExprOrNull should produce same results as CaseExpr
    2197
    0
            let predicate = Arc::new(BinaryExpr::new(
    2198
    0
                make_col("c1", 0),
    2199
    0
                Operator::LtEq,
    2200
    0
                make_lit_i32(250),
    2201
            ));
    2202
    0
            let expr = CaseExpr::try_new(None, vec![(predicate, make_col("c2", 1))], None)?;
    2203
    0
            assert!(matches!(expr.eval_method, EvalMethod::InfallibleExprOrNull));
    2204
    0
            match expr.evaluate(&batch)? {
    2205
    0
                ColumnarValue::Array(array) => {
    2206
    0
                    assert_eq!(1000, array.len());
    2207
    0
                    assert_eq!(785, array.null_count());
    2208
                }
    2209
    0
                _ => unreachable!(),
    2210
            }
    2211
    0
            Ok(())
    2212
    0
        }
    2213
    2214
        #[test]
    2215
    0
        fn test_expr_or_expr_specialization() -> Result<()> {
    2216
    0
            let batch = case_test_batch1()?;
    2217
    0
            let schema = batch.schema();
    2218
    0
            let when = binary(
    2219
    0
                col("a", &schema)?,
    2220
    0
                Operator::LtEq,
    2221
    0
                lit(2i32),
    2222
    0
                &batch.schema(),
    2223
    0
            )?;
    2224
    0
            let then = col("b", &schema)?;
    2225
    0
            let else_expr = col("c", &schema)?;
    2226
    0
            let expr = CaseExpr::try_new(None, vec![(when, then)], Some(else_expr))?;
    2227
    0
            assert!(matches!(
    2228
    0
                expr.eval_method,
    2229
                EvalMethod::ExpressionOrExpression(_)
    2230
            ));
    2231
    0
            let result = expr
    2232
    0
                .evaluate(&batch)?
    2233
    0
                .into_array(batch.num_rows())
    2234
    0
                .expect("Failed to convert to array");
    2235
    0
            let result = as_int32_array(&result).expect("failed to downcast to Int32Array");
    2236
    2237
    0
            let expected = &Int32Array::from(vec![Some(3), None, Some(777), None]);
    2238
    2239
    0
            assert_eq!(expected, result);
    2240
    0
            Ok(())
    2241
    0
        }
    2242
    2243
    0
        fn make_col(name: &str, index: usize) -> Arc<dyn PhysicalExpr> {
    2244
    0
            Arc::new(Column::new(name, index))
    2245
    0
        }
    2246
    2247
    0
        fn make_lit_i32(n: i32) -> Arc<dyn PhysicalExpr> {
    2248
    0
            Arc::new(Literal::new(ScalarValue::Int32(Some(n))))
    2249
    0
        }
    2250
    2251
    0
        fn generate_case_when_with_type_coercion(
    2252
    0
            expr: Option<Arc<dyn PhysicalExpr>>,
    2253
    0
            when_thens: Vec<WhenThen>,
    2254
    0
            else_expr: Option<Arc<dyn PhysicalExpr>>,
    2255
    0
            input_schema: &Schema,
    2256
    0
        ) -> Result<Arc<dyn PhysicalExpr>> {
    2257
    0
            let coerce_type =
    2258
    0
                get_case_common_type(&when_thens, else_expr.clone(), input_schema);
    2259
    0
            let (when_thens, else_expr) = match coerce_type {
    2260
    0
                None => plan_err!(
    2261
    0
                    "Can't get a common type for then {when_thens:?} and else {else_expr:?} expression"
    2262
                ),
    2263
    0
                Some(data_type) => {
    2264
                    // cast then expr
    2265
    0
                    let left = when_thens
    2266
    0
                        .into_iter()
    2267
    0
                        .map(|(when, then)| {
    2268
    0
                            let then = try_cast(then, input_schema, data_type.clone())?;
    2269
    0
                            Ok((when, then))
    2270
    0
                        })
    2271
    0
                        .collect::<Result<Vec<_>>>()?;
    2272
    0
                    let right = match else_expr {
    2273
    0
                        None => None,
    2274
    0
                        Some(expr) => Some(try_cast(expr, input_schema, data_type.clone())?),
    2275
                    };
    2276
    2277
    0
                    Ok((left, right))
    2278
                }
    2279
    0
            }?;
    2280
    0
            case(expr, when_thens, else_expr)
    2281
    0
        }
    2282
    2283
    0
        fn get_case_common_type(
    2284
    0
            when_thens: &[WhenThen],
    2285
    0
            else_expr: Option<Arc<dyn PhysicalExpr>>,
    2286
    0
            input_schema: &Schema,
    2287
    0
        ) -> Option<DataType> {
    2288
    0
            let thens_type = when_thens
    2289
    0
                .iter()
    2290
    0
                .map(|when_then| {
    2291
    0
                    let data_type = &when_then.1.data_type(input_schema).unwrap();
    2292
    0
                    data_type.clone()
    2293
    0
                })
    2294
    0
                .collect::<Vec<_>>();
    2295
    0
            let else_type = match else_expr {
    2296
                None => {
    2297
                    // case when then exprs must have one then value
    2298
    0
                    thens_type[0].clone()
    2299
                }
    2300
    0
                Some(else_phy_expr) => else_phy_expr.data_type(input_schema).unwrap(),
    2301
            };
    2302
    0
            thens_type
    2303
    0
                .iter()
    2304
    0
                .try_fold(else_type, |left_type, right_type| {
    2305
                    // TODO: now just use the `equal` coercion rule for case when. If find the issue, and
    2306
                    // refactor again.
    2307
    0
                    comparison_coercion(&left_type, right_type)
    2308
    0
                })
    2309
    0
        }
    2310
    2311
        #[test]
    2312
    0
        fn test_fmt_sql() -> Result<()> {
    2313
    0
            let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
    2314
    2315
            // CASE WHEN a = 'foo' THEN 123.3 ELSE 999 END
    2316
    0
            let when = binary(col("a", &schema)?, Operator::Eq, lit("foo"), &schema)?;
    2317
    0
            let then = lit(123.3f64);
    2318
    0
            let else_value = lit(999i32);
    2319
    2320
    0
            let expr = generate_case_when_with_type_coercion(
    2321
    0
                None,
    2322
    0
                vec![(when, then)],
    2323
    0
                Some(else_value),
    2324
    0
                &schema,
    2325
    0
            )?;
    2326
    2327
    0
            let display_string = expr.to_string();
    2328
    0
            assert_eq!(
    2329
                display_string,
    2330
                "CASE WHEN a@0 = foo THEN 123.3 ELSE TRY_CAST(999 AS Float64) END"
    2331
            );
    2332
    2333
    0
            let sql_string = fmt_sql(expr.as_ref()).to_string();
    2334
    0
            assert_eq!(
    2335
                sql_string,
    2336
                "CASE WHEN a = foo THEN 123.3 ELSE TRY_CAST(999 AS Float64) END"
    2337
            );
    2338
    2339
    0
            Ok(())
    2340
    0
        }
    2341
    2342
        #[test]
    2343
    0
        fn test_merge_n() {
    2344
    0
            let a1 = StringArray::from(vec![Some("A")]).to_data();
    2345
    0
            let a2 = StringArray::from(vec![Some("B")]).to_data();
    2346
    0
            let a3 = StringArray::from(vec![Some("C"), Some("D")]).to_data();
    2347
    2348
    0
            let indices = vec![
    2349
    0
                PartialResultIndex::none(),
    2350
    0
                PartialResultIndex::try_new(1).unwrap(),
    2351
    0
                PartialResultIndex::try_new(0).unwrap(),
    2352
    0
                PartialResultIndex::none(),
    2353
    0
                PartialResultIndex::try_new(2).unwrap(),
    2354
    0
                PartialResultIndex::try_new(2).unwrap(),
    2355
            ];
    2356
    2357
    0
            let merged = merge_n(&[a1, a2, a3], &indices).unwrap();
    2358
    0
            let merged = merged.as_string::<i32>();
    2359
    2360
    0
            assert_eq!(merged.len(), indices.len());
    2361
    0
            assert!(!merged.is_valid(0));
    2362
    0
            assert!(merged.is_valid(1));
    2363
    0
            assert_eq!(merged.value(1), "B");
    2364
    0
            assert!(merged.is_valid(2));
    2365
    0
            assert_eq!(merged.value(2), "A");
    2366
    0
            assert!(!merged.is_valid(3));
    2367
    0
            assert!(merged.is_valid(4));
    2368
    0
            assert_eq!(merged.value(4), "C");
    2369
    0
            assert!(merged.is_valid(5));
    2370
    0
            assert_eq!(merged.value(5), "D");
    2371
    0
        }
    2372
    2373
        #[test]
    2374
    0
        fn test_merge() {
    2375
    0
            let a1 = Arc::new(StringArray::from(vec![Some("A"), Some("C")]));
    2376
    0
            let a2 = Arc::new(StringArray::from(vec![Some("B")]));
    2377
    2378
    0
            let mask = BooleanArray::from(vec![true, false, true]);
    2379
    2380
    0
            let merged =
    2381
    0
                merge(&mask, ColumnarValue::Array(a1), ColumnarValue::Array(a2)).unwrap();
    2382
    0
            let merged = merged.as_string::<i32>();
    2383
    2384
    0
            assert_eq!(merged.len(), mask.len());
    2385
    0
            assert!(merged.is_valid(0));
    2386
    0
            assert_eq!(merged.value(0), "A");
    2387
    0
            assert!(merged.is_valid(1));
    2388
    0
            assert_eq!(merged.value(1), "B");
    2389
    0
            assert!(merged.is_valid(2));
    2390
    0
            assert_eq!(merged.value(2), "C");
    2391
    0
        }
    2392
    2393
    0
        fn when_then_else(
    2394
    0
            when: &Arc<dyn PhysicalExpr>,
    2395
    0
            then: &Arc<dyn PhysicalExpr>,
    2396
    0
            els: &Arc<dyn PhysicalExpr>,
    2397
    0
        ) -> Result<Arc<dyn PhysicalExpr>> {
    2398
    0
            let case = CaseExpr::try_new(
    2399
    0
                None,
    2400
    0
                vec![(Arc::clone(when), Arc::clone(then))],
    2401
    0
                Some(Arc::clone(els)),
    2402
    0
            )?;
    2403
    0
            Ok(Arc::new(case))
    2404
    0
        }
    2405
    2406
        #[test]
    2407
    0
        fn test_case_expression_nullability_with_nullable_column() -> Result<()> {
    2408
    0
            case_expression_nullability(true)
    2409
    0
        }
    2410
    2411
        #[test]
    2412
    0
        fn test_case_expression_nullability_with_not_nullable_column() -> Result<()> {
    2413
    0
            case_expression_nullability(false)
    2414
    0
        }
    2415
    2416
    0
        fn case_expression_nullability(col_is_nullable: bool) -> Result<()> {
    2417
    0
            let schema =
    2418
    0
                Schema::new(vec![Field::new("foo", DataType::Int32, col_is_nullable)]);
    2419
    2420
    0
            let foo = col("foo", &schema)?;
    2421
    0
            let foo_is_not_null = is_not_null(Arc::clone(&foo))?;
    2422
    0
            let foo_is_null = expressions::is_null(Arc::clone(&foo))?;
    2423
    0
            let not_foo_is_null = expressions::not(Arc::clone(&foo_is_null))?;
    2424
    0
            let zero = lit(0);
    2425
    0
            let foo_eq_zero =
    2426
    0
                binary(Arc::clone(&foo), Operator::Eq, Arc::clone(&zero), &schema)?;
    2427
    2428
    0
            assert_not_nullable(when_then_else(&foo_is_not_null, &foo, &zero)?, &schema);
    2429
    0
            assert_not_nullable(when_then_else(&not_foo_is_null, &foo, &zero)?, &schema);
    2430
    0
            assert_not_nullable(when_then_else(&foo_eq_zero, &foo, &zero)?, &schema);
    2431
    2432
    0
            assert_not_nullable(
    2433
    0
                when_then_else(
    2434
    0
                    &binary(
    2435
    0
                        Arc::clone(&foo_is_not_null),
    2436
    0
                        Operator::And,
    2437
    0
                        Arc::clone(&foo_eq_zero),
    2438
    0
                        &schema,
    2439
    0
                    )?,
    2440
    0
                    &foo,
    2441
    0
                    &zero,
    2442
    0
                )?,
    2443
    0
                &schema,
    2444
            );
    2445
    2446
    0
            assert_not_nullable(
    2447
    0
                when_then_else(
    2448
    0
                    &binary(
    2449
    0
                        Arc::clone(&foo_eq_zero),
    2450
    0
                        Operator::And,
    2451
    0
                        Arc::clone(&foo_is_not_null),
    2452
    0
                        &schema,
    2453
    0
                    )?,
    2454
    0
                    &foo,
    2455
    0
                    &zero,
    2456
    0
                )?,
    2457
    0
                &schema,
    2458
            );
    2459
    2460
    0
            assert_not_nullable(
    2461
    0
                when_then_else(
    2462
    0
                    &binary(
    2463
    0
                        Arc::clone(&foo_is_not_null),
    2464
    0
                        Operator::Or,
    2465
    0
                        Arc::clone(&foo_eq_zero),
    2466
    0
                        &schema,
    2467
    0
                    )?,
    2468
    0
                    &foo,
    2469
    0
                    &zero,
    2470
    0
                )?,
    2471
    0
                &schema,
    2472
            );
    2473
    2474
    0
            assert_not_nullable(
    2475
    0
                when_then_else(
    2476
    0
                    &binary(
    2477
    0
                        Arc::clone(&foo_eq_zero),
    2478
    0
                        Operator::Or,
    2479
    0
                        Arc::clone(&foo_is_not_null),
    2480
    0
                        &schema,
    2481
    0
                    )?,
    2482
    0
                    &foo,
    2483
    0
                    &zero,
    2484
    0
                )?,
    2485
    0
                &schema,
    2486
            );
    2487
    2488
    0
            assert_nullability(
    2489
    0
                when_then_else(
    2490
    0
                    &binary(
    2491
    0
                        Arc::clone(&foo_is_null),
    2492
    0
                        Operator::Or,
    2493
    0
                        Arc::clone(&foo_eq_zero),
    2494
    0
                        &schema,
    2495
    0
                    )?,
    2496
    0
                    &foo,
    2497
    0
                    &zero,
    2498
    0
                )?,
    2499
    0
                &schema,
    2500
    0
                col_is_nullable,
    2501
            );
    2502
    2503
    0
            assert_nullability(
    2504
    0
                when_then_else(
    2505
    0
                    &binary(
    2506
    0
                        binary(Arc::clone(&foo), Operator::Eq, Arc::clone(&zero), &schema)?,
    2507
    0
                        Operator::Or,
    2508
    0
                        Arc::clone(&foo_is_null),
    2509
    0
                        &schema,
    2510
    0
                    )?,
    2511
    0
                    &foo,
    2512
    0
                    &zero,
    2513
    0
                )?,
    2514
    0
                &schema,
    2515
    0
                col_is_nullable,
    2516
            );
    2517
    2518
    0
            assert_not_nullable(
    2519
    0
                when_then_else(
    2520
    0
                    &binary(
    2521
    0
                        binary(
    2522
    0
                            binary(
    2523
    0
                                Arc::clone(&foo),
    2524
    0
                                Operator::Eq,
    2525
    0
                                Arc::clone(&zero),
    2526
    0
                                &schema,
    2527
    0
                            )?,
    2528
    0
                            Operator::And,
    2529
    0
                            Arc::clone(&foo_is_not_null),
    2530
    0
                            &schema,
    2531
    0
                        )?,
    2532
    0
                        Operator::Or,
    2533
    0
                        binary(
    2534
    0
                            binary(
    2535
    0
                                Arc::clone(&foo),
    2536
    0
                                Operator::Eq,
    2537
    0
                                Arc::clone(&foo),
    2538
    0
                                &schema,
    2539
    0
                            )?,
    2540
    0
                            Operator::And,
    2541
    0
                            Arc::clone(&foo_is_not_null),
    2542
    0
                            &schema,
    2543
    0
                        )?,
    2544
    0
                        &schema,
    2545
    0
                    )?,
    2546
    0
                    &foo,
    2547
    0
                    &zero,
    2548
    0
                )?,
    2549
    0
                &schema,
    2550
            );
    2551
    2552
    0
            Ok(())
    2553
    0
        }
    2554
    2555
    0
        fn assert_not_nullable(expr: Arc<dyn PhysicalExpr>, schema: &Schema) {
    2556
    0
            assert!(!expr.nullable(schema).unwrap());
    2557
    0
        }
    2558
    2559
    0
        fn assert_nullable(expr: Arc<dyn PhysicalExpr>, schema: &Schema) {
    2560
    0
            assert!(expr.nullable(schema).unwrap());
    2561
    0
        }
    2562
    2563
    0
        fn assert_nullability(expr: Arc<dyn PhysicalExpr>, schema: &Schema, nullable: bool) {
    2564
    0
            if nullable {
    2565
    0
                assert_nullable(expr, schema);
    2566
    0
            } else {
    2567
    0
                assert_not_nullable(expr, schema);
    2568
    0
            }
    2569
    0
        }
    2570
    }