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,
64    /// fragment_size])` - markup search results around matched terms.
65    UQAHighlight,
66    /// `uqa_facets(field [, field2, ...])` - facet counts over the
67    /// posting list, computed against the current row context.
68    UQAFacets,
69    /// `traverse_match(graph, start, label, max_hops)` - BFS traversal
70    /// emitting `(doc_id, score)` weighted by hop distance.
71    TraverseMatch,
72    /// `temporal_traverse(graph, start, label, max_hops, t_min, t_max)`
73    /// - `traverse_match` filtered by edge `valid_from`/`valid_to`.
74    TemporalTraverse,
75    /// `rpq(expr, start [, graph])` - evaluate a Regular Path Query
76    /// (Definition 5.1.2) and emit endpoint vertex ids reachable from
77    /// `start` along paths matching `expr`.
78    RPQ,
79    /// `graph_create(graph_name)` - register a new in-memory graph.
80    GraphCreate,
81    /// `graph_drop(graph_name)` - drop a registered graph.
82    GraphDrop,
83    /// `graph_exists(graph_name)` - AGE agtype boolean graph probe.
84    GraphExists,
85    /// `create_vlabel(graph_name, label_name)` /
86    /// `create_elabel(graph_name, label_name)` - register an AGE label.
87    GraphLabelCreate,
88    /// `drop_label(graph_name, label_name [, force])` - drop an AGE label
89    /// together with its entities.
90    GraphLabelDrop,
91    /// `alter_graph(graph_name, operation, new_value)` - AGE graph
92    /// alteration (`RENAME`).
93    GraphAlter,
94    /// `graph_edges(graph_name [, label])` - emit every edge in the
95    /// graph as `(source, target, label, weight)` rows.
96    GraphEdges,
97    /// `attention(signal_1, signal_2, ...)` - multi-signal attention
98    /// fusion (single-head).
99    AttentionFusion,
100    /// `learned_fusion(model, signal_1, ...)` - learned per-feature
101    /// weight fusion using a saved `LearnedFusion` model.
102    LearnedFusion,
103    /// `calibrated_vector_match(field, vector, k [, threshold])` -
104    /// KNN with calibrated cosine probabilities (Paper 5).
105    CalibratedVectorMatch,
106    /// `sparse_threshold(signal, threshold)` - drop scores at or below
107    /// the threshold and subtract it from survivors.
108    SparseThreshold,
109    /// `score_bm25([field,] query)` - projection helper exposing the
110    /// current match score.
111    ScoreBM25,
112    /// `score_bayesian_bm25([field,] query)` - projection helper
113    /// exposing the current Bayesian BM25 match score.
114    ScoreBayesianBM25,
115    /// `deep_learn(model, training_set)` - kick off analytical
116    /// training (Paper 4) for the named deep-fusion model.
117    DeepLearn,
118    /// Deep-fusion construction helpers used inside `deep_learn` /
119    /// `deep_predict` argument expressions:
120    Convolve,
121    Pool,
122    Flatten,
123    Dense,
124    Softmax,
125    Layer,
126    Model,
127}
128
129fn registry() -> &'static BTreeMap<&'static str, FunctionKind> {
130    static R: OnceLock<BTreeMap<&'static str, FunctionKind>> = OnceLock::new();
131    R.get_or_init(|| {
132        let mut m = BTreeMap::new();
133        m.insert("text_match", FunctionKind::TextMatch);
134        m.insert("fts_match", FunctionKind::FTSMatch);
135        m.insert("bayesian_match", FunctionKind::BayesianMatch);
136        m.insert(
137            "bayesian_match_with_prior",
138            FunctionKind::BayesianMatchWithPrior,
139        );
140        m.insert("knn_match", FunctionKind::KNNMatch);
141        m.insert("fuse_log_odds", FunctionKind::FuseLogOdds);
142        m.insert("pool_positive_evidence", FunctionKind::PositiveEvidencePool);
143        m.insert(
144            "fuse_bayesian_evidence",
145            FunctionKind::BayesianEvidenceFusion,
146        );
147        m.insert("graph_pagerank", FunctionKind::GraphPagerank);
148        m.insert("pagerank", FunctionKind::GraphPagerank);
149        m.insert("graph_hits", FunctionKind::GraphHits);
150        m.insert("hits", FunctionKind::GraphHits);
151        m.insert("graph_betweenness", FunctionKind::GraphBetweenness);
152        m.insert("betweenness", FunctionKind::GraphBetweenness);
153        m.insert("graph_traverse", FunctionKind::GraphTraverse);
154        m.insert("graph_neighbors", FunctionKind::GraphNeighbors);
155        m.insert("multi_field_match", FunctionKind::MultiFieldMatch);
156        m.insert("staged_retrieval", FunctionKind::StagedRetrieval);
157        m.insert("deep_predict", FunctionKind::DeepPredict);
158        m.insert("uqa_highlight", FunctionKind::UQAHighlight);
159        m.insert("uqa_facets", FunctionKind::UQAFacets);
160        m.insert("traverse_match", FunctionKind::TraverseMatch);
161        m.insert("temporal_traverse", FunctionKind::TemporalTraverse);
162        m.insert("rpq", FunctionKind::RPQ);
163        m.insert("graph_create", FunctionKind::GraphCreate);
164        m.insert("create_graph", FunctionKind::GraphCreate);
165        m.insert("graph_drop", FunctionKind::GraphDrop);
166        m.insert("drop_graph", FunctionKind::GraphDrop);
167        m.insert("graph_exists", FunctionKind::GraphExists);
168        m.insert("create_vlabel", FunctionKind::GraphLabelCreate);
169        m.insert("create_elabel", FunctionKind::GraphLabelCreate);
170        m.insert("drop_label", FunctionKind::GraphLabelDrop);
171        m.insert("alter_graph", FunctionKind::GraphAlter);
172        m.insert("graph_edges", FunctionKind::GraphEdges);
173        m.insert("attention", FunctionKind::AttentionFusion);
174        m.insert("fuse_attention", FunctionKind::AttentionFusion);
175        m.insert("fuse_multihead", FunctionKind::AttentionFusion);
176        m.insert("learned_fusion", FunctionKind::LearnedFusion);
177        m.insert("fuse_learned", FunctionKind::LearnedFusion);
178        m.insert(
179            "calibrated_vector_match",
180            FunctionKind::CalibratedVectorMatch,
181        );
182        m.insert("sparse_threshold", FunctionKind::SparseThreshold);
183        m.insert("score_bm25", FunctionKind::ScoreBM25);
184        m.insert("score_bayesian_bm25", FunctionKind::ScoreBayesianBM25);
185        m.insert("deep_learn", FunctionKind::DeepLearn);
186        m.insert("convolve", FunctionKind::Convolve);
187        m.insert("pool", FunctionKind::Pool);
188        m.insert("flatten", FunctionKind::Flatten);
189        m.insert("dense", FunctionKind::Dense);
190        m.insert("softmax", FunctionKind::Softmax);
191        m.insert("layer", FunctionKind::Layer);
192        m.insert("model", FunctionKind::Model);
193        m
194    })
195}
196
197pub fn lookup(name: &str) -> Option<FunctionKind> {
198    if name.bytes().any(|byte| byte.is_ascii_uppercase()) {
199        registry().get(name.to_ascii_lowercase().as_str()).copied()
200    } else {
201        registry().get(name).copied()
202    }
203}
204
205pub fn is_registered(name: &str) -> bool {
206    lookup(name).is_some()
207}
208
209/// Whether a row-producing operator function owns an explicit relation
210/// argument. These functions are compiled specially so their first SQL
211/// argument is a relation identifier rather than a scalar value.
212pub fn is_operator_join_table_function(name: &str) -> bool {
213    if name.contains('.') {
214        return false;
215    }
216    matches!(
217        name.to_ascii_lowercase().as_str(),
218        "text_similarity_join"
219            | "vector_similarity_join"
220            | "graph_join"
221            | "hybrid_join"
222            | "cross_paradigm_join"
223    )
224}
225
226/// Sorted list of registered SQL function names. CLI completion and
227/// documentation generators should consume this instead of duplicating
228/// function names.
229pub fn registered_names() -> Vec<&'static str> {
230    registry().keys().copied().collect()
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    #[test]
238    fn known_names_resolve() {
239        assert_eq!(lookup("text_match"), Some(FunctionKind::TextMatch));
240        assert_eq!(lookup("KNN_MATCH"), Some(FunctionKind::KNNMatch));
241        assert_eq!(lookup("fuse_log_odds"), Some(FunctionKind::FuseLogOdds));
242        assert_eq!(
243            lookup("pool_positive_evidence"),
244            Some(FunctionKind::PositiveEvidencePool)
245        );
246        assert_eq!(
247            lookup("fuse_bayesian_evidence"),
248            Some(FunctionKind::BayesianEvidenceFusion)
249        );
250        assert!(registered_names().contains(&"deep_predict"));
251    }
252
253    #[test]
254    fn unknown_returns_none() {
255        assert_eq!(lookup("does_not_exist"), None);
256    }
257
258    #[test]
259    fn operator_join_table_functions_are_classified_separately() {
260        assert!(is_operator_join_table_function("vector_similarity_join"));
261        assert!(is_operator_join_table_function("GRAPH_JOIN"));
262        assert!(!is_operator_join_table_function(
263            "app.vector_similarity_join"
264        ));
265        assert!(!is_operator_join_table_function("knn_match"));
266    }
267}