Skip to main content

reddb_server/storage/unified/dsl/builders/
vector.rs

1//! Vector similarity query builder
2//!
3//! Builder for semantic/vector similarity searches.
4
5use std::sync::Arc;
6
7use crate::storage::query::unified::ExecutionError;
8
9use super::super::super::entity::RefType;
10use super::super::super::store::UnifiedStore;
11use super::super::execution::execute_vector_query;
12use super::super::filters::{Filter, FilterAcceptor, WhereClause};
13use super::super::types::QueryResult;
14
15/// Builder for vector similarity queries
16#[derive(Debug, Clone)]
17pub struct VectorQueryBuilder {
18    pub(crate) vector: Vec<f32>,
19    pub(crate) k: usize,
20    pub(crate) collections: Option<Vec<String>>,
21    pub(crate) filters: Vec<Filter>,
22    pub(crate) min_similarity: f32,
23    pub(crate) include_embeddings: bool,
24    pub(crate) embedding_slot: Option<String>,
25    pub(crate) expand_refs: Option<RefType>,
26    pub(crate) expand_depth: u32,
27}
28
29impl VectorQueryBuilder {
30    pub fn new(vector: Vec<f32>, k: usize) -> Self {
31        Self {
32            vector,
33            k,
34            collections: None,
35            filters: Vec::new(),
36            min_similarity: 0.0,
37            include_embeddings: true,
38            embedding_slot: None,
39            expand_refs: None,
40            expand_depth: 1,
41        }
42    }
43
44    /// Limit search to specific collection(s)
45    pub fn in_collection(mut self, name: impl Into<String>) -> Self {
46        self.collections
47            .get_or_insert_with(Vec::new)
48            .push(name.into());
49        self
50    }
51
52    /// Limit search to multiple collections
53    pub fn in_collections(mut self, names: &[&str]) -> Self {
54        let cols = self.collections.get_or_insert_with(Vec::new);
55        for name in names {
56            cols.push((*name).to_string());
57        }
58        self
59    }
60
61    /// Add a filter condition (returns WhereClause for chaining)
62    pub fn where_(self, field: impl Into<String>) -> WhereClause<Self> {
63        WhereClause::new(self, field.into())
64    }
65
66    /// Set minimum similarity threshold
67    pub fn min_similarity(mut self, threshold: f32) -> Self {
68        self.min_similarity = threshold;
69        self
70    }
71
72    /// Search in a specific embedding slot
73    pub fn in_slot(mut self, slot: impl Into<String>) -> Self {
74        self.embedding_slot = Some(slot.into());
75        self
76    }
77
78    /// Skip entity embeddings (only search dedicated vector entities)
79    pub fn vectors_only(mut self) -> Self {
80        self.include_embeddings = false;
81        self
82    }
83
84    /// Expand results by following cross-references
85    pub fn expand_via(mut self, ref_type: RefType) -> Self {
86        self.expand_refs = Some(ref_type);
87        self
88    }
89
90    /// Set expansion depth
91    pub fn depth(mut self, depth: u32) -> Self {
92        self.expand_depth = depth;
93        self
94    }
95
96    /// Execute the query
97    pub fn execute(self, store: &Arc<UnifiedStore>) -> Result<QueryResult, ExecutionError> {
98        execute_vector_query(self, store)
99    }
100}
101
102impl FilterAcceptor for VectorQueryBuilder {
103    fn add_filter(&mut self, filter: Filter) {
104        self.filters.push(filter);
105    }
106}