common_function/
function_registry.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! functions registry
16use std::collections::HashMap;
17use std::sync::{Arc, LazyLock, RwLock};
18
19use datafusion::catalog::TableFunction;
20use datafusion_expr::AggregateUDF;
21
22use crate::admin::AdminFunction;
23use crate::aggrs::aggr_wrapper::StateMergeHelper;
24use crate::aggrs::approximate::ApproximateFunction;
25use crate::aggrs::count_hash::CountHash;
26use crate::aggrs::vector::VectorFunction as VectorAggrFunction;
27use crate::function::{Function, FunctionRef};
28use crate::function_factory::ScalarFunctionFactory;
29use crate::scalars::date::DateFunction;
30use crate::scalars::expression::ExpressionFunction;
31use crate::scalars::hll_count::HllCalcFunction;
32use crate::scalars::ip::IpFunctions;
33use crate::scalars::json::JsonFunction;
34use crate::scalars::matches::MatchesFunction;
35use crate::scalars::matches_term::MatchesTermFunction;
36use crate::scalars::math::MathFunction;
37use crate::scalars::string::register_string_functions;
38use crate::scalars::timestamp::TimestampFunction;
39use crate::scalars::uddsketch_calc::UddSketchCalcFunction;
40use crate::scalars::vector::VectorFunction as VectorScalarFunction;
41use crate::system::SystemFunction;
42
43#[derive(Default)]
44pub struct FunctionRegistry {
45    functions: RwLock<HashMap<String, ScalarFunctionFactory>>,
46    aggregate_functions: RwLock<HashMap<String, AggregateUDF>>,
47    table_functions: RwLock<HashMap<String, Arc<TableFunction>>>,
48}
49
50impl FunctionRegistry {
51    /// Register a function in the registry by converting it into a `ScalarFunctionFactory`.
52    ///
53    /// # Arguments
54    ///
55    /// * `func` - An object that can be converted into a `ScalarFunctionFactory`.
56    ///
57    /// The function is inserted into the internal function map, keyed by its name.
58    /// If a function with the same name already exists, it will be replaced.
59    pub fn register(&self, func: impl Into<ScalarFunctionFactory>) {
60        let func = func.into();
61        let _ = self
62            .functions
63            .write()
64            .unwrap()
65            .insert(func.name().to_string(), func);
66    }
67
68    /// Register a scalar function in the registry.
69    pub fn register_scalar(&self, func: impl Function + 'static) {
70        let func = Arc::new(func) as FunctionRef;
71
72        for alias in func.aliases() {
73            let func: ScalarFunctionFactory = func.clone().into();
74            let alias = ScalarFunctionFactory {
75                name: alias.clone(),
76                ..func
77            };
78            self.register(alias);
79        }
80
81        self.register(func)
82    }
83
84    /// Register an aggregate function in the registry.
85    pub fn register_aggr(&self, func: AggregateUDF) {
86        let _ = self
87            .aggregate_functions
88            .write()
89            .unwrap()
90            .insert(func.name().to_string(), func);
91    }
92
93    /// Register a table function
94    pub fn register_table_function(&self, func: TableFunction) {
95        let _ = self
96            .table_functions
97            .write()
98            .unwrap()
99            .insert(func.name().to_string(), Arc::new(func));
100    }
101
102    pub fn get_function(&self, name: &str) -> Option<ScalarFunctionFactory> {
103        self.functions.read().unwrap().get(name).cloned()
104    }
105
106    /// Returns a list of all scalar functions registered in the registry.
107    pub fn scalar_functions(&self) -> Vec<ScalarFunctionFactory> {
108        self.functions.read().unwrap().values().cloned().collect()
109    }
110
111    /// Returns a list of all aggregate functions registered in the registry.
112    pub fn aggregate_functions(&self) -> Vec<AggregateUDF> {
113        self.aggregate_functions
114            .read()
115            .unwrap()
116            .values()
117            .cloned()
118            .collect()
119    }
120
121    pub fn table_functions(&self) -> Vec<Arc<TableFunction>> {
122        self.table_functions
123            .read()
124            .unwrap()
125            .values()
126            .cloned()
127            .collect()
128    }
129
130    /// Returns true if an aggregate function with the given name exists in the registry.
131    pub fn is_aggr_func_exist(&self, name: &str) -> bool {
132        self.aggregate_functions.read().unwrap().contains_key(name)
133    }
134}
135
136pub static FUNCTION_REGISTRY: LazyLock<Arc<FunctionRegistry>> = LazyLock::new(|| {
137    let function_registry = FunctionRegistry::default();
138
139    // Utility functions
140    MathFunction::register(&function_registry);
141    TimestampFunction::register(&function_registry);
142    DateFunction::register(&function_registry);
143    ExpressionFunction::register(&function_registry);
144    UddSketchCalcFunction::register(&function_registry);
145    HllCalcFunction::register(&function_registry);
146
147    // Full text search function
148    MatchesFunction::register(&function_registry);
149    MatchesTermFunction::register(&function_registry);
150
151    // System and administration functions
152    SystemFunction::register(&function_registry);
153    AdminFunction::register(&function_registry);
154
155    // Json related functions
156    JsonFunction::register(&function_registry);
157
158    // String related functions
159    register_string_functions(&function_registry);
160
161    // Vector related functions
162    VectorScalarFunction::register(&function_registry);
163    VectorAggrFunction::register(&function_registry);
164
165    // Geo functions
166    #[cfg(feature = "geo")]
167    crate::scalars::geo::GeoFunctions::register(&function_registry);
168    #[cfg(feature = "geo")]
169    crate::aggrs::geo::GeoFunction::register(&function_registry);
170
171    // Ip functions
172    IpFunctions::register(&function_registry);
173
174    // Approximate functions
175    ApproximateFunction::register(&function_registry);
176
177    // CountHash function
178    CountHash::register(&function_registry);
179
180    // state function of supported aggregate functions
181    StateMergeHelper::register(&function_registry);
182
183    Arc::new(function_registry)
184});
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189    use crate::scalars::test::TestAndFunction;
190
191    #[test]
192    fn test_function_registry() {
193        let registry = FunctionRegistry::default();
194
195        assert!(registry.get_function("test_and").is_none());
196        assert!(registry.scalar_functions().is_empty());
197        registry.register_scalar(TestAndFunction::default());
198        let _ = registry.get_function("test_and").unwrap();
199        assert_eq!(1, registry.scalar_functions().len());
200    }
201}
common_function/function_registry.rs

common_function/
function_registry.rs
