reddb_server/storage/unified/dsl/builders/
vector.rs1use 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#[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 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 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 pub fn where_(self, field: impl Into<String>) -> WhereClause<Self> {
63 WhereClause::new(self, field.into())
64 }
65
66 pub fn min_similarity(mut self, threshold: f32) -> Self {
68 self.min_similarity = threshold;
69 self
70 }
71
72 pub fn in_slot(mut self, slot: impl Into<String>) -> Self {
74 self.embedding_slot = Some(slot.into());
75 self
76 }
77
78 pub fn vectors_only(mut self) -> Self {
80 self.include_embeddings = false;
81 self
82 }
83
84 pub fn expand_via(mut self, ref_type: RefType) -> Self {
86 self.expand_refs = Some(ref_type);
87 self
88 }
89
90 pub fn depth(mut self, depth: u32) -> Self {
92 self.expand_depth = depth;
93 self
94 }
95
96 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}