weavatrix-search-vector 0.2.0

use crate::{
    MappedVectorIndex, MutableVectorIndex, ScalarQuantizedIndex, SearchError, SearchHit,
    VectorIndex,
};
use std::collections::BTreeMap;
use std::sync::Arc;

/// Transport-neutral vector-search shard.
///
/// Implementations may be local indexes or remote clients. Remote adapters
/// should translate transport failures into [`SearchError::ShardFailed`].
pub trait SearchShard: Send + Sync {
    fn dimensions(&self) -> usize;

    /// Searches this shard.
    ///
    /// # Errors
    ///
    /// Returns a query, allocation, worker, transport, or provider error.
    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError>;
}

impl SearchShard for VectorIndex {
    fn dimensions(&self) -> usize {
        self.dimensions()
    }

    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.search(query, count)
    }
}

impl SearchShard for MappedVectorIndex {
    fn dimensions(&self) -> usize {
        self.dimensions()
    }

    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.search(query, count)
    }
}

impl SearchShard for ScalarQuantizedIndex {
    fn dimensions(&self) -> usize {
        self.dimensions()
    }

    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.search(query, count)
    }
}

impl SearchShard for MutableVectorIndex {
    fn dimensions(&self) -> usize {
        self.config().dimensions
    }

    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.search(query, count)
    }
}

/// Bounded coordinator policy for shard fan-out and stable result merging.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DistributedConfig {
    /// Total local worker budget shared by query and shard parallelism.
    pub worker_budget: usize,
    /// Per-shard candidate count is `top_k * candidate_multiplier`.
    pub candidate_multiplier: usize,
}

impl DistributedConfig {
    #[must_use]
    pub const fn new(worker_budget: usize) -> Self {
        Self {
            worker_budget,
            candidate_multiplier: 2,
        }
    }

    fn validate(self) -> Result<Self, SearchError> {
        if self.worker_budget == 0 {
            return Err(SearchError::InvalidConfig(
                "distributed worker_budget must be non-zero",
            ));
        }
        if self.candidate_multiplier == 0 {
            return Err(SearchError::InvalidConfig(
                "distributed candidate_multiplier must be non-zero",
            ));
        }
        Ok(self)
    }
}

/// Heterogeneous local/remote shard coordinator.
pub struct DistributedIndex {
    dimensions: usize,
    config: DistributedConfig,
    shards: Vec<Arc<dyn SearchShard>>,
}

impl std::fmt::Debug for DistributedIndex {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("DistributedIndex")
            .field("dimensions", &self.dimensions)
            .field("config", &self.config)
            .field("shards", &self.shards.len())
            .finish()
    }
}

impl DistributedIndex {
    /// Creates an empty coordinator for fixed-dimensional shards.
    ///
    /// # Errors
    ///
    /// Returns a typed config error.
    pub fn new(dimensions: usize, config: DistributedConfig) -> Result<Self, SearchError> {
        if dimensions == 0 {
            return Err(SearchError::InvalidConfig(
                "distributed dimensions must be non-zero",
            ));
        }
        Ok(Self {
            dimensions,
            config: config.validate()?,
            shards: Vec::new(),
        })
    }

    /// Adds one local or remote shard.
    ///
    /// # Errors
    ///
    /// Returns a dimension mismatch when shard and coordinator differ.
    pub fn push(&mut self, shard: Arc<dyn SearchShard>) -> Result<(), SearchError> {
        if shard.dimensions() != self.dimensions {
            return Err(SearchError::DimensionMismatch {
                expected: self.dimensions,
                actual: shard.dimensions(),
                vector: None,
            });
        }
        self.shards.push(shard);
        Ok(())
    }

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

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

    /// Fans one query out with bounded workers and keeps the best occurrence
    /// of duplicate keys.
    ///
    /// # Errors
    ///
    /// Returns a query, worker, allocation, or shard error.
    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.search_with_workers(query, count, self.config.worker_budget)
    }

    /// Searches independent queries while keeping total local threads within
    /// `worker_budget`.
    ///
    /// # Errors
    ///
    /// Returns the first query or shard error in input order.
    pub fn search_batch(
        &self,
        queries: &[&[f32]],
        count: usize,
    ) -> Result<Vec<Vec<SearchHit>>, SearchError> {
        if queries.is_empty() {
            return Ok(Vec::new());
        }
        let query_workers = self.config.worker_budget.min(queries.len()).max(1);
        let shard_workers = (self.config.worker_budget / query_workers).max(1);
        crate::parallel::search_batch(queries, query_workers, |query| {
            self.search_with_workers(query, count, shard_workers)
        })
    }

    fn search_with_workers(
        &self,
        query: &[f32],
        count: usize,
        workers: usize,
    ) -> Result<Vec<SearchHit>, SearchError> {
        if query.len() != self.dimensions {
            return Err(SearchError::DimensionMismatch {
                expected: self.dimensions,
                actual: query.len(),
                vector: None,
            });
        }
        if count == 0 || self.shards.is_empty() {
            return Ok(Vec::new());
        }
        let per_shard = count
            .checked_mul(self.config.candidate_multiplier)
            .ok_or(SearchError::CapacityOverflow)?;
        let workers = workers.min(self.shards.len()).max(1);
        let chunk_size = self.shards.len().div_ceil(workers);
        let results = std::thread::scope(|scope| {
            let handles = self
                .shards
                .chunks(chunk_size)
                .enumerate()
                .map(|(chunk_index, chunk)| {
                    let first = chunk_index * chunk_size;
                    scope.spawn(move || {
                        chunk
                            .iter()
                            .enumerate()
                            .map(|(offset, shard)| {
                                let shard_index = first + offset;
                                shard.search_shard(query, per_shard).map_err(|error| {
                                    SearchError::ShardFailed {
                                        shard: shard_index,
                                        message: error.to_string(),
                                    }
                                })
                            })
                            .collect::<Vec<_>>()
                    })
                })
                .collect::<Vec<_>>();
            handles
                .into_iter()
                .map(|handle| handle.join().map_err(|_| SearchError::WorkerPanic))
                .collect::<Result<Vec<_>, _>>()
        })?;
        let mut best_by_key = BTreeMap::<u64, SearchHit>::new();
        for result in results.into_iter().flatten() {
            for hit in result? {
                best_by_key
                    .entry(hit.key)
                    .and_modify(|current| {
                        if hit.distance.total_cmp(&current.distance).is_lt() {
                            *current = hit;
                        }
                    })
                    .or_insert(hit);
            }
        }
        let mut hits = best_by_key.into_values().collect::<Vec<_>>();
        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)
    }
}