nodedb_sql/ddl_ast/graph_parse/
fusion_params.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Shared parameter extraction for graph-vector fusion SQL surfaces.
4//!
5//! Two syntaxes reach the same `GraphOp::RagFusion` executor:
6//!
7//! - `GRAPH RAG FUSION ON <col> QUERY ARRAY[...] ...` (DSL form)
8//! - `SEARCH <col> USING FUSION(ARRAY[...] ...)` (wrapped form)
9//!
10//! They use different keyword aliases for the same parameters
11//! (`EXPANSION_DEPTH` vs `DEPTH`, `EDGE_LABEL` vs `LABEL`, `FINAL_TOP_K`
12//! vs `TOP`). Both read into the same typed bag through one left-to-right
13//! clause reader, so an unknown or misspelled parameter is refused by name.
14
15use super::super::statement::GraphDirection;
16use super::cursor::{Cursor, set_once};
17use super::tokenizer::{Tok, blank_comments, tokenize};
18use crate::error::SqlError;
19
20/// Keyword aliases for the shared fusion parameters.
21///
22/// Each fusion SQL surface picks one of the `*_KEYWORDS` constants below.
23/// New fusion variants add their own constant rather than editing the
24/// extractor.
25pub struct FusionKeywords {
26    pub vector_top_k: &'static str,
27    pub expansion_depth: &'static str,
28    pub edge_label: &'static str,
29    pub final_top_k: &'static str,
30    pub rrf_k: &'static str,
31    pub vector_field: &'static str,
32    pub direction: &'static str,
33    pub max_visited: &'static str,
34    /// Keyword that precedes `ARRAY[...]` in raw SQL (e.g. `QUERY` or
35    /// `ARRAY` itself when there is no leading keyword).
36    pub query_anchor: &'static str,
37    /// Keyword that precedes the BM25 query string for three-source fusion.
38    /// Empty string disables BM25 parsing for surfaces that do not support it.
39    pub bm25_query: &'static str,
40    /// Keyword that precedes the BM25 field name in three-source fusion.
41    pub bm25_field: &'static str,
42}
43
44/// Keywords used by `GRAPH RAG FUSION ON ...`.
45pub const RAG_FUSION_KEYWORDS: FusionKeywords = FusionKeywords {
46    vector_top_k: "VECTOR_TOP_K",
47    expansion_depth: "EXPANSION_DEPTH",
48    edge_label: "EDGE_LABEL",
49    final_top_k: "FINAL_TOP_K",
50    rrf_k: "RRF_K",
51    vector_field: "VECTOR_FIELD",
52    direction: "DIRECTION",
53    max_visited: "MAX_VISITED",
54    query_anchor: "QUERY",
55    bm25_query: "BM25",
56    bm25_field: "ON",
57};
58
59/// Keywords used by `SEARCH ... USING FUSION(...)`.
60pub const SEARCH_FUSION_KEYWORDS: FusionKeywords = FusionKeywords {
61    vector_top_k: "VECTOR_TOP_K",
62    expansion_depth: "DEPTH",
63    edge_label: "LABEL",
64    final_top_k: "TOP",
65    rrf_k: "RRF_K",
66    vector_field: "VECTOR_FIELD",
67    direction: "DIRECTION",
68    max_visited: "MAX_VISITED",
69    query_anchor: "ARRAY",
70    bm25_query: "BM25",
71    bm25_field: "ON",
72};
73
74/// Typed parameter bag for every graph-vector fusion SQL surface.
75///
76/// All fields are optional at parse time — bounds, caps, and
77/// "absent but required" errors are enforced at the pgwire boundary.
78///
79/// Three-source fusion (vector + text + graph) is enabled by populating
80/// `bm25_query` and `bm25_field` together with `rrf_k_triple`. When only
81/// `rrf_k` is set (two values), behaviour is unchanged from the two-source
82/// form. When `rrf_k_triple` is set it takes precedence and the BM25 leg
83/// participates in the fusion.
84#[derive(Debug, Clone, Default, PartialEq)]
85pub struct FusionParams {
86    pub query_vector: Option<Vec<f32>>,
87    pub vector_top_k: Option<usize>,
88    pub expansion_depth: Option<usize>,
89    pub edge_label: Option<String>,
90    pub final_top_k: Option<usize>,
91    /// Two-source RRF k constants: `(vector_k, graph_k)`. Used when no
92    /// `bm25_query` is present (backwards-compatible two-source form).
93    pub rrf_k: Option<(f64, f64)>,
94    /// Three-source RRF k constants: `(vector_k, text_k, graph_k)`. Set
95    /// when `RRF_K (kv, kt, kg)` is parsed and three values are found.
96    pub rrf_k_triple: Option<(f64, f64, f64)>,
97    pub vector_field: Option<String>,
98    pub direction: Option<GraphDirection>,
99    pub max_visited: Option<usize>,
100    /// BM25 query string for the text leg of three-source fusion. Parsed
101    /// from `BM25 'query string' ON 'field_name'` in the FUSION DSL.
102    pub bm25_query: Option<String>,
103    /// Document field on which BM25 scoring is applied in three-source fusion.
104    pub bm25_field: Option<String>,
105}
106
107impl FusionParams {
108    /// Read the fusion clauses left to right until the end of the statement.
109    ///
110    /// A keyword outside `kw` is refused by name, a clause given twice is
111    /// refused, and a value that does not read as its type is refused.
112    pub(super) fn parse(c: &mut Cursor<'_, '_>, kw: &FusionKeywords) -> Result<Self, SqlError> {
113        let stmt = c.statement();
114        let mut p = Self::default();
115        while let Some(clause) = c.next_clause()? {
116            let k = clause.as_str();
117            if k == kw.query_anchor {
118                if kw.query_anchor != "ARRAY" && !c.eat_keyword("ARRAY") {
119                    return Err(SqlError::Parse {
120                        detail: format!("{k} needs ARRAY[...] in {stmt}"),
121                    });
122                }
123                let vector = parse_vector(c.list(k)?, k)?;
124                set_once(&mut p.query_vector, vector, k, stmt)?;
125            } else if k == kw.vector_top_k {
126                set_once(&mut p.vector_top_k, c.usize(k)?, k, stmt)?;
127            } else if k == kw.expansion_depth {
128                set_once(&mut p.expansion_depth, c.usize(k)?, k, stmt)?;
129            } else if k == kw.edge_label {
130                set_once(&mut p.edge_label, c.value(k)?, k, stmt)?;
131            } else if k == kw.final_top_k {
132                set_once(&mut p.final_top_k, c.usize(k)?, k, stmt)?;
133            } else if k == kw.rrf_k {
134                if p.rrf_k.is_some() || p.rrf_k_triple.is_some() {
135                    return Err(SqlError::Parse {
136                        detail: format!("{k} given more than once in {stmt}"),
137                    });
138                }
139                match (c.number_word(), c.number_word(), c.number_word()) {
140                    (Some(vector_k), Some(graph_k), None) => p.rrf_k = Some((vector_k, graph_k)),
141                    (Some(vector_k), Some(text_k), Some(graph_k)) => {
142                        p.rrf_k_triple = Some((vector_k, text_k, graph_k));
143                    }
144                    _ => {
145                        return Err(SqlError::Parse {
146                            detail: format!("{k} needs two or three numbers, as {k} (60.0, 35.0)"),
147                        });
148                    }
149                }
150            } else if k == kw.vector_field {
151                set_once(&mut p.vector_field, c.value(k)?, k, stmt)?;
152            } else if k == kw.direction {
153                set_once(&mut p.direction, c.direction()?, k, stmt)?;
154            } else if k == kw.max_visited {
155                set_once(&mut p.max_visited, c.usize(k)?, k, stmt)?;
156            } else if !kw.bm25_query.is_empty() && k == kw.bm25_query {
157                set_once(&mut p.bm25_query, c.value(k)?, k, stmt)?;
158                // The field clause belongs to the BM25 clause only when it
159                // follows it directly.
160                if c.eat_keyword(kw.bm25_field) {
161                    set_once(
162                        &mut p.bm25_field,
163                        c.value(kw.bm25_field)?,
164                        kw.bm25_field,
165                        stmt,
166                    )?;
167                }
168            } else {
169                return Err(c.unknown_clause());
170            }
171        }
172        Ok(p)
173    }
174}
175
176/// Read the comma-separated numbers of an `ARRAY[...]` literal.
177fn parse_vector(inner: &str, clause: &str) -> Result<Vec<f32>, SqlError> {
178    if inner.trim().is_empty() {
179        return Err(SqlError::Parse {
180            detail: format!("{clause} ARRAY[...] needs one or more numbers"),
181        });
182    }
183    inner
184        .split(',')
185        .map(|part| {
186            let text = part.trim();
187            text.parse::<f32>().map_err(|_| SqlError::Parse {
188                detail: format!("{clause} ARRAY element must be a number — found '{text}'"),
189            })
190        })
191        .collect()
192}
193
194/// Parse `SEARCH <collection> USING FUSION(...)` into its collection name
195/// and a typed [`FusionParams`].
196///
197/// `None` means the SQL does not start with `SEARCH <collection> USING
198/// FUSION`. `Some(Err)` means it does and a clause is malformed. The body
199/// goes through the same quote- and bracket-aware tokenizer and the same
200/// clause reader as `GRAPH RAG FUSION`, so a keyword-shaped string literal
201/// (e.g. a label value `'TOP'`) cannot shadow a real parameter keyword.
202pub fn parse_search_using_fusion(sql: &str) -> Option<Result<(String, FusionParams), SqlError>> {
203    let sql = blank_comments(sql);
204    let toks = tokenize(&sql);
205    let collection = match toks.as_slice() {
206        [Tok::Word(s), Tok::Word(c), Tok::Word(u), Tok::Word(f), ..]
207            if s.eq_ignore_ascii_case("SEARCH")
208                && u.eq_ignore_ascii_case("USING")
209                && f.eq_ignore_ascii_case("FUSION") =>
210        {
211            (*c).to_string()
212        }
213        _ => return None,
214    };
215    let mut cursor = Cursor::new(&toks, "SEARCH USING FUSION", 4);
216    Some(
217        FusionParams::parse(&mut cursor, &SEARCH_FUSION_KEYWORDS)
218            .map(|params| (collection, params)),
219    )
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225
226    #[test]
227    fn search_fusion_full_surface_parses() {
228        let (col, p) = parse_search_using_fusion(
229            "SEARCH mycol USING FUSION(ARRAY[0.1, 0.2] VECTOR_TOP_K 5 DEPTH 2 \
230             LABEL 'related' TOP 10 RRF_K (60.0, 35.0))",
231        )
232        .expect("SEARCH USING FUSION form")
233        .expect("well-formed fusion clauses");
234        assert_eq!(col, "mycol");
235        assert_eq!(p.query_vector.as_deref().map(<[f32]>::len), Some(2));
236        assert_eq!(p.vector_top_k, Some(5));
237        assert_eq!(p.expansion_depth, Some(2));
238        assert_eq!(p.edge_label.as_deref(), Some("related"));
239        assert_eq!(p.final_top_k, Some(10));
240        assert_eq!(p.rrf_k, Some((60.0, 35.0)));
241        assert_eq!(p.rrf_k_triple, None);
242    }
243
244    #[test]
245    fn search_fusion_three_source_parses() {
246        let (col, p) = parse_search_using_fusion(
247            "SEARCH entities USING FUSION(ARRAY[0.1, 0.3] VECTOR_FIELD 'embedding' \
248             VECTOR_TOP_K 50 BM25 'transformer attention' ON 'body' \
249             DEPTH 2 LABEL 'related_to' TOP 10 RRF_K (60.0, 35.0, 50.0))",
250        )
251        .expect("SEARCH USING FUSION form")
252        .expect("well-formed fusion clauses");
253        assert_eq!(col, "entities");
254        assert_eq!(p.rrf_k, None);
255        assert_eq!(p.rrf_k_triple, Some((60.0, 35.0, 50.0)));
256        assert_eq!(p.bm25_query.as_deref(), Some("transformer attention"));
257        assert_eq!(p.bm25_field.as_deref(), Some("body"));
258        assert_eq!(p.expansion_depth, Some(2));
259        assert_eq!(p.edge_label.as_deref(), Some("related_to"));
260        assert_eq!(p.final_top_k, Some(10));
261    }
262
263    #[test]
264    fn search_fusion_label_literal_that_shadows_top_keyword() {
265        // A quoted label value containing the `TOP` keyword must not be
266        // misread as the `TOP` numeric parameter — the tokenizer keeps
267        // quoted strings whole, so `TOP 10` is the real parameter.
268        let (_, p) =
269            parse_search_using_fusion("SEARCH c USING FUSION(ARRAY[0.5] LABEL 'TOP_SECRET' TOP 7)")
270                .expect("SEARCH USING FUSION form")
271                .expect("well-formed fusion clauses");
272        assert_eq!(p.edge_label.as_deref(), Some("TOP_SECRET"));
273        assert_eq!(p.final_top_k, Some(7));
274    }
275
276    #[test]
277    fn search_fusion_rejects_wrong_prefix() {
278        assert!(parse_search_using_fusion("INSERT INTO x VALUES (1)").is_none());
279        assert!(parse_search_using_fusion("SEARCH x USING VECTOR(ARRAY[1.0])").is_none());
280    }
281
282    fn search_error(sql: &str) -> String {
283        parse_search_using_fusion(sql)
284            .expect("SEARCH USING FUSION form")
285            .expect_err("malformed fusion clauses must be refused")
286            .to_string()
287    }
288
289    #[test]
290    fn search_fusion_refuses_misspelled_parameter() {
291        let error = search_error("SEARCH c USING FUSION(ARRAY[0.5] TOPP 7)");
292        assert!(
293            error.contains("unexpected token 'TOPP' in SEARCH USING FUSION"),
294            "{error}"
295        );
296    }
297
298    #[test]
299    fn search_fusion_refuses_non_numeric_vector_element() {
300        let error = search_error("SEARCH c USING FUSION(ARRAY[0.1, x] TOP 3)");
301        assert!(error.contains("'x'"), "{error}");
302    }
303
304    #[test]
305    fn search_fusion_refuses_unknown_direction() {
306        let error = search_error("SEARCH c USING FUSION(ARRAY[0.1] DIRECTION sideways)");
307        assert!(error.contains("sideways"), "{error}");
308    }
309
310    #[test]
311    fn search_fusion_refuses_repeated_parameter() {
312        let error = search_error("SEARCH c USING FUSION(ARRAY[0.1] TOP 3 TOP 4)");
313        assert!(error.contains("TOP given more than once"), "{error}");
314    }
315}
nodedb_sql/ddl_ast/graph_parse/fusion_params.rs

nodedb_sql/ddl_ast/graph_parse/
fusion_params.rs
