Skip to main content

uni_query/query/executor/
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
10use std::collections::HashMap;
11use std::sync::Arc;
12
13use uni_common::{Result, Value};
14
15/// Type alias for a custom scalar function.
16///
17/// The function receives evaluated arguments and returns a single value.
18pub type CustomScalarFn = Arc<dyn Fn(&[Value]) -> Result<Value> + Send + Sync>;
19
20/// Registry of user-defined scalar functions.
21///
22/// Functions are stored with uppercased names for case-insensitive lookup.
23#[derive(Default, Clone)]
24pub struct CustomFunctionRegistry {
25    functions: HashMap<String, CustomScalarFn>,
26}
27
28impl CustomFunctionRegistry {
29    /// Create an empty registry.
30    pub fn new() -> Self {
31        Self::default()
32    }
33
34    /// Register a custom scalar function.
35    ///
36    /// If a function with the same name already exists, it is replaced.
37    pub fn register(&mut self, name: String, func: CustomScalarFn) {
38        self.functions.insert(name.to_uppercase(), func);
39    }
40
41    /// Look up a custom function by name (case-insensitive).
42    pub fn get(&self, name: &str) -> Option<&CustomScalarFn> {
43        self.functions.get(&name.to_uppercase())
44    }
45
46    /// Iterate over all registered functions.
47    pub fn iter(&self) -> impl Iterator<Item = (&str, &CustomScalarFn)> {
48        self.functions.iter().map(|(k, v)| (k.as_str(), v))
49    }
50
51    /// Remove a custom function by name. Returns true if it existed.
52    pub fn remove(&mut self, name: &str) -> bool {
53        self.functions.remove(&name.to_uppercase()).is_some()
54    }
55
56    /// Returns `true` if no functions are registered.
57    pub fn is_empty(&self) -> bool {
58        self.functions.is_empty()
59    }
60}
61
62impl std::fmt::Debug for CustomFunctionRegistry {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        f.debug_struct("CustomFunctionRegistry")
65            .field("functions", &self.functions.keys().collect::<Vec<_>>())
66            .finish()
67    }
68}