reddb_server/storage/query/executors/
vector.rs

1//! Vector Query Executor
2//!
3//! Executes VECTOR SEARCH queries using HNSW approximate nearest neighbor search.
4//! Supports metadata filtering, multiple distance metrics, and cross-references.
5
6use std::collections::HashMap;
7use std::sync::Arc;
8
9use crate::storage::engine::distance::{distance, DistanceMetric};
10use crate::storage::engine::hnsw::{HnswConfig, HnswIndex};
11use crate::storage::engine::unified_index::UnifiedIndex;
12use crate::storage::engine::vector_metadata::{MetadataFilter, MetadataValue};
13use crate::storage::engine::vector_store::VectorStore;
14use crate::storage::query::ast::{QueryExpr, VectorQuery, VectorSource};
15use crate::storage::query::sql_lowering::effective_vector_filter;
16use crate::storage::query::unified::{
17    ExecutionError, QueryStats, UnifiedRecord, UnifiedResult, VectorSearchResult,
18};
19use crate::storage::schema::Value;
20
21/// Vector query executor using HNSW index
22pub struct VectorExecutor {
23    /// Vector store for segment management
24    vector_store: Arc<VectorStore>,
25    /// Cross-reference index for linking vectors to nodes/rows
26    unified_index: Option<Arc<UnifiedIndex>>,
27}
28
29impl VectorExecutor {
30    /// Create a new vector executor
31    pub fn new(vector_store: Arc<VectorStore>) -> Self {
32        Self {
33            vector_store,
34            unified_index: None,
35        }
36    }
37
38    /// Create with cross-reference support
39    pub fn with_unified_index(mut self, index: Arc<UnifiedIndex>) -> Self {
40        self.unified_index = Some(index);
41        self
42    }
43
44    /// Execute a vector search query
45    pub fn execute(&self, query: &VectorQuery) -> Result<UnifiedResult, ExecutionError> {
46        let start = std::time::Instant::now();
47
48        // Resolve the query vector
49        let query_vector = self.resolve_vector_source(&query.query_vector)?;
50
51        // Get the collection
52        let collection = self.vector_store.get(&query.collection).ok_or_else(|| {
53            ExecutionError::new(format!("Vector collection not found: {}", query.collection))
54        })?;
55
56        // Search the vector store with filter
57        let search_results = collection.search_with_filter(
58            &query_vector,
59            query.k,
60            effective_vector_filter(query).as_ref(),
61        );
62
63        // Build result
64        let mut result = UnifiedResult::with_columns(vec![
65            "id".to_string(),
66            "distance".to_string(),
67            "collection".to_string(),
68        ]);
69
70        if query.include_vectors {
71            result.columns.push("vector".to_string());
72        }
73        if query.include_metadata {
74            result.columns.push("metadata".to_string());
75        }
76
77        // Convert search results to unified records
78        for sr in search_results {
79            // Apply threshold filter if specified
80            if let Some(threshold) = query.threshold {
81                if sr.distance > threshold {
82                    continue;
83                }
84            }
85
86            let mut record = UnifiedRecord::new();
87
88            // Build vector search result
89            let mut vsr = VectorSearchResult::new(sr.id, &query.collection, sr.distance);
90
91            // Include vector data if requested and available
92            if query.include_vectors {
93                if let Some(vec_data) = sr.vector {
94                    vsr = vsr.with_vector(vec_data);
95                }
96            }
97
98            // Include metadata if requested and available
99            if query.include_metadata {
100                if let Some(ref meta_entry) = sr.metadata {
101                    // Convert MetadataEntry to HashMap<String, Value>
102                    let mut meta_map: HashMap<String, Value> = HashMap::new();
103                    for (k, v) in &meta_entry.strings {
104                        meta_map.insert(k.clone(), Value::text(v.clone()));
105                    }
106                    for (k, v) in &meta_entry.integers {
107                        meta_map.insert(k.clone(), Value::Integer(*v));
108                    }
109                    for (k, v) in &meta_entry.floats {
110                        meta_map.insert(k.clone(), Value::Float(*v));
111                    }
112                    for (k, v) in &meta_entry.bools {
113                        meta_map.insert(k.clone(), Value::Boolean(*v));
114                    }
115                    vsr = vsr.with_metadata(meta_map);
116                }
117            }
118
119            // Add cross-references if available
120            if let Some(ref unified) = self.unified_index {
121                // Check for linked node
122                if let Some(node_id) = unified.get_vector_node(&query.collection, sr.id) {
123                    vsr = vsr.with_linked_node(node_id);
124                }
125
126                // Check for linked row
127                if let Some(row_key) = unified.get_vector_row(&query.collection, sr.id) {
128                    vsr = vsr.with_linked_row(&row_key.table, row_key.row_id);
129                }
130            }
131
132            // Add basic values to record
133            record.set_arc(Arc::from("id"), Value::Integer(sr.id as i64));
134            record.set_arc(Arc::from("distance"), Value::Float(sr.distance as f64));
135            record.set_arc(
136                Arc::from("collection"),
137                Value::text(query.collection.clone()),
138            );
139
140            record.vector_results.push(vsr);
141            result.push(record);
142        }
143
144        // Update stats
145        result.stats = QueryStats {
146            nodes_scanned: 0,
147            edges_scanned: 0,
148            rows_scanned: result.len() as u64,
149            exec_time_us: start.elapsed().as_micros() as u64,
150        };
151
152        Ok(result)
153    }
154
155    /// Resolve vector source to actual vector data
156    fn resolve_vector_source(&self, source: &VectorSource) -> Result<Vec<f32>, ExecutionError> {
157        match source {
158            VectorSource::Literal(vec) => Ok(vec.clone()),
159
160            VectorSource::Text(text) => {
161                // Text embedding would require an embedding model
162                // For now, return an error indicating this needs external embedding
163                Err(ExecutionError::new(format!(
164                    "Text embedding not yet implemented. Provide a literal vector or use an embedding service for: '{}'",
165                    text
166                )))
167            }
168
169            VectorSource::Reference {
170                collection,
171                vector_id,
172            } => {
173                if let Some(coll) = self.vector_store.get(collection) {
174                    coll.get(*vector_id).cloned().ok_or_else(|| {
175                        ExecutionError::new(format!(
176                            "Reference vector not found: {}:{}",
177                            collection, vector_id
178                        ))
179                    })
180                } else {
181                    Err(ExecutionError::new(format!(
182                        "Vector collection not found: {}",
183                        collection
184                    )))
185                }
186            }
187
188            VectorSource::Subquery(expr) => self.resolve_subquery_vector(expr.as_ref()),
189        }
190    }
191
192    fn resolve_subquery_vector(&self, expr: &QueryExpr) -> Result<Vec<f32>, ExecutionError> {
193        match expr {
194            QueryExpr::Vector(query) => {
195                let result = self.execute(query)?;
196                let (collection, vector_id) =
197                    vector_subquery_reference(&result.records, &query.collection)?;
198                self.resolve_vector_source(&VectorSource::Reference {
199                    collection,
200                    vector_id,
201                })
202            }
203            other => Err(ExecutionError::new(format!(
204                "Vector subqueries currently support only nested VECTOR SEARCH expressions, got {}",
205                query_expr_name(other)
206            ))),
207        }
208    }
209}
210
211/// Convert MetadataValue to Value for unified results
212fn metadata_value_to_value(mv: MetadataValue) -> Value {
213    match mv {
214        MetadataValue::String(s) => Value::text(s),
215        MetadataValue::Integer(i) => Value::Integer(i),
216        MetadataValue::Float(f) => Value::Float(f),
217        MetadataValue::Bool(b) => Value::Boolean(b),
218        MetadataValue::Null => Value::Null,
219    }
220}
221
222// ============================================================================
223// In-Memory Executor for Testing
224// ============================================================================
225
226/// Simple in-memory vector executor for testing without full VectorStore
227pub struct InMemoryVectorExecutor {
228    /// Vectors indexed by (collection, id)
229    vectors: HashMap<(String, u64), Vec<f32>>,
230    /// Metadata indexed by (collection, id)
231    metadata: HashMap<(String, u64), HashMap<String, MetadataValue>>,
232    /// HNSW indexes by collection
233    indexes: HashMap<String, HnswIndex>,
234    /// Cross-reference index
235    unified_index: Option<Arc<UnifiedIndex>>,
236}
237
238impl InMemoryVectorExecutor {
239    /// Create a new in-memory executor
240    pub fn new() -> Self {
241        Self {
242            vectors: HashMap::new(),
243            metadata: HashMap::new(),
244            indexes: HashMap::new(),
245            unified_index: None,
246        }
247    }
248
249    /// Add cross-reference support
250    pub fn with_unified_index(mut self, index: Arc<UnifiedIndex>) -> Self {
251        self.unified_index = Some(index);
252        self
253    }
254
255    /// Add a vector to a collection
256    pub fn add_vector(
257        &mut self,
258        collection: &str,
259        id: u64,
260        vector: Vec<f32>,
261        meta: Option<HashMap<String, MetadataValue>>,
262    ) {
263        let dim = vector.len();
264
265        // Store vector
266        self.vectors
267            .insert((collection.to_string(), id), vector.clone());
268
269        // Store metadata
270        if let Some(m) = meta {
271            self.metadata.insert((collection.to_string(), id), m);
272        }
273
274        // Add to HNSW index
275        let index = self
276            .indexes
277            .entry(collection.to_string())
278            .or_insert_with(|| {
279                let config = HnswConfig {
280                    m: 16,
281                    m_max0: 32,
282                    ef_construction: 200,
283                    ef_search: 50,
284                    ml: 1.0 / (16.0_f64).ln(),
285                    metric: DistanceMetric::L2,
286                };
287                HnswIndex::new(dim, config)
288            });
289
290        index.insert_with_id(id, vector.clone());
291    }
292
293    /// Execute a vector query
294    pub fn execute(&self, query: &VectorQuery) -> Result<UnifiedResult, ExecutionError> {
295        let start = std::time::Instant::now();
296
297        // Resolve query vector
298        let query_vector = match &query.query_vector {
299            VectorSource::Literal(v) => v.clone(),
300            VectorSource::Reference {
301                collection,
302                vector_id,
303            } => self
304                .vectors
305                .get(&(collection.clone(), *vector_id))
306                .cloned()
307                .ok_or_else(|| ExecutionError::new("Reference vector not found"))?,
308            VectorSource::Text(t) => {
309                return Err(ExecutionError::new(format!(
310                    "Text embedding not implemented: '{}'",
311                    t
312                )));
313            }
314            VectorSource::Subquery(expr) => self.resolve_subquery_vector(expr.as_ref())?,
315        };
316
317        let metric = query.metric.unwrap_or(DistanceMetric::L2);
318
319        // Get or create result
320        let mut result = UnifiedResult::with_columns(vec![
321            "id".to_string(),
322            "distance".to_string(),
323            "collection".to_string(),
324        ]);
325
326        // Search using HNSW if available, otherwise brute force
327        let search_results: Vec<(u64, f32)> =
328            if let Some(index) = self.indexes.get(&query.collection) {
329                // HNSW search returns DistanceResult with id and distance
330                let mut results: Vec<_> = index
331                    .search(&query_vector, query.k)
332                    .into_iter()
333                    .map(|r| (r.id, r.distance))
334                    .collect();
335                results.sort_by(|a, b| {
336                    match a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal) {
337                        std::cmp::Ordering::Equal => a.0.cmp(&b.0),
338                        ordering => ordering,
339                    }
340                });
341                results
342            } else {
343                // Brute force search
344                self.brute_force_search(&query.collection, &query_vector, query.k, metric)
345            };
346
347        for (vector_id, dist) in search_results {
348            // Apply threshold
349            if let Some(threshold) = query.threshold {
350                if dist > threshold {
351                    continue;
352                }
353            }
354
355            // Apply metadata filter
356            if let Some(ref filter) = query.filter {
357                let key = (query.collection.clone(), vector_id);
358                if let Some(meta) = self.metadata.get(&key) {
359                    if !evaluate_filter(filter, meta) {
360                        continue;
361                    }
362                } else {
363                    continue; // No metadata, filter fails
364                }
365            }
366
367            let mut record = UnifiedRecord::new();
368            let mut vsr = VectorSearchResult::new(vector_id, &query.collection, dist);
369
370            if query.include_vectors {
371                if let Some(vec) = self.vectors.get(&(query.collection.clone(), vector_id)) {
372                    vsr = vsr.with_vector(vec.clone());
373                }
374            }
375
376            if query.include_metadata {
377                if let Some(meta) = self.metadata.get(&(query.collection.clone(), vector_id)) {
378                    let meta_map: HashMap<String, Value> = meta
379                        .iter()
380                        .map(|(k, v)| (k.clone(), metadata_value_to_value(v.clone())))
381                        .collect();
382                    vsr = vsr.with_metadata(meta_map);
383                }
384            }
385
386            // Add cross-references
387            if let Some(ref unified) = self.unified_index {
388                if let Some(node_id) = unified.get_vector_node(&query.collection, vector_id) {
389                    vsr = vsr.with_linked_node(node_id);
390                }
391
392                if let Some(row_key) = unified.get_vector_row(&query.collection, vector_id) {
393                    vsr = vsr.with_linked_row(&row_key.table, row_key.row_id);
394                }
395            }
396
397            record.set_arc(Arc::from("id"), Value::Integer(vector_id as i64));
398            record.set_arc(Arc::from("distance"), Value::Float(dist as f64));
399            record.set_arc(
400                Arc::from("collection"),
401                Value::text(query.collection.clone()),
402            );
403            record.vector_results.push(vsr);
404            result.push(record);
405        }
406
407        result.stats = QueryStats {
408            nodes_scanned: 0,
409            edges_scanned: 0,
410            rows_scanned: self.vectors.len() as u64,
411            exec_time_us: start.elapsed().as_micros() as u64,
412        };
413
414        Ok(result)
415    }
416
417    fn resolve_subquery_vector(&self, expr: &QueryExpr) -> Result<Vec<f32>, ExecutionError> {
418        match expr {
419            QueryExpr::Vector(query) => {
420                let result = self.execute(query)?;
421                let (collection, vector_id) =
422                    vector_subquery_reference(&result.records, &query.collection)?;
423                self.vectors
424                    .get(&(collection, vector_id))
425                    .cloned()
426                    .ok_or_else(|| ExecutionError::new("Subquery reference vector not found"))
427            }
428            other => Err(ExecutionError::new(format!(
429                "Vector subqueries currently support only nested VECTOR SEARCH expressions, got {}",
430                query_expr_name(other)
431            ))),
432        }
433    }
434
435    /// Brute force search when no index is available
436    fn brute_force_search(
437        &self,
438        collection: &str,
439        query: &[f32],
440        k: usize,
441        metric: DistanceMetric,
442    ) -> Vec<(u64, f32)> {
443        let mut results: Vec<(u64, f32)> = self
444            .vectors
445            .iter()
446            .filter(|((c, _), _)| c == collection)
447            .map(|((_, id), vec)| {
448                let dist = distance(query, vec, metric);
449                (*id, dist)
450            })
451            .collect();
452
453        results.sort_by(
454            |a, b| match a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal) {
455                std::cmp::Ordering::Equal => a.0.cmp(&b.0),
456                ordering => ordering,
457            },
458        );
459        results.truncate(k);
460        results
461    }
462}
463
464impl Default for InMemoryVectorExecutor {
465    fn default() -> Self {
466        Self::new()
467    }
468}
469
470/// Evaluate a metadata filter against metadata values
471fn evaluate_filter(filter: &MetadataFilter, metadata: &HashMap<String, MetadataValue>) -> bool {
472    match filter {
473        MetadataFilter::Eq(field, value) => metadata
474            .get(field)
475            .map(|candidate| candidate.matches_eq(value))
476            .unwrap_or(false),
477        MetadataFilter::Ne(field, value) => metadata
478            .get(field)
479            .map(|candidate| !candidate.matches_eq(value))
480            .unwrap_or(true),
481        MetadataFilter::Lt(field, value) => metadata
482            .get(field)
483            .and_then(|candidate| candidate.compare(value))
484            .map(|ord| ord == std::cmp::Ordering::Less)
485            .unwrap_or(false),
486        MetadataFilter::Lte(field, value) => metadata
487            .get(field)
488            .and_then(|candidate| candidate.compare(value))
489            .map(|ord| ord != std::cmp::Ordering::Greater)
490            .unwrap_or(false),
491        MetadataFilter::Gt(field, value) => metadata
492            .get(field)
493            .and_then(|candidate| candidate.compare(value))
494            .map(|ord| ord == std::cmp::Ordering::Greater)
495            .unwrap_or(false),
496        MetadataFilter::Gte(field, value) => metadata
497            .get(field)
498            .and_then(|candidate| candidate.compare(value))
499            .map(|ord| ord != std::cmp::Ordering::Less)
500            .unwrap_or(false),
501        MetadataFilter::In(field, values) => metadata
502            .get(field)
503            .map(|candidate| values.iter().any(|value| candidate.matches_eq(value)))
504            .unwrap_or(false),
505        MetadataFilter::NotIn(field, values) => metadata
506            .get(field)
507            .map(|candidate| !values.iter().any(|value| candidate.matches_eq(value)))
508            .unwrap_or(true),
509        MetadataFilter::Contains(field, substring) => {
510            if let Some(MetadataValue::String(s)) = metadata.get(field) {
511                s.contains(substring)
512            } else {
513                false
514            }
515        }
516        MetadataFilter::And(filters) => filters.iter().all(|f| evaluate_filter(f, metadata)),
517        MetadataFilter::Or(filters) => filters.iter().any(|f| evaluate_filter(f, metadata)),
518        MetadataFilter::Not(inner) => !evaluate_filter(inner, metadata),
519        MetadataFilter::StartsWith(field, prefix) => {
520            if let Some(MetadataValue::String(s)) = metadata.get(field) {
521                s.starts_with(prefix)
522            } else {
523                false
524            }
525        }
526        MetadataFilter::EndsWith(field, suffix) => {
527            if let Some(MetadataValue::String(s)) = metadata.get(field) {
528                s.ends_with(suffix)
529            } else {
530                false
531            }
532        }
533        MetadataFilter::GeoRadius { .. } => false,
534        MetadataFilter::Exists(field) => metadata.contains_key(field),
535        MetadataFilter::NotExists(field) => !metadata.contains_key(field),
536    }
537}
538
539fn vector_subquery_reference(
540    records: &[UnifiedRecord],
541    default_collection: &str,
542) -> Result<(String, u64), ExecutionError> {
543    let record = records
544        .first()
545        .ok_or_else(|| ExecutionError::new("Vector subquery returned no rows"))?;
546
547    let collection: String = match record.get("collection") {
548        Some(Value::Text(collection)) => collection.to_string(),
549        _ => default_collection.to_string(),
550    };
551
552    let vector_id = match record.get("id") {
553        Some(Value::Integer(id)) if *id >= 0 => *id as u64,
554        Some(Value::UnsignedInteger(id)) => *id,
555        other => {
556            return Err(ExecutionError::new(format!(
557                "Vector subquery must expose an integer id column, got {other:?}"
558            )));
559        }
560    };
561
562    Ok((collection, vector_id))
563}
564
565fn query_expr_name(expr: &QueryExpr) -> &'static str {
566    match expr {
567        QueryExpr::Table(_) => "table",
568        QueryExpr::Graph(_) => "graph",
569        QueryExpr::Join(_) => "join",
570        QueryExpr::Path(_) => "path",
571        QueryExpr::Vector(_) => "vector",
572        QueryExpr::Hybrid(_) => "hybrid",
573        QueryExpr::Insert(_) => "insert",
574        QueryExpr::Update(_) => "update",
575        QueryExpr::Delete(_) => "delete",
576        QueryExpr::CreateTable(_) => "create_table",
577        QueryExpr::CreateCollection(_) => "create_collection",
578        QueryExpr::CreateVector(_) => "create_vector",
579        QueryExpr::DropTable(_) => "drop_table",
580        QueryExpr::DropGraph(_) => "drop_graph",
581        QueryExpr::DropVector(_) => "drop_vector",
582        QueryExpr::DropDocument(_) => "drop_document",
583        QueryExpr::DropKv(_) => "drop_kv",
584        QueryExpr::DropCollection(_) => "drop_collection",
585        QueryExpr::Truncate(_) => "truncate",
586        QueryExpr::AlterTable(_) => "alter_table",
587        QueryExpr::CreateVcsRef(_) => "create_vcs_ref",
588        QueryExpr::DropVcsRef(_) => "drop_vcs_ref",
589        QueryExpr::ForkStore(_) => "fork_store",
590        QueryExpr::PromoteFork(_) => "promote_fork",
591        QueryExpr::DropFork(_) => "drop_fork",
592        QueryExpr::VcsCommand(_) => "vcs_command",
593        QueryExpr::GraphCommand(_) => "graph_command",
594        QueryExpr::SearchCommand(_) => "search_command",
595        QueryExpr::Ask(_) => "ask",
596        QueryExpr::CreateIndex(_) => "create_index",
597        QueryExpr::DropIndex(_) => "drop_index",
598        QueryExpr::ProbabilisticCommand(_) => "probabilistic_command",
599        QueryExpr::CreateTimeSeries(_) => "create_timeseries",
600        QueryExpr::CreateMetric(_) => "create_metric",
601        QueryExpr::AlterMetric(_) => "alter_metric",
602        QueryExpr::CreateSlo(_) => "create_slo",
603        QueryExpr::DropTimeSeries(_) => "drop_timeseries",
604        QueryExpr::CreateQueue(_) => "create_queue",
605        QueryExpr::AlterQueue(_) => "alter_queue",
606        QueryExpr::DropQueue(_) => "drop_queue",
607        QueryExpr::QueueSelect(_) => "queue_select",
608        QueryExpr::QueueCommand(_) => "queue_command",
609        QueryExpr::KvCommand(_) => "kv_command",
610        QueryExpr::ConfigCommand(_) => "config_command",
611        QueryExpr::CreateTree(_) => "create_tree",
612        QueryExpr::DropTree(_) => "drop_tree",
613        QueryExpr::TreeCommand(_) => "tree_command",
614        QueryExpr::SetConfig { .. } => "set_config",
615        QueryExpr::ShowConfig { .. } => "show_config",
616        QueryExpr::Scrub { .. } => "scrub",
617        QueryExpr::SetSecret { .. } => "set_secret",
618        QueryExpr::DeleteSecret { .. } => "delete_secret",
619        QueryExpr::SetKv { .. } => "set_kv",
620        QueryExpr::DeleteKv { .. } => "delete_kv",
621        QueryExpr::ShowSecrets { .. } => "show_secrets",
622        QueryExpr::SetTenant(_) => "set_tenant",
623        QueryExpr::ShowTenant => "show_tenant",
624        QueryExpr::ExplainAlter(_) => "explain_alter",
625        QueryExpr::TransactionControl(_) => "transaction_control",
626        QueryExpr::MaintenanceCommand(_) => "maintenance_command",
627        QueryExpr::CreateSchema(_) => "create_schema",
628        QueryExpr::DropSchema(_) => "drop_schema",
629        QueryExpr::CreateSequence(_) => "create_sequence",
630        QueryExpr::DropSequence(_) => "drop_sequence",
631        QueryExpr::CopyFrom(_) => "copy_from",
632        QueryExpr::CreateView(_) => "create_view",
633        QueryExpr::DropView(_) => "drop_view",
634        QueryExpr::RefreshMaterializedView(_) => "refresh_materialized_view",
635        QueryExpr::CreatePolicy(_) => "create_policy",
636        QueryExpr::DropPolicy(_) => "drop_policy",
637        QueryExpr::CreateServer(_) => "create_server",
638        QueryExpr::DropServer(_) => "drop_server",
639        QueryExpr::CreateForeignTable(_) => "create_foreign_table",
640        QueryExpr::DropForeignTable(_) => "drop_foreign_table",
641        QueryExpr::Grant(_) => "grant",
642        QueryExpr::Revoke(_) => "revoke",
643        QueryExpr::AlterUser(_) => "alter_user",
644        QueryExpr::CreateUser(_) => "create_user",
645        QueryExpr::CreateIamPolicy { .. } => "create_iam_policy",
646        QueryExpr::DropIamPolicy { .. } => "drop_iam_policy",
647        QueryExpr::AttachPolicy { .. } => "attach_policy",
648        QueryExpr::DetachPolicy { .. } => "detach_policy",
649        QueryExpr::ShowPolicies { .. } => "show_policies",
650        QueryExpr::ShowEffectivePermissions { .. } => "show_effective_permissions",
651        QueryExpr::RankOf(_) => "rank_of",
652        QueryExpr::ApproxRankOf(_) => "approx_rank_of",
653        QueryExpr::RankRange(_) => "rank_range",
654        QueryExpr::SimulatePolicy { .. } => "simulate_policy",
655        QueryExpr::LintPolicy { .. } => "lint_policy",
656        QueryExpr::MigratePolicyMode { .. } => "migrate_policy_mode",
657        QueryExpr::CreateMigration(_) => "create_migration",
658        QueryExpr::ApplyMigration(_) => "apply_migration",
659        QueryExpr::RollbackMigration(_) => "rollback_migration",
660        QueryExpr::ExplainMigration(_) => "explain_migration",
661        QueryExpr::Explain(_) => "explain",
662        QueryExpr::EventsBackfill(_) => "events_backfill",
663        QueryExpr::EventsBackfillStatus { .. } => "events_backfill_status",
664    }
665}
666
667// ============================================================================
668// Tests
669// ============================================================================
670
671#[cfg(test)]
672mod tests {
673    use super::*;
674
675    #[test]
676    fn test_in_memory_vector_search() {
677        let mut executor = InMemoryVectorExecutor::new();
678
679        // Add some vectors
680        executor.add_vector("test", 1, vec![1.0, 0.0, 0.0], None);
681        executor.add_vector("test", 2, vec![0.0, 1.0, 0.0], None);
682        executor.add_vector("test", 3, vec![0.0, 0.0, 1.0], None);
683        executor.add_vector("test", 4, vec![0.9, 0.1, 0.0], None);
684
685        let query = VectorQuery {
686            alias: None,
687            collection: "test".to_string(),
688            query_vector: VectorSource::Literal(vec![1.0, 0.0, 0.0]),
689            k: 2,
690            filter: None,
691            metric: Some(DistanceMetric::L2),
692            include_vectors: false,
693            include_metadata: false,
694            threshold: None,
695        };
696
697        let result = executor.execute(&query).unwrap();
698        assert_eq!(result.len(), 2);
699
700        // First result should be vector 1 (exact match)
701        let first = &result.records[0];
702        assert_eq!(first.get("id"), Some(&Value::Integer(1)));
703    }
704
705    #[test]
706    fn test_vector_search_with_metadata_filter() {
707        let mut executor = InMemoryVectorExecutor::new();
708
709        let mut meta1 = HashMap::new();
710        meta1.insert("type".to_string(), MetadataValue::String("cve".to_string()));
711        meta1.insert("severity".to_string(), MetadataValue::Integer(9));
712
713        let mut meta2 = HashMap::new();
714        meta2.insert("type".to_string(), MetadataValue::String("cve".to_string()));
715        meta2.insert("severity".to_string(), MetadataValue::Integer(5));
716
717        let mut meta3 = HashMap::new();
718        meta3.insert(
719            "type".to_string(),
720            MetadataValue::String("advisory".to_string()),
721        );
722        meta3.insert("severity".to_string(), MetadataValue::Integer(8));
723
724        executor.add_vector("vulns", 1, vec![1.0, 0.0], Some(meta1));
725        executor.add_vector("vulns", 2, vec![0.9, 0.1], Some(meta2));
726        executor.add_vector("vulns", 3, vec![0.8, 0.2], Some(meta3));
727
728        // Search with filter: type = 'cve' AND severity >= 7
729        let query = VectorQuery {
730            alias: None,
731            collection: "vulns".to_string(),
732            query_vector: VectorSource::Literal(vec![1.0, 0.0]),
733            k: 10,
734            filter: Some(MetadataFilter::And(vec![
735                MetadataFilter::Eq("type".to_string(), MetadataValue::String("cve".to_string())),
736                MetadataFilter::Gte("severity".to_string(), MetadataValue::Integer(7)),
737            ])),
738            metric: Some(DistanceMetric::L2),
739            include_vectors: false,
740            include_metadata: true,
741            threshold: None,
742        };
743
744        let result = executor.execute(&query).unwrap();
745
746        // Only vector 1 matches (type=cve, severity=9)
747        assert_eq!(result.len(), 1);
748        assert_eq!(result.records[0].get("id"), Some(&Value::Integer(1)));
749    }
750
751    #[test]
752    fn test_vector_search_with_threshold() {
753        let mut executor = InMemoryVectorExecutor::new();
754
755        executor.add_vector("test", 1, vec![1.0, 0.0], None);
756        executor.add_vector("test", 2, vec![0.0, 1.0], None); // Far from query
757
758        let query = VectorQuery {
759            alias: None,
760            collection: "test".to_string(),
761            query_vector: VectorSource::Literal(vec![1.0, 0.0]),
762            k: 10,
763            filter: None,
764            metric: Some(DistanceMetric::L2),
765            include_vectors: false,
766            include_metadata: false,
767            threshold: Some(0.5), // Only include close matches
768        };
769
770        let result = executor.execute(&query).unwrap();
771
772        // Only vector 1 is within threshold
773        assert_eq!(result.len(), 1);
774    }
775
776    #[test]
777    fn test_vector_search_include_vectors() {
778        let mut executor = InMemoryVectorExecutor::new();
779
780        executor.add_vector("test", 1, vec![1.0, 2.0, 3.0], None);
781
782        let query = VectorQuery {
783            alias: None,
784            collection: "test".to_string(),
785            query_vector: VectorSource::Literal(vec![1.0, 2.0, 3.0]),
786            k: 1,
787            filter: None,
788            metric: Some(DistanceMetric::L2),
789            include_vectors: true,
790            include_metadata: false,
791            threshold: None,
792        };
793
794        let result = executor.execute(&query).unwrap();
795        assert_eq!(result.len(), 1);
796
797        let vsr = &result.records[0].vector_results[0];
798        assert!(vsr.vector.is_some());
799        assert_eq!(vsr.vector.as_ref().unwrap(), &vec![1.0, 2.0, 3.0]);
800    }
801
802    #[test]
803    fn test_vector_executor_reference_source() {
804        let mut store = VectorStore::new();
805        let collection = store.create_collection("refs", 2);
806        let ref_id = collection.insert(vec![1.0, 0.0], None).unwrap();
807        collection.insert(vec![0.0, 1.0], None).unwrap();
808
809        let executor = VectorExecutor::new(Arc::new(store));
810        let query = VectorQuery {
811            alias: None,
812            collection: "refs".to_string(),
813            query_vector: VectorSource::Reference {
814                collection: "refs".to_string(),
815                vector_id: ref_id,
816            },
817            k: 1,
818            filter: None,
819            metric: Some(DistanceMetric::L2),
820            include_vectors: false,
821            include_metadata: false,
822            threshold: None,
823        };
824
825        let result = executor.execute(&query).unwrap();
826        assert_eq!(result.len(), 1);
827        assert_eq!(result.records[0].get("id"), Some(&Value::Integer(0)));
828    }
829
830    #[test]
831    fn test_vector_executor_subquery_source() {
832        let mut store = VectorStore::new();
833        let collection = store.create_collection("refs", 2);
834        collection.insert(vec![1.0, 0.0], None).unwrap();
835        collection.insert(vec![0.0, 1.0], None).unwrap();
836
837        let executor = VectorExecutor::new(Arc::new(store));
838        let inner = VectorQuery {
839            alias: None,
840            collection: "refs".to_string(),
841            query_vector: VectorSource::Literal(vec![1.0, 0.0]),
842            k: 1,
843            filter: None,
844            metric: Some(DistanceMetric::L2),
845            include_vectors: false,
846            include_metadata: false,
847            threshold: None,
848        };
849        let query = VectorQuery {
850            alias: None,
851            collection: "refs".to_string(),
852            query_vector: VectorSource::Subquery(Box::new(QueryExpr::Vector(inner))),
853            k: 1,
854            filter: None,
855            metric: Some(DistanceMetric::L2),
856            include_vectors: false,
857            include_metadata: false,
858            threshold: None,
859        };
860
861        let result = executor.execute(&query).unwrap();
862        assert_eq!(result.len(), 1);
863        assert_eq!(result.records[0].get("id"), Some(&Value::Integer(0)));
864    }
865}
reddb_server/storage/query/executors/vector.rs

reddb_server/storage/query/executors/
vector.rs
