1use std::cmp::Ordering;
2use std::sync::Arc;
3use std::thread;
4use std::time::{SystemTime, UNIX_EPOCH};
5
6use crate::storage::Store;
7use feox_ann::{AnnConfig, AnnQuery};
8use serde_json::{Map, Value};
9
10use crate::ann::{AnnRegistry, AnnSnapshot};
11use crate::codec::{decode_ann_record, decode_record, encode_record};
12use crate::filter::{metadata_matches_filter, MetadataFilterIndex};
13use crate::keys::{record_key, record_prefix};
14use crate::validation::{
15 ann_candidate_limit, ann_ef_search, query_limit, validate_metadata, validate_segment,
16 validate_upsert_len, validate_values,
17};
18use crate::{
19 Result, VectorDeleteInput, VectorDeleteResult, VectorMatch, VectorQueryInput, VectorQueryMode,
20 VectorQueryResult, VectorRecord, VectorUpsertInput, VectorUpsertResult,
21};
22
23const QUERY_SCAN_BATCH: usize = 512;
24const FILTER_EXACT_MAX_CANDIDATES: usize = 4_096;
25
26#[derive(Clone)]
27pub struct VectorStore {
28 store: Store,
29 indexes: Arc<AnnRegistry>,
30}
31
32impl VectorStore {
33 pub fn new(store: Store) -> Self {
34 Self {
35 store,
36 indexes: Arc::new(AnnRegistry::default()),
37 }
38 }
39
40 pub fn upsert(
41 &self,
42 namespace: &str,
43 index: &str,
44 partition: &str,
45 input: VectorUpsertInput,
46 ) -> Result<VectorUpsertResult> {
47 validate_scope(namespace, index, partition)?;
48 validate_upsert_len(input.records.len())?;
49
50 let now = now_ms();
51 let scope = scope_key(namespace, index, partition);
52 let mut upserted = 0;
53 for record in input.records {
54 validate_segment(&record.id, "vector_id")?;
55 validate_values(&record.values)?;
56 validate_metadata(&Value::Object(Map::from_iter(record.metadata.clone())))?;
57
58 let key = record_key(namespace, index, partition, &record.id);
59 let existing = self.get_record(&key)?;
60 let created_at_ms = existing
61 .as_ref()
62 .map(|record| record.created_at_ms)
63 .unwrap_or(now);
64 let encoded = encode_record(&record.values, &record.metadata, created_at_ms, now)?;
65 self.store.put_bytes(&key, &encoded)?;
66 self.indexes.mark_dirty_if_present(&scope);
67 upserted += 1;
68 }
69
70 Ok(VectorUpsertResult { upserted })
71 }
72
73 pub fn query(
74 &self,
75 namespace: &str,
76 index: &str,
77 partition: &str,
78 input: VectorQueryInput,
79 ) -> Result<VectorQueryResult> {
80 validate_scope(namespace, index, partition)?;
81 validate_values(&input.vector)?;
82 validate_metadata(&Value::Object(Map::from_iter(input.filter.clone())))?;
83
84 let limit = query_limit(input.top_k)?;
85 if input
86 .min_score
87 .map(|score| !score.is_finite())
88 .unwrap_or(false)
89 {
90 return Err(crate::VectorError::Invalid(
91 "min_score must be a finite number".to_string(),
92 ));
93 }
94 let min_score = input.min_score.unwrap_or(f32::NEG_INFINITY);
95
96 if input.mode == Some(VectorQueryMode::Ann) {
97 if let Some(result) =
98 self.query_ann(namespace, index, partition, &input, limit, min_score)?
99 {
100 return Ok(result);
101 }
102 }
103
104 self.query_exact(namespace, index, partition, &input, limit, min_score)
105 }
106
107 pub fn delete(
108 &self,
109 namespace: &str,
110 index: &str,
111 partition: &str,
112 input: VectorDeleteInput,
113 ) -> Result<VectorDeleteResult> {
114 validate_scope(namespace, index, partition)?;
115
116 let scope = scope_key(namespace, index, partition);
117 let mut deleted = 0;
118 for id in input.ids {
119 validate_segment(&id, "vector_id")?;
120 if self
121 .store
122 .delete(&record_key(namespace, index, partition, &id))?
123 {
124 self.indexes.mark_dirty_if_present(&scope);
125 deleted += 1;
126 }
127 }
128
129 Ok(VectorDeleteResult { deleted })
130 }
131
132 pub fn flush(&self) -> Result<()> {
133 self.store.flush()?;
134 Ok(())
135 }
136
137 fn query_exact(
138 &self,
139 namespace: &str,
140 index: &str,
141 partition: &str,
142 input: &VectorQueryInput,
143 limit: usize,
144 min_score: f32,
145 ) -> Result<VectorQueryResult> {
146 let mut matches = Vec::new();
147 let prefix = record_prefix(namespace, index, partition);
148 let mut after = None;
149 loop {
150 let rows = self
151 .store
152 .list_prefix_after(&prefix, after.as_deref(), QUERY_SCAN_BATCH)?;
153 if rows.is_empty() {
154 break;
155 }
156 after = rows.last().map(|row| row.key.clone());
157
158 for row in rows {
159 let record = decode_record(&row.key, &row.value)?;
160 if record.values.len() != input.vector.len() {
161 continue;
162 }
163 if !metadata_matches_filter(&record.metadata, &input.filter) {
164 continue;
165 }
166
167 let score = cosine_similarity(&input.vector, &record.values);
168 if score < min_score {
169 continue;
170 }
171
172 matches.push((
173 record.updated_at_ms,
174 VectorMatch {
175 id: record.id,
176 score,
177 metadata: record.metadata,
178 },
179 ));
180 }
181 }
182
183 Ok(rank_matches(matches, limit))
184 }
185
186 #[allow(clippy::too_many_arguments)]
187 fn query_exact_ids(
188 &self,
189 namespace: &str,
190 index: &str,
191 partition: &str,
192 ids: &[&str],
193 input: &VectorQueryInput,
194 limit: usize,
195 min_score: f32,
196 ) -> Result<VectorQueryResult> {
197 let mut matches = Vec::new();
198 for id in ids {
199 let key = record_key(namespace, index, partition, id);
200 let Some(record) = self.get_record(&key)? else {
201 continue;
202 };
203 if record.values.len() != input.vector.len() {
204 continue;
205 }
206 if !metadata_matches_filter(&record.metadata, &input.filter) {
207 continue;
208 }
209
210 let score = cosine_similarity(&input.vector, &record.values);
211 if score < min_score {
212 continue;
213 }
214
215 matches.push((
216 record.updated_at_ms,
217 VectorMatch {
218 id: record.id,
219 score,
220 metadata: record.metadata,
221 },
222 ));
223 }
224
225 Ok(rank_matches(matches, limit))
226 }
227
228 fn query_ann(
229 &self,
230 namespace: &str,
231 index: &str,
232 partition: &str,
233 input: &VectorQueryInput,
234 limit: usize,
235 min_score: f32,
236 ) -> Result<Option<VectorQueryResult>> {
237 let requested_candidate_limit = input
238 .candidate_limit
239 .or_else(|| (!input.filter.is_empty()).then_some((limit * 32).max(256)));
240 let candidate_limit = ann_candidate_limit(requested_candidate_limit, limit)?;
241 let ef_search = ann_ef_search(input.ef_search, candidate_limit)?;
242 let scope = scope_key(namespace, index, partition);
243 let ann_scope = self.indexes.scope(&scope);
244 let Some(snapshot) = ann_scope.query_snapshot() else {
245 self.schedule_ann_rebuild(
246 namespace.to_string(),
247 index.to_string(),
248 partition.to_string(),
249 scope,
250 input.vector.len(),
251 );
252 return Ok(None);
253 };
254 if snapshot.index.dimensions() != input.vector.len() {
255 self.schedule_ann_rebuild(
256 namespace.to_string(),
257 index.to_string(),
258 partition.to_string(),
259 scope,
260 input.vector.len(),
261 );
262 return Ok(None);
263 }
264 let compiled_filter = snapshot.filters.compile(&input.filter);
265 if compiled_filter.is_none() {
266 return Ok(Some(VectorQueryResult {
267 matches: Vec::new(),
268 }));
269 }
270 if !input.filter.is_empty() {
271 if let Some(ids) = compiled_filter.candidate_ids(FILTER_EXACT_MAX_CANDIDATES) {
272 return Ok(Some(self.query_exact_ids(
273 namespace, index, partition, &ids, input, limit, min_score,
274 )?));
275 }
276 }
277 let filter = |id: &str| compiled_filter.accepts(id);
278 let query_filter = if input.filter.is_empty() || compiled_filter.is_all() {
279 None
280 } else {
281 Some(&filter as &dyn feox_ann::AnnFilter)
282 };
283
284 let candidates = snapshot.index.query(AnnQuery {
285 vector: &input.vector,
286 top_k: candidate_limit,
287 ef_search: Some(ef_search),
288 filter: query_filter,
289 })?;
290
291 let mut matches = Vec::new();
292 for candidate in candidates {
293 let key = record_key(namespace, index, partition, &candidate.id);
294 let Some(record) = self.get_record(&key)? else {
295 continue;
296 };
297 if record.values.len() != input.vector.len() {
298 continue;
299 }
300 if !metadata_matches_filter(&record.metadata, &input.filter) {
301 continue;
302 }
303 let score = cosine_similarity(&input.vector, &record.values);
304 if score < min_score {
305 continue;
306 }
307 matches.push((
308 record.updated_at_ms,
309 VectorMatch {
310 id: record.id,
311 score,
312 metadata: record.metadata,
313 },
314 ));
315 }
316
317 Ok(Some(rank_matches(matches, limit)))
318 }
319
320 pub fn rebuild_ann(
321 &self,
322 namespace: &str,
323 index: &str,
324 partition: &str,
325 dimensions: usize,
326 ) -> Result<usize> {
327 self.rebuild_ann_with_config(
328 namespace,
329 index,
330 partition,
331 AnnConfig::for_dimensions(dimensions),
332 )
333 }
334
335 pub fn rebuild_ann_with_config(
336 &self,
337 namespace: &str,
338 index: &str,
339 partition: &str,
340 config: AnnConfig,
341 ) -> Result<usize> {
342 validate_scope(namespace, index, partition)?;
343 let scope = scope_key(namespace, index, partition);
344 self.rebuild_ann_for_scope(namespace, index, partition, &scope, config)
345 }
346
347 fn rebuild_ann_index(
348 &self,
349 namespace: &str,
350 index: &str,
351 partition: &str,
352 config: AnnConfig,
353 ) -> Result<AnnSnapshot> {
354 let dimensions = config.dimensions;
355 let mut ann = feox_ann::AnnIndex::new(config)?;
356 let mut filters = MetadataFilterIndex::default();
357 let prefix = record_prefix(namespace, index, partition);
358 let mut after = None;
359 let mut cursor = ann.insert_cursor();
360 loop {
361 let rows = self
362 .store
363 .list_prefix_after(&prefix, after.as_deref(), QUERY_SCAN_BATCH)?;
364 if rows.is_empty() {
365 break;
366 }
367 after = rows.last().map(|row| row.key.clone());
368 cursor.reserve(rows.len());
369 for row in rows {
370 if let Some(record) = decode_ann_record(&row.key, &row.value, dimensions)? {
371 filters.insert(&record.id, &record.metadata);
372 cursor.upsert_owned(record.id, record.values)?;
373 }
374 }
375 }
376 drop(cursor);
377 Ok(AnnSnapshot {
378 index: ann,
379 filters,
380 })
381 }
382
383 fn rebuild_ann_for_scope(
384 &self,
385 namespace: &str,
386 index: &str,
387 partition: &str,
388 scope_key: &str,
389 config: AnnConfig,
390 ) -> Result<usize> {
391 let scope = self.indexes.scope(scope_key);
392 if !scope.begin_rebuild() {
393 return Ok(scope
394 .snapshot()
395 .map(|snapshot| snapshot.index.len())
396 .unwrap_or(0));
397 }
398
399 match self.rebuild_ann_index(namespace, index, partition, config) {
400 Ok(snapshot) => {
401 let len = snapshot.index.len();
402 scope.publish(snapshot);
403 Ok(len)
404 }
405 Err(error) => {
406 scope.finish_failed_rebuild();
407 Err(error)
408 }
409 }
410 }
411
412 fn schedule_ann_rebuild(
413 &self,
414 namespace: String,
415 index: String,
416 partition: String,
417 scope_key: String,
418 dimensions: usize,
419 ) {
420 let scope = self.indexes.scope(&scope_key);
421 if !scope.begin_rebuild() {
422 return;
423 }
424
425 let store = self.clone();
426 let rebuild_scope = scope.clone();
427 let rebuild_scope_key = scope_key.clone();
428 let spawn_result = thread::Builder::new()
429 .name("feox-ann-rebuild".to_string())
430 .spawn(move || {
431 let result = store.rebuild_ann_index(
432 &namespace,
433 &index,
434 &partition,
435 AnnConfig::for_dimensions(dimensions),
436 );
437 match result {
438 Ok(index) => rebuild_scope.publish(index),
439 Err(error) => {
440 rebuild_scope.finish_failed_rebuild();
441 eprintln!(
442 "[feox-vector] ANN rebuild failed for {rebuild_scope_key}: {error}"
443 );
444 }
445 }
446 });
447 if let Err(error) = spawn_result {
448 scope.finish_failed_rebuild();
449 eprintln!("[feox-vector] failed to start ANN rebuild for {scope_key}: {error}");
450 }
451 }
452
453 fn get_record(&self, key: &str) -> Result<Option<VectorRecord>> {
454 let Some(bytes) = self.store.get_bytes(key)? else {
455 return Ok(None);
456 };
457 Ok(Some(decode_record(key, &bytes)?))
458 }
459}
460
461fn validate_scope(namespace: &str, index: &str, partition: &str) -> Result<()> {
462 validate_segment(namespace, "namespace")?;
463 validate_segment(index, "index")?;
464 validate_segment(partition, "partition")?;
465 Ok(())
466}
467
468fn cosine_similarity(left: &[f32], right: &[f32]) -> f32 {
469 let len = left.len().min(right.len());
470 let (left, right) = (&left[..len], &right[..len]);
471 let left_norm = feox_ann::dot(left, left);
472 let right_norm = feox_ann::dot(right, right);
473 if left_norm == 0.0 || right_norm == 0.0 {
474 return 0.0;
475 }
476 feox_ann::dot(left, right) / (left_norm.sqrt() * right_norm.sqrt())
477}
478
479fn rank_matches(mut matches: Vec<(u64, VectorMatch)>, limit: usize) -> VectorQueryResult {
480 matches.sort_by(|left, right| {
481 right
482 .1
483 .score
484 .partial_cmp(&left.1.score)
485 .unwrap_or(Ordering::Equal)
486 .then_with(|| right.0.cmp(&left.0))
487 .then_with(|| left.1.id.cmp(&right.1.id))
488 });
489 matches.truncate(limit);
490
491 VectorQueryResult {
492 matches: matches.into_iter().map(|(_, item)| item).collect(),
493 }
494}
495
496fn scope_key(namespace: &str, index: &str, partition: &str) -> String {
497 format!("{namespace}/{index}/{partition}")
498}
499
500fn now_ms() -> u64 {
501 SystemTime::now()
502 .duration_since(UNIX_EPOCH)
503 .expect("system clock before unix epoch")
504 .as_millis() as u64
505}