linkrs_query/executor/expression/functions/
registry.rs

1//! Function Registry
2//!
3//! Provide functions for registration, lookup, and execution.
4//! The specific implementation of the function is located in the builtin submodule.
5
6mod builtin;
7mod custom;
8mod table;
9mod udf;
10
11#[cfg(test)]
12mod tests;
13
14pub use udf::DynamicUdfInfo;
15
16use super::{BuiltinFunction, CustomFunction, TableFunction};
17use crate::executor::expression::evaluation_context::graph_storage::GraphStorageRef;
18use crate::executor::expression::{ExpressionError, ExpressionErrorType};
19use linkrs_core::DataType;
20use linkrs_core::Value;
21use std::collections::HashMap;
22use std::sync::Arc;
23
24use udf::LoadedExtension;
25
26/// Registry entry carrying the function and its static return type.
27#[derive(Debug, Clone)]
28pub struct RegistryEntry {
29    pub function: BuiltinFunction,
30    pub return_type: DataType,
31}
32
33/// Function Registry
34///
35/// Using a static distribution mechanism, functions are called directly through the BuiltinFunction and CustomFunction enumerations.
36/// The overhead associated with dynamic distribution (dyn) was avoided.
37#[derive(Debug)]
38pub struct FunctionRegistry {
39    /// Built-in function mapping (function name -> RegistryEntry)
40    builtin_functions: HashMap<String, RegistryEntry>,
41    /// Custom function mapping (function name -> CustomFunction)
42    custom_functions: HashMap<String, CustomFunction>,
43    /// Table function mapping (function name -> Box<dyn TableFunction>)
44    table_functions: HashMap<String, Box<dyn TableFunction>>,
45    /// Dynamically loaded UDF libraries (upper-cased name -> extension metadata)
46    dynamic_libraries: HashMap<String, LoadedExtension>,
47}
48
49impl Default for FunctionRegistry {
50    fn default() -> Self {
51        Self::new()
52    }
53}
54
55impl FunctionRegistry {
56    pub fn new() -> Self {
57        let mut registry = Self {
58            builtin_functions: HashMap::new(),
59            custom_functions: HashMap::new(),
60            table_functions: HashMap::new(),
61            dynamic_libraries: HashMap::new(),
62        };
63        registry.register_all_builtin_functions();
64        registry.register_builtin_table_functions();
65        registry
66    }
67
68    fn resolve_entry(&self, name: &str) -> Option<&RegistryEntry> {
69        let upper = name.to_uppercase();
70        self.builtin_functions.get(&upper)
71    }
72
73    /// Check whether the function exists.
74    pub fn contains(&self, name: &str) -> bool {
75        let upper_name = name.to_uppercase();
76        self.builtin_functions.contains_key(&upper_name)
77            || self.custom_functions.contains_key(&upper_name)
78    }
79
80    /// Obtain all function names
81    pub fn function_names(&self) -> Vec<&str> {
82        let mut names: Vec<&str> = self.builtin_functions.keys().map(|s| s.as_str()).collect();
83        names.extend(self.custom_functions.keys().map(|s| s.as_str()));
84        names
85    }
86
87    /// Execute a function (based on its name)
88    pub fn execute(&self, name: &str, args: &[Value]) -> Result<Value, ExpressionError> {
89        let upper_name = name.to_uppercase();
90        if let Some(entry) = self.builtin_functions.get(&upper_name) {
91            return entry.function.execute(args);
92        }
93        if let Some(func) = self.custom_functions.get(&upper_name) {
94            return func.execute(args);
95        }
96
97        Err(ExpressionError::new(
98            ExpressionErrorType::UndefinedFunction,
99            format!("Undefined function: {}", name),
100        ))
101    }
102
103    /// Execute a function with graph storage access
104    pub fn execute_with_storage(
105        &self,
106        name: &str,
107        args: &[Value],
108        storage: &GraphStorageRef,
109    ) -> Result<Value, ExpressionError> {
110        let upper_name = name.to_uppercase();
111        if let Some(entry) = self.builtin_functions.get(&upper_name) {
112            return entry.function.execute_with_storage(args, storage);
113        }
114        if let Some(func) = self.custom_functions.get(&upper_name) {
115            return func.execute(args);
116        }
117        Err(ExpressionError::new(
118            ExpressionErrorType::UndefinedFunction,
119            format!("Undefined function: {}", name),
120        ))
121    }
122}
123
124/// Global function registry instance
125pub fn global_registry() -> Arc<FunctionRegistry> {
126    use std::sync::OnceLock;
127    static REGISTRY: OnceLock<Arc<FunctionRegistry>> = OnceLock::new();
128    REGISTRY
129        .get_or_init(|| Arc::new(FunctionRegistry::new()))
130        .clone()
131}
132
133/// Obtain a static reference to the global function registry.
134///
135/// Used in scenarios where it is necessary to retrieve a function reference (such as in ExpressionContext::get_function).
136pub fn global_registry_ref() -> &'static FunctionRegistry {
137    use std::sync::OnceLock;
138    static REGISTRY: OnceLock<FunctionRegistry> = OnceLock::new();
139    REGISTRY.get_or_init(FunctionRegistry::new)
140}
linkrs_query/executor/expression/functions/registry.rs

linkrs_query/executor/expression/functions/
registry.rs
