Skip to main content

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

1//! Hybrid query builder
2//!
3//! Builder for complex hybrid queries combining vector, graph, and filter operations.
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_hybrid_query;
12use super::super::filters::{Filter, FilterAcceptor, WhereClause};
13use super::super::types::QueryResult;
14
15/// Builder for complex hybrid queries
16#[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    /// Add vector similarity component
67    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    /// Add graph pattern component
73    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    /// Limit to collections
83    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    /// Add filter
91    pub fn where_(self, field: impl Into<String>) -> WhereClause<Self> {
92        WhereClause::new(self, field.into())
93    }
94
95    /// Set scoring weights
96    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    /// Set minimum score threshold
106    pub fn min_score(mut self, score: f32) -> Self {
107        self.min_score = score;
108        self
109    }
110
111    /// Limit results
112    pub fn limit(mut self, n: usize) -> Self {
113        self.limit = Some(n);
114        self
115    }
116
117    /// Expand via cross-references
118    pub fn expand_via(mut self, ref_type: RefType) -> Self {
119        self.expand_refs = Some(ref_type);
120        self
121    }
122
123    /// Execute the query
124    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}