weavatrix-search-vector 0.3.1

use super::{SearchPolicy, SearchScratch, VectorIndex};
use crate::error::SearchError;
use crate::hit::SearchHit;

impl VectorIndex {
    /// Searches independent queries using bounded scoped workers. Each worker
    /// reuses its visited-set and heaps across queries.
    ///
    /// # Errors
    ///
    /// Returns the first query error in input order or a worker-panic error.
    pub fn search_batch(
        &self,
        queries: &[&[f32]],
        count: usize,
    ) -> Result<Vec<Vec<SearchHit>>, SearchError> {
        self.search_batch_with_policy(
            queries,
            count,
            SearchPolicy::new(self.config.expansion_query),
        )
    }

    /// Searches independent queries with a shared per-query recall policy.
    ///
    /// # Errors
    ///
    /// Returns the first query error in input order or a worker-panic error.
    pub fn search_batch_with_policy(
        &self,
        queries: &[&[f32]],
        count: usize,
        policy: SearchPolicy,
    ) -> Result<Vec<Vec<SearchHit>>, SearchError> {
        let policy = policy.validate()?;
        if queries.is_empty() {
            return Ok(Vec::new());
        }
        let workers = self.config.query_threads.min(queries.len()).max(1);
        let chunk_size = queries.len().div_ceil(workers);
        let mut output = Vec::new();
        output
            .try_reserve_exact(queries.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        output.extend(std::iter::repeat_with(|| None).take(queries.len()));
        let panicked = std::thread::scope(|scope| {
            let handles = queries
                .chunks(chunk_size)
                .enumerate()
                .map(|(chunk_index, chunk)| {
                    let start = chunk_index * chunk_size;
                    scope.spawn(move || {
                        let mut scratch = SearchScratch::new(self.len());
                        let mut local = Vec::with_capacity(chunk.len());
                        for (offset, query) in chunk.iter().enumerate() {
                            let result = match &mut scratch {
                                Ok(scratch) => {
                                    self.search_with_scratch(query, count, policy, scratch)
                                }
                                Err(error) => Err(error.clone()),
                            };
                            local.push((start + offset, result));
                        }
                        local
                    })
                })
                .collect::<Vec<_>>();
            let mut panicked = false;
            for handle in handles {
                match handle.join() {
                    Ok(local) => {
                        for (index, result) in local {
                            output[index] = Some(result);
                        }
                    }
                    Err(_) => panicked = true,
                }
            }
            panicked
        });
        if panicked {
            return Err(SearchError::WorkerPanic);
        }
        output
            .into_iter()
            .map(|slot| slot.ok_or(SearchError::WorkerPanic)?)
            .collect()
    }
}