weavatrix-search-vector 0.3.1

use super::graph_helpers::from_node;
use super::{FilterSearchPolicy, SearchPolicy, SearchScratch, VectorIndex};
use crate::error::SearchError;
use crate::hit::SearchHit;
use crate::vector::Candidate;

impl VectorIndex {
    /// Returns approximate top-K hits ordered by exact distance and key.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query or scratch allocation.
    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.search_with_policy(query, count, SearchPolicy::new(self.config.expansion_query))
    }

    /// Returns approximate top-K hits with a per-query recall policy.
    ///
    /// This separates graph expansion from deterministic routing recovery so
    /// callers can raise recall without rebuilding the index.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid policy, query, or scratch
    /// allocation.
    pub fn search_with_policy(
        &self,
        query: &[f32],
        count: usize,
        policy: SearchPolicy,
    ) -> Result<Vec<SearchHit>, SearchError> {
        let policy = policy.validate()?;
        let mut scratch = SearchScratch::new(self.len())?;
        self.search_with_scratch(query, count, policy, &mut scratch)
    }

    /// Returns exact top-K results over the same normalized vectors.
    ///
    /// This is intended as a correctness and recall oracle.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query.
    pub fn search_exact(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.vectors.exact(query, count)
    }

    /// Returns exact top-K results accepted by `filter`.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query.
    pub fn search_exact_filtered<F>(
        &self,
        query: &[f32],
        count: usize,
        filter: F,
    ) -> Result<Vec<SearchHit>, SearchError>
    where
        F: Fn(u64) -> bool,
    {
        self.vectors.exact_filtered(query, count, filter)
    }

    /// Searches HNSW candidates accepted by `filter`, with an exact fallback
    /// when selective filters do not yield enough approximate candidates.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query or scratch allocation.
    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 with the predicate applied during graph traversal.
    ///
    /// `Traversal` avoids a full scan and may return fewer than `count` hits.
    /// `ExactFallback` preserves completeness when the graph cannot find
    /// enough accepted candidates.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query or scratch allocation.
    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.vectors.query_squared_norm(query)?;
            return Ok(Vec::new());
        }
        let mut scratch = SearchScratch::new(self.len())?;
        let hits = self.search_with_scratch_where(
            query,
            requested,
            SearchPolicy::new(self.config.expansion_query),
            &mut scratch,
            &filter,
        )?;
        if hits.len() == requested || policy == FilterSearchPolicy::Traversal {
            return Ok(hits);
        }
        self.vectors.exact_filtered(query, requested, filter)
    }

    pub(super) fn search_with_scratch(
        &self,
        query: &[f32],
        count: usize,
        policy: SearchPolicy,
        scratch: &mut SearchScratch,
    ) -> Result<Vec<SearchHit>, SearchError> {
        self.search_with_scratch_where(query, count, policy, scratch, &|_| true)
    }

    fn search_with_scratch_where<F>(
        &self,
        query: &[f32],
        count: usize,
        policy: SearchPolicy,
        scratch: &mut SearchScratch,
        accepts: &F,
    ) -> Result<Vec<SearchHit>, SearchError>
    where
        F: Fn(u64) -> bool,
    {
        let query_norm = self.vectors.query_squared_norm(query)?;
        let limit = count.min(self.len());
        if limit == 0 {
            return Ok(Vec::new());
        }
        let expansion = policy.expansion.max(limit);
        let mut merged = std::mem::take(&mut scratch.merged);
        merged.clear();
        merged
            .try_reserve(self.graphs.len().saturating_mul(expansion))
            .map_err(|_| SearchError::AllocationFailed)?;
        for graph in &self.graphs {
            graph.search_into(
                &self.vectors,
                query,
                query_norm,
                expansion,
                scratch,
                &mut merged,
                accepts,
            );
        }
        self.append_routing_candidates(query, query_norm, policy, scratch, accepts, &mut merged)?;
        merged.sort_unstable();
        let hits = collect_unique_hits(&self.vectors, &merged, limit)?;
        merged.clear();
        scratch.merged = merged;
        Ok(hits)
    }

    fn append_routing_candidates<F>(
        &self,
        query: &[f32],
        query_norm: f32,
        policy: SearchPolicy,
        scratch: &mut SearchScratch,
        accepts: &F,
        merged: &mut Vec<Candidate>,
    ) -> Result<(), SearchError>
    where
        F: Fn(u64) -> bool,
    {
        let mut routing_nodes = std::mem::take(&mut scratch.routing_nodes);
        routing_nodes.clear();
        self.vectors.routing_probes_into(
            query,
            policy.routing_probes,
            &mut scratch.routing_probes,
            &mut scratch.routing_probe_heap,
        )?;
        self.routing
            .append_candidates(&scratch.routing_probes, &mut routing_nodes);
        merged
            .try_reserve(routing_nodes.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        merged.extend(
            routing_nodes
                .iter()
                .copied()
                .map(from_node)
                .filter(|index| accepts(self.vectors.key(*index)))
                .map(|index| {
                    Candidate::new(self.vectors.distance_query(index, query, query_norm), index)
                }),
        );
        routing_nodes.clear();
        scratch.routing_nodes = routing_nodes;
        Ok(())
    }
}

fn collect_unique_hits(
    vectors: &crate::vector::VectorStore,
    candidates: &[Candidate],
    limit: usize,
) -> Result<Vec<SearchHit>, SearchError> {
    let mut hits = Vec::new();
    hits.try_reserve_exact(limit)
        .map_err(|_| SearchError::AllocationFailed)?;
    for candidate in candidates.iter().copied() {
        let key = vectors.key(candidate.index());
        if hits.iter().any(|hit: &SearchHit| hit.key == key) {
            continue;
        }
        hits.push(SearchHit {
            key,
            distance: candidate.distance,
        });
        if hits.len() == limit {
            break;
        }
    }
    Ok(hits)
}