|
| 1 | +//! Built-in function registry for SQL planning. |
| 2 | +//! |
| 3 | +//! Tracks known functions, their categories, and whether they trigger |
| 4 | +//! special engine routing (e.g., vector_distance → VectorSearch). |
| 5 | +
|
| 6 | +/// Function category. |
| 7 | +#[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 8 | +pub enum FunctionCategory { |
| 9 | + Scalar, |
| 10 | + Aggregate, |
| 11 | + Window, |
| 12 | +} |
| 13 | + |
| 14 | +/// Whether a function triggers special engine routing. |
| 15 | +#[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 16 | +pub enum SearchTrigger { |
| 17 | + None, |
| 18 | + VectorSearch, |
| 19 | + MultiVectorSearch, |
| 20 | + TextSearch, |
| 21 | + HybridSearch, |
| 22 | + TextMatch, |
| 23 | + SpatialDWithin, |
| 24 | + SpatialContains, |
| 25 | + SpatialIntersects, |
| 26 | + SpatialWithin, |
| 27 | + TimeBucket, |
| 28 | +} |
| 29 | + |
| 30 | +/// Metadata about a known function. |
| 31 | +#[derive(Debug, Clone)] |
| 32 | +pub struct FunctionMeta { |
| 33 | + pub name: &'static str, |
| 34 | + pub category: FunctionCategory, |
| 35 | + pub min_args: usize, |
| 36 | + pub max_args: usize, |
| 37 | + pub search_trigger: SearchTrigger, |
| 38 | +} |
| 39 | + |
| 40 | +/// The function registry. |
| 41 | +pub struct FunctionRegistry { |
| 42 | + functions: Vec<FunctionMeta>, |
| 43 | +} |
| 44 | + |
| 45 | +impl FunctionRegistry { |
| 46 | + /// Create the default registry with all built-in functions. |
| 47 | + pub fn new() -> Self { |
| 48 | + Self { |
| 49 | + functions: builtin_functions(), |
| 50 | + } |
| 51 | + } |
| 52 | + |
| 53 | + /// Look up a function by name (case-insensitive). |
| 54 | + pub fn lookup(&self, name: &str) -> Option<&FunctionMeta> { |
| 55 | + let lower = name.to_lowercase(); |
| 56 | + self.functions.iter().find(|f| f.name == lower) |
| 57 | + } |
| 58 | + |
| 59 | + /// Check if a function triggers special search routing. |
| 60 | + pub fn search_trigger(&self, name: &str) -> SearchTrigger { |
| 61 | + self.lookup(name) |
| 62 | + .map(|f| f.search_trigger) |
| 63 | + .unwrap_or(SearchTrigger::None) |
| 64 | + } |
| 65 | + |
| 66 | + /// Check if a function is an aggregate. |
| 67 | + pub fn is_aggregate(&self, name: &str) -> bool { |
| 68 | + self.lookup(name) |
| 69 | + .is_some_and(|f| f.category == FunctionCategory::Aggregate) |
| 70 | + } |
| 71 | + |
| 72 | + /// Check if a function is a window function. |
| 73 | + pub fn is_window(&self, name: &str) -> bool { |
| 74 | + self.lookup(name) |
| 75 | + .is_some_and(|f| f.category == FunctionCategory::Window) |
| 76 | + } |
| 77 | +} |
| 78 | + |
| 79 | +impl Default for FunctionRegistry { |
| 80 | + fn default() -> Self { |
| 81 | + Self::new() |
| 82 | + } |
| 83 | +} |
| 84 | + |
| 85 | +fn s( |
| 86 | + name: &'static str, |
| 87 | + cat: FunctionCategory, |
| 88 | + min: usize, |
| 89 | + max: usize, |
| 90 | + trigger: SearchTrigger, |
| 91 | +) -> FunctionMeta { |
| 92 | + FunctionMeta { |
| 93 | + name, |
| 94 | + category: cat, |
| 95 | + min_args: min, |
| 96 | + max_args: max, |
| 97 | + search_trigger: trigger, |
| 98 | + } |
| 99 | +} |
| 100 | + |
| 101 | +fn builtin_functions() -> Vec<FunctionMeta> { |
| 102 | + use FunctionCategory::*; |
| 103 | + use SearchTrigger::*; |
| 104 | + |
| 105 | + vec![ |
| 106 | + // ── Standard aggregates ── |
| 107 | + s("count", Aggregate, 0, 1, None), |
| 108 | + s("sum", Aggregate, 1, 1, None), |
| 109 | + s("avg", Aggregate, 1, 1, None), |
| 110 | + s("min", Aggregate, 1, 1, None), |
| 111 | + s("max", Aggregate, 1, 1, None), |
| 112 | + // ── Standard window ── |
| 113 | + s("row_number", Window, 0, 0, None), |
| 114 | + s("rank", Window, 0, 0, None), |
| 115 | + s("dense_rank", Window, 0, 0, None), |
| 116 | + s("lag", Window, 1, 3, None), |
| 117 | + s("lead", Window, 1, 3, None), |
| 118 | + s("first_value", Window, 1, 1, None), |
| 119 | + s("last_value", Window, 1, 1, None), |
| 120 | + s("nth_value", Window, 2, 2, None), |
| 121 | + // ── Vector search ── |
| 122 | + s("vector_distance", Scalar, 2, 3, VectorSearch), |
| 123 | + s("multi_vector_search", Scalar, 1, 2, MultiVectorSearch), |
| 124 | + s("multi_vector_score", Scalar, 3, 3, None), |
| 125 | + s("sparse_score", Scalar, 3, 3, None), |
| 126 | + // ── Text search ── |
| 127 | + s("bm25_score", Scalar, 2, 2, TextSearch), |
| 128 | + s("text_match", Scalar, 2, 2, TextMatch), |
| 129 | + // ── Hybrid search ── |
| 130 | + s("rrf_score", Scalar, 2, 4, HybridSearch), |
| 131 | + // ── Spatial ── |
| 132 | + s("st_dwithin", Scalar, 3, 3, SpatialDWithin), |
| 133 | + s("st_contains", Scalar, 2, 2, SpatialContains), |
| 134 | + s("st_intersects", Scalar, 2, 2, SpatialIntersects), |
| 135 | + s("st_within", Scalar, 2, 2, SpatialWithin), |
| 136 | + s("st_distance", Scalar, 2, 2, None), |
| 137 | + s("st_point", Scalar, 2, 2, None), |
| 138 | + // ── Timeseries ── |
| 139 | + s("time_bucket", Scalar, 2, 2, TimeBucket), |
| 140 | + // ── Timeseries aggregates ── |
| 141 | + s("ts_percentile", Aggregate, 2, 2, None), |
| 142 | + s("ts_stddev", Aggregate, 1, 1, None), |
| 143 | + s("ts_correlate", Aggregate, 2, 2, None), |
| 144 | + // ── Timeseries window ── |
| 145 | + s("ts_rate", Window, 1, 1, None), |
| 146 | + s("ts_derivative", Window, 1, 1, None), |
| 147 | + s("ts_moving_avg", Window, 2, 2, None), |
| 148 | + s("ts_ema", Window, 2, 2, None), |
| 149 | + s("ts_delta", Window, 1, 1, None), |
| 150 | + s("ts_interpolate", Window, 1, 1, None), |
| 151 | + s("ts_lag", Window, 1, 3, None), |
| 152 | + s("ts_lead", Window, 1, 3, None), |
| 153 | + s("ts_rank", Window, 0, 0, None), |
| 154 | + // ── Approximate aggregates ── |
| 155 | + s("approx_count_distinct", Aggregate, 1, 1, None), |
| 156 | + s("approx_percentile", Aggregate, 2, 2, None), |
| 157 | + s("approx_topk", Aggregate, 2, 2, None), |
| 158 | + s("approx_count", Aggregate, 1, 1, None), |
| 159 | + // ── Document helpers ── |
| 160 | + s("doc_get", Scalar, 2, 3, None), |
| 161 | + s("doc_exists", Scalar, 2, 2, None), |
| 162 | + s("doc_array_contains", Scalar, 3, 3, None), |
| 163 | + s("nav", Scalar, 2, 2, None), |
| 164 | + // ── Utility ── |
| 165 | + s("chunk_text", Scalar, 2, 3, None), |
| 166 | + s("currency", Scalar, 1, 2, None), |
| 167 | + s("distribute", Scalar, 2, 3, None), |
| 168 | + s("allocate", Scalar, 2, 3, None), |
| 169 | + s("resolve_permission", Scalar, 2, 3, None), |
| 170 | + // ── Standard scalar ── |
| 171 | + s("coalesce", Scalar, 1, 255, None), |
| 172 | + s("nullif", Scalar, 2, 2, None), |
| 173 | + s("abs", Scalar, 1, 1, None), |
| 174 | + s("ceil", Scalar, 1, 1, None), |
| 175 | + s("floor", Scalar, 1, 1, None), |
| 176 | + s("round", Scalar, 1, 2, None), |
| 177 | + s("lower", Scalar, 1, 1, None), |
| 178 | + s("upper", Scalar, 1, 1, None), |
| 179 | + s("length", Scalar, 1, 1, None), |
| 180 | + s("trim", Scalar, 1, 1, None), |
| 181 | + s("substring", Scalar, 2, 3, None), |
| 182 | + s("concat", Scalar, 1, 255, None), |
| 183 | + s("replace", Scalar, 3, 3, None), |
| 184 | + s("now", Scalar, 0, 0, None), |
| 185 | + s("current_timestamp", Scalar, 0, 0, None), |
| 186 | + s("make_array", Scalar, 0, 255, None), |
| 187 | + ] |
| 188 | +} |
| 189 | + |
| 190 | +#[cfg(test)] |
| 191 | +mod tests { |
| 192 | + use super::*; |
| 193 | + |
| 194 | + #[test] |
| 195 | + fn lookup_builtin() { |
| 196 | + let reg = FunctionRegistry::new(); |
| 197 | + assert!(reg.is_aggregate("COUNT")); |
| 198 | + assert!(reg.is_aggregate("sum")); |
| 199 | + assert!(!reg.is_aggregate("vector_distance")); |
| 200 | + assert!(reg.is_window("row_number")); |
| 201 | + assert_eq!( |
| 202 | + reg.search_trigger("vector_distance"), |
| 203 | + SearchTrigger::VectorSearch |
| 204 | + ); |
| 205 | + assert_eq!( |
| 206 | + reg.search_trigger("st_dwithin"), |
| 207 | + SearchTrigger::SpatialDWithin |
| 208 | + ); |
| 209 | + assert_eq!(reg.search_trigger("time_bucket"), SearchTrigger::TimeBucket); |
| 210 | + } |
| 211 | +} |
0 commit comments