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::collections::hash_map::Entry;
18use std::sync::{Arc, LazyLock, RwLock};
19
20use datafusion::catalog::TableFunction;
21use datafusion_expr::expr_rewriter::FunctionRewrite;
22use datafusion_expr::{AggregateUDF, WindowUDF};
23
24use crate::admin::AdminFunction;
25use crate::aggrs::aggr_wrapper::StateMergeHelper;
26use crate::aggrs::approximate::ApproximateFunction;
27use crate::aggrs::count_hash::CountHash;
28use crate::aggrs::vector::VectorFunction as VectorAggrFunction;
29use crate::function::{Function, FunctionRef};
30use crate::function_factory::ScalarFunctionFactory;
31use crate::scalars::anomaly::AnomalyFunction;
32use crate::scalars::date::DateFunction;
33use crate::scalars::expression::ExpressionFunction;
34use crate::scalars::hll_count::HllCalcFunction;
35use crate::scalars::ip::IpFunctions;
36use crate::scalars::json::JsonFunction;
37use crate::scalars::matches::MatchesFunction;
38use crate::scalars::matches_term::MatchesTermFunction;
39use crate::scalars::math::MathFunction;
40use crate::scalars::primary_key::DecodePrimaryKeyFunction;
41use crate::scalars::string::register_string_functions;
42use crate::scalars::timestamp::TimestampFunction;
43use crate::scalars::uddsketch_calc::UddSketchCalcFunction;
44use crate::scalars::uddsketch_rank::UddSketchRankFunction;
45use crate::scalars::vector::VectorFunction as VectorScalarFunction;
46use crate::scalars::welford_stddev::WelfordStddevFunction;
47use crate::system::SystemFunction;
48
49#[derive(Default)]
50pub struct FunctionRegistry {
51    functions: RwLock<HashMap<String, ScalarFunctionFactory>>,
52    aggregate_functions: RwLock<HashMap<String, AggregateUDF>>,
53    table_functions: RwLock<HashMap<String, Arc<TableFunction>>>,
54    function_rewrites: RwLock<Vec<Arc<dyn FunctionRewrite + Send + Sync>>>,
55    window_functions: RwLock<HashMap<String, WindowUDF>>,
56}
57
58/// The result of registering a function.
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60pub enum FunctionRegistrationResult {
61    /// The function was newly registered.
62    Registered,
63    /// A function with the same name was already registered and was kept.
64    AlreadyExists,
65}
66
67impl FunctionRegistry {
68    /// Register a function in the registry by converting it into a `ScalarFunctionFactory`.
69    ///
70    /// # Arguments
71    ///
72    /// * `func` - An object that can be converted into a `ScalarFunctionFactory`.
73    ///
74    /// The function is inserted into the internal function map, keyed by its name.
75    /// If a function with the same name already exists, it will be replaced.
76    pub fn register(&self, func: impl Into<ScalarFunctionFactory>) {
77        let func = func.into();
78        let _ = self
79            .functions
80            .write()
81            .unwrap()
82            .insert(func.name().to_string(), func);
83    }
84
85    /// Register a function only if no function with the same name exists.
86    ///
87    /// The duplicate check and the insert happen atomically under the same
88    /// write lock of the functions map. If a function with the same name
89    /// already exists, it is kept unchanged and
90    /// [`FunctionRegistrationResult::AlreadyExists`] is returned; otherwise the
91    /// function is registered and [`FunctionRegistrationResult::Registered`] is
92    /// returned.
93    pub fn register_if_absent(
94        &self,
95        func: impl Into<ScalarFunctionFactory>,
96    ) -> FunctionRegistrationResult {
97        let func = func.into();
98        let mut functions = self.functions.write().unwrap();
99        match functions.entry(func.name().to_string()) {
100            Entry::Occupied(_) => FunctionRegistrationResult::AlreadyExists,
101            Entry::Vacant(entry) => {
102                entry.insert(func);
103                FunctionRegistrationResult::Registered
104            }
105        }
106    }
107
108    /// Register a scalar function in the registry.
109    pub fn register_scalar(&self, func: impl Function + 'static) {
110        let func = Arc::new(func) as FunctionRef;
111
112        for alias in func.aliases() {
113            let func: ScalarFunctionFactory = func.clone().into();
114            let alias = ScalarFunctionFactory {
115                name: alias.clone(),
116                ..func
117            };
118            self.register(alias);
119        }
120
121        self.register(func)
122    }
123
124    /// Register an aggregate function in the registry.
125    pub fn register_aggr(&self, func: AggregateUDF) {
126        let _ = self
127            .aggregate_functions
128            .write()
129            .unwrap()
130            .insert(func.name().to_string(), func);
131    }
132
133    /// Register a table function
134    pub fn register_table_function(&self, func: TableFunction) {
135        let _ = self
136            .table_functions
137            .write()
138            .unwrap()
139            .insert(func.name().to_string(), Arc::new(func));
140    }
141
142    /// Register a function rewrite rule.
143    pub fn register_function_rewrite(&self, func: impl FunctionRewrite + Send + Sync + 'static) {
144        self.function_rewrites.write().unwrap().push(Arc::new(func));
145    }
146
147    /// Register a window function (UDWF).
148    pub fn register_window(&self, func: WindowUDF) {
149        let _ = self
150            .window_functions
151            .write()
152            .unwrap()
153            .insert(func.name().to_string(), func);
154    }
155
156    pub fn get_function(&self, name: &str) -> Option<ScalarFunctionFactory> {
157        self.functions.read().unwrap().get(name).cloned()
158    }
159
160    /// Returns a list of all scalar functions registered in the registry.
161    pub fn scalar_functions(&self) -> Vec<ScalarFunctionFactory> {
162        self.functions.read().unwrap().values().cloned().collect()
163    }
164
165    /// Returns a list of all aggregate functions registered in the registry.
166    pub fn aggregate_functions(&self) -> Vec<AggregateUDF> {
167        self.aggregate_functions
168            .read()
169            .unwrap()
170            .values()
171            .cloned()
172            .collect()
173    }
174
175    pub fn table_functions(&self) -> Vec<Arc<TableFunction>> {
176        self.table_functions
177            .read()
178            .unwrap()
179            .values()
180            .cloned()
181            .collect()
182    }
183
184    /// Returns a list of all window functions registered in the registry.
185    pub fn window_functions(&self) -> Vec<WindowUDF> {
186        self.window_functions
187            .read()
188            .unwrap()
189            .values()
190            .cloned()
191            .collect()
192    }
193
194    /// Returns a registered aggregate function by name.
195    pub fn get_aggr_func(&self, name: &str) -> Option<AggregateUDF> {
196        self.aggregate_functions.read().unwrap().get(name).cloned()
197    }
198
199    /// Returns true if an aggregate function with the given name exists in the registry.
200    pub fn is_aggr_func_exist(&self, name: &str) -> bool {
201        self.aggregate_functions.read().unwrap().contains_key(name)
202    }
203
204    /// Returns a list of all function rewrite rules registered in the registry.
205    pub fn function_rewrites(&self) -> Vec<Arc<dyn FunctionRewrite + Send + Sync>> {
206        self.function_rewrites.read().unwrap().clone()
207    }
208}
209
210pub static FUNCTION_REGISTRY: LazyLock<Arc<FunctionRegistry>> = LazyLock::new(|| {
211    let function_registry = FunctionRegistry::default();
212
213    // Utility functions
214    MathFunction::register(&function_registry);
215    TimestampFunction::register(&function_registry);
216    DateFunction::register(&function_registry);
217    ExpressionFunction::register(&function_registry);
218    UddSketchCalcFunction::register(&function_registry);
219    UddSketchRankFunction::register(&function_registry);
220    HllCalcFunction::register(&function_registry);
221    WelfordStddevFunction::register(&function_registry);
222    DecodePrimaryKeyFunction::register(&function_registry);
223
224    // Full text search function
225    MatchesFunction::register(&function_registry);
226    MatchesTermFunction::register(&function_registry);
227
228    // System and administration functions
229    SystemFunction::register(&function_registry);
230    AdminFunction::register(&function_registry);
231
232    // Json related functions
233    JsonFunction::register(&function_registry);
234
235    // String related functions
236    register_string_functions(&function_registry);
237
238    // Vector related functions
239    VectorScalarFunction::register(&function_registry);
240    VectorAggrFunction::register(&function_registry);
241
242    // Geo functions
243    #[cfg(feature = "geo")]
244    crate::scalars::geo::GeoFunctions::register(&function_registry);
245    #[cfg(feature = "geo")]
246    crate::aggrs::geo::GeoFunction::register(&function_registry);
247
248    // Ip functions
249    IpFunctions::register(&function_registry);
250
251    // Approximate functions
252    ApproximateFunction::register(&function_registry);
253
254    // CountHash function
255    CountHash::register(&function_registry);
256
257    // state function of supported aggregate functions
258    StateMergeHelper::register(&function_registry);
259
260    // Anomaly detection window functions
261    AnomalyFunction::register(&function_registry);
262
263    Arc::new(function_registry)
264});
265
266static ADMIN_FUNCTION_REGISTRY: LazyLock<FunctionRegistry> = LazyLock::new(|| {
267    let registry = FunctionRegistry::default();
268    AdminFunction::register_admin_only(&registry);
269    registry
270});
271
272/// Returns a function that is only available to the ADMIN statement executor.
273pub fn get_admin_function(name: &str) -> Option<ScalarFunctionFactory> {
274    ADMIN_FUNCTION_REGISTRY.get_function(name)
275}
276
277/// Register a function that is only available to the ADMIN statement executor.
278///
279/// If a function with the same name is already registered in the ADMIN
280/// registry, the existing one is kept and
281/// [`FunctionRegistrationResult::AlreadyExists`] is returned. A name that
282/// already exists in the normal [`FUNCTION_REGISTRY`] when this call
283/// linearizes is also rejected: the ADMIN executor resolves admin-only
284/// functions before falling back to the normal registry, so inserting such a
285/// name here would shadow the built-in. Otherwise the function is registered
286/// and [`FunctionRegistrationResult::Registered`] is returned.
287///
288/// The enforced contract is one-way: it only guards the ADMIN registration
289/// against names already present in the normal registry. A later ordinary
290/// [`FunctionRegistry::register`] may still install the same name in the
291/// normal registry because the normal registry keeps its legacy replace
292/// semantics.
293pub fn register_admin_function(
294    func: impl Into<ScalarFunctionFactory>,
295) -> FunctionRegistrationResult {
296    register_admin_function_in(&ADMIN_FUNCTION_REGISTRY, &FUNCTION_REGISTRY, func)
297}
298
299/// Core implementation of [`register_admin_function`] against a pair of
300/// registries, parameterized so tests can exercise it with local registries.
301///
302/// Locking: the ADMIN-registry write lock is acquired first, then a read lock
303/// on the normal registry, and the normal-registry guard (bound to
304/// `normal_functions`) is kept alive through both the normal-name check and
305/// the ADMIN insertion below. This is the only code path that holds both
306/// registries' locks, so the ADMIN -> FUNCTION acquisition order is
307/// consistent and a concurrent normal-registry registration cannot slip in
308/// between the check and the ADMIN insert and be shadowed.
309///
310/// The enforced contract is one-way: it only guards the ADMIN registration
311/// against names already present in the normal registry. A later ordinary
312/// [`FunctionRegistry::register`] may still install the same name in the
313/// normal registry because the normal registry keeps its legacy replace
314/// semantics.
315fn register_admin_function_in(
316    admin_registry: &FunctionRegistry,
317    normal_registry: &FunctionRegistry,
318    func: impl Into<ScalarFunctionFactory>,
319) -> FunctionRegistrationResult {
320    let func = func.into();
321    let mut admin_functions = admin_registry.functions.write().unwrap();
322    // The normal-registry guard is a read lock: it is held across the
323    // normal-name check and the ADMIN insertion below, and while it is alive
324    // no writer can acquire the normal-registry write lock, so a concurrent
325    // normal-registry registration cannot slip in between the check and the
326    // ADMIN insert and be shadowed.
327    let normal_functions = normal_registry.functions.read().unwrap();
328    if normal_functions.contains_key(func.name()) {
329        drop(normal_functions);
330        return FunctionRegistrationResult::AlreadyExists;
331    }
332    let result = match admin_functions.entry(func.name().to_string()) {
333        Entry::Occupied(_) => FunctionRegistrationResult::AlreadyExists,
334        Entry::Vacant(entry) => {
335            entry.insert(func);
336            FunctionRegistrationResult::Registered
337        }
338    };
339    // Drop the read guard only after the ADMIN insertion, so writers to the
340    // normal registry stay blocked until the check-and-insert is complete.
341    drop(normal_functions);
342    result
343}
344
345#[cfg(test)]
346mod tests {
347    use std::sync::{Arc, Barrier};
348    use std::thread;
349
350    use super::*;
351    use crate::scalars::test::TestAndFunction;
352    use crate::scalars::udf::create_udf;
353
354    /// Creates a [`ScalarFunctionFactory`] with the given name. Each call
355    /// allocates a distinct factory closure, so factories can be told apart by
356    /// [`Arc::ptr_eq`] on their `factory` field even when names are identical.
357    fn named_factory(name: &str) -> ScalarFunctionFactory {
358        ScalarFunctionFactory {
359            name: name.to_string(),
360            factory: Arc::new(|_ctx| create_udf(Arc::new(TestAndFunction::default()))),
361        }
362    }
363
364    #[test]
365    fn test_function_registry() {
366        let registry = FunctionRegistry::default();
367
368        assert!(registry.get_function("test_and").is_none());
369        assert!(registry.scalar_functions().is_empty());
370        registry.register_scalar(TestAndFunction::default());
371        let _ = registry.get_function("test_and").unwrap();
372        assert_eq!(1, registry.scalar_functions().len());
373    }
374
375    #[test]
376    fn test_uddsketch_rank_registered() {
377        assert!(FUNCTION_REGISTRY.get_function("uddsketch_rank").is_some());
378    }
379
380    #[test]
381    fn test_register_if_absent_registers_new_function() {
382        let registry = FunctionRegistry::default();
383        let name = "pr3_register_if_absent_new";
384        let factory = named_factory(name);
385
386        assert_eq!(
387            registry.register_if_absent(factory.clone()),
388            FunctionRegistrationResult::Registered
389        );
390        let registered = registry
391            .get_function(name)
392            .expect("function should be registered");
393        assert!(Arc::ptr_eq(&registered.factory, &factory.factory));
394    }
395
396    #[test]
397    fn test_register_if_absent_first_registration_wins() {
398        let registry = FunctionRegistry::default();
399        let name = "pr3_register_if_absent_duplicate";
400        let first = named_factory(name);
401        let second = named_factory(name);
402
403        assert_eq!(
404            registry.register_if_absent(first.clone()),
405            FunctionRegistrationResult::Registered
406        );
407        assert_eq!(
408            registry.register_if_absent(second.clone()),
409            FunctionRegistrationResult::AlreadyExists
410        );
411
412        let stored = registry
413            .get_function(name)
414            .expect("function should be registered");
415        assert!(Arc::ptr_eq(&stored.factory, &first.factory));
416        assert!(!Arc::ptr_eq(&stored.factory, &second.factory));
417    }
418
419    #[test]
420    fn test_register_replaces_existing_function() {
421        // Regression test: `register` keeps its replace semantics.
422        let registry = FunctionRegistry::default();
423        let name = "pr3_register_replaces";
424        let first = named_factory(name);
425        let second = named_factory(name);
426
427        registry.register(first.clone());
428        registry.register(second.clone());
429
430        let stored = registry
431            .get_function(name)
432            .expect("function should be registered");
433        assert!(Arc::ptr_eq(&stored.factory, &second.factory));
434        assert!(!Arc::ptr_eq(&stored.factory, &first.factory));
435    }
436
437    #[test]
438    fn test_concurrent_register_if_absent_same_name() {
439        const THREADS: usize = 8;
440        let registry = Arc::new(FunctionRegistry::default());
441        let name = "pr3_concurrent_same_name";
442        let barrier = Arc::new(Barrier::new(THREADS));
443
444        let handles: Vec<_> = (0..THREADS)
445            .map(|_| {
446                let registry = Arc::clone(&registry);
447                let barrier = Arc::clone(&barrier);
448                thread::spawn(move || {
449                    let factory = named_factory(name);
450                    // Synchronize so every thread attempts registration at the
451                    // same time; only one may win the write lock.
452                    barrier.wait();
453                    let result = registry.register_if_absent(factory.clone());
454                    (result, factory)
455                })
456            })
457            .collect();
458
459        let mut results: Vec<(FunctionRegistrationResult, ScalarFunctionFactory)> =
460            Vec::with_capacity(THREADS);
461        for handle in handles {
462            results.push(handle.join().expect("thread should not panic"));
463        }
464
465        let registered = results
466            .iter()
467            .filter(|(result, _)| *result == FunctionRegistrationResult::Registered)
468            .count();
469        let already_exists = results
470            .iter()
471            .filter(|(result, _)| *result == FunctionRegistrationResult::AlreadyExists)
472            .count();
473        assert_eq!(registered, 1);
474        assert_eq!(already_exists, THREADS - 1);
475
476        let winner = results
477            .iter()
478            .find(|(result, _)| *result == FunctionRegistrationResult::Registered)
479            .map(|(_, factory)| factory)
480            .expect("exactly one registration must win");
481
482        let stored = registry
483            .get_function(name)
484            .expect("function should be registered");
485        assert!(
486            Arc::ptr_eq(&stored.factory, &winner.factory),
487            "the stored factory must be the factory of the winning registration"
488        );
489    }
490
491    #[test]
492    fn test_register_admin_function_first_wins() {
493        // Tests touching the global registry must use unique names.
494        let name = "pr3_admin_runtime_register";
495        let first = named_factory(name);
496        let second = named_factory(name);
497
498        assert_eq!(
499            register_admin_function(first.clone()),
500            FunctionRegistrationResult::Registered
501        );
502        assert_eq!(
503            register_admin_function(second.clone()),
504            FunctionRegistrationResult::AlreadyExists
505        );
506
507        let stored = get_admin_function(name).expect("admin function should be queryable");
508        assert!(Arc::ptr_eq(&stored.factory, &first.factory));
509        assert!(!Arc::ptr_eq(&stored.factory, &second.factory));
510    }
511
512    #[test]
513    fn test_register_admin_function_duplicate_same_factory() {
514        // The exact same factory clone registered twice: the first registration
515        // wins, the duplicate is rejected, and the stored factory is
516        // pointer-equal to the original. Tests touching the global registry
517        // must use unique names.
518        let name = "pr3_admin_runtime_register_same_factory";
519        let factory = named_factory(name);
520
521        assert_eq!(
522            register_admin_function(factory.clone()),
523            FunctionRegistrationResult::Registered
524        );
525        assert_eq!(
526            register_admin_function(factory.clone()),
527            FunctionRegistrationResult::AlreadyExists
528        );
529
530        let stored = get_admin_function(name).expect("admin function should be queryable");
531        assert!(Arc::ptr_eq(&stored.factory, &factory.factory));
532    }
533
534    #[test]
535    fn test_register_admin_function_in_is_one_way_later_normal_register_allowed() {
536        // The guard is one-way: after the ADMIN registration completes, an
537        // ordinary normal-registry registration of the same name still
538        // succeeds because the normal registry keeps its legacy replace
539        // semantics.
540        let admin = FunctionRegistry::default();
541        let normal = FunctionRegistry::default();
542        let name = "pr3_admin_in_one_way";
543        let admin_factory = named_factory(name);
544        let normal_factory = named_factory(name);
545
546        assert_eq!(
547            register_admin_function_in(&admin, &normal, admin_factory.clone()),
548            FunctionRegistrationResult::Registered
549        );
550
551        normal.register(normal_factory.clone());
552
553        let stored_admin = admin
554            .get_function(name)
555            .expect("the ADMIN registration must be kept");
556        assert!(Arc::ptr_eq(&stored_admin.factory, &admin_factory.factory));
557        let stored_normal = normal
558            .get_function(name)
559            .expect("the later normal registration must succeed");
560        assert!(Arc::ptr_eq(&stored_normal.factory, &normal_factory.factory));
561    }
562
563    #[test]
564    fn test_register_admin_function_rejects_normal_registry_builtin_name() {
565        // Regression test: the ADMIN executor resolves `get_admin_function`
566        // before falling back to `FUNCTION_REGISTRY`, so registering a function
567        // whose name already exists in the normal registry would shadow the
568        // ADMIN-invocable built-in (e.g. `flush_table`). Such registrations
569        // must be rejected with
570        // [`FunctionRegistrationResult::AlreadyExists`] and must not be
571        // inserted into the ADMIN registry.
572        let factory = named_factory("flush_table");
573
574        assert_eq!(
575            register_admin_function(factory.clone()),
576            FunctionRegistrationResult::AlreadyExists
577        );
578        assert!(
579            get_admin_function("flush_table").is_none(),
580            "a normal-registry built-in must not be shadowed into the ADMIN registry"
581        );
582        assert!(
583            FUNCTION_REGISTRY.get_function("flush_table").is_some(),
584            "the normal-registry built-in must remain registered"
585        );
586    }
587
588    #[test]
589    fn test_builtin_admin_functions_remain_queryable() {
590        // Built-in admin-only functions registered at startup stay queryable
591        // through the same global registry used for runtime registrations.
592        #[cfg(feature = "enterprise")]
593        {
594            assert!(get_admin_function("purge_table").is_some());
595        }
596        #[cfg(not(feature = "enterprise"))]
597        {
598            assert!(get_admin_function("purge_table").is_none());
599        }
600    }
601}
common_function/function_registry.rs

common_function/
function_registry.rs
