## 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/scalar_function.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
    //! Declaration of built-in (scalar) functions.
    19
    //! This module contains built-in functions' enumeration and metadata.
    20
    //!
    21
    //! Generally, a function has:
    22
    //! * a signature
    23
    //! * a return type, that is a function of the incoming argument's types
    24
    //! * the computation, that must accept each valid signature
    25
    //!
    26
    //! * Signature: see `Signature`
    27
    //! * Return type: a function `(arg_types) -> return_type`. E.g. for sqrt, ([f32]) -> f32, ([f64]) -> f64.
    28
    //!
    29
    //! This module also has a set of coercion rules to improve user experience: if an argument i32 is passed
    30
    //! to a function that supports f64, it is coerced to f64.
    31
    32
    use std::any::Any;
    33
    use std::fmt::{self, Debug, Formatter};
    34
    use std::hash::{Hash, Hasher};
    35
    use std::sync::Arc;
    36
    37
    use crate::expressions::Literal;
    38
    use crate::PhysicalExpr;
    39
    40
    use arrow::array::{Array, RecordBatch};
    41
    use arrow::datatypes::{DataType, FieldRef, Schema};
    42
    use datafusion_common::config::{ConfigEntry, ConfigOptions};
    43
    use datafusion_common::{internal_err, Result, ScalarValue};
    44
    use datafusion_expr::interval_arithmetic::Interval;
    45
    use datafusion_expr::sort_properties::ExprProperties;
    46
    use datafusion_expr::type_coercion::functions::data_types_with_scalar_udf;
    47
    use datafusion_expr::{
    48
        expr_vec_fmt, ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF,
    49
        Volatility,
    50
    };
    51
    52
    /// Physical expression of a scalar function
    53
    pub struct ScalarFunctionExpr {
    54
        fun: Arc<ScalarUDF>,
    55
        name: String,
    56
        args: Vec<Arc<dyn PhysicalExpr>>,
    57
        return_field: FieldRef,
    58
        config_options: Arc<ConfigOptions>,
    59
    }
    60
    61
    impl Debug for ScalarFunctionExpr {
    62
    0
        fn fmt(&self, f: &mut Formatter) -> fmt::Result {
    63
    0
            f.debug_struct("ScalarFunctionExpr")
    64
    0
                .field("fun", &"<FUNC>")
    65
    0
                .field("name", &self.name)
    66
    0
                .field("args", &self.args)
    67
    0
                .field("return_field", &self.return_field)
    68
    0
                .finish()
    69
    0
        }
    70
    }
    71
    72
    impl ScalarFunctionExpr {
    73
        /// Create a new Scalar function
    74
    0
        pub fn new(
    75
    0
            name: &str,
    76
    0
            fun: Arc<ScalarUDF>,
    77
    0
            args: Vec<Arc<dyn PhysicalExpr>>,
    78
    0
            return_field: FieldRef,
    79
    0
            config_options: Arc<ConfigOptions>,
    80
    0
        ) -> Self {
    81
    0
            Self {
    82
    0
                fun,
    83
    0
                name: name.to_owned(),
    84
    0
                args,
    85
    0
                return_field,
    86
    0
                config_options,
    87
    0
            }
    88
    0
        }
    89
    90
        /// Create a new Scalar function
    91
    0
        pub fn try_new(
    92
    0
            fun: Arc<ScalarUDF>,
    93
    0
            args: Vec<Arc<dyn PhysicalExpr>>,
    94
    0
            schema: &Schema,
    95
    0
            config_options: Arc<ConfigOptions>,
    96
    0
        ) -> Result<Self> {
    97
    0
            let name = fun.name().to_string();
    98
    0
            let arg_fields = args
    99
    0
                .iter()
    100
    0
                .map(|e| e.return_field(schema))
    101
    0
                .collect::<Result<Vec<_>>>()?;
    102
    103
            // verify that input data types is consistent with function's `TypeSignature`
    104
    0
            let arg_types = arg_fields
    105
    0
                .iter()
    106
    0
                .map(|f| f.data_type().clone())
    107
    0
                .collect::<Vec<_>>();
    108
    0
            data_types_with_scalar_udf(&arg_types, &fun)?;
    109
    110
    0
            let arguments = args
    111
    0
                .iter()
    112
    0
                .map(|e| {
    113
    0
                    e.as_any()
    114
    0
                        .downcast_ref::<Literal>()
    115
    0
                        .map(|literal| literal.value())
    116
    0
                })
    117
    0
                .collect::<Vec<_>>();
    118
    0
            let ret_args = ReturnFieldArgs {
    119
    0
                arg_fields: &arg_fields,
    120
    0
                scalar_arguments: &arguments,
    121
    0
            };
    122
    0
            let return_field = fun.return_field_from_args(ret_args)?;
    123
    0
            Ok(Self {
    124
    0
                fun,
    125
    0
                name,
    126
    0
                args,
    127
    0
                return_field,
    128
    0
                config_options,
    129
    0
            })
    130
    0
        }
    131
    132
        /// Get the scalar function implementation
    133
    0
        pub fn fun(&self) -> &ScalarUDF {
    134
    0
            &self.fun
    135
    0
        }
    136
    137
        /// The name for this expression
    138
    0
        pub fn name(&self) -> &str {
    139
    0
            &self.name
    140
    0
        }
    141
    142
        /// Input arguments
    143
    0
        pub fn args(&self) -> &[Arc<dyn PhysicalExpr>] {
    144
    0
            &self.args
    145
    0
        }
    146
    147
        /// Data type produced by this expression
    148
    0
        pub fn return_type(&self) -> &DataType {
    149
    0
            self.return_field.data_type()
    150
    0
        }
    151
    152
    0
        pub fn with_nullable(mut self, nullable: bool) -> Self {
    153
    0
            self.return_field = self
    154
    0
                .return_field
    155
    0
                .as_ref()
    156
    0
                .clone()
    157
    0
                .with_nullable(nullable)
    158
    0
                .into();
    159
    0
            self
    160
    0
        }
    161
    162
    0
        pub fn nullable(&self) -> bool {
    163
    0
            self.return_field.is_nullable()
    164
    0
        }
    165
    166
    0
        pub fn config_options(&self) -> &ConfigOptions {
    167
    0
            &self.config_options
    168
    0
        }
    169
    170
        /// Given an arbitrary PhysicalExpr attempt to downcast it to a ScalarFunctionExpr
    171
        /// and verify that its inner function is of type T.
    172
        /// If the downcast fails, or the function is not of type T, returns `None`.
    173
        /// Otherwise returns `Some(ScalarFunctionExpr)`.
    174
    0
        pub fn try_downcast_func<T>(expr: &dyn PhysicalExpr) -> Option<&ScalarFunctionExpr>
    175
    0
        where
    176
    0
            T: 'static,
    177
        {
    178
    0
            match expr.as_any().downcast_ref::<ScalarFunctionExpr>() {
    179
    0
                Some(scalar_expr)
    180
    0
                    if scalar_expr
    181
    0
                        .fun()
    182
    0
                        .inner()
    183
    0
                        .as_any()
    184
    0
                        .downcast_ref::<T>()
    185
    0
                        .is_some() =>
    186
                {
    187
    0
                    Some(scalar_expr)
    188
                }
    189
    0
                _ => None,
    190
            }
    191
    0
        }
    192
    }
    193
    194
    impl fmt::Display for ScalarFunctionExpr {
    195
    0
        fn fmt(&self, f: &mut Formatter) -> fmt::Result {
    196
    0
            write!(f, "{}({})", self.name, expr_vec_fmt!(self.args))
    197
    0
        }
    198
    }
    199
    200
    impl PartialEq for ScalarFunctionExpr {
    201
    0
        fn eq(&self, o: &Self) -> bool {
    202
    0
            if std::ptr::eq(self, o) {
    203
                // The equality implementation is somewhat expensive, so let's short-circuit when possible.
    204
    0
                return true;
    205
    0
            }
    206
            let Self {
    207
    0
                fun,
    208
    0
                name,
    209
    0
                args,
    210
    0
                return_field,
    211
    0
                config_options,
    212
    0
            } = self;
    213
    0
            fun.eq(&o.fun)
    214
    0
                && name.eq(&o.name)
    215
    0
                && args.eq(&o.args)
    216
    0
                && return_field.eq(&o.return_field)
    217
    0
                && (Arc::ptr_eq(config_options, &o.config_options)
    218
    0
                    || sorted_config_entries(config_options)
    219
    0
                        == sorted_config_entries(&o.config_options))
    220
    0
        }
    221
    }
    222
    impl Eq for ScalarFunctionExpr {}
    223
    impl Hash for ScalarFunctionExpr {
    224
    0
        fn hash<H: Hasher>(&self, state: &mut H) {
    225
            let Self {
    226
    0
                fun,
    227
    0
                name,
    228
    0
                args,
    229
    0
                return_field,
    230
                config_options: _, // expensive to hash, and often equal
    231
    0
            } = self;
    232
    0
            fun.hash(state);
    233
    0
            name.hash(state);
    234
    0
            args.hash(state);
    235
    0
            return_field.hash(state);
    236
    0
        }
    237
    }
    238
    239
    0
    fn sorted_config_entries(config_options: &ConfigOptions) -> Vec<ConfigEntry> {
    240
    0
        let mut entries = config_options.entries();
    241
    0
        entries.sort_by(|l, r| l.key.cmp(&r.key));
    242
    0
        entries
    243
    0
    }
    244
    245
    impl PhysicalExpr for ScalarFunctionExpr {
    246
        /// Return a reference to Any that can be used for downcasting
    247
    0
        fn as_any(&self) -> &dyn Any {
    248
    0
            self
    249
    0
        }
    250
    251
    0
        fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
    252
    0
            Ok(self.return_field.data_type().clone())
    253
    0
        }
    254
    255
    0
        fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
    256
    0
            Ok(self.return_field.is_nullable())
    257
    0
        }
    258
    259
    0
        fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
    260
    0
            let args = self
    261
    0
                .args
    262
    0
                .iter()
    263
    0
                .map(|e| e.evaluate(batch))
    264
    0
                .collect::<Result<Vec<_>>>()?;
    265
    266
    0
            let arg_fields = self
    267
    0
                .args
    268
    0
                .iter()
    269
    0
                .map(|e| e.return_field(batch.schema_ref()))
    270
    0
                .collect::<Result<Vec<_>>>()?;
    271
    272
    0
            let input_empty = args.is_empty();
    273
    0
            let input_all_scalar = args
    274
    0
                .iter()
    275
    0
                .all(|arg| matches!(arg, ColumnarValue::Scalar(_)));
    276
    277
            // evaluate the function
    278
    0
            let output = self.fun.invoke_with_args(ScalarFunctionArgs {
    279
    0
                args,
    280
    0
                arg_fields,
    281
    0
                number_rows: batch.num_rows(),
    282
    0
                return_field: Arc::clone(&self.return_field),
    283
    0
                config_options: Arc::clone(&self.config_options),
    284
    0
            })?;
    285
    286
    0
            if let ColumnarValue::Array(array) = &output {
    287
    0
                if array.len() != batch.num_rows() {
    288
                    // If the arguments are a non-empty slice of scalar values, we can assume that
    289
                    // returning a one-element array is equivalent to returning a scalar.
    290
    0
                    let preserve_scalar =
    291
    0
                        array.len() == 1 && !input_empty && input_all_scalar;
    292
    0
                    return if preserve_scalar {
    293
    0
                        ScalarValue::try_from_array(array, 0).map(ColumnarValue::Scalar)
    294
                    } else {
    295
    0
                        internal_err!("UDF {} returned a different number of rows than expected. Expected: {}, Got: {}",
    296
    0
                                self.name, batch.num_rows(), array.len())
    297
                    };
    298
    0
                }
    299
    0
            }
    300
    0
            Ok(output)
    301
    0
        }
    302
    303
    0
        fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
    304
    0
            Ok(Arc::clone(&self.return_field))
    305
    0
        }
    306
    307
    0
        fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
    308
    0
            self.args.iter().collect()
    309
    0
        }
    310
    311
    0
        fn with_new_children(
    312
    0
            self: Arc<Self>,
    313
    0
            children: Vec<Arc<dyn PhysicalExpr>>,
    314
    0
        ) -> Result<Arc<dyn PhysicalExpr>> {
    315
    0
            Ok(Arc::new(ScalarFunctionExpr::new(
    316
    0
                &self.name,
    317
    0
                Arc::clone(&self.fun),
    318
    0
                children,
    319
    0
                Arc::clone(&self.return_field),
    320
    0
                Arc::clone(&self.config_options),
    321
    0
            )))
    322
    0
        }
    323
    324
    0
        fn evaluate_bounds(&self, children: &[&Interval]) -> Result<Interval> {
    325
    0
            self.fun.evaluate_bounds(children)
    326
    0
        }
    327
    328
    0
        fn propagate_constraints(
    329
    0
            &self,
    330
    0
            interval: &Interval,
    331
    0
            children: &[&Interval],
    332
    0
        ) -> Result<Option<Vec<Interval>>> {
    333
    0
            self.fun.propagate_constraints(interval, children)
    334
    0
        }
    335
    336
    0
        fn get_properties(&self, children: &[ExprProperties]) -> Result<ExprProperties> {
    337
    0
            let sort_properties = self.fun.output_ordering(children)?;
    338
    0
            let preserves_lex_ordering = self.fun.preserves_lex_ordering(children)?;
    339
    0
            let children_range = children
    340
    0
                .iter()
    341
    0
                .map(|props| &props.range)
    342
    0
                .collect::<Vec<_>>();
    343
    0
            let range = self.fun().evaluate_bounds(&children_range)?;
    344
    345
    0
            Ok(ExprProperties {
    346
    0
                sort_properties,
    347
    0
                range,
    348
    0
                preserves_lex_ordering,
    349
    0
            })
    350
    0
        }
    351
    352
    0
        fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result {
    353
    0
            write!(f, "{}(", self.name)?;
    354
    0
            for (i, expr) in self.args.iter().enumerate() {
    355
    0
                if i > 0 {
    356
    0
                    write!(f, ", ")?;
    357
    0
                }
    358
    0
                expr.fmt_sql(f)?;
    359
            }
    360
    0
            write!(f, ")")
    361
    0
        }
    362
    363
    0
        fn is_volatile_node(&self) -> bool {
    364
    0
            self.fun.signature().volatility == Volatility::Volatile
    365
    0
        }
    366
    }
    367
    368
    #[cfg(test)]
    369
    mod tests {
    370
        use super::*;
    371
        use crate::expressions::Column;
    372
        use arrow::datatypes::{DataType, Field, Schema};
    373
        use datafusion_expr::{ScalarUDF, ScalarUDFImpl, Signature};
    374
        use datafusion_physical_expr_common::physical_expr::is_volatile;
    375
        use std::any::Any;
    376
    377
        /// Test helper to create a mock UDF with a specific volatility
    378
        #[derive(Debug, PartialEq, Eq, Hash)]
    379
        struct MockScalarUDF {
    380
            signature: Signature,
    381
        }
    382
    383
        impl ScalarUDFImpl for MockScalarUDF {
    384
    0
            fn as_any(&self) -> &dyn Any {
    385
    0
                self
    386
    0
            }
    387
    388
    0
            fn name(&self) -> &str {
    389
    0
                "mock_function"
    390
    0
            }
    391
    392
    0
            fn signature(&self) -> &Signature {
    393
    0
                &self.signature
    394
    0
            }
    395
    396
    0
            fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
    397
    0
                Ok(DataType::Int32)
    398
    0
            }
    399
    400
    0
            fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    401
    0
                Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(42))))
    402
    0
            }
    403
        }
    404
    405
        #[test]
    406
    0
        fn test_scalar_function_volatile_node() {
    407
            // Create a volatile UDF
    408
    0
            let volatile_udf = Arc::new(ScalarUDF::from(MockScalarUDF {
    409
    0
                signature: Signature::uniform(
    410
    0
                    1,
    411
    0
                    vec![DataType::Float32],
    412
    0
                    Volatility::Volatile,
    413
    0
                ),
    414
    0
            }));
    415
    416
            // Create a non-volatile UDF
    417
    0
            let stable_udf = Arc::new(ScalarUDF::from(MockScalarUDF {
    418
    0
                signature: Signature::uniform(1, vec![DataType::Float32], Volatility::Stable),
    419
    0
            }));
    420
    421
    0
            let schema = Schema::new(vec![Field::new("a", DataType::Float32, false)]);
    422
    0
            let args = vec![Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>];
    423
    0
            let config_options = Arc::new(ConfigOptions::new());
    424
    425
            // Test volatile function
    426
    0
            let volatile_expr = ScalarFunctionExpr::try_new(
    427
    0
                volatile_udf,
    428
    0
                args.clone(),
    429
    0
                &schema,
    430
    0
                Arc::clone(&config_options),
    431
            )
    432
    0
            .unwrap();
    433
    434
    0
            assert!(volatile_expr.is_volatile_node());
    435
    0
            let volatile_arc: Arc<dyn PhysicalExpr> = Arc::new(volatile_expr);
    436
    0
            assert!(is_volatile(&volatile_arc));
    437
    438
            // Test non-volatile function
    439
    0
            let stable_expr =
    440
    0
                ScalarFunctionExpr::try_new(stable_udf, args, &schema, config_options)
    441
    0
                    .unwrap();
    442
    443
    0
            assert!(!stable_expr.is_volatile_node());
    444
    0
            let stable_arc: Arc<dyn PhysicalExpr> = Arc::new(stable_expr);
    445
    0
            assert!(!is_volatile(&stable_arc));
    446
    0
        }
    447
    }