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

common_function/
function_registry.rs
