reddb_rql/parser/
vector.rs

1//! Vector query parsing (VECTOR SEARCH ... SIMILAR TO ...)
2
3use super::error::ParseError;
4use super::Parser;
5use crate::ast::{QueryExpr, VectorQuery, VectorSource};
6use crate::lexer::Token;
7use reddb_types::distance::DistanceMetric;
8use reddb_types::vector_metadata::{MetadataFilter, MetadataValue};
9
10impl<'a> Parser<'a> {
11    /// Parse VECTOR SEARCH ... SIMILAR TO ... query
12    ///
13    /// Syntax:
14    /// ```text
15    /// VECTOR SEARCH collection
16    /// SIMILAR TO [0.1, 0.2, ...] | 'text query' | (subquery)
17    /// [WHERE metadata conditions]
18    /// [METRIC L2|COSINE|INNER_PRODUCT]
19    /// [THRESHOLD 0.5]
20    /// [INCLUDE VECTORS] [INCLUDE METADATA]
21    /// [LIMIT k]
22    /// ```
23    pub fn parse_vector_query(&mut self) -> Result<QueryExpr, ParseError> {
24        self.expect(Token::Vector)?;
25        self.expect(Token::Search)?;
26
27        // Collection name
28        let collection = self.expect_ident()?;
29
30        // SIMILAR TO clause
31        self.expect(Token::Similar)?;
32        self.expect(Token::To)?;
33
34        let query_vector = self.parse_vector_source()?;
35
36        // Parse optional clauses
37        let mut filter: Option<MetadataFilter> = None;
38        let mut metric: Option<DistanceMetric> = None;
39        let mut threshold: Option<f32> = None;
40        let mut include_vectors = false;
41        let mut include_metadata = false;
42        let mut k: usize = 10; // Default
43
44        // Parse optional clauses in any order
45        loop {
46            if self.consume(&Token::Where)? {
47                filter = Some(self.parse_metadata_filter()?);
48            } else if self.consume(&Token::Metric)? {
49                metric = Some(self.parse_distance_metric()?);
50            } else if self.consume(&Token::Threshold)? {
51                threshold = Some(self.parse_float()? as f32);
52            } else if self.consume(&Token::Include)? {
53                if self.consume(&Token::Vectors)? {
54                    include_vectors = true;
55                } else if self.consume(&Token::Metadata)? {
56                    include_metadata = true;
57                } else {
58                    return Err(ParseError::expected(
59                        vec!["VECTORS", "METADATA"],
60                        self.peek(),
61                        self.position(),
62                    ));
63                }
64            } else if self.consume(&Token::Limit)? {
65                k = self.parse_integer()? as usize;
66            } else if self.consume(&Token::K)? {
67                // Alternative: K = 10
68                self.expect(Token::Eq)?;
69                k = self.parse_integer()? as usize;
70            } else {
71                break;
72            }
73        }
74
75        Ok(QueryExpr::Vector(VectorQuery {
76            alias: None,
77            collection,
78            query_vector,
79            k,
80            filter,
81            metric,
82            include_vectors,
83            include_metadata,
84            threshold,
85        }))
86    }
87
88    /// Parse vector source: literal array, text, reference, or subquery
89    pub fn parse_vector_source(&mut self) -> Result<VectorSource, ParseError> {
90        match self.peek() {
91            // Literal vector: [0.1, 0.2, 0.3]
92            Token::LBracket => {
93                self.advance()?;
94                let mut values = Vec::new();
95                loop {
96                    let value = self.parse_float()?;
97                    values.push(value as f32);
98                    if !self.consume(&Token::Comma)? {
99                        break;
100                    }
101                }
102                self.expect(Token::RBracket)?;
103                Ok(VectorSource::Literal(values))
104            }
105            // Text query: 'find similar vulnerabilities'
106            Token::String(_) => {
107                let text = self.parse_string()?;
108                Ok(VectorSource::Text(text))
109            }
110            // Parenthesized source: subquery or (collection, vector_id) reference
111            Token::LParen => {
112                self.advance()?;
113                if self.vector_source_starts_subquery() {
114                    let expr = self.parse_query_expr()?;
115                    self.expect(Token::RParen)?;
116                    Ok(VectorSource::Subquery(Box::new(expr)))
117                } else {
118                    // Reference: (collection, vector_id)
119                    let collection = self.expect_ident()?;
120                    self.expect(Token::Comma)?;
121                    let vector_id = self.parse_integer()? as u64;
122                    self.expect(Token::RParen)?;
123                    Ok(VectorSource::Reference {
124                        collection,
125                        vector_id,
126                    })
127                }
128            }
129            // Reference by name: embedding_name
130            Token::Ident(_) => {
131                let name = self.expect_ident()?;
132                // Check for (collection, id) format
133                if self.consume(&Token::LParen)? {
134                    let vector_id = self.parse_integer()? as u64;
135                    self.expect(Token::RParen)?;
136                    Ok(VectorSource::Reference {
137                        collection: name,
138                        vector_id,
139                    })
140                } else {
141                    // Just a name reference, treat as text
142                    Ok(VectorSource::Text(name))
143                }
144            }
145            other => Err(ParseError::expected(
146                vec!["vector literal [...]", "string", "reference"],
147                other,
148                self.position(),
149            )),
150        }
151    }
152
153    fn vector_source_starts_subquery(&self) -> bool {
154        matches!(
155            self.peek(),
156            Token::Select
157                | Token::Match
158                | Token::Path
159                | Token::From
160                | Token::Vector
161                | Token::Hybrid
162        )
163    }
164
165    /// Parse metadata filter for vector queries
166    pub fn parse_metadata_filter(&mut self) -> Result<MetadataFilter, ParseError> {
167        self.parse_metadata_or_expr()
168    }
169
170    /// Parse OR expression in metadata filter
171    fn parse_metadata_or_expr(&mut self) -> Result<MetadataFilter, ParseError> {
172        let mut left = self.parse_metadata_and_expr()?;
173
174        while self.consume(&Token::Or)? {
175            let right = self.parse_metadata_and_expr()?;
176            left = MetadataFilter::Or(vec![left, right]);
177        }
178
179        Ok(left)
180    }
181
182    /// Parse AND expression in metadata filter
183    fn parse_metadata_and_expr(&mut self) -> Result<MetadataFilter, ParseError> {
184        let mut left = self.parse_metadata_primary()?;
185
186        while self.consume(&Token::And)? {
187            let right = self.parse_metadata_primary()?;
188            left = MetadataFilter::And(vec![left, right]);
189        }
190
191        Ok(left)
192    }
193
194    /// Parse primary metadata filter
195    fn parse_metadata_primary(&mut self) -> Result<MetadataFilter, ParseError> {
196        // Parenthesized expression
197        if self.consume(&Token::LParen)? {
198            let expr = self.parse_metadata_filter()?;
199            self.expect(Token::RParen)?;
200            return Ok(expr);
201        }
202
203        if matches!(self.peek(), Token::Ident(name) if is_vector_geo_distance_function(name)) {
204            return self.parse_metadata_geo_radius();
205        }
206
207        // field op value
208        let field = self.expect_ident()?;
209
210        // Handle different operators
211        if self.consume(&Token::Eq)? {
212            let value = self.parse_metadata_value()?;
213            Ok(MetadataFilter::Eq(field, value))
214        } else if self.consume(&Token::Ne)? {
215            let value = self.parse_metadata_value()?;
216            Ok(MetadataFilter::Ne(field, value))
217        } else if self.consume(&Token::Lt)? {
218            let value = self.parse_metadata_value()?;
219            Ok(MetadataFilter::Lt(field, value))
220        } else if self.consume(&Token::Le)? {
221            let value = self.parse_metadata_value()?;
222            Ok(MetadataFilter::Lte(field, value))
223        } else if self.consume(&Token::Gt)? {
224            let value = self.parse_metadata_value()?;
225            Ok(MetadataFilter::Gt(field, value))
226        } else if self.consume(&Token::Ge)? {
227            let value = self.parse_metadata_value()?;
228            Ok(MetadataFilter::Gte(field, value))
229        } else if self.consume(&Token::In)? {
230            self.expect(Token::LParen)?;
231            let values = self.parse_metadata_value_list()?;
232            self.expect(Token::RParen)?;
233            Ok(MetadataFilter::In(field, values))
234        } else if self.consume(&Token::Not)? {
235            self.expect(Token::In)?;
236            self.expect(Token::LParen)?;
237            let values = self.parse_metadata_value_list()?;
238            self.expect(Token::RParen)?;
239            Ok(MetadataFilter::NotIn(field, values))
240        } else if self.consume(&Token::Contains)? {
241            let value = self.parse_string()?;
242            Ok(MetadataFilter::Contains(field, value))
243        } else {
244            Err(ParseError::expected(
245                vec!["=", "<>", "<", "<=", ">", ">=", "IN", "NOT IN", "CONTAINS"],
246                self.peek(),
247                self.position(),
248            ))
249        }
250    }
251
252    fn parse_metadata_geo_radius(&mut self) -> Result<MetadataFilter, ParseError> {
253        let function = self.expect_ident()?;
254        self.expect(Token::LParen)?;
255        let field = self.expect_ident()?;
256        self.expect(Token::Comma)?;
257        let center_lat = self.parse_float()?;
258        self.expect(Token::Comma)?;
259        let center_lon = self.parse_float()?;
260        self.expect(Token::RParen)?;
261        if !self.consume(&Token::Le)? && !self.consume(&Token::Lt)? {
262            return Err(ParseError::expected(
263                vec!["<", "<="],
264                self.peek(),
265                self.position(),
266            ));
267        }
268        let radius_km = self.parse_float()?;
269        if radius_km.partial_cmp(&0.0) != Some(std::cmp::Ordering::Greater) {
270            return Err(ParseError::new(
271                format!("{function} radius must be > 0"),
272                self.position(),
273            ));
274        }
275        Ok(MetadataFilter::GeoRadius {
276            key: field,
277            center_lat,
278            center_lon,
279            radius_km,
280        })
281    }
282
283    /// Parse metadata value
284    fn parse_metadata_value(&mut self) -> Result<MetadataValue, ParseError> {
285        match self.peek() {
286            Token::String(_) => {
287                let s = self.parse_string()?;
288                Ok(MetadataValue::String(s))
289            }
290            Token::Integer(_) => {
291                let n = self.parse_integer()?;
292                Ok(MetadataValue::Integer(n))
293            }
294            Token::Float(_) => {
295                let n = self.parse_float()?;
296                Ok(MetadataValue::Float(n))
297            }
298            Token::True => {
299                self.advance()?;
300                Ok(MetadataValue::Bool(true))
301            }
302            Token::False => {
303                self.advance()?;
304                Ok(MetadataValue::Bool(false))
305            }
306            other => Err(ParseError::expected(
307                vec!["string", "number", "true", "false"],
308                other,
309                self.position(),
310            )),
311        }
312    }
313
314    /// Parse list of metadata values
315    fn parse_metadata_value_list(&mut self) -> Result<Vec<MetadataValue>, ParseError> {
316        let mut values = Vec::new();
317        loop {
318            values.push(self.parse_metadata_value()?);
319            if !self.consume(&Token::Comma)? {
320                break;
321            }
322        }
323        Ok(values)
324    }
325
326    /// Parse distance metric
327    pub fn parse_distance_metric(&mut self) -> Result<DistanceMetric, ParseError> {
328        match self.peek() {
329            Token::L2 => {
330                self.advance()?;
331                Ok(DistanceMetric::L2)
332            }
333            Token::Cosine => {
334                self.advance()?;
335                Ok(DistanceMetric::Cosine)
336            }
337            Token::InnerProduct => {
338                self.advance()?;
339                Ok(DistanceMetric::InnerProduct)
340            }
341            Token::Ident(name) => {
342                let name_upper = name.to_uppercase();
343                let name_clone = name.clone();
344                self.advance()?;
345                match name_upper.as_str() {
346                    "L2" | "EUCLIDEAN" => Ok(DistanceMetric::L2),
347                    "COSINE" | "COS" => Ok(DistanceMetric::Cosine),
348                    "INNER_PRODUCT" | "IP" | "DOT" => Ok(DistanceMetric::InnerProduct),
349                    _ => Err(ParseError::new(
350                        format!(
351                            "Unknown distance metric: {}. Valid: L2, COSINE, INNER_PRODUCT",
352                            name_clone
353                        ),
354                        self.position(),
355                    )),
356                }
357            }
358            other => Err(ParseError::expected(
359                vec!["L2", "COSINE", "INNER_PRODUCT"],
360                other,
361                self.position(),
362            )),
363        }
364    }
365}
366
367fn is_vector_geo_distance_function(name: &str) -> bool {
368    name.eq_ignore_ascii_case("GEO_DISTANCE") || name.eq_ignore_ascii_case("HAVERSINE")
369}
370
371#[cfg(test)]
372mod tests {
373    use super::*;
374
375    fn parse_query(input: &str) -> Result<QueryExpr, ParseError> {
376        crate::parser::parse(input).map(|query| query.query)
377    }
378
379    #[test]
380    fn vector_query_uses_defaults_for_bare_identifier_source() {
381        let query = parse_query("VECTOR SEARCH embeddings SIMILAR TO nearest_neighbor").unwrap();
382
383        let QueryExpr::Vector(vector) = query else {
384            panic!("expected vector query");
385        };
386        assert_eq!(vector.collection, "embeddings");
387        assert_eq!(vector.k, 10);
388        assert!(vector.filter.is_none());
389        assert_eq!(vector.metric, None);
390        assert_eq!(vector.threshold, None);
391        assert!(!vector.include_vectors);
392        assert!(!vector.include_metadata);
393        assert!(matches!(
394            vector.query_vector,
395            VectorSource::Text(text) if text == "nearest_neighbor"
396        ));
397    }
398
399    #[test]
400    fn vector_query_parses_reference_sources_and_k_alias() {
401        let query =
402            parse_query("VECTOR SEARCH embeddings SIMILAR TO docs(42) INCLUDE METADATA K = 7")
403                .unwrap();
404        let QueryExpr::Vector(vector) = query else {
405            panic!("expected vector query");
406        };
407        assert_eq!(vector.k, 7);
408        assert!(vector.include_metadata);
409        assert!(matches!(
410            vector.query_vector,
411            VectorSource::Reference {
412                collection,
413                vector_id,
414            } if collection == "docs" && vector_id == 42
415        ));
416
417        let query =
418            parse_query("VECTOR SEARCH embeddings SIMILAR TO (archive, 99) LIMIT 4").unwrap();
419        let QueryExpr::Vector(vector) = query else {
420            panic!("expected vector query");
421        };
422        assert_eq!(vector.k, 4);
423        assert!(matches!(
424            vector.query_vector,
425            VectorSource::Reference {
426                collection,
427                vector_id,
428            } if collection == "archive" && vector_id == 99
429        ));
430    }
431
432    #[test]
433    fn vector_query_parses_subquery_source() {
434        let query =
435            parse_query("VECTOR SEARCH docs SIMILAR TO (SELECT id FROM seeds) LIMIT 2").unwrap();
436
437        let QueryExpr::Vector(vector) = query else {
438            panic!("expected vector query");
439        };
440        assert_eq!(vector.collection, "docs");
441        assert_eq!(vector.k, 2);
442        match vector.query_vector {
443            VectorSource::Subquery(expr) => match *expr {
444                QueryExpr::Table(table) => assert_eq!(table.table, "seeds"),
445                other => panic!("expected table subquery, got {other:?}"),
446            },
447            other => panic!("expected subquery source, got {other:?}"),
448        }
449    }
450
451    #[test]
452    fn vector_query_parses_filter_sets_metric_threshold_and_includes() {
453        let query = parse_query(
454            "VECTOR SEARCH docs SIMILAR TO [0.1, 0.2] \
455             WHERE (source IN ('nmap', 'nessus') OR severity NOT IN (1, 2)) \
456             AND archived = false METRIC DOT THRESHOLD 0.25 INCLUDE VECTORS LIMIT 3",
457        )
458        .unwrap();
459
460        let QueryExpr::Vector(vector) = query else {
461            panic!("expected vector query");
462        };
463        assert_eq!(vector.k, 3);
464        assert_eq!(vector.metric, Some(DistanceMetric::InnerProduct));
465        assert_eq!(vector.threshold, Some(0.25));
466        assert!(vector.include_vectors);
467        assert!(
468            matches!(vector.query_vector, VectorSource::Literal(values) if values == vec![0.1, 0.2])
469        );
470
471        let Some(MetadataFilter::And(and_parts)) = vector.filter else {
472            panic!("expected AND filter");
473        };
474        assert_eq!(and_parts.len(), 2);
475        match &and_parts[0] {
476            MetadataFilter::Or(or_parts) => {
477                assert_eq!(or_parts.len(), 2);
478                assert!(matches!(
479                    &or_parts[0],
480                    MetadataFilter::In(field, values)
481                        if field == "source"
482                            && values == &vec![
483                                MetadataValue::String("nmap".to_string()),
484                                MetadataValue::String("nessus".to_string())
485                            ]
486                ));
487                assert!(matches!(
488                    &or_parts[1],
489                    MetadataFilter::NotIn(field, values)
490                        if field == "severity"
491                            && values == &vec![MetadataValue::Integer(1), MetadataValue::Integer(2)]
492                ));
493            }
494            other => panic!("expected OR filter, got {other:?}"),
495        }
496        assert!(matches!(
497            &and_parts[1],
498            MetadataFilter::Eq(field, MetadataValue::Bool(false)) if field == "archived"
499        ));
500    }
501
502    #[test]
503    fn metadata_filter_parses_comparisons_and_contains() {
504        let query = parse_query(
505            "VECTOR SEARCH docs SIMILAR TO [0.3] \
506             WHERE score < 0.7 OR rank >= 10 AND title CONTAINS 'redis'",
507        )
508        .unwrap();
509
510        let QueryExpr::Vector(vector) = query else {
511            panic!("expected vector query");
512        };
513        let Some(MetadataFilter::Or(or_parts)) = vector.filter else {
514            panic!("expected OR filter");
515        };
516        assert_eq!(or_parts.len(), 2);
517        assert!(matches!(
518            &or_parts[0],
519            MetadataFilter::Lt(field, MetadataValue::Float(value))
520                if field == "score" && (*value - 0.7).abs() < f64::EPSILON
521        ));
522        match &or_parts[1] {
523            MetadataFilter::And(and_parts) => {
524                assert_eq!(and_parts.len(), 2);
525                assert!(matches!(
526                    &and_parts[0],
527                    MetadataFilter::Gte(field, MetadataValue::Integer(10)) if field == "rank"
528                ));
529                assert!(matches!(
530                    &and_parts[1],
531                    MetadataFilter::Contains(field, value)
532                        if field == "title" && value == "redis"
533                ));
534            }
535            other => panic!("expected AND filter, got {other:?}"),
536        }
537    }
538
539    #[test]
540    fn metadata_filter_parses_geo_radius() {
541        let query = parse_query(
542            "VECTOR SEARCH restaurants SIMILAR TO [1.0, 0.0] \
543             WHERE GEO_DISTANCE(location, 48.8566, 2.3522) <= 5.0 AND cuisine = 'bistro' \
544             LIMIT 3",
545        )
546        .unwrap();
547
548        let QueryExpr::Vector(vector) = query else {
549            panic!("expected vector query");
550        };
551        let Some(MetadataFilter::And(and_parts)) = vector.filter else {
552            panic!("expected AND filter");
553        };
554        assert!(matches!(
555            &and_parts[0],
556            MetadataFilter::GeoRadius {
557                key,
558                center_lat,
559                center_lon,
560                radius_km,
561            } if key == "location"
562                && (*center_lat - 48.8566).abs() < f64::EPSILON
563                && (*center_lon - 2.3522).abs() < f64::EPSILON
564                && (*radius_km - 5.0).abs() < f64::EPSILON
565        ));
566        assert!(matches!(
567            &and_parts[1],
568            MetadataFilter::Eq(field, MetadataValue::String(value))
569                if field == "cuisine" && value == "bistro"
570        ));
571    }
572
573    #[test]
574    fn vector_parser_reports_malformed_queries() {
575        for sql in [
576            "VECTOR SEARCH docs SIMILAR TO []",
577            "VECTOR SEARCH docs SIMILAR TO [0.1] INCLUDE SCORES",
578            "VECTOR SEARCH docs SIMILAR TO [0.1] METRIC MANHATTAN",
579            "VECTOR SEARCH docs SIMILAR TO [0.1] WHERE source",
580            "VECTOR SEARCH docs SIMILAR TO (docs)",
581        ] {
582            assert!(parse_query(sql).is_err(), "{sql} should not parse");
583        }
584    }
585}
reddb_rql/parser/vector.rs

reddb_rql/parser/
vector.rs
