reddb_server/storage/unified/dsl/builders/
hybrid.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_hybrid_query;
12use super::super::filters::{Filter, FilterAcceptor, WhereClause};
13use super::super::types::QueryResult;
14
15#[derive(Debug, Clone)]
17pub struct HybridQueryBuilder {
18 pub(crate) vector_query: Option<(Vec<f32>, usize)>,
19 pub(crate) graph_pattern: Option<GraphPatternDsl>,
20 pub(crate) filters: Vec<Filter>,
21 pub(crate) collections: Option<Vec<String>>,
22 pub(crate) weights: QueryWeights,
23 pub(crate) min_score: f32,
24 pub(crate) limit: Option<usize>,
25 pub(crate) expand_refs: Option<RefType>,
26}
27
28#[derive(Debug, Clone)]
29pub struct GraphPatternDsl {
30 pub node_label: Option<String>,
31 pub node_type: Option<String>,
32 pub edge_labels: Vec<String>,
33}
34
35#[derive(Debug, Clone)]
36pub struct QueryWeights {
37 pub vector: f32,
38 pub graph: f32,
39 pub filter: f32,
40}
41
42impl Default for QueryWeights {
43 fn default() -> Self {
44 Self {
45 vector: 0.5,
46 graph: 0.3,
47 filter: 0.2,
48 }
49 }
50}
51
52impl HybridQueryBuilder {
53 pub fn new() -> Self {
54 Self {
55 vector_query: None,
56 graph_pattern: None,
57 filters: Vec::new(),
58 collections: None,
59 weights: QueryWeights::default(),
60 min_score: 0.1,
61 limit: None,
62 expand_refs: None,
63 }
64 }
65
66 pub fn similar_to(mut self, vector: &[f32], k: usize) -> Self {
68 self.vector_query = Some((vector.to_vec(), k));
69 self
70 }
71
72 pub fn matching_nodes(mut self, label: impl Into<String>) -> Self {
74 self.graph_pattern = Some(GraphPatternDsl {
75 node_label: Some(label.into()),
76 node_type: None,
77 edge_labels: Vec::new(),
78 });
79 self
80 }
81
82 pub fn in_collection(mut self, name: impl Into<String>) -> Self {
84 self.collections
85 .get_or_insert_with(Vec::new)
86 .push(name.into());
87 self
88 }
89
90 pub fn where_(self, field: impl Into<String>) -> WhereClause<Self> {
92 WhereClause::new(self, field.into())
93 }
94
95 pub fn with_weights(mut self, vector: f32, graph: f32, filter: f32) -> Self {
97 self.weights = QueryWeights {
98 vector,
99 graph,
100 filter,
101 };
102 self
103 }
104
105 pub fn min_score(mut self, score: f32) -> Self {
107 self.min_score = score;
108 self
109 }
110
111 pub fn limit(mut self, n: usize) -> Self {
113 self.limit = Some(n);
114 self
115 }
116
117 pub fn expand_via(mut self, ref_type: RefType) -> Self {
119 self.expand_refs = Some(ref_type);
120 self
121 }
122
123 pub fn execute(self, store: &Arc<UnifiedStore>) -> Result<QueryResult, ExecutionError> {
125 execute_hybrid_query(self, store)
126 }
127}
128
129impl Default for HybridQueryBuilder {
130 fn default() -> Self {
131 Self::new()
132 }
133}
134
135impl FilterAcceptor for HybridQueryBuilder {
136 fn add_filter(&mut self, filter: Filter) {
137 self.filters.push(filter);
138 }
139}