weavatrix-search-vector 0.3.1

use super::storage::{QuantizedStorage, binary_word};
use crate::config::{DistanceMetric, IndexConfig};
use crate::error::SearchError;
use crate::hit::SearchHit;
use crate::hnsw::VectorIndex;
use crate::parallel;
use crate::vector::{Candidate, squared_norm};
use std::collections::BinaryHeap;

/// Compact vector representation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum QuantizationKind {
    BFloat16,
    Float16,
    Float8E4M3,
    Int8,
    Binary,
}

/// Exact scan over a selected compact representation.
///
/// `Int8` and `Binary` require cosine distance. Floating-point formats support
/// every [`DistanceMetric`]. Returned compact distances can be re-scored with
/// [`Self::search_rerank`].
#[derive(Debug)]
pub struct QuantizedIndex {
    pub(crate) config: IndexConfig,
    pub(crate) kind: QuantizationKind,
    pub(crate) keys: Vec<u64>,
    pub(crate) storage: QuantizedStorage,
    pub(crate) squared_norms: Vec<f32>,
    pub(crate) words_per_vector: usize,
}

impl QuantizedIndex {
    /// Builds a compact exact-scan index.
    ///
    /// # Errors
    ///
    /// Returns typed config, vector, key, capacity, or allocation errors.
    pub fn build(
        config: IndexConfig,
        kind: QuantizationKind,
        vectors: &[(u64, &[f32])],
    ) -> Result<Self, SearchError> {
        config.validate()?;
        if matches!(kind, QuantizationKind::Int8 | QuantizationKind::Binary)
            && config.metric != DistanceMetric::Cosine
        {
            return Err(SearchError::InvalidConfig(
                "int8 and binary quantization require cosine distance",
            ));
        }
        let mut order = (0..vectors.len()).collect::<Vec<_>>();
        order.sort_unstable_by_key(|index| vectors[*index].0);
        for pair in order.windows(2) {
            if vectors[pair[0]].0 == vectors[pair[1]].0 {
                return Err(SearchError::DuplicateKey(vectors[pair[0]].0));
            }
        }
        let elements = config
            .dimensions
            .checked_mul(vectors.len())
            .ok_or(SearchError::CapacityOverflow)?;
        let words_per_vector = config.dimensions.div_ceil(64);
        let mut keys = Vec::new();
        keys.try_reserve_exact(vectors.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        let mut squared_norms = Vec::new();
        squared_norms
            .try_reserve_exact(vectors.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        let mut storage = QuantizedStorage::with_capacity(
            kind,
            elements,
            words_per_vector
                .checked_mul(vectors.len())
                .ok_or(SearchError::CapacityOverflow)?,
        )?;
        for source in order {
            let (key, vector) = vectors[source];
            if vector.len() != config.dimensions {
                return Err(SearchError::DimensionMismatch {
                    expected: config.dimensions,
                    actual: vector.len(),
                    vector: Some(source),
                });
            }
            let norm_squared = squared_norm(vector, Some(source))?;
            if config.metric == DistanceMetric::Cosine && norm_squared == 0.0 {
                return Err(SearchError::ZeroVector {
                    vector: Some(source),
                });
            }
            let inverse = if config.metric == DistanceMetric::Cosine {
                norm_squared.sqrt().recip()
            } else {
                1.0
            };
            keys.push(key);
            let quantized_norm = storage.push_vector(vector, inverse, words_per_vector)?;
            squared_norms.push(quantized_norm);
        }
        Ok(Self {
            config,
            kind,
            keys,
            storage,
            squared_norms,
            words_per_vector,
        })
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.keys.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.keys.is_empty()
    }

    #[must_use]
    pub const fn dimensions(&self) -> usize {
        self.config.dimensions
    }

    #[must_use]
    pub const fn kind(&self) -> QuantizationKind {
        self.kind
    }

    #[must_use]
    pub const fn config(&self) -> &IndexConfig {
        &self.config
    }

    #[must_use]
    pub fn estimated_memory_bytes(&self) -> usize {
        self.keys
            .capacity()
            .saturating_mul(std::mem::size_of::<u64>())
            .saturating_add(
                self.squared_norms
                    .capacity()
                    .saturating_mul(std::mem::size_of::<f32>()),
            )
            .saturating_add(self.storage.estimated_bytes())
    }

    /// Searches all compact vectors.
    ///
    /// # Errors
    ///
    /// Returns a typed query validation or allocation error.
    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        if query.len() != self.dimensions() {
            return Err(SearchError::DimensionMismatch {
                expected: self.dimensions(),
                actual: query.len(),
                vector: None,
            });
        }
        let query_squared_norm = squared_norm(query, None)?;
        if self.config.metric == DistanceMetric::Cosine && query_squared_norm == 0.0 {
            return Err(SearchError::ZeroVector { vector: None });
        }
        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 distance = self.compact_distance(index, query, query_squared_norm);
            let candidate = Candidate::new(distance, 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.keys[candidate.index()],
                distance: candidate.distance,
            })
            .collect())
    }

    /// Searches independent queries with bounded workers.
    ///
    /// # Errors
    ///
    /// Returns the first input or worker error.
    pub fn search_batch(
        &self,
        queries: &[&[f32]],
        count: usize,
    ) -> Result<Vec<Vec<SearchHit>>, SearchError> {
        parallel::search_batch(queries, self.config.query_threads, |query| {
            self.search(query, count)
        })
    }

    /// Retrieves compact candidates and re-scores them against an f32 index.
    ///
    /// # Errors
    ///
    /// Returns a typed dimension, query, key, or allocation error.
    pub fn search_rerank(
        &self,
        oracle: &VectorIndex,
        query: &[f32],
        count: usize,
        candidates: usize,
    ) -> Result<Vec<SearchHit>, SearchError> {
        if oracle.dimensions() != self.dimensions() || oracle.config().metric != self.config.metric
        {
            return Err(SearchError::InvalidConfig(
                "rerank oracle config differs from quantized index",
            ));
        }
        let candidate_count = candidates.max(count).min(self.len());
        let compact = self.search(query, candidate_count)?;
        oracle.search_exact_filtered(query, count, |key| compact.iter().any(|hit| hit.key == key))
    }

    #[allow(clippy::cast_precision_loss)]
    fn compact_distance(&self, index: usize, query: &[f32], query_squared_norm: f32) -> f32 {
        match &self.storage {
            QuantizedStorage::Binary(words) => {
                let start = index * self.words_per_vector;
                let mut different = 0_u32;
                for word_index in 0..self.words_per_vector {
                    let query_word = binary_word(query, word_index);
                    different = different
                        .saturating_add((words[start + word_index] ^ query_word).count_ones());
                }
                2.0 * different as f32 / self.dimensions() as f32
            }
            storage => {
                let start = index * self.dimensions();
                let mut dot = 0.0_f32;
                for (dimension, query_value) in query.iter().copied().enumerate() {
                    dot += storage.value(start + dimension) * query_value;
                }
                match self.config.metric {
                    DistanceMetric::Cosine => {
                        let denominator = (self.squared_norms[index] * query_squared_norm).sqrt();
                        (1.0 - dot / denominator).clamp(0.0, 2.0)
                    }
                    DistanceMetric::Dot => -dot,
                    DistanceMetric::SquaredEuclidean => self.squared_norms[index]
                        .mul_add(1.0, query_squared_norm - 2.0 * dot)
                        .max(0.0),
                }
            }
        }
    }
}