uni_query_functions/
custom_functions.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2024-2026 Dragonscale Team
3
4//! User-defined scalar function registry.
5//!
6//! Custom scalar functions can be registered at the database level and are
7//! available in both the DataFusion columnar execution path and the direct
8//! expression evaluator.
9//!
10//! This is a dependency-light leaf type: it stores `(name, fn)` pairs and
11//! nothing else. The plugin-framework shadow (`uni_plugin::PluginRegistry`)
12//! that the DataFusion adapter consumes is built on demand by
13//! `uni_query::query::df_udfs_plugin::plugin_registry_for_custom_functions`,
14//! which lives in `uni-query` because it depends on `uni-plugin`.
15
16use std::collections::HashMap;
17use std::sync::Arc;
18
19use uni_common::{Result, Value};
20
21/// Type alias for a custom scalar function.
22///
23/// The function receives evaluated arguments and returns a single value.
24pub type CustomScalarFn = Arc<dyn Fn(&[Value]) -> Result<Value> + Send + Sync>;
25
26/// Reserved plugin id used for legacy `CustomFunctionRegistry` registrations
27/// when they are mirrored into a `uni_plugin::PluginRegistry` shadow.
28///
29/// All legacy `register(name, fn)` calls land under this id in the shadow
30/// registry built by the `uni-query` plugin adapter. User plugins may not
31/// use this id. Defined here (a plain `&str` constant, no plugin dependency)
32/// so both the leaf crate and the host crate share the single source of
33/// truth.
34pub const LEGACY_USER_PLUGIN_ID: &str = "user.legacy";
35
36/// Registry of user-defined scalar functions.
37///
38/// Functions are stored with uppercased names for case-insensitive lookup.
39#[derive(Default, Clone)]
40pub struct CustomFunctionRegistry {
41    functions: HashMap<String, CustomScalarFn>,
42}
43
44impl CustomFunctionRegistry {
45    /// Create an empty registry.
46    pub fn new() -> Self {
47        Self::default()
48    }
49
50    /// Register a custom scalar function.
51    ///
52    /// If a function with the same name already exists, it is replaced.
53    pub fn register(&mut self, name: String, func: CustomScalarFn) {
54        self.functions.insert(name.to_uppercase(), func);
55    }
56
57    /// Look up a custom function by name (case-insensitive).
58    pub fn get(&self, name: &str) -> Option<&CustomScalarFn> {
59        self.functions.get(&name.to_uppercase())
60    }
61
62    /// Iterate over all registered functions.
63    pub fn iter(&self) -> impl Iterator<Item = (&str, &CustomScalarFn)> {
64        self.functions.iter().map(|(k, v)| (k.as_str(), v))
65    }
66
67    /// Remove a custom function by name. Returns true if it existed.
68    pub fn remove(&mut self, name: &str) -> bool {
69        self.functions.remove(&name.to_uppercase()).is_some()
70    }
71
72    /// Returns `true` if no functions are registered.
73    pub fn is_empty(&self) -> bool {
74        self.functions.is_empty()
75    }
76}
77
78impl std::fmt::Debug for CustomFunctionRegistry {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        f.debug_struct("CustomFunctionRegistry")
81            .field("functions", &self.functions.keys().collect::<Vec<_>>())
82            .finish()
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn legacy_register_and_get_roundtrip() {
92        let mut reg = CustomFunctionRegistry::new();
93        let f: CustomScalarFn = Arc::new(|_args: &[Value]| Ok(Value::String("ok".to_owned())));
94        reg.register("MYFN".to_owned(), f);
95
96        // Case-insensitive lookup.
97        assert!(reg.get("myfn").is_some());
98        assert!(reg.get("MYFN").is_some());
99    }
100
101    #[test]
102    fn legacy_remove_clears_entry() {
103        let mut reg = CustomFunctionRegistry::new();
104        let f: CustomScalarFn = Arc::new(|_| Ok(Value::Null));
105        reg.register("MYFN".to_owned(), f);
106        assert!(reg.remove("myfn"));
107        assert!(reg.get("MYFN").is_none());
108    }
109
110    #[test]
111    fn legacy_replace_updates_entry() {
112        let mut reg = CustomFunctionRegistry::new();
113        reg.register(
114            "MYFN".to_owned(),
115            Arc::new(|_| Ok(Value::String("first".to_owned()))),
116        );
117        reg.register(
118            "MYFN".to_owned(),
119            Arc::new(|_| Ok(Value::String("second".to_owned()))),
120        );
121
122        let v = (reg.get("MYFN").unwrap())(&[]).unwrap();
123        assert_eq!(v, Value::String("second".to_owned()));
124    }
125}
uni_query_functions/custom_functions.rs

uni_query_functions/
custom_functions.rs
