weavatrix_search_vector/
multi.rs

1use crate::config::IndexConfig;
2use crate::error::SearchError;
3use crate::hit::SearchHit;
4use crate::hnsw::VectorIndex;
5use std::collections::BTreeSet;
6
7/// Stable caller identity for one vector among several attached to a key.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
9pub struct MultiVectorKey {
10    pub key: u64,
11    pub vector_id: u64,
12}
13
14/// Borrowed multi-vector build record.
15#[derive(Debug, Clone, Copy)]
16pub struct MultiVectorRef<'a> {
17    pub key: u64,
18    pub vector_id: u64,
19    pub vector: &'a [f32],
20}
21
22/// One nearest-neighbor result retaining both caller identifiers.
23#[derive(Debug, Clone, Copy, PartialEq)]
24pub struct MultiSearchHit {
25    pub key: u64,
26    pub vector_id: u64,
27    pub distance: f32,
28}
29
30/// Deterministic vector index allowing multiple vectors per caller key.
31#[derive(Debug)]
32pub struct MultiVectorIndex {
33    pub(crate) index: VectorIndex,
34    pub(crate) identities: Vec<MultiVectorKey>,
35}
36
37impl MultiVectorIndex {
38    /// Builds one HNSW node per unique `(key, vector_id)` pair.
39    ///
40    /// Input order does not affect internal IDs or graph construction.
41    ///
42    /// # Errors
43    ///
44    /// Returns typed config, vector, duplicate identity, capacity, or
45    /// allocation errors.
46    pub fn build(config: IndexConfig, vectors: &[MultiVectorRef<'_>]) -> Result<Self, SearchError> {
47        if vectors.len() > u32::MAX as usize {
48            return Err(SearchError::CapacityOverflow);
49        }
50        let mut order = (0..vectors.len()).collect::<Vec<_>>();
51        order.sort_unstable_by_key(|index| (vectors[*index].key, vectors[*index].vector_id));
52        for pair in order.windows(2) {
53            let left = vectors[pair[0]];
54            let right = vectors[pair[1]];
55            if (left.key, left.vector_id) == (right.key, right.vector_id) {
56                return Err(SearchError::DuplicateVectorId {
57                    key: left.key,
58                    vector_id: left.vector_id,
59                });
60            }
61        }
62        let mut identities = Vec::new();
63        identities
64            .try_reserve_exact(vectors.len())
65            .map_err(|_| SearchError::AllocationFailed)?;
66        let mut internal = Vec::new();
67        internal
68            .try_reserve_exact(vectors.len())
69            .map_err(|_| SearchError::AllocationFailed)?;
70        for (internal_id, source) in order.into_iter().enumerate() {
71            let record = vectors[source];
72            identities.push(MultiVectorKey {
73                key: record.key,
74                vector_id: record.vector_id,
75            });
76            internal.push((
77                u64::try_from(internal_id).map_err(|_| SearchError::CapacityOverflow)?,
78                record.vector,
79            ));
80        }
81        Ok(Self {
82            index: VectorIndex::build(config, &internal)?,
83            identities,
84        })
85    }
86
87    #[must_use]
88    pub fn len(&self) -> usize {
89        self.identities.len()
90    }
91
92    #[must_use]
93    pub fn is_empty(&self) -> bool {
94        self.identities.is_empty()
95    }
96
97    #[must_use]
98    pub const fn dimensions(&self) -> usize {
99        self.index.dimensions()
100    }
101
102    #[must_use]
103    pub fn config(&self) -> &IndexConfig {
104        self.index.config()
105    }
106
107    #[must_use]
108    pub fn estimated_memory_bytes(&self) -> usize {
109        self.index.estimated_memory_bytes().saturating_add(
110            self.identities
111                .capacity()
112                .saturating_mul(std::mem::size_of::<MultiVectorKey>()),
113        )
114    }
115
116    /// Returns top vectors, allowing repeated caller keys.
117    ///
118    /// # Errors
119    ///
120    /// Returns a typed query or allocation error.
121    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<MultiSearchHit>, SearchError> {
122        self.map_hits(self.index.search(query, count)?)
123    }
124
125    /// Returns exact top vectors, allowing repeated caller keys.
126    ///
127    /// # Errors
128    ///
129    /// Returns a typed query or allocation error.
130    pub fn search_exact(
131        &self,
132        query: &[f32],
133        count: usize,
134    ) -> Result<Vec<MultiSearchHit>, SearchError> {
135        self.map_hits(self.index.search_exact(query, count)?)
136    }
137
138    /// Returns at most one best vector for each caller key.
139    ///
140    /// The approximate pass widens candidates; an exact fallback fills sparse
141    /// key distributions.
142    ///
143    /// # Errors
144    ///
145    /// Returns a typed query or allocation error.
146    pub fn search_unique_keys(
147        &self,
148        query: &[f32],
149        count: usize,
150    ) -> Result<Vec<MultiSearchHit>, SearchError> {
151        let requested = count.min(self.distinct_key_count());
152        if requested == 0 {
153            self.index.search(query, 0)?;
154            return Ok(Vec::new());
155        }
156        let candidate_count = requested.saturating_mul(8).min(self.len());
157        let mut hits = Self::unique(self.search(query, candidate_count)?, requested);
158        if hits.len() < requested {
159            hits = Self::unique(self.search_exact(query, self.len())?, requested);
160        }
161        Ok(hits)
162    }
163
164    /// Searches only identities accepted by `filter`.
165    ///
166    /// # Errors
167    ///
168    /// Returns a typed query or allocation error.
169    pub fn search_filtered<F>(
170        &self,
171        query: &[f32],
172        count: usize,
173        filter: F,
174    ) -> Result<Vec<MultiSearchHit>, SearchError>
175    where
176        F: Fn(MultiVectorKey) -> bool,
177    {
178        let hits = self.index.search_filtered(query, count, |internal| {
179            usize::try_from(internal)
180                .ok()
181                .and_then(|index| self.identities.get(index))
182                .is_some_and(|identity| filter(*identity))
183        })?;
184        self.map_hits(hits)
185    }
186
187    /// Searches independent queries with bounded workers.
188    ///
189    /// # Errors
190    ///
191    /// Returns the first query or worker error.
192    pub fn search_batch(
193        &self,
194        queries: &[&[f32]],
195        count: usize,
196    ) -> Result<Vec<Vec<MultiSearchHit>>, SearchError> {
197        self.index
198            .search_batch(queries, count)?
199            .into_iter()
200            .map(|hits| self.map_hits(hits))
201            .collect()
202    }
203
204    fn map_hits(&self, hits: Vec<SearchHit>) -> Result<Vec<MultiSearchHit>, SearchError> {
205        let mut mapped = Vec::new();
206        mapped
207            .try_reserve_exact(hits.len())
208            .map_err(|_| SearchError::AllocationFailed)?;
209        for hit in hits {
210            let index = usize::try_from(hit.key).map_err(|_| SearchError::CapacityOverflow)?;
211            let identity = self
212                .identities
213                .get(index)
214                .ok_or(SearchError::CorruptSnapshot(
215                    "multi-vector internal key is outside identity table",
216                ))?;
217            mapped.push(MultiSearchHit {
218                key: identity.key,
219                vector_id: identity.vector_id,
220                distance: hit.distance,
221            });
222        }
223        Ok(mapped)
224    }
225
226    fn unique(hits: Vec<MultiSearchHit>, count: usize) -> Vec<MultiSearchHit> {
227        let mut seen = BTreeSet::new();
228        hits.into_iter()
229            .filter(|hit| seen.insert(hit.key))
230            .take(count)
231            .collect()
232    }
233
234    fn distinct_key_count(&self) -> usize {
235        self.identities
236            .iter()
237            .map(|identity| identity.key)
238            .collect::<BTreeSet<_>>()
239            .len()
240    }
241}
weavatrix_search_vector/multi.rs

weavatrix_search_vector/
multi.rs
