weavatrix_search_vector/storage/
mapped_exact.rs

1use super::MappedVectorIndex;
2use crate::error::SearchError;
3use crate::hit::SearchHit;
4use crate::hnsw::{FilterSearchPolicy, SearchPolicy};
5use crate::vector::Candidate;
6use std::collections::BinaryHeap;
7
8impl MappedVectorIndex {
9    /// Returns exact top-K hits directly from mapped vectors.
10    ///
11    /// # Errors
12    ///
13    /// Returns a typed error for an invalid query.
14    pub fn search_exact(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
15        self.search_exact_filtered(query, count, |_| true)
16    }
17
18    /// Returns exact top-K mapped hits accepted by `filter`.
19    ///
20    /// # Errors
21    ///
22    /// Returns a typed error for an invalid query.
23    pub fn search_exact_filtered<F>(
24        &self,
25        query: &[f32],
26        count: usize,
27        mut filter: F,
28    ) -> Result<Vec<SearchHit>, SearchError>
29    where
30        F: FnMut(u64) -> bool,
31    {
32        let query_norm = self.query_squared_norm(query)?;
33        let limit = count.min(self.len());
34        if limit == 0 {
35            return Ok(Vec::new());
36        }
37        let mut best = BinaryHeap::with_capacity(limit);
38        for index in 0..self.len() {
39            let key = self.key_slice()[index];
40            if !filter(key) {
41                continue;
42            }
43            let candidate = Candidate::new(self.distance_query(index, query, query_norm), index);
44            if best.len() < limit {
45                best.push(candidate);
46            } else if best
47                .peek()
48                .is_some_and(|worst| candidate.cmp(worst).is_lt())
49            {
50                best.pop();
51                best.push(candidate);
52            }
53        }
54        let mut candidates = best.into_vec();
55        candidates.sort_unstable();
56        Ok(candidates
57            .into_iter()
58            .map(|candidate| SearchHit {
59                key: self.key_slice()[candidate.index()],
60                distance: candidate.distance,
61            })
62            .collect())
63    }
64
65    /// Searches mapped HNSW candidates accepted by `filter`, with an exact
66    /// fallback for selective filters.
67    ///
68    /// # Errors
69    ///
70    /// Returns a typed query or allocation error.
71    pub fn search_filtered<F>(
72        &self,
73        query: &[f32],
74        count: usize,
75        filter: F,
76    ) -> Result<Vec<SearchHit>, SearchError>
77    where
78        F: Fn(u64) -> bool,
79    {
80        self.search_filtered_with_policy(query, count, filter, FilterSearchPolicy::ExactFallback)
81    }
82
83    /// Searches a mapped graph with the predicate applied during traversal.
84    ///
85    /// # Errors
86    ///
87    /// Returns a typed query or allocation error.
88    pub fn search_filtered_with_policy<F>(
89        &self,
90        query: &[f32],
91        count: usize,
92        filter: F,
93        policy: FilterSearchPolicy,
94    ) -> Result<Vec<SearchHit>, SearchError>
95    where
96        F: Fn(u64) -> bool,
97    {
98        let requested = count.min(self.len());
99        if requested == 0 {
100            self.query_squared_norm(query)?;
101            return Ok(Vec::new());
102        }
103        let hits = self.search_where(
104            query,
105            requested,
106            SearchPolicy::new(self.header.config.expansion_query),
107            &filter,
108        )?;
109        if hits.len() == requested || policy == FilterSearchPolicy::Traversal {
110            return Ok(hits);
111        }
112        self.search_exact_filtered(query, requested, filter)
113    }
114}
weavatrix_search_vector/storage/mapped_exact.rs

weavatrix_search_vector/storage/
mapped_exact.rs
