## 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/core/arrow_cast.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
    //! [`ArrowCastFunc`]: Implementation of the `arrow_cast`
    19
    20
    use arrow::datatypes::{DataType, Field, FieldRef};
    21
    use arrow::error::ArrowError;
    22
    use datafusion_common::{
    23
        arrow_datafusion_err, exec_err, internal_err, Result, ScalarValue,
    24
    };
    25
    use datafusion_common::{
    26
        exec_datafusion_err, utils::take_function_args, DataFusionError,
    27
    };
    28
    use std::any::Any;
    29
    30
    use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
    31
    use datafusion_expr::{
    32
        ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs,
    33
        ScalarUDFImpl, Signature, Volatility,
    34
    };
    35
    use datafusion_macros::user_doc;
    36
    37
    /// Implements casting to arbitrary arrow types (rather than SQL types)
    38
    ///
    39
    /// Note that the `arrow_cast` function is somewhat special in that its
    40
    /// return depends only on the *value* of its second argument (not its type)
    41
    ///
    42
    /// It is implemented by calling the same underlying arrow `cast` kernel as
    43
    /// normal SQL casts.
    44
    ///
    45
    /// For example to cast to `int` using SQL  (which is then mapped to the arrow
    46
    /// type `Int32`)
    47
    ///
    48
    /// ```sql
    49
    /// select cast(column_x as int) ...
    50
    /// ```
    51
    ///
    52
    /// Use the `arrow_cast` function to cast to a specific arrow type
    53
    ///
    54
    /// For example
    55
    /// ```sql
    56
    /// select arrow_cast(column_x, 'Float64')
    57
    /// ```
    58
    #[user_doc(
    59
        doc_section(label = "Other Functions"),
    60
        description = "Casts a value to a specific Arrow data type.",
    61
        syntax_example = "arrow_cast(expression, datatype)",
    62
        sql_example = r#"```sql
    63
    > select
    64
      arrow_cast(-5,    'Int8') as a,
    65
      arrow_cast('foo', 'Dictionary(Int32, Utf8)') as b,
    66
      arrow_cast('bar', 'LargeUtf8') as c;
    67
    68
    +----+-----+-----+
    69
    | a  | b   | c   |
    70
    +----+-----+-----+
    71
    | -5 | foo | bar |
    72
    +----+-----+-----+
    73
    74
    > select
    75
      arrow_cast('2023-01-02T12:53:02', 'Timestamp(µs, "+08:00")') as d,
    76
      arrow_cast('2023-01-02T12:53:02', 'Timestamp(µs)') as e;
    77
    78
    +---------------------------+---------------------+
    79
    | d                         | e                   |
    80
    +---------------------------+---------------------+
    81
    | 2023-01-02T12:53:02+08:00 | 2023-01-02T12:53:02 |
    82
    +---------------------------+---------------------+
    83
    ```"#,
    84
        argument(
    85
            name = "expression",
    86
            description = "Expression to cast. The expression can be a constant, column, or function, and any combination of operators."
    87
        ),
    88
        argument(
    89
            name = "datatype",
    90
            description = "[Arrow data type](https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html) name to cast to, as a string. The format is the same as that returned by [`arrow_typeof`]"
    91
        )
    92
    )]
    93
    #[derive(Debug, PartialEq, Eq, Hash)]
    94
    pub struct ArrowCastFunc {
    95
        signature: Signature,
    96
    }
    97
    98
    impl Default for ArrowCastFunc {
    99
    0
        fn default() -> Self {
    100
    0
            Self::new()
    101
    0
        }
    102
    }
    103
    104
    impl ArrowCastFunc {
    105
    0
        pub fn new() -> Self {
    106
    0
            Self {
    107
    0
                signature: Signature::any(2, Volatility::Immutable),
    108
    0
            }
    109
    0
        }
    110
    }
    111
    112
    impl ScalarUDFImpl for ArrowCastFunc {
    113
    0
        fn as_any(&self) -> &dyn Any {
    114
    0
            self
    115
    0
        }
    116
    117
    0
        fn name(&self) -> &str {
    118
    0
            "arrow_cast"
    119
    0
        }
    120
    121
    0
        fn signature(&self) -> &Signature {
    122
    0
            &self.signature
    123
    0
        }
    124
    125
    0
        fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
    126
    0
            internal_err!("return_field_from_args should be called instead")
    127
    0
        }
    128
    129
    0
        fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
    130
    0
            let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
    131
    132
    0
            let [_, type_arg] = take_function_args(self.name(), args.scalar_arguments)?;
    133
    134
    0
            type_arg
    135
    0
                .and_then(|sv| sv.try_as_str().flatten().filter(|s| !s.is_empty()))
    136
    0
                .map_or_else(
    137
    0
                    || {
    138
    0
                        exec_err!(
    139
    0
                            "{} requires its second argument to be a non-empty constant string",
    140
    0
                            self.name()
    141
                        )
    142
    0
                    },
    143
    0
                    |casted_type| match casted_type.parse::<DataType>() {
    144
    0
                        Ok(data_type) => Ok(Field::new(self.name(), data_type, nullable).into()),
    145
    0
                        Err(ArrowError::ParseError(e)) => Err(exec_datafusion_err!("{e}")),
    146
    0
                        Err(e) => Err(arrow_datafusion_err!(e)),
    147
    0
                    },
    148
                )
    149
    0
        }
    150
    151
    0
        fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    152
    0
            internal_err!("arrow_cast should have been simplified to cast")
    153
    0
        }
    154
    155
    0
        fn simplify(
    156
    0
            &self,
    157
    0
            mut args: Vec<Expr>,
    158
    0
            info: &dyn SimplifyInfo,
    159
    0
        ) -> Result<ExprSimplifyResult> {
    160
            // convert this into a real cast
    161
    0
            let target_type = data_type_from_args(&args)?;
    162
            // remove second (type) argument
    163
    0
            args.pop().unwrap();
    164
    0
            let arg = args.pop().unwrap();
    165
    166
    0
            let source_type = info.get_data_type(&arg)?;
    167
    0
            let new_expr = if source_type == target_type {
    168
                // the argument's data type is already the correct type
    169
    0
                arg
    170
            } else {
    171
                // Use an actual cast to get the correct type
    172
    0
                Expr::Cast(datafusion_expr::Cast {
    173
    0
                    expr: Box::new(arg),
    174
    0
                    data_type: target_type,
    175
    0
                })
    176
            };
    177
            // return the newly written argument to DataFusion
    178
    0
            Ok(ExprSimplifyResult::Simplified(new_expr))
    179
    0
        }
    180
    181
    0
        fn documentation(&self) -> Option<&Documentation> {
    182
    0
            self.doc()
    183
    0
        }
    184
    }
    185
    186
    /// Returns the requested type from the arguments
    187
    0
    fn data_type_from_args(args: &[Expr]) -> Result<DataType> {
    188
    0
        let [_, type_arg] = take_function_args("arrow_cast", args)?;
    189
    190
    0
        let Expr::Literal(ScalarValue::Utf8(Some(val)), _) = type_arg else {
    191
    0
            return exec_err!(
    192
    0
                "arrow_cast requires its second argument to be a constant string, got {:?}",
    193
                type_arg
    194
            );
    195
        };
    196
    197
    0
        val.parse().map_err(|e| match e {
    198
            // If the data type cannot be parsed, return a Plan error to signal an
    199
            // error in the input rather than a more general ArrowError
    200
    0
            ArrowError::ParseError(e) => exec_datafusion_err!("{e}"),
    201
    0
            e => arrow_datafusion_err!(e),
    202
    0
        })
    203
    0
    }