weavatrix_search_vector/storage/
mapped_search.rs

1use super::MappedVectorIndex;
2use crate::error::SearchError;
3use crate::hit::SearchHit;
4use crate::hnsw::SearchPolicy;
5use crate::vector::{Candidate, RoutingCandidate, routing_probes_from_signs};
6use std::cmp::Reverse;
7use std::collections::BinaryHeap;
8
9pub(super) struct MappedScratch {
10    pub(super) marks: Vec<u32>,
11    pub(super) generation: u32,
12    pub(super) candidates: BinaryHeap<Reverse<Candidate>>,
13    pub(super) results: BinaryHeap<Candidate>,
14    pub(super) routing_probes: Vec<u16>,
15    pub(super) routing_probe_heap: BinaryHeap<Reverse<RoutingCandidate>>,
16}
17
18impl MappedScratch {
19    pub(super) fn new(len: usize) -> Result<Self, SearchError> {
20        let mut marks = Vec::new();
21        marks
22            .try_reserve_exact(len)
23            .map_err(|_| SearchError::AllocationFailed)?;
24        marks.resize(len, 0);
25        Ok(Self {
26            marks,
27            generation: 0,
28            candidates: BinaryHeap::new(),
29            results: BinaryHeap::new(),
30            routing_probes: Vec::new(),
31            routing_probe_heap: BinaryHeap::new(),
32        })
33    }
34
35    pub(super) fn begin(&mut self) {
36        self.candidates.clear();
37        self.results.clear();
38        self.generation = self.generation.wrapping_add(1);
39        if self.generation == 0 {
40            self.marks.fill(0);
41            self.generation = 1;
42        }
43    }
44
45    pub(super) fn mark(&mut self, index: usize) -> bool {
46        if self.marks[index] == self.generation {
47            return false;
48        }
49        self.marks[index] = self.generation;
50        true
51    }
52}
53
54impl MappedVectorIndex {
55    /// Returns approximate top-K hits ordered by exact distance and key.
56    ///
57    /// # Errors
58    ///
59    /// Returns a typed error for an invalid query or scratch allocation.
60    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
61        self.search_with_policy(
62            query,
63            count,
64            SearchPolicy::new(self.header.config.expansion_query),
65        )
66    }
67
68    /// Returns approximate top-K hits with a per-query recall policy.
69    ///
70    /// # Errors
71    ///
72    /// Returns a typed error for an invalid policy, query, or scratch
73    /// allocation.
74    pub fn search_with_policy(
75        &self,
76        query: &[f32],
77        count: usize,
78        policy: SearchPolicy,
79    ) -> Result<Vec<SearchHit>, SearchError> {
80        self.search_where(query, count, policy.validate()?, &|_| true)
81    }
82
83    pub(super) fn search_where<F>(
84        &self,
85        query: &[f32],
86        count: usize,
87        policy: SearchPolicy,
88        accepts: &F,
89    ) -> Result<Vec<SearchHit>, SearchError>
90    where
91        F: Fn(u64) -> bool,
92    {
93        let query_norm = self.query_squared_norm(query)?;
94        let limit = count.min(self.len());
95        if limit == 0 {
96            return Ok(Vec::new());
97        }
98        let expansion = policy.expansion.max(limit);
99        let mut scratch = MappedScratch::new(self.len())?;
100        let mut merged = Vec::new();
101        merged
102            .try_reserve(self.graphs.len().saturating_mul(expansion))
103            .map_err(|_| SearchError::AllocationFailed)?;
104        for graph in &self.graphs {
105            graph.search_into(
106                self,
107                query,
108                query_norm,
109                expansion,
110                &mut scratch,
111                &mut merged,
112                accepts,
113            );
114        }
115        self.append_routing(
116            query,
117            query_norm,
118            policy,
119            accepts,
120            &mut scratch,
121            &mut merged,
122        )?;
123        merged.sort_unstable();
124        collect_hits(self, merged, limit)
125    }
126
127    fn append_routing<F>(
128        &self,
129        query: &[f32],
130        query_norm: f32,
131        policy: SearchPolicy,
132        accepts: &F,
133        scratch: &mut MappedScratch,
134        merged: &mut Vec<Candidate>,
135    ) -> Result<(), SearchError>
136    where
137        F: Fn(u64) -> bool,
138    {
139        routing_probes_from_signs(
140            query,
141            self.routing_signs(),
142            policy.routing_probes,
143            &mut scratch.routing_probes,
144            &mut scratch.routing_probe_heap,
145        )?;
146        let codes = self.routing_codes();
147        let nodes = self.routing_nodes();
148        for code in scratch.routing_probes.iter().copied() {
149            let start = codes.partition_point(|entry| *entry < code);
150            let end = codes.partition_point(|entry| *entry <= code);
151            merged
152                .try_reserve(end.saturating_sub(start))
153                .map_err(|_| SearchError::AllocationFailed)?;
154            merged.extend(
155                nodes[start..end]
156                    .iter()
157                    .copied()
158                    .map(|node| node as usize)
159                    .filter(|index| accepts(self.key_slice()[*index]))
160                    .map(|index| {
161                        Candidate::new(self.distance_query(index, query, query_norm), index)
162                    }),
163            );
164        }
165        Ok(())
166    }
167}
168
169fn collect_hits(
170    index: &MappedVectorIndex,
171    candidates: Vec<Candidate>,
172    limit: usize,
173) -> Result<Vec<SearchHit>, SearchError> {
174    let mut hits = Vec::new();
175    hits.try_reserve_exact(limit)
176        .map_err(|_| SearchError::AllocationFailed)?;
177    for candidate in candidates {
178        let key = index.key_slice()[candidate.index()];
179        if hits.iter().any(|hit: &SearchHit| hit.key == key) {
180            continue;
181        }
182        hits.push(SearchHit {
183            key,
184            distance: candidate.distance,
185        });
186        if hits.len() == limit {
187            break;
188        }
189    }
190    Ok(hits)
191}
weavatrix_search_vector/storage/mapped_search.rs

weavatrix_search_vector/storage/
mapped_search.rs
