weavatrix-search-vector 0.2.0

use crate::parallel;
use crate::vector::{Candidate, inverse_norm};
use crate::{IndexConfig, SearchError, SearchHit, VectorIndex};
use std::collections::BinaryHeap;

const SCALE: f32 = 127.0;

/// Compact scalar-int8 cosine index.
///
/// Stored vectors use one byte per component plus one norm scalar. Search
/// distances are quantized approximations unless [`Self::search_rerank`] is
/// used with an owned f32 index.
#[derive(Debug)]
pub struct ScalarQuantizedIndex {
    config: IndexConfig,
    keys: Vec<u64>,
    values: Vec<i8>,
    inverse_norms: Vec<f32>,
}

impl ScalarQuantizedIndex {
    /// Validates, normalizes, and quantizes dense vectors.
    ///
    /// # Errors
    ///
    /// Returns typed config, vector, duplicate-key, capacity, or allocation
    /// errors.
    pub fn build(config: IndexConfig, vectors: &[(u64, &[f32])]) -> Result<Self, SearchError> {
        config.validate()?;
        let elements = config
            .dimensions
            .checked_mul(vectors.len())
            .ok_or(SearchError::CapacityOverflow)?;
        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 mut keys = Vec::new();
        keys.try_reserve_exact(vectors.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        let mut values = Vec::new();
        values
            .try_reserve_exact(elements)
            .map_err(|_| SearchError::AllocationFailed)?;
        let mut inverse_norms = Vec::new();
        inverse_norms
            .try_reserve_exact(vectors.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        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 inverse = inverse_norm(vector, Some(source))?;
            let start = values.len();
            values.extend(
                vector
                    .iter()
                    .map(|value| quantize_component(value * inverse)),
            );
            let integer_norm = values[start..]
                .iter()
                .map(|value| {
                    let value = f32::from(*value);
                    value * value
                })
                .sum::<f32>()
                .sqrt();
            if integer_norm == 0.0 {
                return Err(SearchError::ZeroVector {
                    vector: Some(source),
                });
            }
            keys.push(key);
            inverse_norms.push(integer_norm.recip());
        }
        Ok(Self {
            config,
            keys,
            values,
            inverse_norms,
        })
    }

    #[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 config(&self) -> &IndexConfig {
        &self.config
    }

    /// Allocation-based retained size estimate.
    #[must_use]
    pub fn estimated_memory_bytes(&self) -> usize {
        self.keys
            .capacity()
            .saturating_mul(std::mem::size_of::<u64>())
            .saturating_add(
                self.values
                    .capacity()
                    .saturating_mul(std::mem::size_of::<i8>()),
            )
            .saturating_add(
                self.inverse_norms
                    .capacity()
                    .saturating_mul(std::mem::size_of::<f32>()),
            )
    }

    /// Returns scalar-int8 approximate cosine hits.
    ///
    /// # Errors
    ///
    /// Returns a typed query validation or allocation error.
    #[allow(clippy::cast_precision_loss)]
    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        let (quantized_query, query_inverse_norm) = self.quantize_query(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 dot = self
                .vector(index)
                .iter()
                .zip(&quantized_query)
                .map(|(left, right)| i64::from(*left) * i64::from(*right))
                .sum::<i64>();
            let similarity = dot as f32 * self.inverse_norms[index] * query_inverse_norm;
            let candidate = Candidate::new((1.0 - similarity).clamp(0.0, 2.0), 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())
    }

    /// Retrieves a wider quantized candidate set and re-scores it with exact
    /// normalized f32 vectors from `oracle`.
    ///
    /// # Errors
    ///
    /// Returns a typed query, missing-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() {
            return Err(SearchError::DimensionMismatch {
                expected: self.dimensions(),
                actual: oracle.dimensions(),
                vector: None,
            });
        }
        let query_inverse = inverse_norm(query, None)?;
        let candidate_count = candidates.max(count).min(self.len());
        let mut hits = self.search(query, candidate_count)?;
        for hit in &mut hits {
            let vector = oracle
                .vector(hit.key)
                .ok_or(SearchError::MissingKey(hit.key))?;
            let dot = vector
                .iter()
                .zip(query)
                .map(|(left, right)| left * right)
                .sum::<f32>();
            hit.distance = (1.0 - dot * query_inverse).clamp(0.0, 2.0);
        }
        hits.sort_unstable_by(|left, right| {
            left.distance
                .total_cmp(&right.distance)
                .then_with(|| left.key.cmp(&right.key))
        });
        hits.truncate(count.min(hits.len()));
        Ok(hits)
    }

    /// Searches independent queries with bounded workers.
    ///
    /// # 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> {
        parallel::search_batch(queries, self.config.query_threads, |query| {
            self.search(query, count)
        })
    }

    fn vector(&self, index: usize) -> &[i8] {
        let start = index * self.dimensions();
        &self.values[start..start + self.dimensions()]
    }

    fn quantize_query(&self, query: &[f32]) -> Result<(Vec<i8>, f32), SearchError> {
        if query.len() != self.dimensions() {
            return Err(SearchError::DimensionMismatch {
                expected: self.dimensions(),
                actual: query.len(),
                vector: None,
            });
        }
        let inverse = inverse_norm(query, None)?;
        let mut quantized = Vec::new();
        quantized
            .try_reserve_exact(query.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        quantized.extend(
            query
                .iter()
                .map(|value| quantize_component(value * inverse)),
        );
        let norm = quantized
            .iter()
            .map(|value| {
                let value = f32::from(*value);
                value * value
            })
            .sum::<f32>()
            .sqrt();
        if norm == 0.0 {
            return Err(SearchError::ZeroVector { vector: None });
        }
        Ok((quantized, norm.recip()))
    }
}

#[allow(clippy::cast_possible_truncation)]
fn quantize_component(value: f32) -> i8 {
    (value * SCALE).round().clamp(-SCALE, SCALE) as i8
}