weavatrix_search_vector/mutable/
search.rs

1use super::support::current_len;
2use super::{MutableState, MutableVectorIndex};
3use crate::config::{DistanceMetric, IndexConfig};
4use crate::error::SearchError;
5use crate::hit::SearchHit;
6use crate::metadata::MetadataFilter;
7use crate::vector::{distance, squared_norm};
8
9impl MutableVectorIndex {
10    /// Searches the immutable base and exact delta, then deterministically
11    /// merges equal keys.
12    ///
13    /// # Errors
14    ///
15    /// Returns a typed query or allocation error.
16    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
17        self.search_where(query, count, |_| true)
18    }
19
20    /// Searches only records matching `filter`.
21    ///
22    /// # Errors
23    ///
24    /// Returns a typed query or allocation error.
25    pub fn search_filtered(
26        &self,
27        query: &[f32],
28        count: usize,
29        filter: &MetadataFilter,
30    ) -> Result<Vec<SearchHit>, SearchError> {
31        let state = self
32            .state
33            .read()
34            .unwrap_or_else(std::sync::PoisonError::into_inner);
35        self.search_locked(&state, query, count, |key| {
36            state.metadata.matches(key, filter)
37        })
38    }
39
40    fn search_where<F>(
41        &self,
42        query: &[f32],
43        count: usize,
44        accepts: F,
45    ) -> Result<Vec<SearchHit>, SearchError>
46    where
47        F: Fn(u64) -> bool,
48    {
49        let state = self
50            .state
51            .read()
52            .unwrap_or_else(std::sync::PoisonError::into_inner);
53        self.search_locked(&state, query, count, accepts)
54    }
55
56    fn search_locked<F>(
57        &self,
58        state: &MutableState,
59        query: &[f32],
60        count: usize,
61        accepts: F,
62    ) -> Result<Vec<SearchHit>, SearchError>
63    where
64        F: Fn(u64) -> bool,
65    {
66        validate_query(&self.config, query)?;
67        let query_squared_norm = squared_norm(query, None)?;
68        let limit = count.min(current_len(state));
69        if limit == 0 {
70            return Ok(Vec::new());
71        }
72        let mut hits = state.base.search_filtered(query, limit, |key| {
73            !state.deleted.contains(&key)
74                && !state.pending.contains_key(&key)
75                && state
76                    .sealed
77                    .as_ref()
78                    .is_none_or(|index| index.vector(key).is_none())
79                && accepts(key)
80        })?;
81        if let Some(sealed) = &state.sealed {
82            let mut sealed_hits = sealed.search_filtered(query, limit, |key| {
83                !state.deleted.contains(&key) && !state.pending.contains_key(&key) && accepts(key)
84            })?;
85            hits.try_reserve(sealed_hits.len())
86                .map_err(|_| SearchError::AllocationFailed)?;
87            hits.append(&mut sealed_hits);
88        }
89        hits.try_reserve(state.pending.len())
90            .map_err(|_| SearchError::AllocationFailed)?;
91        hits.extend(
92            state
93                .pending
94                .iter()
95                .filter(|(key, _)| accepts(**key))
96                .map(|(key, vector)| SearchHit {
97                    key: *key,
98                    distance: distance(
99                        self.distance_kernel,
100                        self.config.metric,
101                        vector,
102                        pending_norm(self.config.metric, vector),
103                        query,
104                        query_squared_norm,
105                    ),
106                }),
107        );
108        hits.sort_unstable_by(|left, right| {
109            left.distance
110                .total_cmp(&right.distance)
111                .then_with(|| left.key.cmp(&right.key))
112        });
113        hits.dedup_by_key(|hit| hit.key);
114        hits.truncate(limit);
115        Ok(hits)
116    }
117}
118
119fn validate_query(config: &IndexConfig, query: &[f32]) -> Result<(), SearchError> {
120    if query.len() != config.dimensions {
121        return Err(SearchError::DimensionMismatch {
122            expected: config.dimensions,
123            actual: query.len(),
124            vector: None,
125        });
126    }
127    let query_squared_norm = squared_norm(query, None)?;
128    if config.metric == DistanceMetric::Cosine && query_squared_norm == 0.0 {
129        return Err(SearchError::ZeroVector { vector: None });
130    }
131    Ok(())
132}
133
134fn pending_norm(metric: DistanceMetric, vector: &[f32]) -> f32 {
135    if metric == DistanceMetric::Cosine {
136        1.0
137    } else {
138        squared_norm(vector, None).expect("validated pending vector")
139    }
140}
weavatrix_search_vector/mutable/search.rs

weavatrix_search_vector/mutable/
search.rs
