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

common_function/
function_registry.rs
