Skip to main content

uqa_sql/
registry.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Function registry: maps SQL function names called inside `WHERE` /
8//! projections to UQA-side semantics (text match, vector knn, hybrid
9//! fusion, ...).
10//!
11//! The registry only **classifies** a function by name; the compiler
12//! dispatches the actual operator construction once the call signature
13//! is bound to its arguments.
14
15use std::collections::BTreeMap;
16use std::sync::OnceLock;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum FunctionKind {
20    /// `text_match(field, query_string)` - BM25 text retrieval.
21    TextMatch,
22    /// `field @@ query` - full-text query-string parser over text and
23    /// vector signals.
24    FTSMatch,
25    /// `bayesian_match(field, query_string)` - Bayesian BM25 retrieval.
26    BayesianMatch,
27    /// `bayesian_match_with_prior(field, query, prior_field, mode)` -
28    /// Bayesian BM25 adjusted by a document-level external prior.
29    BayesianMatchWithPrior,
30    /// `knn_match(field, query_vector, k)` - top-k cosine KNN.
31    KNNMatch,
32    /// `fuse_log_odds(signal_1, signal_2, ...)` - exact signed
33    /// log-likelihood-ratio addition with one relevance prior.
34    FuseLogOdds,
35    /// `pool_positive_evidence(signal_1, signal_2, ...)` - gated,
36    /// confidence-scaled retrieval pooling without a calibration theorem.
37    PositiveEvidencePool,
38    /// `fuse_bayesian_evidence(signal_1, signal_2, ...)` - exact signed
39    /// log-likelihood-ratio addition with one relevance prior.
40    BayesianEvidenceFusion,
41    /// `graph_pagerank([graph_name])` - `PageRank` over a named graph.
42    GraphPagerank,
43    /// `graph_hits([graph_name])` - `HITS` over a named graph.
44    GraphHits,
45    /// `graph_betweenness([graph_name])` - betweenness centrality over a
46    /// named graph.
47    GraphBetweenness,
48    /// `graph_traverse(graph_name, start_vertex, label, max_hops)` -
49    /// BFS traversal scoring.
50    GraphTraverse,
51    /// `graph_neighbors(graph_name, vertex_id, label, direction)` -
52    /// 1-hop neighbor expansion.
53    GraphNeighbors,
54    /// `multi_field_match(field_1, query_1, field_2, query_2, ...)` -
55    /// per-field Bayesian BM25 probabilities fused with log-odds conjunction.
56    MultiFieldMatch,
57    /// `staged_retrieval(field_1, query_1, top_k_1, field_2, query_2,
58    /// top_k_2, ...)` - cascading BM25 `text_match`: each stage filters the
59    /// candidate set from the previous stage and keeps top-k.
60    StagedRetrieval,
61    /// `deep_predict(model_name)` - runs the saved deep-fusion model.
62    DeepPredict,
63    /// `uqa_highlight(field, query [, start_tag, end_tag, max_fragments, fragment_size, analyzer])` - markup search results around matched terms.
64    UQAHighlight,
65    /// `uqa_facets(field [, field2, ...])` - facet counts over the
66    /// posting list, computed against the current row context.
67    UQAFacets,
68    /// `traverse_match(graph, start, label, max_hops)` - BFS traversal
69    /// emitting `(doc_id, score)` weighted by hop distance.
70    TraverseMatch,
71    /// `temporal_traverse(graph, start, label, max_hops, t_min, t_max)`
72    /// - `traverse_match` filtered by edge `valid_from`/`valid_to`.
73    TemporalTraverse,
74    /// `rpq(expr, start [, graph])` - evaluate a Regular Path Query
75    /// (Definition 5.1.2) and emit endpoint vertex ids reachable from
76    /// `start` along paths matching `expr`.
77    RPQ,
78    /// `graph_create(graph_name)` - register a new in-memory graph.
79    GraphCreate,
80    /// `graph_drop(graph_name)` - drop a registered graph.
81    GraphDrop,
82    /// `graph_exists(graph_name)` - AGE agtype boolean graph probe.
83    GraphExists,
84    /// `create_vlabel(graph_name, label_name)` /
85    /// `create_elabel(graph_name, label_name)` - register an AGE label.
86    GraphLabelCreate,
87    /// `drop_label(graph_name, label_name [, force])` - drop an AGE label
88    /// together with its entities.
89    GraphLabelDrop,
90    /// `alter_graph(graph_name, operation, new_value)` - AGE graph
91    /// alteration (`RENAME`).
92    GraphAlter,
93    /// `graph_edges(graph_name [, label])` - emit every edge in the
94    /// graph as `(source, target, label, weight)` rows.
95    GraphEdges,
96    /// `attention(signal_1, signal_2, ...)` - multi-signal attention
97    /// fusion (single-head).
98    AttentionFusion,
99    /// `learned_fusion(model, signal_1, ...)` - learned per-feature
100    /// weight fusion using a saved `LearnedFusion` model.
101    LearnedFusion,
102    /// `calibrated_vector_match(field, vector, k [, threshold])` -
103    /// KNN with calibrated cosine probabilities (Paper 5).
104    CalibratedVectorMatch,
105    /// `sparse_threshold(signal, threshold)` - drop scores at or below
106    /// the threshold and subtract it from survivors.
107    SparseThreshold,
108    /// `score_bm25([field,] query)` - projection helper exposing the
109    /// current match score.
110    ScoreBM25,
111    /// `score_bayesian_bm25([field,] query)` - projection helper
112    /// exposing the current Bayesian BM25 match score.
113    ScoreBayesianBM25,
114    /// `deep_learn(model, training_set)` - kick off analytical
115    /// training (Paper 4) for the named deep-fusion model.
116    DeepLearn,
117    /// Deep-fusion construction helpers used inside `deep_learn` /
118    /// `deep_predict` argument expressions:
119    Convolve,
120    Pool,
121    Flatten,
122    Dense,
123    Softmax,
124    Layer,
125    Model,
126}
127
128fn registry() -> &'static BTreeMap<&'static str, FunctionKind> {
129    static R: OnceLock<BTreeMap<&'static str, FunctionKind>> = OnceLock::new();
130    R.get_or_init(|| {
131        let mut m = BTreeMap::new();
132        m.insert("text_match", FunctionKind::TextMatch);
133        m.insert("fts_match", FunctionKind::FTSMatch);
134        m.insert("bayesian_match", FunctionKind::BayesianMatch);
135        m.insert(
136            "bayesian_match_with_prior",
137            FunctionKind::BayesianMatchWithPrior,
138        );
139        m.insert("knn_match", FunctionKind::KNNMatch);
140        m.insert("fuse_log_odds", FunctionKind::FuseLogOdds);
141        m.insert("pool_positive_evidence", FunctionKind::PositiveEvidencePool);
142        m.insert(
143            "fuse_bayesian_evidence",
144            FunctionKind::BayesianEvidenceFusion,
145        );
146        m.insert("graph_pagerank", FunctionKind::GraphPagerank);
147        m.insert("pagerank", FunctionKind::GraphPagerank);
148        m.insert("graph_hits", FunctionKind::GraphHits);
149        m.insert("hits", FunctionKind::GraphHits);
150        m.insert("graph_betweenness", FunctionKind::GraphBetweenness);
151        m.insert("betweenness", FunctionKind::GraphBetweenness);
152        m.insert("graph_traverse", FunctionKind::GraphTraverse);
153        m.insert("graph_neighbors", FunctionKind::GraphNeighbors);
154        m.insert("multi_field_match", FunctionKind::MultiFieldMatch);
155        m.insert("staged_retrieval", FunctionKind::StagedRetrieval);
156        m.insert("deep_predict", FunctionKind::DeepPredict);
157        m.insert("uqa_highlight", FunctionKind::UQAHighlight);
158        m.insert("uqa_facets", FunctionKind::UQAFacets);
159        m.insert("traverse_match", FunctionKind::TraverseMatch);
160        m.insert("temporal_traverse", FunctionKind::TemporalTraverse);
161        m.insert("rpq", FunctionKind::RPQ);
162        m.insert("graph_create", FunctionKind::GraphCreate);
163        m.insert("create_graph", FunctionKind::GraphCreate);
164        m.insert("graph_drop", FunctionKind::GraphDrop);
165        m.insert("drop_graph", FunctionKind::GraphDrop);
166        m.insert("graph_exists", FunctionKind::GraphExists);
167        m.insert("create_vlabel", FunctionKind::GraphLabelCreate);
168        m.insert("create_elabel", FunctionKind::GraphLabelCreate);
169        m.insert("drop_label", FunctionKind::GraphLabelDrop);
170        m.insert("alter_graph", FunctionKind::GraphAlter);
171        m.insert("graph_edges", FunctionKind::GraphEdges);
172        m.insert("attention", FunctionKind::AttentionFusion);
173        m.insert("fuse_attention", FunctionKind::AttentionFusion);
174        m.insert("fuse_multihead", FunctionKind::AttentionFusion);
175        m.insert("learned_fusion", FunctionKind::LearnedFusion);
176        m.insert("fuse_learned", FunctionKind::LearnedFusion);
177        m.insert(
178            "calibrated_vector_match",
179            FunctionKind::CalibratedVectorMatch,
180        );
181        m.insert("sparse_threshold", FunctionKind::SparseThreshold);
182        m.insert("score_bm25", FunctionKind::ScoreBM25);
183        m.insert("score_bayesian_bm25", FunctionKind::ScoreBayesianBM25);
184        m.insert("deep_learn", FunctionKind::DeepLearn);
185        m.insert("convolve", FunctionKind::Convolve);
186        m.insert("pool", FunctionKind::Pool);
187        m.insert("flatten", FunctionKind::Flatten);
188        m.insert("dense", FunctionKind::Dense);
189        m.insert("softmax", FunctionKind::Softmax);
190        m.insert("layer", FunctionKind::Layer);
191        m.insert("model", FunctionKind::Model);
192        m
193    })
194}
195
196pub fn lookup(name: &str) -> Option<FunctionKind> {
197    if name.bytes().any(|byte| byte.is_ascii_uppercase()) {
198        registry().get(name.to_ascii_lowercase().as_str()).copied()
199    } else {
200        registry().get(name).copied()
201    }
202}
203
204pub fn is_registered(name: &str) -> bool {
205    lookup(name).is_some()
206}
207
208/// Whether a row-producing operator function owns an explicit relation
209/// argument. These functions are compiled specially so their first SQL
210/// argument is a relation identifier rather than a scalar value.
211pub fn is_operator_join_table_function(name: &str) -> bool {
212    if name.contains('.') {
213        return false;
214    }
215    matches!(
216        name.to_ascii_lowercase().as_str(),
217        "text_similarity_join"
218            | "vector_similarity_join"
219            | "graph_join"
220            | "hybrid_join"
221            | "cross_paradigm_join"
222    )
223}
224
225/// Sorted list of registered SQL function names. CLI completion and
226/// documentation generators should consume this instead of duplicating
227/// function names.
228pub fn registered_names() -> Vec<&'static str> {
229    registry().keys().copied().collect()
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235
236    #[test]
237    fn known_names_resolve() {
238        assert_eq!(lookup("text_match"), Some(FunctionKind::TextMatch));
239        assert_eq!(lookup("KNN_MATCH"), Some(FunctionKind::KNNMatch));
240        assert_eq!(lookup("fuse_log_odds"), Some(FunctionKind::FuseLogOdds));
241        assert_eq!(
242            lookup("pool_positive_evidence"),
243            Some(FunctionKind::PositiveEvidencePool)
244        );
245        assert_eq!(
246            lookup("fuse_bayesian_evidence"),
247            Some(FunctionKind::BayesianEvidenceFusion)
248        );
249        assert!(registered_names().contains(&"deep_predict"));
250    }
251
252    #[test]
253    fn unknown_returns_none() {
254        assert_eq!(lookup("does_not_exist"), None);
255    }
256
257    #[test]
258    fn operator_join_table_functions_are_classified_separately() {
259        assert!(is_operator_join_table_function("vector_similarity_join"));
260        assert!(is_operator_join_table_function("GRAPH_JOIN"));
261        assert!(!is_operator_join_table_function(
262            "app.vector_similarity_join"
263        ));
264        assert!(!is_operator_join_table_function("knn_match"));
265    }
266}