nodedb_sql/functions/
registry.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Built-in function registry for SQL planning.
4//!
5//! Tracks known functions, their categories, arg specs, return types,
6//! and whether they trigger special engine routing (e.g., vector_distance
7//! → VectorSearch).
8
9use nodedb_types::columnar::ColumnType;
10
11use super::builtins::builtin_functions;
12
13/// Semantic version for tracking when a function was introduced.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
15pub struct Version {
16    pub major: u8,
17    pub minor: u8,
18    pub patch: u8,
19}
20
21impl Version {
22    pub const fn new(major: u8, minor: u8, patch: u8) -> Self {
23        Self {
24            major,
25            minor,
26            patch,
27        }
28    }
29}
30
31/// Specification for a single function argument position.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub struct ArgTypeSpec {
34    /// Argument name, for documentation and error messages.
35    pub name: &'static str,
36    /// Accepted column types. Empty slice means any type is accepted (wildcard).
37    pub accepted: &'static [ColumnType],
38    /// If true on the last argument, this argument may repeat zero or more times.
39    pub variadic: bool,
40}
41
42/// Function category.
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub enum FunctionCategory {
45    Scalar,
46    Aggregate,
47    Window,
48}
49
50/// Whether a function triggers special engine routing.
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub enum SearchTrigger {
53    None,
54    VectorSearch,
55    MultiVectorSearch,
56    /// Sparse-vector inverted-index search: `sparse_score(field, '{dim: weight, ...}')`
57    /// in ORDER BY routes to `SqlPlan::SparseSearch` and lowers to
58    /// `VectorOp::SparseSearch` (dot-product top-k over the sparse index).
59    SparseSearch,
60    TextSearch,
61    HybridSearch,
62    TextMatch,
63    SpatialDWithin,
64    SpatialContains,
65    SpatialIntersects,
66    SpatialWithin,
67    TimeBucket,
68    /// Array engine table-valued read: `ARRAY_SLICE(name, slice_obj, attrs?, limit?)`.
69    ArraySlice,
70    /// Array engine table-valued read: `ARRAY_PROJECT(name, attrs)`.
71    ArrayProject,
72    /// Array engine table-valued aggregate:
73    /// `ARRAY_AGG(name, attr, reducer, group_by_dim?)`.
74    ArrayAgg,
75    /// Array engine table-valued elementwise:
76    /// `ARRAY_ELEMENTWISE(left, right, op, attr)`.
77    ArrayElementwise,
78    /// Array engine maintenance scalar (returns BOOL): `ARRAY_FLUSH(name)`.
79    ArrayFlush,
80    /// Array engine maintenance scalar (returns BOOL): `ARRAY_COMPACT(name)`.
81    ArrayCompact,
82    /// Planner-intercepted graph-distance marker for three-source RRF fusion.
83    /// `graph_score(node_id_col, seed_id, depth => N, label => 'edge_label')`
84    /// is never evaluated per-row; the hybrid planner extracts its arguments
85    /// and builds a graph BFS spec in the physical plan.
86    GraphSearch,
87}
88
89/// Metadata about a known function.
90#[derive(Debug, Clone)]
91pub struct FunctionMeta {
92    pub name: &'static str,
93    pub category: FunctionCategory,
94    pub min_args: usize,
95    pub max_args: usize,
96    pub search_trigger: SearchTrigger,
97    /// Static return type, when known at plan time. `None` means the type
98    /// is context-dependent or unknown (resolved at runtime).
99    pub return_type: Option<ColumnType>,
100    /// Per-position argument type specifications.
101    pub arg_types: &'static [ArgTypeSpec],
102    /// Version in which this function was introduced.
103    pub since: Version,
104}
105
106/// The function registry.
107pub struct FunctionRegistry {
108    functions: Vec<FunctionMeta>,
109}
110
111impl FunctionRegistry {
112    /// Create the default registry with all built-in functions.
113    pub fn new() -> Self {
114        Self {
115            functions: builtin_functions(),
116        }
117    }
118
119    /// Look up a function by name (case-insensitive).
120    pub fn lookup(&self, name: &str) -> Option<&FunctionMeta> {
121        let lower = name.to_lowercase();
122        self.functions.iter().find(|f| f.name == lower)
123    }
124
125    /// Check if a function triggers special search routing.
126    pub fn search_trigger(&self, name: &str) -> SearchTrigger {
127        self.lookup(name)
128            .map(|f| f.search_trigger)
129            .unwrap_or(SearchTrigger::None)
130    }
131
132    /// Check if a function is an aggregate.
133    pub fn is_aggregate(&self, name: &str) -> bool {
134        self.lookup(name)
135            .is_some_and(|f| f.category == FunctionCategory::Aggregate)
136    }
137
138    /// Check if a function is a window function.
139    pub fn is_window(&self, name: &str) -> bool {
140        self.lookup(name)
141            .is_some_and(|f| f.category == FunctionCategory::Window)
142    }
143}
144
145impl Default for FunctionRegistry {
146    fn default() -> Self {
147        Self::new()
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[test]
156    fn lookup_builtin() {
157        let reg = FunctionRegistry::new();
158        assert!(reg.is_aggregate("COUNT"));
159        assert!(reg.is_aggregate("sum"));
160        assert!(!reg.is_aggregate("vector_distance"));
161        assert!(reg.is_window("row_number"));
162        assert_eq!(
163            reg.search_trigger("vector_distance"),
164            SearchTrigger::VectorSearch
165        );
166        assert_eq!(
167            reg.search_trigger("st_dwithin"),
168            SearchTrigger::SpatialDWithin
169        );
170        assert_eq!(reg.search_trigger("time_bucket"), SearchTrigger::TimeBucket);
171    }
172
173    #[test]
174    fn chunk_text_renamed_to_ndb_chunk_text() {
175        let reg = FunctionRegistry::new();
176        assert!(
177            reg.lookup("ndb_chunk_text").is_some(),
178            "ndb_chunk_text must be registered"
179        );
180        assert!(
181            reg.lookup("chunk_text").is_none(),
182            "old name chunk_text must not exist"
183        );
184    }
185
186    #[test]
187    fn unimplemented_functions_removed() {
188        let reg = FunctionRegistry::new();
189        assert!(reg.lookup("currency").is_none());
190        assert!(reg.lookup("distribute").is_none());
191        assert!(reg.lookup("allocate").is_none());
192        assert!(reg.lookup("resolve_permission").is_none());
193        assert!(reg.lookup("convert_currency").is_none());
194    }
195
196    #[test]
197    fn all_builtins_have_since_set() {
198        let reg = FunctionRegistry::new();
199        let v0_1_0 = Version::new(0, 1, 0);
200        for f in &reg.functions {
201            assert_eq!(
202                f.since, v0_1_0,
203                "function '{}' must have since = Version::new(0, 1, 0)",
204                f.name
205            );
206        }
207    }
208
209    #[test]
210    fn all_builtins_arg_counts_consistent() {
211        let reg = FunctionRegistry::new();
212        for f in &reg.functions {
213            let n = f.arg_types.len();
214            let last_variadic = f.arg_types.last().is_some_and(|a| a.variadic);
215            if last_variadic {
216                // min_args must be <= arg_types.len()
217                assert!(
218                    f.min_args <= n,
219                    "function '{}': min_args ({}) > arg_types.len() ({}) with variadic last arg",
220                    f.name,
221                    f.min_args,
222                    n
223                );
224            } else {
225                // min_args <= n <= max_args
226                assert!(
227                    f.min_args <= n,
228                    "function '{}': min_args ({}) > arg_types.len() ({})",
229                    f.name,
230                    f.min_args,
231                    n
232                );
233                assert!(
234                    n <= f.max_args,
235                    "function '{}': arg_types.len() ({}) > max_args ({})",
236                    f.name,
237                    n,
238                    f.max_args
239                );
240            }
241        }
242    }
243
244    #[test]
245    fn return_type_spot_checks() {
246        let reg = FunctionRegistry::new();
247        assert_eq!(
248            reg.lookup("now").and_then(|f| f.return_type),
249            Some(ColumnType::Timestamptz),
250            "now() must return Timestamptz"
251        );
252        assert_eq!(
253            reg.lookup("count").and_then(|f| f.return_type),
254            Some(ColumnType::Int64),
255            "count must return Int64"
256        );
257        assert_eq!(
258            reg.lookup("doc_exists").and_then(|f| f.return_type),
259            Some(ColumnType::Bool),
260            "doc_exists must return Bool"
261        );
262        assert_eq!(
263            reg.lookup("st_contains").and_then(|f| f.return_type),
264            Some(ColumnType::Bool),
265            "st_contains must return Bool"
266        );
267        assert_eq!(
268            reg.lookup("pg_fts_match").and_then(|f| f.return_type),
269            Some(ColumnType::Bool),
270            "pg_fts_match must return Bool"
271        );
272        assert_eq!(
273            reg.lookup("lower").and_then(|f| f.return_type),
274            Some(ColumnType::String),
275            "lower must return String"
276        );
277    }
278}
nodedb_sql/functions/registry.rs

nodedb_sql/functions/
registry.rs
