reddb_server/runtime/
hybrid_search.rs

1//! Hybrid ranking and one-hop expansion share one authorized snapshot.
2use super::*;
3use crate::storage::unified::dsl::{calculate_entity_similarity, extract_searchable_text};
4
5struct Candidate {
6    collection: String,
7    entity: UnifiedEntity,
8    text: f32,
9    vector: f32,
10}
11
12fn matches_node(entity: &UnifiedEntity, pattern: &RuntimeGraphPattern) -> bool {
13    matches!(&entity.kind, EntityKind::GraphNode(node)
14        if pattern.node_label.as_ref().is_none_or(|label| *label == node.label)
15            && pattern.node_type.as_ref().is_none_or(|kind| *kind == node.node_type))
16}
17
18impl RedDBRuntime {
19    pub fn search_hybrid(
20        &self,
21        vector: Option<Vec<f32>>,
22        query: Option<String>,
23        k: Option<usize>,
24        collections: Option<Vec<String>>,
25        entity_types: Option<Vec<String>>,
26        capabilities: Option<Vec<String>>,
27        graph_pattern: Option<RuntimeGraphPattern>,
28        filters: Vec<RuntimeFilter>,
29        weights: Option<RuntimeQueryWeights>,
30        min_score: Option<f32>,
31        limit: Option<usize>,
32    ) -> RedDBResult<DslQueryResult> {
33        let started = std::time::Instant::now();
34        let query = query
35            .map(|text| text.trim().to_lowercase())
36            .filter(|text| !text.is_empty());
37        if vector.is_none() && query.is_none() {
38            return Err(RedDBError::Query(
39                "field 'query' or 'vector' is required for hybrid search".into(),
40            ));
41        }
42        if vector.as_ref().is_some_and(|values| {
43            values.is_empty() || values.iter().any(|value| !value.is_finite())
44        }) {
45            return Err(RedDBError::Query(
46                "hybrid search vector must contain finite values".into(),
47            ));
48        }
49        let weights = weights.unwrap_or(RuntimeQueryWeights {
50            vector: 0.5,
51            graph: 0.3,
52            filter: 0.2,
53        });
54        if [weights.vector, weights.graph, weights.filter]
55            .iter()
56            .any(|weight| !weight.is_finite() || *weight < 0.0)
57        {
58            return Err(RedDBError::Query(
59                "hybrid search weights must be finite and nonnegative".into(),
60            ));
61        }
62        let filters = filters
63            .into_iter()
64            .map(runtime_filter_to_dsl)
65            .collect::<RedDBResult<Vec<_>>>()?;
66        let limit = limit.or(k).unwrap_or(10).max(1);
67        let min_score = min_score
68            .filter(|score| score.is_finite())
69            .unwrap_or(0.0)
70            .max(0.0);
71        let relation_labels = graph_pattern
72            .as_ref()
73            .map(|pattern| pattern.edge_labels.as_slice())
74            .unwrap_or_default();
75        let relations = !relation_labels.is_empty();
76        let explicit_scope = collections.as_ref().is_some_and(|names| !names.is_empty());
77        let collections =
78            runtime_search_collections(&self.inner.db, collections).unwrap_or_default();
79        let (snapshot, _retention) = self.native_read_snapshot()?;
80        let _snapshot = super::execution_context::CurrentSnapshotGuard::install(snapshot.clone());
81        self.refresh_memory_accounting();
82        let store = self.db().store();
83        let mut candidates = Vec::new();
84        let mut reservations = Vec::new();
85        let mut policies = HashMap::new();
86        let terms: Vec<_> = query
87            .as_deref()
88            .unwrap_or_default()
89            .split_whitespace()
90            .collect();
91        for collection in collections {
92            if super::impl_ddl::is_system_schema_name(&collection) || collection.starts_with("red_")
93            {
94                continue;
95            }
96            if let Err(error) =
97                self.check_collection_privilege(reddb_rql::ast::PolicyAction::Select, &collection)
98            {
99                if explicit_scope {
100                    return Err(error);
101                }
102                continue;
103            }
104            let Some(manager) = store.get_collection(&collection) else {
105                continue;
106            };
107            let (_scan_memory, read_bytes_max) = self.reserve_index_scan(&manager)?;
108            let mut rows = Vec::new();
109            super::function_budget::try_scan_result(
110                |visit| manager.try_for_each_entity(read_bytes_max, visit),
111                |entity| {
112                    if super::execution_context::entity_visible_with_context(
113                        Some(&snapshot),
114                        entity,
115                    ) {
116                        // Own candidate payloads, ranking copies and adjacency keys. Reserve
117                        // from the sampled ledger here; never evict under a storage lock.
118                        let bytes = crate::storage::unified::memory_size::entity_bytes(entity)
119                            .saturating_mul(3)
120                            .saturating_add(collection.len().saturating_mul(4))
121                            .saturating_add(1024);
122                        let reservation = self
123                            .try_reserve_query_buffer_from_sample(bytes)
124                            .ok_or_else(|| {
125                                RedDBError::Query(
126                                    "hybrid search exceeds the query memory budget".into(),
127                                )
128                            })?;
129                        reservations.push(reservation);
130                        rows.push(entity.clone());
131                    }
132                    Ok(true)
133                },
134            )?;
135            // RLS may read metadata or query storage; evaluate outside segment locks.
136            for entity in rows {
137                super::function_budget::charge(1)?;
138                if !self.search_entity_allowed(&collection, &entity, Some(&snapshot), &mut policies)
139                {
140                    continue;
141                }
142                let text = if terms.is_empty() {
143                    0.0
144                } else {
145                    let body = extract_searchable_text(&entity).to_lowercase();
146                    terms.iter().filter(|term| body.contains(**term)).count() as f32
147                        / terms.len() as f32
148                };
149                let similarity = vector.as_ref().map_or(0.0, |vector| {
150                    calculate_entity_similarity(&entity, vector, &None)
151                });
152                candidates.push(Candidate {
153                    collection: collection.clone(),
154                    entity,
155                    text,
156                    vector: similarity,
157                });
158            }
159        }
160        let mut matches: HashMap<EntityId, ScoredMatch> = HashMap::new();
161        for candidate in &candidates {
162            let graph = if !relations
163                && vector.is_some()
164                && graph_pattern
165                    .as_ref()
166                    .is_some_and(|pattern| matches_node(&candidate.entity, pattern))
167            {
168                1.0
169            } else {
170                0.0
171            };
172            let score = if vector.is_none() && !relations {
173                candidate.text
174            } else {
175                candidate.vector * weights.vector
176                    + candidate.text * weights.filter
177                    + graph * weights.graph
178            };
179            if score > 0.0 {
180                matches.insert(
181                    candidate.entity.id,
182                    ScoredMatch {
183                        entity: candidate.entity.clone(),
184                        score,
185                        path: None,
186                        components: MatchComponents {
187                            vector_similarity: (candidate.vector > 0.0).then_some(candidate.vector),
188                            text_relevance: (candidate.text > 0.0).then_some(candidate.text),
189                            graph_match: (graph > 0.0).then_some(graph),
190                            ..Default::default()
191                        },
192                    },
193                );
194            }
195        }
196        if relations {
197            expand_relations(
198                &candidates,
199                relation_labels,
200                k.unwrap_or(limit).max(1),
201                weights.graph,
202                &mut matches,
203            )?;
204        }
205        let mut result = DslQueryResult {
206            matches: matches
207                .into_values()
208                .filter(|item| {
209                    item.score >= min_score
210                        && apply_filters(&item.entity, &filters)
211                        && graph_pattern.as_ref().is_none_or(|pattern| {
212                            // Legacy vector patterns boost nodes; relation patterns and
213                            // text-only patterns constrain the returned node types.
214                            (!relations && vector.is_some()) || matches_node(&item.entity, pattern)
215                        })
216                })
217                .collect(),
218            scanned: candidates.len(),
219            execution_time_us: started.elapsed().as_micros() as u64,
220            explanation: if relations {
221                "Hybrid search with authorized one-hop relations"
222            } else {
223                "Hybrid multi-modal query"
224            }
225            .into(),
226        };
227        for item in &mut result.matches {
228            item.components.filter_match = true;
229            item.components.structured_match = Some(1.0);
230            item.components.final_score = Some(item.score);
231        }
232        normalize_runtime_dsl_result_scores(&mut result);
233        runtime_filter_dsl_result(&mut result, entity_types, capabilities);
234        result.matches.truncate(limit);
235        Ok(result)
236    }
237}
238
239fn expand_relations(
240    candidates: &[Candidate],
241    labels: &[String],
242    seed_limit: usize,
243    weight: f32,
244    matches: &mut HashMap<EntityId, ScoredMatch>,
245) -> RedDBResult<()> {
246    let mut seeds: Vec<_> = candidates
247        .iter()
248        .filter(|candidate| matches!(candidate.entity.kind, EntityKind::GraphNode(_)))
249        .map(|candidate| (candidate.entity.id, candidate.text.max(candidate.vector)))
250        .filter(|(_, score)| *score > 0.0)
251        .collect();
252    seeds.sort_by(|left, right| {
253        right
254            .1
255            .total_cmp(&left.1)
256            .then_with(|| left.0.cmp(&right.0))
257    });
258    seeds.truncate(seed_limit);
259    let seeds: HashMap<_, _> = seeds.into_iter().collect();
260    if seeds.is_empty() || weight == 0.0 {
261        return Ok(());
262    }
263    let mut nodes: HashMap<(String, String), Vec<usize>> = HashMap::new();
264    for (index, candidate) in candidates.iter().enumerate() {
265        if let EntityKind::GraphNode(node) = &candidate.entity.kind {
266            for alias in [
267                node.label.clone(),
268                candidate.entity.id.raw().to_string(),
269                candidate.entity.logical_id().raw().to_string(),
270            ] {
271                let entries = nodes
272                    .entry((candidate.collection.clone(), alias))
273                    .or_default();
274                if entries.last() != Some(&index) {
275                    entries.push(index);
276                }
277            }
278        }
279    }
280    for candidate in candidates {
281        let EntityKind::GraphEdge(edge) = &candidate.entity.kind else {
282            continue;
283        };
284        if !labels.contains(&edge.label) {
285            continue;
286        }
287        let Some(from) = nodes.get(&(candidate.collection.clone(), edge.from_node.clone())) else {
288            continue;
289        };
290        let Some(to) = nodes.get(&(candidate.collection.clone(), edge.to_node.clone())) else {
291            continue;
292        };
293        for &source in from {
294            for &target in to {
295                super::function_budget::charge(1)?;
296                for (source, target) in [(source, target), (target, source)] {
297                    let source = &candidates[source].entity;
298                    let target = &candidates[target].entity;
299                    if source.id == target.id {
300                        continue;
301                    }
302                    let Some(&strength) = seeds.get(&source.id) else {
303                        continue;
304                    };
305                    let path = vec![source.id, candidate.entity.id, target.id];
306                    let item = matches.entry(target.id).or_insert_with(|| ScoredMatch {
307                        entity: target.clone(),
308                        score: 0.0,
309                        components: MatchComponents::default(),
310                        path: None,
311                    });
312                    let previous = item.components.graph_match.unwrap_or(0.0);
313                    if strength > previous
314                        || (strength == previous
315                            && item.path.as_ref().is_none_or(|old| path < *old))
316                    {
317                        item.score += (strength - previous) * weight;
318                        item.components.graph_match = Some(strength);
319                        item.components.hop_distance = Some(1);
320                        item.path = Some(path);
321                    }
322                }
323            }
324        }
325    }
326    Ok(())
327}
reddb_server/runtime/hybrid_search.rs

reddb_server/runtime/
hybrid_search.rs
