## Coverage Report

#### Created: 2024-10-13 08:39

    /Users/andrewlamb/Software/datafusion/datafusion/physical-expr/src/expressions/try_cast.rs

    Line
    Count
    Source (jump to first uncovered line)
    1
    // Licensed to the Apache Software Foundation (ASF) under one
    2
    // or more contributor license agreements.  See the NOTICE file
    3
    // distributed with this work for additional information
    4
    // regarding copyright ownership.  The ASF licenses this file
    5
    // to you under the Apache License, Version 2.0 (the
    6
    // "License"); you may not use this file except in compliance
    7
    // with the License.  You may obtain a copy of the License at
    8
    //
    9
    //   http://www.apache.org/licenses/LICENSE-2.0
    10
    //
    11
    // Unless required by applicable law or agreed to in writing,
    12
    // software distributed under the License is distributed on an
    13
    // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    14
    // KIND, either express or implied.  See the License for the
    15
    // specific language governing permissions and limitations
    16
    // under the License.
    17
    18
    use std::any::Any;
    19
    use std::fmt;
    20
    use std::hash::{Hash, Hasher};
    21
    use std::sync::Arc;
    22
    23
    use crate::physical_expr::down_cast_any_ref;
    24
    use crate::PhysicalExpr;
    25
    use arrow::compute;
    26
    use arrow::compute::{cast_with_options, CastOptions};
    27
    use arrow::datatypes::{DataType, Schema};
    28
    use arrow::record_batch::RecordBatch;
    29
    use compute::can_cast_types;
    30
    use datafusion_common::format::DEFAULT_FORMAT_OPTIONS;
    31
    use datafusion_common::{not_impl_err, Result, ScalarValue};
    32
    use datafusion_expr::ColumnarValue;
    33
    34
    /// TRY_CAST expression casts an expression to a specific data type and returns NULL on invalid cast
    35
    #[derive(Debug, Hash)]
    36
    pub struct TryCastExpr {
    37
        /// The expression to cast
    38
        expr: Arc<dyn PhysicalExpr>,
    39
        /// The data type to cast to
    40
        cast_type: DataType,
    41
    }
    42
    43
    impl TryCastExpr {
    44
        /// Create a new CastExpr
    45
    0
        pub fn new(expr: Arc<dyn PhysicalExpr>, cast_type: DataType) -> Self {
    46
    0
            Self { expr, cast_type }
    47
    0
        }
    48
    49
        /// The expression to cast
    50
    0
        pub fn expr(&self) -> &Arc<dyn PhysicalExpr> {
    51
    0
            &self.expr
    52
    0
        }
    53
    54
        /// The data type to cast to
    55
    0
        pub fn cast_type(&self) -> &DataType {
    56
    0
            &self.cast_type
    57
    0
        }
    58
    }
    59
    60
    impl fmt::Display for TryCastExpr {
    61
    0
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    62
    0
            write!(f, "TRY_CAST({} AS {:?})", self.expr, self.cast_type)
    63
    0
        }
    64
    }
    65
    66
    impl PhysicalExpr for TryCastExpr {
    67
        /// Return a reference to Any that can be used for downcasting
    68
    0
        fn as_any(&self) -> &dyn Any {
    69
    0
            self
    70
    0
        }
    71
    72
    0
        fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
    73
    0
            Ok(self.cast_type.clone())
    74
    0
        }
    75
    76
    0
        fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
    77
    0
            Ok(true)
    78
    0
        }
    79
    80
    0
        fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
    81
    0
            let value = self.expr.evaluate(batch)?;
    82
    0
            let options = CastOptions {
    83
    0
                safe: true,
    84
    0
                format_options: DEFAULT_FORMAT_OPTIONS,
    85
    0
            };
    86
    0
            match value {
    87
    0
                ColumnarValue::Array(array) => {
    88
    0
                    let cast = cast_with_options(&array, &self.cast_type, &options)?;
    89
    0
                    Ok(ColumnarValue::Array(cast))
    90
                }
    91
    0
                ColumnarValue::Scalar(scalar) => {
    92
    0
                    let array = scalar.to_array()?;
    93
    0
                    let cast_array = cast_with_options(&array, &self.cast_type, &options)?;
    94
    0
                    let cast_scalar = ScalarValue::try_from_array(&cast_array, 0)?;
    95
    0
                    Ok(ColumnarValue::Scalar(cast_scalar))
    96
                }
    97
            }
    98
    0
        }
    99
    100
    0
        fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
    101
    0
            vec![&self.expr]
    102
    0
        }
    103
    104
    0
        fn with_new_children(
    105
    0
            self: Arc<Self>,
    106
    0
            children: Vec<Arc<dyn PhysicalExpr>>,
    107
    0
        ) -> Result<Arc<dyn PhysicalExpr>> {
    108
    0
            Ok(Arc::new(TryCastExpr::new(
    109
    0
                Arc::clone(&children[0]),
    110
    0
                self.cast_type.clone(),
    111
    0
            )))
    112
    0
        }
    113
    114
    0
        fn dyn_hash(&self, state: &mut dyn Hasher) {
    115
    0
            let mut s = state;
    116
    0
            self.hash(&mut s);
    117
    0
        }
    118
    }
    119
    120
    impl PartialEq<dyn Any> for TryCastExpr {
    121
    0
        fn eq(&self, other: &dyn Any) -> bool {
    122
    0
            down_cast_any_ref(other)
    123
    0
                .downcast_ref::<Self>()
    124
    0
                .map(|x| self.expr.eq(&x.expr) && self.cast_type == x.cast_type)
    125
    0
                .unwrap_or(false)
    126
    0
        }
    127
    }
    128
    129
    /// Return a PhysicalExpression representing `expr` casted to
    130
    /// `cast_type`, if any casting is needed.
    131
    ///
    132
    /// Note that such casts may lose type information
    133
    0
    pub fn try_cast(
    134
    0
        expr: Arc<dyn PhysicalExpr>,
    135
    0
        input_schema: &Schema,
    136
    0
        cast_type: DataType,
    137
    0
    ) -> Result<Arc<dyn PhysicalExpr>> {
    138
    0
        let expr_type = expr.data_type(input_schema)?;
    139
    0
        if expr_type == cast_type {
    140
    0
            Ok(Arc::clone(&expr))
    141
    0
        } else if can_cast_types(&expr_type, &cast_type) {
    142
    0
            Ok(Arc::new(TryCastExpr::new(expr, cast_type)))
    143
        } else {
    144
    0
            not_impl_err!("Unsupported TRY_CAST from {expr_type:?} to {cast_type:?}")
    145
        }
    146
    0
    }
    147
    148
    #[cfg(test)]
    149
    mod tests {
    150
        use super::*;
    151
        use crate::expressions::col;
    152
        use arrow::array::{
    153
            Decimal128Array, Decimal128Builder, StringArray, Time64NanosecondArray,
    154
        };
    155
        use arrow::{
    156
            array::{
    157
                Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array,
    158
                Int8Array, TimestampNanosecondArray, UInt32Array,
    159
            },
    160
            datatypes::*,
    161
        };
    162
    163
        // runs an end-to-end test of physical type cast
    164
        // 1. construct a record batch with a column "a" of type A
    165
        // 2. construct a physical expression of TRY_CAST(a AS B)
    166
        // 3. evaluate the expression
    167
        // 4. verify that the resulting expression is of type B
    168
        // 5. verify that the resulting values are downcastable and correct
    169
        macro_rules! generic_decimal_to_other_test_cast {
    170
            ($DECIMAL_ARRAY:ident, $A_TYPE:expr, $TYPEARRAY:ident, $TYPE:expr, $VEC:expr) => {{
    171
                let schema = Schema::new(vec![Field::new("a", $A_TYPE, true)]);
    172
                let batch = RecordBatch::try_new(
    173
                    Arc::new(schema.clone()),
    174
                    vec![Arc::new($DECIMAL_ARRAY)],
    175
                )?;
    176
                // verify that we can construct the expression
    177
                let expression = try_cast(col("a", &schema)?, &schema, $TYPE)?;
    178
    179
                // verify that its display is correct
    180
                assert_eq!(
    181
                    format!("TRY_CAST(a@0 AS {:?})", $TYPE),
    182
                    format!("{}", expression)
    183
                );
    184
    185
                // verify that the expression's type is correct
    186
                assert_eq!(expression.data_type(&schema)?, $TYPE);
    187
    188
                // compute
    189
                let result = expression
    190
                    .evaluate(&batch)?
    191
                    .into_array(batch.num_rows())
    192
                    .expect("Failed to convert to array");
    193
    194
                // verify that the array's data_type is correct
    195
                assert_eq!(*result.data_type(), $TYPE);
    196
    197
                // verify that the data itself is downcastable
    198
                let result = result
    199
                    .as_any()
    200
                    .downcast_ref::<$TYPEARRAY>()
    201
                    .expect("failed to downcast");
    202
    203
                // verify that the result itself is correct
    204
                for (i, x) in $VEC.iter().enumerate() {
    205
                    match x {
    206
                        Some(x) => assert_eq!(result.value(i), *x),
    207
                        None => assert!(!result.is_valid(i)),
    208
                    }
    209
                }
    210
            }};
    211
        }
    212
    213
        // runs an end-to-end test of physical type cast
    214
        // 1. construct a record batch with a column "a" of type A
    215
        // 2. construct a physical expression of TRY_CAST(a AS B)
    216
        // 3. evaluate the expression
    217
        // 4. verify that the resulting expression is of type B
    218
        // 5. verify that the resulting values are downcastable and correct
    219
        macro_rules! generic_test_cast {
    220
            ($A_ARRAY:ident, $A_TYPE:expr, $A_VEC:expr, $TYPEARRAY:ident, $TYPE:expr, $VEC:expr) => {{
    221
                let schema = Schema::new(vec![Field::new("a", $A_TYPE, true)]);
    222
                let a_vec_len = $A_VEC.len();
    223
                let a = $A_ARRAY::from($A_VEC);
    224
                let batch =
    225
                    RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)])?;
    226
    227
                // verify that we can construct the expression
    228
                let expression = try_cast(col("a", &schema)?, &schema, $TYPE)?;
    229
    230
                // verify that its display is correct
    231
                assert_eq!(
    232
                    format!("TRY_CAST(a@0 AS {:?})", $TYPE),
    233
                    format!("{}", expression)
    234
                );
    235
    236
                // verify that the expression's type is correct
    237
                assert_eq!(expression.data_type(&schema)?, $TYPE);
    238
    239
                // compute
    240
                let result = expression
    241
                    .evaluate(&batch)?
    242
                    .into_array(batch.num_rows())
    243
                    .expect("Failed to convert to array");
    244
    245
                // verify that the array's data_type is correct
    246
                assert_eq!(*result.data_type(), $TYPE);
    247
    248
                // verify that the len is correct
    249
                assert_eq!(result.len(), a_vec_len);
    250
    251
                // verify that the data itself is downcastable
    252
                let result = result
    253
                    .as_any()
    254
                    .downcast_ref::<$TYPEARRAY>()
    255
                    .expect("failed to downcast");
    256
    257
                // verify that the result itself is correct
    258
                for (i, x) in $VEC.iter().enumerate() {
    259
                    match x {
    260
                        Some(x) => assert_eq!(result.value(i), *x),
    261
                        None => assert!(!result.is_valid(i)),
    262
                    }
    263
                }
    264
            }};
    265
        }
    266
    267
        #[test]
    268
        fn test_try_cast_decimal_to_decimal() -> Result<()> {
    269
            // try cast one decimal data type to another decimal data type
    270
            let array: Vec<i128> = vec![1234, 2222, 3, 4000, 5000];
    271
            let decimal_array = create_decimal_array(&array, 10, 3);
    272
            generic_decimal_to_other_test_cast!(
    273
                decimal_array,
    274
                DataType::Decimal128(10, 3),
    275
                Decimal128Array,
    276
                DataType::Decimal128(20, 6),
    277
                [
    278
                    Some(1_234_000),
    279
                    Some(2_222_000),
    280
                    Some(3_000),
    281
                    Some(4_000_000),
    282
                    Some(5_000_000),
    283
                    None
    284
                ]
    285
            );
    286
    287
            let decimal_array = create_decimal_array(&array, 10, 3);
    288
            generic_decimal_to_other_test_cast!(
    289
                decimal_array,
    290
                DataType::Decimal128(10, 3),
    291
                Decimal128Array,
    292
                DataType::Decimal128(10, 2),
    293
                [Some(123), Some(222), Some(0), Some(400), Some(500), None]
    294
            );
    295
    296
            Ok(())
    297
        }
    298
    299
        #[test]
    300
        fn test_try_cast_decimal_to_numeric() -> Result<()> {
    301
            // TODO we should add function to create Decimal128Array with value and metadata
    302
            // https://github.com/apache/arrow-rs/issues/1009
    303
            let array: Vec<i128> = vec![1, 2, 3, 4, 5];
    304
            let decimal_array = create_decimal_array(&array, 10, 0);
    305
            // decimal to i8
    306
            generic_decimal_to_other_test_cast!(
    307
                decimal_array,
    308
                DataType::Decimal128(10, 0),
    309
                Int8Array,
    310
                DataType::Int8,
    311
                [
    312
                    Some(1_i8),
    313
                    Some(2_i8),
    314
                    Some(3_i8),
    315
                    Some(4_i8),
    316
                    Some(5_i8),
    317
                    None
    318
                ]
    319
            );
    320
    321
            // decimal to i16
    322
            let decimal_array = create_decimal_array(&array, 10, 0);
    323
            generic_decimal_to_other_test_cast!(
    324
                decimal_array,
    325
                DataType::Decimal128(10, 0),
    326
                Int16Array,
    327
                DataType::Int16,
    328
                [
    329
                    Some(1_i16),
    330
                    Some(2_i16),
    331
                    Some(3_i16),
    332
                    Some(4_i16),
    333
                    Some(5_i16),
    334
                    None
    335
                ]
    336
            );
    337
    338
            // decimal to i32
    339
            let decimal_array = create_decimal_array(&array, 10, 0);
    340
            generic_decimal_to_other_test_cast!(
    341
                decimal_array,
    342
                DataType::Decimal128(10, 0),
    343
                Int32Array,
    344
                DataType::Int32,
    345
                [
    346
                    Some(1_i32),
    347
                    Some(2_i32),
    348
                    Some(3_i32),
    349
                    Some(4_i32),
    350
                    Some(5_i32),
    351
                    None
    352
                ]
    353
            );
    354
    355
            // decimal to i64
    356
            let decimal_array = create_decimal_array(&array, 10, 0);
    357
            generic_decimal_to_other_test_cast!(
    358
                decimal_array,
    359
                DataType::Decimal128(10, 0),
    360
                Int64Array,
    361
                DataType::Int64,
    362
                [
    363
                    Some(1_i64),
    364
                    Some(2_i64),
    365
                    Some(3_i64),
    366
                    Some(4_i64),
    367
                    Some(5_i64),
    368
                    None
    369
                ]
    370
            );
    371
    372
            // decimal to float32
    373
            let array: Vec<i128> = vec![1234, 2222, 3, 4000, 5000];
    374
            let decimal_array = create_decimal_array(&array, 10, 3);
    375
            generic_decimal_to_other_test_cast!(
    376
                decimal_array,
    377
                DataType::Decimal128(10, 3),
    378
                Float32Array,
    379
                DataType::Float32,
    380
                [
    381
                    Some(1.234_f32),
    382
                    Some(2.222_f32),
    383
                    Some(0.003_f32),
    384
                    Some(4.0_f32),
    385
                    Some(5.0_f32),
    386
                    None
    387
                ]
    388
            );
    389
            // decimal to float64
    390
            let decimal_array = create_decimal_array(&array, 20, 6);
    391
            generic_decimal_to_other_test_cast!(
    392
                decimal_array,
    393
                DataType::Decimal128(20, 6),
    394
                Float64Array,
    395
                DataType::Float64,
    396
                [
    397
                    Some(0.001234_f64),
    398
                    Some(0.002222_f64),
    399
                    Some(0.000003_f64),
    400
                    Some(0.004_f64),
    401
                    Some(0.005_f64),
    402
                    None
    403
                ]
    404
            );
    405
    406
            Ok(())
    407
        }
    408
    409
        #[test]
    410
        fn test_try_cast_numeric_to_decimal() -> Result<()> {
    411
            // int8
    412
            generic_test_cast!(
    413
                Int8Array,
    414
                DataType::Int8,
    415
                vec![1, 2, 3, 4, 5],
    416
                Decimal128Array,
    417
                DataType::Decimal128(3, 0),
    418
                [Some(1), Some(2), Some(3), Some(4), Some(5)]
    419
            );
    420
    421
            // int16
    422
            generic_test_cast!(
    423
                Int16Array,
    424
                DataType::Int16,
    425
                vec![1, 2, 3, 4, 5],
    426
                Decimal128Array,
    427
                DataType::Decimal128(5, 0),
    428
                [Some(1), Some(2), Some(3), Some(4), Some(5)]
    429
            );
    430
    431
            // int32
    432
            generic_test_cast!(
    433
                Int32Array,
    434
                DataType::Int32,
    435
                vec![1, 2, 3, 4, 5],
    436
                Decimal128Array,
    437
                DataType::Decimal128(10, 0),
    438
                [Some(1), Some(2), Some(3), Some(4), Some(5)]
    439
            );
    440
    441
            // int64
    442
            generic_test_cast!(
    443
                Int64Array,
    444
                DataType::Int64,
    445
                vec![1, 2, 3, 4, 5],
    446
                Decimal128Array,
    447
                DataType::Decimal128(20, 0),
    448
                [Some(1), Some(2), Some(3), Some(4), Some(5)]
    449
            );
    450
    451
            // int64 to different scale
    452
            generic_test_cast!(
    453
                Int64Array,
    454
                DataType::Int64,
    455
                vec![1, 2, 3, 4, 5],
    456
                Decimal128Array,
    457
                DataType::Decimal128(20, 2),
    458
                [Some(100), Some(200), Some(300), Some(400), Some(500)]
    459
            );
    460
    461
            // float32
    462
            generic_test_cast!(
    463
                Float32Array,
    464
                DataType::Float32,
    465
                vec![1.5, 2.5, 3.0, 1.123_456_8, 5.50],
    466
                Decimal128Array,
    467
                DataType::Decimal128(10, 2),
    468
                [Some(150), Some(250), Some(300), Some(112), Some(550)]
    469
            );
    470
    471
            // float64
    472
            generic_test_cast!(
    473
                Float64Array,
    474
                DataType::Float64,
    475
                vec![1.5, 2.5, 3.0, 1.123_456_8, 5.50],
    476
                Decimal128Array,
    477
                DataType::Decimal128(20, 4),
    478
                [
    479
                    Some(15000),
    480
                    Some(25000),
    481
                    Some(30000),
    482
                    Some(11235),
    483
                    Some(55000)
    484
                ]
    485
            );
    486
            Ok(())
    487
        }
    488
    489
        #[test]
    490
        fn test_cast_i32_u32() -> Result<()> {
    491
            generic_test_cast!(
    492
                Int32Array,
    493
                DataType::Int32,
    494
                vec![1, 2, 3, 4, 5],
    495
                UInt32Array,
    496
                DataType::UInt32,
    497
                [
    498
                    Some(1_u32),
    499
                    Some(2_u32),
    500
                    Some(3_u32),
    501
                    Some(4_u32),
    502
                    Some(5_u32)
    503
                ]
    504
            );
    505
            Ok(())
    506
        }
    507
    508
        #[test]
    509
        fn test_cast_i32_utf8() -> Result<()> {
    510
            generic_test_cast!(
    511
                Int32Array,
    512
                DataType::Int32,
    513
                vec![1, 2, 3, 4, 5],
    514
                StringArray,
    515
                DataType::Utf8,
    516
                [Some("1"), Some("2"), Some("3"), Some("4"), Some("5")]
    517
            );
    518
            Ok(())
    519
        }
    520
    521
        #[test]
    522
        fn test_try_cast_utf8_i32() -> Result<()> {
    523
            generic_test_cast!(
    524
                StringArray,
    525
                DataType::Utf8,
    526
                vec!["a", "2", "3", "b", "5"],
    527
                Int32Array,
    528
                DataType::Int32,
    529
                [None, Some(2), Some(3), None, Some(5)]
    530
            );
    531
            Ok(())
    532
        }
    533
    534
        #[test]
    535
        fn test_cast_i64_t64() -> Result<()> {
    536
            let original = vec![1, 2, 3, 4, 5];
    537
            let expected: Vec<Option<i64>> = original
    538
                .iter()
    539
                .map(|i| Some(Time64NanosecondArray::from(vec![*i]).value(0)))
    540
                .collect();
    541
            generic_test_cast!(
    542
                Int64Array,
    543
                DataType::Int64,
    544
                original,
    545
                TimestampNanosecondArray,
    546
                DataType::Timestamp(TimeUnit::Nanosecond, None),
    547
                expected
    548
            );
    549
            Ok(())
    550
        }
    551
    552
        #[test]
    553
        fn invalid_cast() {
    554
            // Ensure a useful error happens at plan time if invalid casts are used
    555
            let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
    556
    557
            let result = try_cast(
    558
                col("a", &schema).unwrap(),
    559
                &schema,
    560
                DataType::Interval(IntervalUnit::MonthDayNano),
    561
            );
    562
            result.expect_err("expected Invalid TRY_CAST");
    563
        }
    564
    565
        // create decimal array with the specified precision and scale
    566
        fn create_decimal_array(array: &[i128], precision: u8, scale: i8) -> Decimal128Array {
    567
            let mut decimal_builder = Decimal128Builder::with_capacity(array.len());
    568
            for value in array {
    569
                decimal_builder.append_value(*value);
    570
            }
    571
            decimal_builder.append_null();
    572
            decimal_builder
    573
                .finish()
    574
                .with_precision_and_scale(precision, scale)
    575
                .unwrap()
    576
        }
    577
    }