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
}