## Coverage Report

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

    /Users/andrewlamb/Software/datafusion/datafusion/expr/src/registry.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
    //! FunctionRegistry trait
    19
    20
    use crate::expr_rewriter::FunctionRewrite;
    21
    use crate::planner::ExprPlanner;
    22
    use crate::{AggregateUDF, ScalarUDF, UserDefinedLogicalNode, WindowUDF};
    23
    use datafusion_common::{not_impl_err, plan_datafusion_err, Result};
    24
    use std::collections::{HashMap, HashSet};
    25
    use std::fmt::Debug;
    26
    use std::sync::Arc;
    27
    28
    /// A registry knows how to build logical expressions out of user-defined function' names
    29
    pub trait FunctionRegistry {
    30
        /// Set of all available udfs.
    31
        fn udfs(&self) -> HashSet<String>;
    32
    33
        /// Returns a reference to the user defined scalar function (udf) named
    34
        /// `name`.
    35
        fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>>;
    36
    37
        /// Returns a reference to the user defined aggregate function (udaf) named
    38
        /// `name`.
    39
        fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>>;
    40
    41
        /// Returns a reference to the user defined window function (udwf) named
    42
        /// `name`.
    43
        fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>>;
    44
    45
        /// Registers a new [`ScalarUDF`], returning any previously registered
    46
        /// implementation.
    47
        ///
    48
        /// Returns an error (the default) if the function can not be registered,
    49
        /// for example if the registry is read only.
    50
    0
        fn register_udf(&mut self, _udf: Arc<ScalarUDF>) -> Result<Option<Arc<ScalarUDF>>> {
    51
    0
            not_impl_err!("Registering ScalarUDF")
    52
    0
        }
    53
        /// Registers a new [`AggregateUDF`], returning any previously registered
    54
        /// implementation.
    55
        ///
    56
        /// Returns an error (the default) if the function can not be registered,
    57
        /// for example if the registry is read only.
    58
    0
        fn register_udaf(
    59
    0
            &mut self,
    60
    0
            _udaf: Arc<AggregateUDF>,
    61
    0
        ) -> Result<Option<Arc<AggregateUDF>>> {
    62
    0
            not_impl_err!("Registering AggregateUDF")
    63
    0
        }
    64
        /// Registers a new [`WindowUDF`], returning any previously registered
    65
        /// implementation.
    66
        ///
    67
        /// Returns an error (the default) if the function can not be registered,
    68
        /// for example if the registry is read only.
    69
    0
        fn register_udwf(&mut self, _udaf: Arc<WindowUDF>) -> Result<Option<Arc<WindowUDF>>> {
    70
    0
            not_impl_err!("Registering WindowUDF")
    71
    0
        }
    72
    73
        /// Deregisters a [`ScalarUDF`], returning the implementation that was
    74
        /// deregistered.
    75
        ///
    76
        /// Returns an error (the default) if the function can not be deregistered,
    77
        /// for example if the registry is read only.
    78
    0
        fn deregister_udf(&mut self, _name: &str) -> Result<Option<Arc<ScalarUDF>>> {
    79
    0
            not_impl_err!("Deregistering ScalarUDF")
    80
    0
        }
    81
    82
        /// Deregisters a [`AggregateUDF`], returning the implementation that was
    83
        /// deregistered.
    84
        ///
    85
        /// Returns an error (the default) if the function can not be deregistered,
    86
        /// for example if the registry is read only.
    87
    0
        fn deregister_udaf(&mut self, _name: &str) -> Result<Option<Arc<AggregateUDF>>> {
    88
    0
            not_impl_err!("Deregistering AggregateUDF")
    89
    0
        }
    90
    91
        /// Deregisters a [`WindowUDF`], returning the implementation that was
    92
        /// deregistered.
    93
        ///
    94
        /// Returns an error (the default) if the function can not be deregistered,
    95
        /// for example if the registry is read only.
    96
    0
        fn deregister_udwf(&mut self, _name: &str) -> Result<Option<Arc<WindowUDF>>> {
    97
    0
            not_impl_err!("Deregistering WindowUDF")
    98
    0
        }
    99
    100
        /// Registers a new [`FunctionRewrite`] with the registry.
    101
        ///
    102
        /// `FunctionRewrite` rules are used to rewrite certain / operators in the
    103
        /// logical plan to function calls.  For example `a || b` might be written to
    104
        /// `array_concat(a, b)`.
    105
        ///
    106
        /// This allows the behavior of operators to be customized by the user.
    107
    0
        fn register_function_rewrite(
    108
    0
            &mut self,
    109
    0
            _rewrite: Arc<dyn FunctionRewrite + Send + Sync>,
    110
    0
        ) -> Result<()> {
    111
    0
            not_impl_err!("Registering FunctionRewrite")
    112
    0
        }
    113
    114
        /// Set of all registered [`ExprPlanner`]s
    115
        fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>>;
    116
    117
        /// Registers a new [`ExprPlanner`] with the registry.
    118
    0
        fn register_expr_planner(
    119
    0
            &mut self,
    120
    0
            _expr_planner: Arc<dyn ExprPlanner>,
    121
    0
        ) -> Result<()> {
    122
    0
            not_impl_err!("Registering ExprPlanner")
    123
    0
        }
    124
    }
    125
    126
    /// Serializer and deserializer registry for extensions like [UserDefinedLogicalNode].
    127
    pub trait SerializerRegistry: Debug + Send + Sync {
    128
        /// Serialize this node to a byte array. This serialization should not include
    129
        /// input plans.
    130
        fn serialize_logical_plan(
    131
            &self,
    132
            node: &dyn UserDefinedLogicalNode,
    133
        ) -> Result<Vec<u8>>;
    134
    135
        /// Deserialize user defined logical plan node ([UserDefinedLogicalNode]) from
    136
        /// bytes.
    137
        fn deserialize_logical_plan(
    138
            &self,
    139
            name: &str,
    140
            bytes: &[u8],
    141
        ) -> Result<Arc<dyn UserDefinedLogicalNode>>;
    142
    }
    143
    144
    /// A  [`FunctionRegistry`] that uses in memory [`HashMap`]s
    145
    #[derive(Default, Debug)]
    146
    pub struct MemoryFunctionRegistry {
    147
        /// Scalar Functions
    148
        udfs: HashMap<String, Arc<ScalarUDF>>,
    149
        /// Aggregate Functions
    150
        udafs: HashMap<String, Arc<AggregateUDF>>,
    151
        /// Window Functions
    152
        udwfs: HashMap<String, Arc<WindowUDF>>,
    153
    }
    154
    155
    impl MemoryFunctionRegistry {
    156
        pub fn new() -> Self {
    157
            Self::default()
    158
        }
    159
    }
    160
    161
    impl FunctionRegistry for MemoryFunctionRegistry {
    162
        fn udfs(&self) -> HashSet<String> {
    163
            self.udfs.keys().cloned().collect()
    164
        }
    165
    166
        fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>> {
    167
            self.udfs
    168
                .get(name)
    169
                .cloned()
    170
    0
                .ok_or_else(|| plan_datafusion_err!("Function {name} not found"))
    171
        }
    172
    173
        fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {
    174
            self.udafs
    175
                .get(name)
    176
                .cloned()
    177
    0
                .ok_or_else(|| plan_datafusion_err!("Aggregate Function {name} not found"))
    178
        }
    179
    180
        fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>> {
    181
            self.udwfs
    182
                .get(name)
    183
                .cloned()
    184
    0
                .ok_or_else(|| plan_datafusion_err!("Window Function {name} not found"))
    185
        }
    186
    187
        fn register_udf(&mut self, udf: Arc<ScalarUDF>) -> Result<Option<Arc<ScalarUDF>>> {
    188
            Ok(self.udfs.insert(udf.name().to_string(), udf))
    189
        }
    190
        fn register_udaf(
    191
            &mut self,
    192
            udaf: Arc<AggregateUDF>,
    193
        ) -> Result<Option<Arc<AggregateUDF>>> {
    194
            Ok(self.udafs.insert(udaf.name().into(), udaf))
    195
        }
    196
        fn register_udwf(&mut self, udaf: Arc<WindowUDF>) -> Result<Option<Arc<WindowUDF>>> {
    197
            Ok(self.udwfs.insert(udaf.name().into(), udaf))
    198
        }
    199
    200
        fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
    201
            vec![]
    202
        }
    203
    }