weavatrix-search-vector 0.3.1

use super::support::current_len;
use super::{MutableState, MutableVectorIndex};
use crate::config::{DistanceMetric, IndexConfig};
use crate::error::SearchError;
use crate::hit::SearchHit;
use crate::metadata::MetadataFilter;
use crate::vector::{distance, squared_norm};

impl MutableVectorIndex {
    /// Searches the immutable base and exact delta, then deterministically
    /// merges equal keys.
    ///
    /// # Errors
    ///
    /// Returns a typed query or allocation error.
    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.search_where(query, count, |_| true)
    }

    /// Searches only records matching `filter`.
    ///
    /// # Errors
    ///
    /// Returns a typed query or allocation error.
    pub fn search_filtered(
        &self,
        query: &[f32],
        count: usize,
        filter: &MetadataFilter,
    ) -> Result<Vec<SearchHit>, SearchError> {
        let state = self
            .state
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        self.search_locked(&state, query, count, |key| {
            state.metadata.matches(key, filter)
        })
    }

    fn search_where<F>(
        &self,
        query: &[f32],
        count: usize,
        accepts: F,
    ) -> Result<Vec<SearchHit>, SearchError>
    where
        F: Fn(u64) -> bool,
    {
        let state = self
            .state
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        self.search_locked(&state, query, count, accepts)
    }

    fn search_locked<F>(
        &self,
        state: &MutableState,
        query: &[f32],
        count: usize,
        accepts: F,
    ) -> Result<Vec<SearchHit>, SearchError>
    where
        F: Fn(u64) -> bool,
    {
        validate_query(&self.config, query)?;
        let query_squared_norm = squared_norm(query, None)?;
        let limit = count.min(current_len(state));
        if limit == 0 {
            return Ok(Vec::new());
        }
        let mut hits = state.base.search_filtered(query, limit, |key| {
            !state.deleted.contains(&key)
                && !state.pending.contains_key(&key)
                && state
                    .sealed
                    .as_ref()
                    .is_none_or(|index| index.vector(key).is_none())
                && accepts(key)
        })?;
        if let Some(sealed) = &state.sealed {
            let mut sealed_hits = sealed.search_filtered(query, limit, |key| {
                !state.deleted.contains(&key) && !state.pending.contains_key(&key) && accepts(key)
            })?;
            hits.try_reserve(sealed_hits.len())
                .map_err(|_| SearchError::AllocationFailed)?;
            hits.append(&mut sealed_hits);
        }
        hits.try_reserve(state.pending.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        hits.extend(
            state
                .pending
                .iter()
                .filter(|(key, _)| accepts(**key))
                .map(|(key, vector)| SearchHit {
                    key: *key,
                    distance: distance(
                        self.distance_kernel,
                        self.config.metric,
                        vector,
                        pending_norm(self.config.metric, vector),
                        query,
                        query_squared_norm,
                    ),
                }),
        );
        hits.sort_unstable_by(|left, right| {
            left.distance
                .total_cmp(&right.distance)
                .then_with(|| left.key.cmp(&right.key))
        });
        hits.dedup_by_key(|hit| hit.key);
        hits.truncate(limit);
        Ok(hits)
    }
}

fn validate_query(config: &IndexConfig, query: &[f32]) -> Result<(), SearchError> {
    if query.len() != config.dimensions {
        return Err(SearchError::DimensionMismatch {
            expected: config.dimensions,
            actual: query.len(),
            vector: None,
        });
    }
    let query_squared_norm = squared_norm(query, None)?;
    if config.metric == DistanceMetric::Cosine && query_squared_norm == 0.0 {
        return Err(SearchError::ZeroVector { vector: None });
    }
    Ok(())
}

fn pending_norm(metric: DistanceMetric, vector: &[f32]) -> f32 {
    if metric == DistanceMetric::Cosine {
        1.0
    } else {
        squared_norm(vector, None).expect("validated pending vector")
    }
}