weavatrix_search_vector/hnsw/
batch.rs

1use super::{SearchPolicy, SearchScratch, VectorIndex};
2use crate::error::SearchError;
3use crate::hit::SearchHit;
4
5impl VectorIndex {
6    /// Searches independent queries using bounded scoped workers. Each worker
7    /// reuses its visited-set and heaps across queries.
8    ///
9    /// # Errors
10    ///
11    /// Returns the first query error in input order or a worker-panic error.
12    pub fn search_batch(
13        &self,
14        queries: &[&[f32]],
15        count: usize,
16    ) -> Result<Vec<Vec<SearchHit>>, SearchError> {
17        self.search_batch_with_policy(
18            queries,
19            count,
20            SearchPolicy::new(self.config.expansion_query),
21        )
22    }
23
24    /// Searches independent queries with a shared per-query recall policy.
25    ///
26    /// # Errors
27    ///
28    /// Returns the first query error in input order or a worker-panic error.
29    pub fn search_batch_with_policy(
30        &self,
31        queries: &[&[f32]],
32        count: usize,
33        policy: SearchPolicy,
34    ) -> Result<Vec<Vec<SearchHit>>, SearchError> {
35        let policy = policy.validate()?;
36        if queries.is_empty() {
37            return Ok(Vec::new());
38        }
39        let workers = self.config.query_threads.min(queries.len()).max(1);
40        let chunk_size = queries.len().div_ceil(workers);
41        let mut output = Vec::new();
42        output
43            .try_reserve_exact(queries.len())
44            .map_err(|_| SearchError::AllocationFailed)?;
45        output.extend(std::iter::repeat_with(|| None).take(queries.len()));
46        let panicked = std::thread::scope(|scope| {
47            let handles = queries
48                .chunks(chunk_size)
49                .enumerate()
50                .map(|(chunk_index, chunk)| {
51                    let start = chunk_index * chunk_size;
52                    scope.spawn(move || {
53                        let mut scratch = SearchScratch::new(self.len());
54                        let mut local = Vec::with_capacity(chunk.len());
55                        for (offset, query) in chunk.iter().enumerate() {
56                            let result = match &mut scratch {
57                                Ok(scratch) => {
58                                    self.search_with_scratch(query, count, policy, scratch)
59                                }
60                                Err(error) => Err(error.clone()),
61                            };
62                            local.push((start + offset, result));
63                        }
64                        local
65                    })
66                })
67                .collect::<Vec<_>>();
68            let mut panicked = false;
69            for handle in handles {
70                match handle.join() {
71                    Ok(local) => {
72                        for (index, result) in local {
73                            output[index] = Some(result);
74                        }
75                    }
76                    Err(_) => panicked = true,
77                }
78            }
79            panicked
80        });
81        if panicked {
82            return Err(SearchError::WorkerPanic);
83        }
84        output
85            .into_iter()
86            .map(|slot| slot.ok_or(SearchError::WorkerPanic)?)
87            .collect()
88    }
89}
weavatrix_search_vector/hnsw/batch.rs

weavatrix_search_vector/hnsw/
batch.rs
