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, RwLock};
18
19use once_cell::sync::Lazy;
20
21use crate::admin::AdminFunction;
22use crate::function::{AsyncFunctionRef, FunctionRef};
23use crate::scalars::aggregate::{AggregateFunctionMetaRef, AggregateFunctions};
24use crate::scalars::date::DateFunction;
25use crate::scalars::expression::ExpressionFunction;
26use crate::scalars::hll_count::HllCalcFunction;
27use crate::scalars::ip::IpFunctions;
28use crate::scalars::json::JsonFunction;
29use crate::scalars::matches::MatchesFunction;
30use crate::scalars::matches_term::MatchesTermFunction;
31use crate::scalars::math::MathFunction;
32use crate::scalars::timestamp::TimestampFunction;
33use crate::scalars::uddsketch_calc::UddSketchCalcFunction;
34use crate::scalars::vector::VectorFunction;
35use crate::system::SystemFunction;
36
37#[derive(Default)]
38pub struct FunctionRegistry {
39    functions: RwLock<HashMap<String, FunctionRef>>,
40    async_functions: RwLock<HashMap<String, AsyncFunctionRef>>,
41    aggregate_functions: RwLock<HashMap<String, AggregateFunctionMetaRef>>,
42}
43
44impl FunctionRegistry {
45    pub fn register(&self, func: FunctionRef) {
46        let _ = self
47            .functions
48            .write()
49            .unwrap()
50            .insert(func.name().to_string(), func);
51    }
52
53    pub fn register_async(&self, func: AsyncFunctionRef) {
54        let _ = self
55            .async_functions
56            .write()
57            .unwrap()
58            .insert(func.name().to_string(), func);
59    }
60
61    pub fn get_async_function(&self, name: &str) -> Option<AsyncFunctionRef> {
62        self.async_functions.read().unwrap().get(name).cloned()
63    }
64
65    pub fn async_functions(&self) -> Vec<AsyncFunctionRef> {
66        self.async_functions
67            .read()
68            .unwrap()
69            .values()
70            .cloned()
71            .collect()
72    }
73
74    pub fn register_aggregate_function(&self, func: AggregateFunctionMetaRef) {
75        let _ = self
76            .aggregate_functions
77            .write()
78            .unwrap()
79            .insert(func.name(), func);
80    }
81
82    pub fn get_aggr_function(&self, name: &str) -> Option<AggregateFunctionMetaRef> {
83        self.aggregate_functions.read().unwrap().get(name).cloned()
84    }
85
86    pub fn get_function(&self, name: &str) -> Option<FunctionRef> {
87        self.functions.read().unwrap().get(name).cloned()
88    }
89
90    pub fn functions(&self) -> Vec<FunctionRef> {
91        self.functions.read().unwrap().values().cloned().collect()
92    }
93
94    pub fn aggregate_functions(&self) -> Vec<AggregateFunctionMetaRef> {
95        self.aggregate_functions
96            .read()
97            .unwrap()
98            .values()
99            .cloned()
100            .collect()
101    }
102}
103
104pub static FUNCTION_REGISTRY: Lazy<Arc<FunctionRegistry>> = Lazy::new(|| {
105    let function_registry = FunctionRegistry::default();
106
107    // Utility functions
108    MathFunction::register(&function_registry);
109    TimestampFunction::register(&function_registry);
110    DateFunction::register(&function_registry);
111    ExpressionFunction::register(&function_registry);
112    UddSketchCalcFunction::register(&function_registry);
113    HllCalcFunction::register(&function_registry);
114
115    // Aggregate functions
116    AggregateFunctions::register(&function_registry);
117
118    // Full text search function
119    MatchesFunction::register(&function_registry);
120    MatchesTermFunction::register(&function_registry);
121
122    // System and administration functions
123    SystemFunction::register(&function_registry);
124    AdminFunction::register(&function_registry);
125
126    // Json related functions
127    JsonFunction::register(&function_registry);
128
129    // Vector related functions
130    VectorFunction::register(&function_registry);
131
132    // Geo functions
133    #[cfg(feature = "geo")]
134    crate::scalars::geo::GeoFunctions::register(&function_registry);
135
136    // Ip functions
137    IpFunctions::register(&function_registry);
138
139    Arc::new(function_registry)
140});
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145    use crate::scalars::test::TestAndFunction;
146
147    #[test]
148    fn test_function_registry() {
149        let registry = FunctionRegistry::default();
150        let func = Arc::new(TestAndFunction);
151
152        assert!(registry.get_function("test_and").is_none());
153        assert!(registry.functions().is_empty());
154        registry.register(func);
155        let _ = registry.get_function("test_and").unwrap();
156        assert_eq!(1, registry.functions().len());
157    }
158}
common_function/function_registry.rs

common_function/
function_registry.rs
