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 fn query_exact(
133 &self,
134 namespace: &str,
135 index: &str,
136 partition: &str,
137 input: &VectorQueryInput,
138 limit: usize,
139 min_score: f32,
140 ) -> Result<VectorQueryResult> {
141 let mut matches = Vec::new();
142 let prefix = record_prefix(namespace, index, partition);
143 let mut after = None;
144 loop {
145 let rows = self
146 .store
147 .list_prefix_after(&prefix, after.as_deref(), QUERY_SCAN_BATCH)?;
148 if rows.is_empty() {
149 break;
150 }
151 after = rows.last().map(|row| row.key.clone());
152
153 for row in rows {
154 let record = decode_record(&row.key, &row.value)?;
155 if record.values.len() != input.vector.len() {
156 continue;
157 }
158 if !metadata_matches_filter(&record.metadata, &input.filter) {
159 continue;
160 }
161
162 let score = cosine_similarity(&input.vector, &record.values);
163 if score < min_score {
164 continue;
165 }
166
167 matches.push((
168 record.updated_at_ms,
169 VectorMatch {
170 id: record.id,
171 score,
172 metadata: record.metadata,
173 },
174 ));
175 }
176 }
177
178 Ok(rank_matches(matches, limit))
179 }
180
181 #[allow(clippy::too_many_arguments)]
182 fn query_exact_ids(
183 &self,
184 namespace: &str,
185 index: &str,
186 partition: &str,
187 ids: &[&str],
188 input: &VectorQueryInput,
189 limit: usize,
190 min_score: f32,
191 ) -> Result<VectorQueryResult> {
192 let mut matches = Vec::new();
193 for id in ids {
194 let key = record_key(namespace, index, partition, id);
195 let Some(record) = self.get_record(&key)? else {
196 continue;
197 };
198 if record.values.len() != input.vector.len() {
199 continue;
200 }
201 if !metadata_matches_filter(&record.metadata, &input.filter) {
202 continue;
203 }
204
205 let score = cosine_similarity(&input.vector, &record.values);
206 if score < min_score {
207 continue;
208 }
209
210 matches.push((
211 record.updated_at_ms,
212 VectorMatch {
213 id: record.id,
214 score,
215 metadata: record.metadata,
216 },
217 ));
218 }
219
220 Ok(rank_matches(matches, limit))
221 }
222
223 fn query_ann(
224 &self,
225 namespace: &str,
226 index: &str,
227 partition: &str,
228 input: &VectorQueryInput,
229 limit: usize,
230 min_score: f32,
231 ) -> Result<Option<VectorQueryResult>> {
232 let requested_candidate_limit = input
233 .candidate_limit
234 .or_else(|| (!input.filter.is_empty()).then_some((limit * 32).max(256)));
235 let candidate_limit = ann_candidate_limit(requested_candidate_limit, limit)?;
236 let ef_search = ann_ef_search(input.ef_search, candidate_limit)?;
237 let scope = scope_key(namespace, index, partition);
238 let ann_scope = self.indexes.scope(&scope);
239 let Some(snapshot) = ann_scope.snapshot() else {
240 self.schedule_ann_rebuild(
241 namespace.to_string(),
242 index.to_string(),
243 partition.to_string(),
244 scope,
245 input.vector.len(),
246 );
247 return Ok(None);
248 };
249 if snapshot.index.dimensions() != input.vector.len() {
250 self.schedule_ann_rebuild(
251 namespace.to_string(),
252 index.to_string(),
253 partition.to_string(),
254 scope,
255 input.vector.len(),
256 );
257 return Ok(None);
258 }
259 if ann_scope.is_dirty() {
260 self.schedule_ann_rebuild(
261 namespace.to_string(),
262 index.to_string(),
263 partition.to_string(),
264 scope.clone(),
265 input.vector.len(),
266 );
267 }
268
269 let compiled_filter = snapshot.filters.compile(&input.filter);
270 if compiled_filter.is_none() {
271 return Ok(Some(VectorQueryResult {
272 matches: Vec::new(),
273 }));
274 }
275 if !input.filter.is_empty() {
276 if let Some(ids) = compiled_filter.candidate_ids(FILTER_EXACT_MAX_CANDIDATES) {
277 return Ok(Some(self.query_exact_ids(
278 namespace, index, partition, &ids, input, limit, min_score,
279 )?));
280 }
281 }
282 let filter = |id: &str| compiled_filter.accepts(id);
283 let query_filter = if input.filter.is_empty() || compiled_filter.is_all() {
284 None
285 } else {
286 Some(&filter as &dyn feox_ann::AnnFilter)
287 };
288
289 let candidates = snapshot.index.query(AnnQuery {
290 vector: &input.vector,
291 top_k: candidate_limit,
292 ef_search: Some(ef_search),
293 filter: query_filter,
294 })?;
295
296 let mut matches = Vec::new();
297 for candidate in candidates {
298 let key = record_key(namespace, index, partition, &candidate.id);
299 let Some(record) = self.get_record(&key)? else {
300 continue;
301 };
302 if record.values.len() != input.vector.len() {
303 continue;
304 }
305 if !metadata_matches_filter(&record.metadata, &input.filter) {
306 continue;
307 }
308 let score = cosine_similarity(&input.vector, &record.values);
309 if score < min_score {
310 continue;
311 }
312 matches.push((
313 record.updated_at_ms,
314 VectorMatch {
315 id: record.id,
316 score,
317 metadata: record.metadata,
318 },
319 ));
320 }
321
322 Ok(Some(rank_matches(matches, limit)))
323 }
324
325 pub fn rebuild_ann(
326 &self,
327 namespace: &str,
328 index: &str,
329 partition: &str,
330 dimensions: usize,
331 ) -> Result<usize> {
332 self.rebuild_ann_with_config(
333 namespace,
334 index,
335 partition,
336 AnnConfig::for_dimensions(dimensions),
337 )
338 }
339
340 pub fn rebuild_ann_with_config(
341 &self,
342 namespace: &str,
343 index: &str,
344 partition: &str,
345 config: AnnConfig,
346 ) -> Result<usize> {
347 validate_scope(namespace, index, partition)?;
348 let scope = scope_key(namespace, index, partition);
349 self.rebuild_ann_for_scope(namespace, index, partition, &scope, config)
350 }
351
352 fn rebuild_ann_index(
353 &self,
354 namespace: &str,
355 index: &str,
356 partition: &str,
357 config: AnnConfig,
358 ) -> Result<AnnSnapshot> {
359 let dimensions = config.dimensions;
360 let mut ann = feox_ann::AnnIndex::new(config)?;
361 let mut filters = MetadataFilterIndex::default();
362 let prefix = record_prefix(namespace, index, partition);
363 let mut after = None;
364 let mut cursor = ann.insert_cursor();
365 loop {
366 let rows = self
367 .store
368 .list_prefix_after(&prefix, after.as_deref(), QUERY_SCAN_BATCH)?;
369 if rows.is_empty() {
370 break;
371 }
372 after = rows.last().map(|row| row.key.clone());
373 cursor.reserve(rows.len());
374 for row in rows {
375 if let Some(record) = decode_ann_record(&row.key, &row.value, dimensions)? {
376 filters.insert(&record.id, &record.metadata);
377 cursor.upsert_owned(record.id, record.values)?;
378 }
379 }
380 }
381 drop(cursor);
382 Ok(AnnSnapshot {
383 index: ann,
384 filters,
385 })
386 }
387
388 fn rebuild_ann_for_scope(
389 &self,
390 namespace: &str,
391 index: &str,
392 partition: &str,
393 scope_key: &str,
394 config: AnnConfig,
395 ) -> Result<usize> {
396 let scope = self.indexes.scope(scope_key);
397 if !scope.begin_rebuild() {
398 return Ok(scope
399 .snapshot()
400 .map(|snapshot| snapshot.index.len())
401 .unwrap_or(0));
402 }
403
404 match self.rebuild_ann_index(namespace, index, partition, config) {
405 Ok(snapshot) => {
406 let len = snapshot.index.len();
407 scope.publish(snapshot);
408 Ok(len)
409 }
410 Err(error) => {
411 scope.finish_failed_rebuild();
412 Err(error)
413 }
414 }
415 }
416
417 fn schedule_ann_rebuild(
418 &self,
419 namespace: String,
420 index: String,
421 partition: String,
422 scope_key: String,
423 dimensions: usize,
424 ) {
425 let scope = self.indexes.scope(&scope_key);
426 if !scope.begin_rebuild() {
427 return;
428 }
429
430 let store = self.clone();
431 let _ = thread::Builder::new()
432 .name("feox-ann-rebuild".to_string())
433 .spawn(move || {
434 let result = store.rebuild_ann_index(
435 &namespace,
436 &index,
437 &partition,
438 AnnConfig::for_dimensions(dimensions),
439 );
440 match result {
441 Ok(index) => scope.publish(index),
442 Err(error) => {
443 scope.finish_failed_rebuild();
444 eprintln!("[feox-vector] ANN rebuild failed for {scope_key}: {error}");
445 }
446 }
447 });
448 }
449
450 fn get_record(&self, key: &str) -> Result<Option<VectorRecord>> {
451 let Some(bytes) = self.store.get_bytes(key)? else {
452 return Ok(None);
453 };
454 Ok(Some(decode_record(key, &bytes)?))
455 }
456}
457
458fn validate_scope(namespace: &str, index: &str, partition: &str) -> Result<()> {
459 validate_segment(namespace, "namespace")?;
460 validate_segment(index, "index")?;
461 validate_segment(partition, "partition")?;
462 Ok(())
463}
464
465fn cosine_similarity(left: &[f32], right: &[f32]) -> f32 {
466 let len = left.len().min(right.len());
467 let (left, right) = (&left[..len], &right[..len]);
468 let left_norm = feox_ann::dot(left, left);
469 let right_norm = feox_ann::dot(right, right);
470 if left_norm == 0.0 || right_norm == 0.0 {
471 return 0.0;
472 }
473 feox_ann::dot(left, right) / (left_norm.sqrt() * right_norm.sqrt())
474}
475
476fn rank_matches(mut matches: Vec<(u64, VectorMatch)>, limit: usize) -> VectorQueryResult {
477 matches.sort_by(|left, right| {
478 right
479 .1
480 .score
481 .partial_cmp(&left.1.score)
482 .unwrap_or(Ordering::Equal)
483 .then_with(|| right.0.cmp(&left.0))
484 .then_with(|| left.1.id.cmp(&right.1.id))
485 });
486 matches.truncate(limit);
487
488 VectorQueryResult {
489 matches: matches.into_iter().map(|(_, item)| item).collect(),
490 }
491}
492
493fn scope_key(namespace: &str, index: &str, partition: &str) -> String {
494 format!("{namespace}/{index}/{partition}")
495}
496
497fn now_ms() -> u64 {
498 SystemTime::now()
499 .duration_since(UNIX_EPOCH)
500 .expect("system clock before unix epoch")
501 .as_millis() as u64
502}