weavatrix-search-vector 0.3.1

use super::MappedVectorIndex;
use crate::error::SearchError;
use crate::hit::SearchHit;
use crate::hnsw::{FilterSearchPolicy, SearchPolicy};
use crate::vector::Candidate;
use std::collections::BinaryHeap;

impl MappedVectorIndex {
    /// Returns exact top-K hits directly from mapped vectors.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query.
    pub fn search_exact(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.search_exact_filtered(query, count, |_| true)
    }

    /// Returns exact top-K mapped hits accepted by `filter`.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query.
    pub fn search_exact_filtered<F>(
        &self,
        query: &[f32],
        count: usize,
        mut filter: F,
    ) -> Result<Vec<SearchHit>, SearchError>
    where
        F: FnMut(u64) -> bool,
    {
        let query_norm = self.query_squared_norm(query)?;
        let limit = count.min(self.len());
        if limit == 0 {
            return Ok(Vec::new());
        }
        let mut best = BinaryHeap::with_capacity(limit);
        for index in 0..self.len() {
            let key = self.key_slice()[index];
            if !filter(key) {
                continue;
            }
            let candidate = Candidate::new(self.distance_query(index, query, query_norm), index);
            if best.len() < limit {
                best.push(candidate);
            } else if best
                .peek()
                .is_some_and(|worst| candidate.cmp(worst).is_lt())
            {
                best.pop();
                best.push(candidate);
            }
        }
        let mut candidates = best.into_vec();
        candidates.sort_unstable();
        Ok(candidates
            .into_iter()
            .map(|candidate| SearchHit {
                key: self.key_slice()[candidate.index()],
                distance: candidate.distance,
            })
            .collect())
    }

    /// Searches mapped HNSW candidates accepted by `filter`, with an exact
    /// fallback for selective filters.
    ///
    /// # Errors
    ///
    /// Returns a typed query or allocation error.
    pub fn search_filtered<F>(
        &self,
        query: &[f32],
        count: usize,
        filter: F,
    ) -> Result<Vec<SearchHit>, SearchError>
    where
        F: Fn(u64) -> bool,
    {
        self.search_filtered_with_policy(query, count, filter, FilterSearchPolicy::ExactFallback)
    }

    /// Searches a mapped graph with the predicate applied during traversal.
    ///
    /// # Errors
    ///
    /// Returns a typed query or allocation error.
    pub fn search_filtered_with_policy<F>(
        &self,
        query: &[f32],
        count: usize,
        filter: F,
        policy: FilterSearchPolicy,
    ) -> Result<Vec<SearchHit>, SearchError>
    where
        F: Fn(u64) -> bool,
    {
        let requested = count.min(self.len());
        if requested == 0 {
            self.query_squared_norm(query)?;
            return Ok(Vec::new());
        }
        let hits = self.search_where(
            query,
            requested,
            SearchPolicy::new(self.header.config.expansion_query),
            &filter,
        )?;
        if hits.len() == requested || policy == FilterSearchPolicy::Traversal {
            return Ok(hits);
        }
        self.search_exact_filtered(query, requested, filter)
    }
}