weavatrix_search_vector/
distributed.rs

1use crate::error::SearchError;
2use crate::hit::SearchHit;
3use crate::hnsw::VectorIndex;
4use crate::mutable::MutableVectorIndex;
5use crate::quantized::{QuantizedIndex, ScalarQuantizedIndex};
6use crate::storage::MappedVectorIndex;
7use std::collections::BTreeMap;
8use std::sync::Arc;
9
10/// Transport-neutral vector-search shard.
11///
12/// Implementations may be local indexes or remote clients. Remote adapters
13/// should translate transport failures into [`SearchError::ShardFailed`].
14pub trait SearchShard: Send + Sync {
15    fn dimensions(&self) -> usize;
16
17    /// Searches this shard.
18    ///
19    /// # Errors
20    ///
21    /// Returns a query, allocation, worker, transport, or provider error.
22    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError>;
23}
24
25impl SearchShard for VectorIndex {
26    fn dimensions(&self) -> usize {
27        self.dimensions()
28    }
29
30    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
31        self.search(query, count)
32    }
33}
34
35impl SearchShard for MappedVectorIndex {
36    fn dimensions(&self) -> usize {
37        self.dimensions()
38    }
39
40    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
41        self.search(query, count)
42    }
43}
44
45impl SearchShard for ScalarQuantizedIndex {
46    fn dimensions(&self) -> usize {
47        self.dimensions()
48    }
49
50    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
51        self.search(query, count)
52    }
53}
54
55impl SearchShard for QuantizedIndex {
56    fn dimensions(&self) -> usize {
57        self.dimensions()
58    }
59
60    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
61        self.search(query, count)
62    }
63}
64
65impl SearchShard for MutableVectorIndex {
66    fn dimensions(&self) -> usize {
67        self.config().dimensions
68    }
69
70    fn search_shard(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
71        self.search(query, count)
72    }
73}
74
75/// Bounded coordinator policy for shard fan-out and stable result merging.
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub struct DistributedConfig {
78    /// Total local worker budget shared by query and shard parallelism.
79    pub worker_budget: usize,
80    /// Per-shard candidate count is `top_k * candidate_multiplier`.
81    pub candidate_multiplier: usize,
82}
83
84impl DistributedConfig {
85    #[must_use]
86    pub const fn new(worker_budget: usize) -> Self {
87        Self {
88            worker_budget,
89            candidate_multiplier: 2,
90        }
91    }
92
93    fn validate(self) -> Result<Self, SearchError> {
94        if self.worker_budget == 0 {
95            return Err(SearchError::InvalidConfig(
96                "distributed worker_budget must be non-zero",
97            ));
98        }
99        if self.candidate_multiplier == 0 {
100            return Err(SearchError::InvalidConfig(
101                "distributed candidate_multiplier must be non-zero",
102            ));
103        }
104        Ok(self)
105    }
106}
107
108/// Heterogeneous local/remote shard coordinator.
109pub struct DistributedIndex {
110    dimensions: usize,
111    config: DistributedConfig,
112    shards: Vec<Arc<dyn SearchShard>>,
113}
114
115impl std::fmt::Debug for DistributedIndex {
116    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117        formatter
118            .debug_struct("DistributedIndex")
119            .field("dimensions", &self.dimensions)
120            .field("config", &self.config)
121            .field("shards", &self.shards.len())
122            .finish()
123    }
124}
125
126impl DistributedIndex {
127    /// Creates an empty coordinator for fixed-dimensional shards.
128    ///
129    /// # Errors
130    ///
131    /// Returns a typed config error.
132    pub fn new(dimensions: usize, config: DistributedConfig) -> Result<Self, SearchError> {
133        if dimensions == 0 {
134            return Err(SearchError::InvalidConfig(
135                "distributed dimensions must be non-zero",
136            ));
137        }
138        Ok(Self {
139            dimensions,
140            config: config.validate()?,
141            shards: Vec::new(),
142        })
143    }
144
145    /// Adds one local or remote shard.
146    ///
147    /// # Errors
148    ///
149    /// Returns a dimension mismatch when shard and coordinator differ.
150    pub fn push(&mut self, shard: Arc<dyn SearchShard>) -> Result<(), SearchError> {
151        if shard.dimensions() != self.dimensions {
152            return Err(SearchError::DimensionMismatch {
153                expected: self.dimensions,
154                actual: shard.dimensions(),
155                vector: None,
156            });
157        }
158        self.shards.push(shard);
159        Ok(())
160    }
161
162    #[must_use]
163    pub fn len(&self) -> usize {
164        self.shards.len()
165    }
166
167    #[must_use]
168    pub fn is_empty(&self) -> bool {
169        self.shards.is_empty()
170    }
171
172    /// Fans one query out with bounded workers and keeps the best occurrence
173    /// of duplicate keys.
174    ///
175    /// # Errors
176    ///
177    /// Returns a query, worker, allocation, or shard error.
178    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
179        self.search_with_workers(query, count, self.config.worker_budget)
180    }
181
182    /// Searches independent queries while keeping total local threads within
183    /// `worker_budget`.
184    ///
185    /// # Errors
186    ///
187    /// Returns the first query or shard error in input order.
188    pub fn search_batch(
189        &self,
190        queries: &[&[f32]],
191        count: usize,
192    ) -> Result<Vec<Vec<SearchHit>>, SearchError> {
193        if queries.is_empty() {
194            return Ok(Vec::new());
195        }
196        let query_workers = self.config.worker_budget.min(queries.len()).max(1);
197        let shard_workers = (self.config.worker_budget / query_workers).max(1);
198        crate::parallel::search_batch(queries, query_workers, |query| {
199            self.search_with_workers(query, count, shard_workers)
200        })
201    }
202
203    fn search_with_workers(
204        &self,
205        query: &[f32],
206        count: usize,
207        workers: usize,
208    ) -> Result<Vec<SearchHit>, SearchError> {
209        if query.len() != self.dimensions {
210            return Err(SearchError::DimensionMismatch {
211                expected: self.dimensions,
212                actual: query.len(),
213                vector: None,
214            });
215        }
216        if count == 0 || self.shards.is_empty() {
217            return Ok(Vec::new());
218        }
219        let per_shard = count
220            .checked_mul(self.config.candidate_multiplier)
221            .ok_or(SearchError::CapacityOverflow)?;
222        let workers = workers.min(self.shards.len()).max(1);
223        let chunk_size = self.shards.len().div_ceil(workers);
224        let results = std::thread::scope(|scope| {
225            let handles = self
226                .shards
227                .chunks(chunk_size)
228                .enumerate()
229                .map(|(chunk_index, chunk)| {
230                    let first = chunk_index * chunk_size;
231                    scope.spawn(move || {
232                        chunk
233                            .iter()
234                            .enumerate()
235                            .map(|(offset, shard)| {
236                                let shard_index = first + offset;
237                                shard.search_shard(query, per_shard).map_err(|error| {
238                                    SearchError::ShardFailed {
239                                        shard: shard_index,
240                                        message: error.to_string(),
241                                    }
242                                })
243                            })
244                            .collect::<Vec<_>>()
245                    })
246                })
247                .collect::<Vec<_>>();
248            handles
249                .into_iter()
250                .map(|handle| handle.join().map_err(|_| SearchError::WorkerPanic))
251                .collect::<Result<Vec<_>, _>>()
252        })?;
253        let mut best_by_key = BTreeMap::<u64, SearchHit>::new();
254        for result in results.into_iter().flatten() {
255            for hit in result? {
256                best_by_key
257                    .entry(hit.key)
258                    .and_modify(|current| {
259                        if hit.distance.total_cmp(&current.distance).is_lt() {
260                            *current = hit;
261                        }
262                    })
263                    .or_insert(hit);
264            }
265        }
266        let mut hits = best_by_key.into_values().collect::<Vec<_>>();
267        hits.sort_unstable_by(|left, right| {
268            left.distance
269                .total_cmp(&right.distance)
270                .then_with(|| left.key.cmp(&right.key))
271        });
272        hits.truncate(count.min(hits.len()));
273        Ok(hits)
274    }
275}
weavatrix_search_vector/distributed.rs

weavatrix_search_vector/
distributed.rs
