weavatrix-search-vector 0.3.1

use crate::config::IndexConfig;
use crate::error::SearchError;
use crate::hit::SearchHit;
use crate::hnsw::VectorIndex;
use std::collections::BTreeSet;

/// Stable caller identity for one vector among several attached to a key.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct MultiVectorKey {
    pub key: u64,
    pub vector_id: u64,
}

/// Borrowed multi-vector build record.
#[derive(Debug, Clone, Copy)]
pub struct MultiVectorRef<'a> {
    pub key: u64,
    pub vector_id: u64,
    pub vector: &'a [f32],
}

/// One nearest-neighbor result retaining both caller identifiers.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct MultiSearchHit {
    pub key: u64,
    pub vector_id: u64,
    pub distance: f32,
}

/// Deterministic vector index allowing multiple vectors per caller key.
#[derive(Debug)]
pub struct MultiVectorIndex {
    pub(crate) index: VectorIndex,
    pub(crate) identities: Vec<MultiVectorKey>,
}

impl MultiVectorIndex {
    /// Builds one HNSW node per unique `(key, vector_id)` pair.
    ///
    /// Input order does not affect internal IDs or graph construction.
    ///
    /// # Errors
    ///
    /// Returns typed config, vector, duplicate identity, capacity, or
    /// allocation errors.
    pub fn build(config: IndexConfig, vectors: &[MultiVectorRef<'_>]) -> Result<Self, SearchError> {
        if vectors.len() > u32::MAX as usize {
            return Err(SearchError::CapacityOverflow);
        }
        let mut order = (0..vectors.len()).collect::<Vec<_>>();
        order.sort_unstable_by_key(|index| (vectors[*index].key, vectors[*index].vector_id));
        for pair in order.windows(2) {
            let left = vectors[pair[0]];
            let right = vectors[pair[1]];
            if (left.key, left.vector_id) == (right.key, right.vector_id) {
                return Err(SearchError::DuplicateVectorId {
                    key: left.key,
                    vector_id: left.vector_id,
                });
            }
        }
        let mut identities = Vec::new();
        identities
            .try_reserve_exact(vectors.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        let mut internal = Vec::new();
        internal
            .try_reserve_exact(vectors.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        for (internal_id, source) in order.into_iter().enumerate() {
            let record = vectors[source];
            identities.push(MultiVectorKey {
                key: record.key,
                vector_id: record.vector_id,
            });
            internal.push((
                u64::try_from(internal_id).map_err(|_| SearchError::CapacityOverflow)?,
                record.vector,
            ));
        }
        Ok(Self {
            index: VectorIndex::build(config, &internal)?,
            identities,
        })
    }

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

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

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

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

    #[must_use]
    pub fn estimated_memory_bytes(&self) -> usize {
        self.index.estimated_memory_bytes().saturating_add(
            self.identities
                .capacity()
                .saturating_mul(std::mem::size_of::<MultiVectorKey>()),
        )
    }

    /// Returns top vectors, allowing repeated caller keys.
    ///
    /// # Errors
    ///
    /// Returns a typed query or allocation error.
    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<MultiSearchHit>, SearchError> {
        self.map_hits(self.index.search(query, count)?)
    }

    /// Returns exact top vectors, allowing repeated caller keys.
    ///
    /// # Errors
    ///
    /// Returns a typed query or allocation error.
    pub fn search_exact(
        &self,
        query: &[f32],
        count: usize,
    ) -> Result<Vec<MultiSearchHit>, SearchError> {
        self.map_hits(self.index.search_exact(query, count)?)
    }

    /// Returns at most one best vector for each caller key.
    ///
    /// The approximate pass widens candidates; an exact fallback fills sparse
    /// key distributions.
    ///
    /// # Errors
    ///
    /// Returns a typed query or allocation error.
    pub fn search_unique_keys(
        &self,
        query: &[f32],
        count: usize,
    ) -> Result<Vec<MultiSearchHit>, SearchError> {
        let requested = count.min(self.distinct_key_count());
        if requested == 0 {
            self.index.search(query, 0)?;
            return Ok(Vec::new());
        }
        let candidate_count = requested.saturating_mul(8).min(self.len());
        let mut hits = Self::unique(self.search(query, candidate_count)?, requested);
        if hits.len() < requested {
            hits = Self::unique(self.search_exact(query, self.len())?, requested);
        }
        Ok(hits)
    }

    /// Searches only identities accepted by `filter`.
    ///
    /// # Errors
    ///
    /// Returns a typed query or allocation error.
    pub fn search_filtered<F>(
        &self,
        query: &[f32],
        count: usize,
        filter: F,
    ) -> Result<Vec<MultiSearchHit>, SearchError>
    where
        F: Fn(MultiVectorKey) -> bool,
    {
        let hits = self.index.search_filtered(query, count, |internal| {
            usize::try_from(internal)
                .ok()
                .and_then(|index| self.identities.get(index))
                .is_some_and(|identity| filter(*identity))
        })?;
        self.map_hits(hits)
    }

    /// Searches independent queries with bounded workers.
    ///
    /// # Errors
    ///
    /// Returns the first query or worker error.
    pub fn search_batch(
        &self,
        queries: &[&[f32]],
        count: usize,
    ) -> Result<Vec<Vec<MultiSearchHit>>, SearchError> {
        self.index
            .search_batch(queries, count)?
            .into_iter()
            .map(|hits| self.map_hits(hits))
            .collect()
    }

    fn map_hits(&self, hits: Vec<SearchHit>) -> Result<Vec<MultiSearchHit>, SearchError> {
        let mut mapped = Vec::new();
        mapped
            .try_reserve_exact(hits.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        for hit in hits {
            let index = usize::try_from(hit.key).map_err(|_| SearchError::CapacityOverflow)?;
            let identity = self
                .identities
                .get(index)
                .ok_or(SearchError::CorruptSnapshot(
                    "multi-vector internal key is outside identity table",
                ))?;
            mapped.push(MultiSearchHit {
                key: identity.key,
                vector_id: identity.vector_id,
                distance: hit.distance,
            });
        }
        Ok(mapped)
    }

    fn unique(hits: Vec<MultiSearchHit>, count: usize) -> Vec<MultiSearchHit> {
        let mut seen = BTreeSet::new();
        hits.into_iter()
            .filter(|hit| seen.insert(hit.key))
            .take(count)
            .collect()
    }

    fn distinct_key_count(&self) -> usize {
        self.identities
            .iter()
            .map(|identity| identity.key)
            .collect::<BTreeSet<_>>()
            .len()
    }
}