## 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/functions/src/utils.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 arrow::array::{Array, ArrayRef, ArrowPrimitiveType, AsArray, PrimitiveArray};
    19
    use arrow::compute::try_binary;
    20
    use arrow::datatypes::{DataType, DecimalType};
    21
    use arrow::error::ArrowError;
    22
    use datafusion_common::{not_impl_err, DataFusionError, Result, ScalarValue};
    23
    use datafusion_expr::function::Hint;
    24
    use datafusion_expr::ColumnarValue;
    25
    use std::sync::Arc;
    26
    27
    /// Creates a function to identify the optimal return type of a string function given
    28
    /// the type of its first argument.
    29
    ///
    30
    /// If the input type is `LargeUtf8` or `LargeBinary` the return type is
    31
    /// `$largeUtf8Type`,
    32
    ///
    33
    /// If the input type is `Utf8` or `Binary` the return type is `$utf8Type`,
    34
    ///
    35
    /// If the input type is `Utf8View` the return type is $utf8Type,
    36
    macro_rules! get_optimal_return_type {
    37
        ($FUNC:ident, $largeUtf8Type:expr, $utf8Type:expr) => {
    38
    0
            pub(crate) fn $FUNC(arg_type: &DataType, name: &str) -> Result<DataType> {
    39
    0
                Ok(match arg_type {
    40
                    // LargeBinary inputs are automatically coerced to Utf8
    41
                    DataType::LargeUtf8 | DataType::LargeBinary => $largeUtf8Type,
    42
                    // Binary inputs are automatically coerced to Utf8
    43
                    DataType::Utf8 | DataType::Binary => $utf8Type,
    44
                    // Utf8View max offset size is u32::MAX, the same as UTF8
    45
                    DataType::Utf8View | DataType::BinaryView => $utf8Type,
    46
    0
                    DataType::Null => DataType::Null,
    47
    0
                    DataType::Dictionary(_, value_type) => match **value_type {
    48
                        DataType::LargeUtf8 | DataType::LargeBinary => $largeUtf8Type,
    49
                        DataType::Utf8 | DataType::Binary => $utf8Type,
    50
    0
                        DataType::Null => DataType::Null,
    51
                        _ => {
    52
    0
                            return datafusion_common::exec_err!(
    53
    0
                                "The {} function can only accept strings, but got {:?}.",
    54
    0
                                name.to_uppercase(),
    55
                                **value_type
    56
                            );
    57
                        }
    58
                    },
    59
    0
                    data_type => {
    60
    0
                        return datafusion_common::exec_err!(
    61
    0
                            "The {} function can only accept strings, but got {:?}.",
    62
    0
                            name.to_uppercase(),
    63
                            data_type
    64
                        );
    65
                    }
    66
                })
    67
    0
            }
    68
        };
    69
    }
    70
    71
    // `utf8_to_str_type`: returns either a Utf8 or LargeUtf8 based on the input type size.
    72
    get_optimal_return_type!(utf8_to_str_type, DataType::LargeUtf8, DataType::Utf8);
    73
    74
    // `utf8_to_int_type`: returns either a Int32 or Int64 based on the input type size.
    75
    get_optimal_return_type!(utf8_to_int_type, DataType::Int64, DataType::Int32);
    76
    77
    /// Creates a scalar function implementation for the given function.
    78
    /// * `inner` - the function to be executed
    79
    /// * `hints` - hints to be used when expanding scalars to arrays
    80
    0
    pub fn make_scalar_function<F>(
    81
    0
        inner: F,
    82
    0
        hints: Vec<Hint>,
    83
    0
    ) -> impl Fn(&[ColumnarValue]) -> Result<ColumnarValue>
    84
    0
    where
    85
    0
        F: Fn(&[ArrayRef]) -> Result<ArrayRef>,
    86
    {
    87
    0
        move |args: &[ColumnarValue]| {
    88
            // first, identify if any of the arguments is an Array. If yes, store its `len`,
    89
            // as any scalar will need to be converted to an array of len `len`.
    90
    0
            let len = args
    91
    0
                .iter()
    92
    0
                .fold(Option::<usize>::None, |acc, arg| match arg {
    93
    0
                    ColumnarValue::Scalar(_) => acc,
    94
    0
                    ColumnarValue::Array(a) => Some(a.len()),
    95
    0
                });
    96
    97
    0
            let is_scalar = len.is_none();
    98
    99
    0
            let inferred_length = len.unwrap_or(1);
    100
    0
            let args = args
    101
    0
                .iter()
    102
    0
                .zip(hints.iter().chain(std::iter::repeat(&Hint::Pad)))
    103
    0
                .map(|(arg, hint)| {
    104
                    // Decide on the length to expand this scalar to depending
    105
                    // on the given hints.
    106
    0
                    let expansion_len = match hint {
    107
    0
                        Hint::AcceptsSingular => 1,
    108
    0
                        Hint::Pad => inferred_length,
    109
                    };
    110
    0
                    arg.to_array(expansion_len)
    111
    0
                })
    112
    0
                .collect::<Result<Vec<_>>>()?;
    113
    114
    0
            let result = (inner)(&args);
    115
    0
            if is_scalar {
    116
                // If all inputs are scalar, keeps output as scalar
    117
    0
                let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
    118
    0
                result.map(ColumnarValue::Scalar)
    119
            } else {
    120
    0
                result.map(ColumnarValue::Array)
    121
            }
    122
    0
        }
    123
    0
    }
    124
    125
    /// Computes a binary math function for input arrays using a specified function.
    126
    /// Generic types:
    127
    /// - `L`: Left array primitive type
    128
    /// - `R`: Right array primitive type
    129
    /// - `O`: Output array primitive type
    130
    /// - `F`: Functor computing `fun(l: L, r: R) -> Result<OutputType>`
    131
    0
    pub fn calculate_binary_math<L, R, O, F>(
    132
    0
        left: &dyn Array,
    133
    0
        right: &ColumnarValue,
    134
    0
        fun: F,
    135
    0
    ) -> Result<Arc<PrimitiveArray<O>>>
    136
    0
    where
    137
    0
        L: ArrowPrimitiveType,
    138
    0
        R: ArrowPrimitiveType,
    139
    0
        O: ArrowPrimitiveType,
    140
    0
        F: Fn(L::Native, R::Native) -> Result<O::Native, ArrowError>,
    141
    0
        R::Native: TryFrom<ScalarValue>,
    142
    {
    143
    0
        let left = left.as_primitive::<L>();
    144
    0
        let right = right.cast_to(&R::DATA_TYPE, None)?;
    145
    0
        let result = match right {
    146
    0
            ColumnarValue::Scalar(scalar) => {
    147
    0
                if scalar.is_null() {
    148
                    // Null scalar is castable to any numeric, creating a non-null expression.
    149
                    // Provide null array explicitly to make result null
    150
    0
                    PrimitiveArray::<O>::new_null(1)
    151
                } else {
    152
    0
                    let right = R::Native::try_from(scalar.clone()).map_err(|_| {
    153
    0
                        DataFusionError::NotImplemented(format!(
    154
    0
                            "Cannot convert scalar value {} to {}",
    155
    0
                            &scalar,
    156
    0
                            R::DATA_TYPE
    157
    0
                        ))
    158
    0
                    })?;
    159
    0
                    left.try_unary::<_, O, _>(|lvalue| fun(lvalue, right))?
    160
                }
    161
            }
    162
    0
            ColumnarValue::Array(right) => {
    163
    0
                let right = right.as_primitive::<R>();
    164
    0
                try_binary::<_, _, _, O>(left, right, &fun)?
    165
            }
    166
        };
    167
    0
        Ok(Arc::new(result) as _)
    168
    0
    }
    169
    170
    /// Computes a binary math function for input arrays using a specified function
    171
    /// and apply rescaling to given precision and scale.
    172
    /// Generic types:
    173
    /// - `L`: Left array decimal type
    174
    /// - `R`: Right array primitive type
    175
    /// - `O`: Output array decimal type
    176
    /// - `F`: Functor computing `fun(l: L, r: R) -> Result<OutputType>`
    177
    0
    pub fn calculate_binary_decimal_math<L, R, O, F>(
    178
    0
        left: &dyn Array,
    179
    0
        right: &ColumnarValue,
    180
    0
        fun: F,
    181
    0
        precision: u8,
    182
    0
        scale: i8,
    183
    0
    ) -> Result<Arc<PrimitiveArray<O>>>
    184
    0
    where
    185
    0
        L: DecimalType,
    186
    0
        R: ArrowPrimitiveType,
    187
    0
        O: DecimalType,
    188
    0
        F: Fn(L::Native, R::Native) -> Result<O::Native, ArrowError>,
    189
    0
        R::Native: TryFrom<ScalarValue>,
    190
    {
    191
    0
        let result_array = calculate_binary_math::<L, R, O, F>(left, right, fun)?;
    192
    0
        if scale < 0 {
    193
    0
            not_impl_err!("Negative scale is not supported for power for decimal types")
    194
        } else {
    195
    0
            Ok(Arc::new(
    196
    0
                result_array
    197
    0
                    .as_ref()
    198
    0
                    .clone()
    199
    0
                    .with_precision_and_scale(precision, scale)?,
    200
            ))
    201
        }
    202
    0
    }
    203
    204
    /// Converts Decimal128 components (value and scale) to an unscaled i128
    205
    0
    pub fn decimal128_to_i128(value: i128, scale: i8) -> Result<i128, ArrowError> {
    206
    0
        if scale < 0 {
    207
    0
            Err(ArrowError::ComputeError(
    208
    0
                "Negative scale is not supported".into(),
    209
    0
            ))
    210
    0
        } else if scale == 0 {
    211
    0
            Ok(value)
    212
        } else {
    213
    0
            match i128::from(10).checked_pow(scale as u32) {
    214
    0
                Some(divisor) => Ok(value / divisor),
    215
    0
                None => Err(ArrowError::ComputeError(format!(
    216
    0
                    "Cannot get a power of {scale}"
    217
    0
                ))),
    218
            }
    219
        }
    220
    0
    }
    221
    222
    #[cfg(test)]
    223
    pub mod test {
    224
        /// $FUNC ScalarUDFImpl to test
    225
        /// $ARGS arguments (vec) to pass to function
    226
        /// $EXPECTED a Result<ColumnarValue>
    227
        /// $EXPECTED_TYPE is the expected value type
    228
        /// $EXPECTED_DATA_TYPE is the expected result type
    229
        /// $ARRAY_TYPE is the column type after function applied
    230
        /// $CONFIG_OPTIONS config options to pass to function
    231
        macro_rules! test_function {
    232
        ($FUNC:expr, $ARGS:expr, $EXPECTED:expr, $EXPECTED_TYPE:ty, $EXPECTED_DATA_TYPE:expr, $ARRAY_TYPE:ident, $CONFIG_OPTIONS:expr) => {
    233
            let expected: Result<Option<$EXPECTED_TYPE>> = $EXPECTED;
    234
            let func = $FUNC;
    235
    236
            let data_array = $ARGS.iter().map(|arg| arg.data_type()).collect::<Vec<_>>();
    237
            let cardinality = $ARGS
    238
                .iter()
    239
                .fold(Option::<usize>::None, |acc, arg| match arg {
    240
                    ColumnarValue::Scalar(_) => acc,
    241
                    ColumnarValue::Array(a) => Some(a.len()),
    242
                })
    243
                .unwrap_or(1);
    244
    245
                let scalar_arguments = $ARGS.iter().map(|arg| match arg {
    246
                    ColumnarValue::Scalar(scalar) => Some(scalar.clone()),
    247
                    ColumnarValue::Array(_) => None,
    248
                }).collect::<Vec<_>>();
    249
                let scalar_arguments_refs = scalar_arguments.iter().map(|arg| arg.as_ref()).collect::<Vec<_>>();
    250
    251
                let nullables = $ARGS.iter().map(|arg| match arg {
    252
                    ColumnarValue::Scalar(scalar) => scalar.is_null(),
    253
                    ColumnarValue::Array(a) => a.null_count() > 0,
    254
                }).collect::<Vec<_>>();
    255
    256
                let field_array = data_array.into_iter().zip(nullables).enumerate()
    257
                    .map(|(idx, (data_type, nullable))| arrow::datatypes::Field::new(format!("field_{idx}"), data_type, nullable))
    258
                .map(std::sync::Arc::new)
    259
                .collect::<Vec<_>>();
    260
    261
            let return_field = func.return_field_from_args(datafusion_expr::ReturnFieldArgs {
    262
                arg_fields: &field_array,
    263
                scalar_arguments: &scalar_arguments_refs,
    264
            });
    265
                let arg_fields = $ARGS.iter()
    266
                .enumerate()
    267
                    .map(|(idx, arg)| arrow::datatypes::Field::new(format!("f_{idx}"), arg.data_type(), true).into())
    268
                .collect::<Vec<_>>();
    269
    270
            match expected {
    271
                Ok(expected) => {
    272
                    assert_eq!(return_field.is_ok(), true);
    273
                    let return_field = return_field.unwrap();
    274
                    let return_type = return_field.data_type();
    275
                    assert_eq!(return_type, &$EXPECTED_DATA_TYPE);
    276
    277
                        let result = func.invoke_with_args(datafusion_expr::ScalarFunctionArgs{
    278
                        args: $ARGS,
    279
                        arg_fields,
    280
                        number_rows: cardinality,
    281
                        return_field,
    282
                            config_options: $CONFIG_OPTIONS
    283
                    });
    284
                        assert_eq!(result.is_ok(), true, "function returned an error: {}", result.unwrap_err());
    285
    286
                        let result = result.unwrap().to_array(cardinality).expect("Failed to convert to array");
    287
                        let result = result.as_any().downcast_ref::<$ARRAY_TYPE>().expect("Failed to convert to type");
    288
                    assert_eq!(result.data_type(), &$EXPECTED_DATA_TYPE);
    289
    290
                    // value is correct
    291
                    match expected {
    292
                        Some(v) => assert_eq!(result.value(0), v),
    293
                        None => assert!(result.is_null(0)),
    294
                    };
    295
                }
    296
                Err(expected_error) => {
    297
                    if let Ok(return_field) = return_field {
    298
                        // invoke is expected error - cannot use .expect_err() due to Debug not being implemented
    299
                        match func.invoke_with_args(datafusion_expr::ScalarFunctionArgs {
    300
                            args: $ARGS,
    301
                            arg_fields,
    302
                            number_rows: cardinality,
    303
                            return_field,
    304
                            config_options: $CONFIG_OPTIONS,
    305
                        }) {
    306
                            Ok(_) => assert!(false, "expected error"),
    307
                            Err(error) => {
    308
                                assert!(expected_error
    309
                                    .strip_backtrace()
    310
                                    .starts_with(&error.strip_backtrace()));
    311
                            }
    312
                        }
    313
                    } else if let Err(error) = return_field {
    314
                        datafusion_common::assert_contains!(
    315
                            expected_error.strip_backtrace(),
    316
                            error.strip_backtrace()
    317
                        );
    318
                    }
    319
                }
    320
            };
    321
        };
    322
    323
            ($FUNC:expr, $ARGS:expr, $EXPECTED:expr, $EXPECTED_TYPE:ty, $EXPECTED_DATA_TYPE:expr, $ARRAY_TYPE:ident) => {
    324
                test_function!(
    325
                    $FUNC,
    326
                    $ARGS,
    327
                    $EXPECTED,
    328
                    $EXPECTED_TYPE,
    329
                    $EXPECTED_DATA_TYPE,
    330
                    $ARRAY_TYPE,
    331
                    std::sync::Arc::new(datafusion_common::config::ConfigOptions::default())
    332
                )
    333
            };
    334
        }
    335
    336
        use arrow::datatypes::DataType;
    337
        #[allow(unused_imports)]
    338
        pub(crate) use test_function;
    339
    340
        use super::*;
    341
    342
        #[test]
    343
        fn string_to_int_type() {
    344
            let v = utf8_to_int_type(&DataType::Utf8, "test").unwrap();
    345
            assert_eq!(v, DataType::Int32);
    346
    347
            let v = utf8_to_int_type(&DataType::Utf8View, "test").unwrap();
    348
            assert_eq!(v, DataType::Int32);
    349
    350
            let v = utf8_to_int_type(&DataType::LargeUtf8, "test").unwrap();
    351
            assert_eq!(v, DataType::Int64);
    352
        }
    353
    354
        #[test]
    355
        fn test_decimal128_to_i128() {
    356
            let cases = [
    357
                (123, 0, Some(123)),
    358
                (1230, 1, Some(123)),
    359
                (123000, 3, Some(123)),
    360
                (1, 0, Some(1)),
    361
                (123, -3, None),
    362
                (123, i8::MAX, None),
    363
                (i128::MAX, 0, Some(i128::MAX)),
    364
                (i128::MAX, 3, Some(i128::MAX / 1000)),
    365
            ];
    366
    367
            for (value, scale, expected) in cases {
    368
                match decimal128_to_i128(value, scale) {
    369
                    Ok(actual) => {
    370
                        assert_eq!(
    371
                            actual,
    372
                            expected.expect("Got value but expected none"),
    373
                            "{value} and {scale} vs {expected:?}"
    374
                        );
    375
                    }
    376
                    Err(_) => assert!(expected.is_none()),
    377
                }
    378
            }
    379
        }
    380
    }