oxirs_embed/
vector_store.rs

1//! In-memory vector store with namespace support.
2//!
3//! Provides `VectorStore`: a namespace-partitioned collection of embedding
4//! vectors with cosine-similarity search, upsert/delete, and dimension
5//! enforcement.
6
7use std::collections::HashMap;
8
9// ─────────────────────────────────────────────────────────────────────────────
10// Public types
11// ─────────────────────────────────────────────────────────────────────────────
12
13/// A single vector stored in the vector store.
14#[derive(Debug, Clone)]
15pub struct VectorEntry {
16    /// Unique identifier within its namespace.
17    pub id: String,
18    /// Namespace this entry belongs to.
19    pub namespace: String,
20    /// The embedding vector.
21    pub vector: Vec<f32>,
22    /// Arbitrary key-value metadata.
23    pub metadata: HashMap<String, String>,
24    /// Creation timestamp (arbitrary unit — caller-supplied, e.g. Unix ms).
25    pub created_at: u64,
26}
27
28/// A single result returned by a similarity search.
29#[derive(Debug, Clone)]
30pub struct SearchResult {
31    /// Entry identifier.
32    pub id: String,
33    /// Namespace the entry lives in.
34    pub namespace: String,
35    /// Cosine similarity score in `[-1, 1]` (1 = identical direction).
36    pub score: f32,
37    /// Entry metadata.
38    pub metadata: HashMap<String, String>,
39}
40
41/// Aggregate statistics about the store.
42#[derive(Debug, Clone)]
43pub struct VectorStoreStats {
44    /// Total number of stored vectors across all namespaces.
45    pub total_vectors: usize,
46    /// Number of distinct namespaces that contain at least one vector.
47    pub namespace_count: usize,
48    /// Dimensionality of stored vectors, or `None` when the store is empty or
49    /// no dimension was enforced.
50    pub dimension: Option<usize>,
51}
52
53/// Errors returned by store operations.
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub enum StoreError {
56    /// The supplied vector does not match the store's configured dimensionality.
57    DimensionMismatch { expected: usize, got: usize },
58    /// An empty (zero-length) vector was supplied.
59    EmptyVector,
60}
61
62impl std::fmt::Display for StoreError {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        match self {
65            StoreError::DimensionMismatch { expected, got } => {
66                write!(f, "dimension mismatch: expected {expected}, got {got}")
67            }
68            StoreError::EmptyVector => write!(f, "vector must not be empty"),
69        }
70    }
71}
72
73impl std::error::Error for StoreError {}
74
75// ─────────────────────────────────────────────────────────────────────────────
76// VectorStore
77// ─────────────────────────────────────────────────────────────────────────────
78
79/// In-memory vector store partitioned by namespace.
80pub struct VectorStore {
81    /// namespace → list of entries.
82    namespaces: HashMap<String, Vec<VectorEntry>>,
83    /// Optional fixed dimensionality; enforced on upsert.
84    dimension: Option<usize>,
85}
86
87// ── helpers ───────────────────────────────────────────────────────────────────
88
89/// Compute cosine similarity between two equal-length slices.
90///
91/// Returns `0.0` if either vector has zero magnitude.
92fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
93    let mut dot = 0.0_f32;
94    let mut norm_a = 0.0_f32;
95    let mut norm_b = 0.0_f32;
96    for (&x, &y) in a.iter().zip(b.iter()) {
97        dot += x * y;
98        norm_a += x * x;
99        norm_b += y * y;
100    }
101    let denom = norm_a.sqrt() * norm_b.sqrt();
102    if denom == 0.0 {
103        0.0
104    } else {
105        (dot / denom).clamp(-1.0, 1.0)
106    }
107}
108
109impl Default for VectorStore {
110    fn default() -> Self {
111        Self::new()
112    }
113}
114
115impl VectorStore {
116    /// Create an empty store without a fixed dimension.
117    pub fn new() -> Self {
118        Self {
119            namespaces: HashMap::new(),
120            dimension: None,
121        }
122    }
123
124    /// Create an empty store that enforces a fixed dimension on upsert.
125    pub fn with_dimension(dim: usize) -> Self {
126        Self {
127            namespaces: HashMap::new(),
128            dimension: Some(dim),
129        }
130    }
131
132    // ── validation helper ─────────────────────────────────────────────────────
133
134    fn validate_vector(&self, vector: &[f32]) -> Result<(), StoreError> {
135        if vector.is_empty() {
136            return Err(StoreError::EmptyVector);
137        }
138        if let Some(dim) = self.dimension {
139            if vector.len() != dim {
140                return Err(StoreError::DimensionMismatch {
141                    expected: dim,
142                    got: vector.len(),
143                });
144            }
145        }
146        Ok(())
147    }
148
149    // ── CRUD ──────────────────────────────────────────────────────────────────
150
151    /// Insert or update a vector entry.
152    ///
153    /// Returns `Ok(true)` if the entry was newly inserted, `Ok(false)` if an
154    /// existing entry with the same `(namespace, id)` was replaced.
155    ///
156    /// When the store is dimension-aware (created with [`VectorStore::with_dimension`]),
157    /// vectors of a different length are rejected.
158    pub fn upsert(&mut self, entry: VectorEntry) -> Result<bool, StoreError> {
159        self.validate_vector(&entry.vector)?;
160        let entries = self.namespaces.entry(entry.namespace.clone()).or_default();
161        if let Some(existing) = entries.iter_mut().find(|e| e.id == entry.id) {
162            *existing = entry;
163            return Ok(false);
164        }
165        entries.push(entry);
166        Ok(true)
167    }
168
169    /// Delete the entry with the given `(namespace, id)`.
170    ///
171    /// Returns `true` if the entry existed and was removed.
172    pub fn delete(&mut self, namespace: &str, id: &str) -> bool {
173        match self.namespaces.get_mut(namespace) {
174            Some(entries) => {
175                let before = entries.len();
176                entries.retain(|e| e.id != id);
177                entries.len() < before
178            }
179            None => false,
180        }
181    }
182
183    /// Look up a specific entry by `(namespace, id)`.
184    pub fn get(&self, namespace: &str, id: &str) -> Option<&VectorEntry> {
185        self.namespaces
186            .get(namespace)
187            .and_then(|entries| entries.iter().find(|e| e.id == id))
188    }
189
190    /// Check whether a specific `(namespace, id)` exists.
191    pub fn contains(&self, namespace: &str, id: &str) -> bool {
192        self.get(namespace, id).is_some()
193    }
194
195    // ── search ────────────────────────────────────────────────────────────────
196
197    /// Return the `top_k` most-similar entries in `namespace` to `query`,
198    /// sorted by descending cosine similarity.
199    pub fn search(&self, namespace: &str, query: &[f32], top_k: usize) -> Vec<SearchResult> {
200        let entries = match self.namespaces.get(namespace) {
201            Some(e) => e,
202            None => return Vec::new(),
203        };
204        let mut scored: Vec<(f32, &VectorEntry)> = entries
205            .iter()
206            .filter(|e| e.vector.len() == query.len())
207            .map(|e| (cosine_similarity(&e.vector, query), e))
208            .collect();
209        scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
210        scored
211            .into_iter()
212            .take(top_k)
213            .map(|(score, entry)| SearchResult {
214                id: entry.id.clone(),
215                namespace: entry.namespace.clone(),
216                score,
217                metadata: entry.metadata.clone(),
218            })
219            .collect()
220    }
221
222    /// Return the `top_k` most-similar entries across **all** namespaces.
223    pub fn search_all_namespaces(&self, query: &[f32], top_k: usize) -> Vec<SearchResult> {
224        let mut scored: Vec<(f32, &VectorEntry)> = self
225            .namespaces
226            .values()
227            .flat_map(|entries| entries.iter())
228            .filter(|e| e.vector.len() == query.len())
229            .map(|e| (cosine_similarity(&e.vector, query), e))
230            .collect();
231        scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
232        scored
233            .into_iter()
234            .take(top_k)
235            .map(|(score, entry)| SearchResult {
236                id: entry.id.clone(),
237                namespace: entry.namespace.clone(),
238                score,
239                metadata: entry.metadata.clone(),
240            })
241            .collect()
242    }
243
244    // ── listing ───────────────────────────────────────────────────────────────
245
246    /// List all entries in `namespace`.  Returns an empty slice if the
247    /// namespace does not exist.
248    pub fn list(&self, namespace: &str) -> Vec<&VectorEntry> {
249        match self.namespaces.get(namespace) {
250            Some(entries) => entries.iter().collect(),
251            None => Vec::new(),
252        }
253    }
254
255    // ── namespace management ──────────────────────────────────────────────────
256
257    /// Remove all entries belonging to `namespace`.
258    ///
259    /// Returns the number of entries that were deleted.
260    pub fn delete_namespace(&mut self, namespace: &str) -> usize {
261        match self.namespaces.remove(namespace) {
262            Some(entries) => entries.len(),
263            None => 0,
264        }
265    }
266
267    // ── statistics ────────────────────────────────────────────────────────────
268
269    /// Return aggregate statistics about the store.
270    pub fn stats(&self) -> VectorStoreStats {
271        let total_vectors: usize = self.namespaces.values().map(|v| v.len()).sum();
272        let namespace_count = self.namespaces.values().filter(|v| !v.is_empty()).count();
273        // Derive dimension from first stored vector if not explicitly configured.
274        let dimension = self.dimension.or_else(|| {
275            self.namespaces
276                .values()
277                .flat_map(|v| v.iter())
278                .next()
279                .map(|e| e.vector.len())
280        });
281        VectorStoreStats {
282            total_vectors,
283            namespace_count,
284            dimension,
285        }
286    }
287}
288
289// ─────────────────────────────────────────────────────────────────────────────
290// Tests
291// ─────────────────────────────────────────────────────────────────────────────
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296
297    fn entry(id: &str, ns: &str, v: Vec<f32>) -> VectorEntry {
298        VectorEntry {
299            id: id.to_string(),
300            namespace: ns.to_string(),
301            vector: v,
302            metadata: HashMap::new(),
303            created_at: 0,
304        }
305    }
306
307    fn unit_vec(dim: usize, one_at: usize) -> Vec<f32> {
308        let mut v = vec![0.0_f32; dim];
309        v[one_at] = 1.0;
310        v
311    }
312
313    // ── upsert new / update ───────────────────────────────────────────────────
314
315    #[test]
316    fn test_upsert_new_returns_true() {
317        let mut store = VectorStore::new();
318        let is_new = store
319            .upsert(entry("e1", "ns", vec![1.0, 0.0]))
320            .expect("should succeed");
321        assert!(is_new);
322    }
323
324    #[test]
325    fn test_upsert_update_returns_false() {
326        let mut store = VectorStore::new();
327        store
328            .upsert(entry("e1", "ns", vec![1.0, 0.0]))
329            .expect("should succeed");
330        let is_new = store
331            .upsert(entry("e1", "ns", vec![0.0, 1.0]))
332            .expect("should succeed");
333        assert!(!is_new);
334    }
335
336    #[test]
337    fn test_upsert_update_replaces_vector() {
338        let mut store = VectorStore::new();
339        store
340            .upsert(entry("e1", "ns", vec![1.0, 0.0]))
341            .expect("should succeed");
342        store
343            .upsert(entry("e1", "ns", vec![0.0, 1.0]))
344            .expect("should succeed");
345        let got = store.get("ns", "e1").expect("should succeed");
346        assert_eq!(got.vector, vec![0.0, 1.0]);
347    }
348
349    #[test]
350    fn test_upsert_empty_vector_errors() {
351        let mut store = VectorStore::new();
352        let res = store.upsert(entry("e1", "ns", vec![]));
353        assert_eq!(res, Err(StoreError::EmptyVector));
354    }
355
356    #[test]
357    fn test_upsert_dimension_mismatch_errors() {
358        let mut store = VectorStore::with_dimension(3);
359        let res = store.upsert(entry("e1", "ns", vec![1.0, 2.0]));
360        assert_eq!(
361            res,
362            Err(StoreError::DimensionMismatch {
363                expected: 3,
364                got: 2
365            })
366        );
367    }
368
369    #[test]
370    fn test_upsert_correct_dimension_ok() {
371        let mut store = VectorStore::with_dimension(2);
372        assert!(store.upsert(entry("e1", "ns", vec![1.0, 0.0])).is_ok());
373    }
374
375    // ── delete ────────────────────────────────────────────────────────────────
376
377    #[test]
378    fn test_delete_existing() {
379        let mut store = VectorStore::new();
380        store
381            .upsert(entry("e1", "ns", vec![1.0]))
382            .expect("should succeed");
383        assert!(store.delete("ns", "e1"));
384        assert!(store.get("ns", "e1").is_none());
385    }
386
387    #[test]
388    fn test_delete_nonexistent_returns_false() {
389        let mut store = VectorStore::new();
390        assert!(!store.delete("ns", "ghost"));
391    }
392
393    // ── get ───────────────────────────────────────────────────────────────────
394
395    #[test]
396    fn test_get_existing() {
397        let mut store = VectorStore::new();
398        store
399            .upsert(entry("e1", "ns", vec![1.0, 2.0]))
400            .expect("should succeed");
401        assert!(store.get("ns", "e1").is_some());
402    }
403
404    #[test]
405    fn test_get_nonexistent_none() {
406        let store = VectorStore::new();
407        assert!(store.get("ns", "missing").is_none());
408    }
409
410    // ── contains ─────────────────────────────────────────────────────────────
411
412    #[test]
413    fn test_contains_true() {
414        let mut store = VectorStore::new();
415        store
416            .upsert(entry("x", "ns", vec![1.0]))
417            .expect("should succeed");
418        assert!(store.contains("ns", "x"));
419    }
420
421    #[test]
422    fn test_contains_false() {
423        let store = VectorStore::new();
424        assert!(!store.contains("ns", "x"));
425    }
426
427    // ── search sorted by score ────────────────────────────────────────────────
428
429    #[test]
430    fn test_search_sorted_by_score_descending() {
431        let mut store = VectorStore::new();
432        // e1 aligned with query [1,0,0]
433        store
434            .upsert(entry("e1", "ns", unit_vec(3, 0)))
435            .expect("should succeed");
436        // e2 aligned with [0,1,0] — low similarity to query
437        store
438            .upsert(entry("e2", "ns", unit_vec(3, 1)))
439            .expect("should succeed");
440        let query = unit_vec(3, 0);
441        let results = store.search("ns", &query, 2);
442        assert_eq!(results.len(), 2);
443        assert!(results[0].score >= results[1].score);
444        assert_eq!(results[0].id, "e1");
445    }
446
447    #[test]
448    fn test_search_top_k_limit() {
449        let mut store = VectorStore::new();
450        for i in 0..10 {
451            store
452                .upsert(entry(&i.to_string(), "ns", vec![i as f32]))
453                .expect("should succeed");
454        }
455        let results = store.search("ns", &[5.0_f32], 3);
456        assert_eq!(results.len(), 3);
457    }
458
459    #[test]
460    fn test_search_same_vector_max_score() {
461        let mut store = VectorStore::new();
462        let v = vec![1.0_f32, 1.0, 1.0];
463        store
464            .upsert(entry("e", "ns", v.clone()))
465            .expect("should succeed");
466        let results = store.search("ns", &v, 1);
467        assert_eq!(results.len(), 1);
468        assert!((results[0].score - 1.0).abs() < 1e-6);
469    }
470
471    #[test]
472    fn test_search_empty_namespace_returns_empty() {
473        let store = VectorStore::new();
474        let results = store.search("missing-ns", &[1.0_f32], 5);
475        assert!(results.is_empty());
476    }
477
478    // ── search_all_namespaces ─────────────────────────────────────────────────
479
480    #[test]
481    fn test_search_all_namespaces_cross_namespace() {
482        let mut store = VectorStore::new();
483        store
484            .upsert(entry("a", "ns1", vec![1.0_f32, 0.0]))
485            .expect("should succeed");
486        store
487            .upsert(entry("b", "ns2", vec![0.0_f32, 1.0]))
488            .expect("should succeed");
489        let results = store.search_all_namespaces(&[1.0_f32, 0.0], 2);
490        assert_eq!(results.len(), 2);
491        assert_eq!(results[0].id, "a"); // aligned with query
492    }
493
494    #[test]
495    fn test_search_all_namespaces_top_k() {
496        let mut store = VectorStore::new();
497        for i in 0..5 {
498            store
499                .upsert(entry(&format!("e{i}"), &format!("ns{i}"), vec![i as f32]))
500                .expect("should succeed");
501        }
502        let results = store.search_all_namespaces(&[2.0_f32], 2);
503        assert_eq!(results.len(), 2);
504    }
505
506    // ── list ──────────────────────────────────────────────────────────────────
507
508    #[test]
509    fn test_list_all_in_namespace() {
510        let mut store = VectorStore::new();
511        store
512            .upsert(entry("a", "ns", vec![1.0]))
513            .expect("should succeed");
514        store
515            .upsert(entry("b", "ns", vec![2.0]))
516            .expect("should succeed");
517        let listed = store.list("ns");
518        assert_eq!(listed.len(), 2);
519    }
520
521    #[test]
522    fn test_list_nonexistent_namespace_empty() {
523        let store = VectorStore::new();
524        assert!(store.list("ghost").is_empty());
525    }
526
527    // ── delete_namespace ──────────────────────────────────────────────────────
528
529    #[test]
530    fn test_delete_namespace_returns_count() {
531        let mut store = VectorStore::new();
532        store
533            .upsert(entry("a", "ns", vec![1.0]))
534            .expect("should succeed");
535        store
536            .upsert(entry("b", "ns", vec![2.0]))
537            .expect("should succeed");
538        assert_eq!(store.delete_namespace("ns"), 2);
539    }
540
541    #[test]
542    fn test_delete_namespace_nonexistent_returns_zero() {
543        let mut store = VectorStore::new();
544        assert_eq!(store.delete_namespace("ghost"), 0);
545    }
546
547    #[test]
548    fn test_delete_namespace_removes_entries() {
549        let mut store = VectorStore::new();
550        store
551            .upsert(entry("a", "ns", vec![1.0]))
552            .expect("should succeed");
553        store.delete_namespace("ns");
554        assert!(store.list("ns").is_empty());
555    }
556
557    // ── stats ──────────────────────────────────────────────────────────────────
558
559    #[test]
560    fn test_stats_empty_store() {
561        let store = VectorStore::new();
562        let s = store.stats();
563        assert_eq!(s.total_vectors, 0);
564        assert_eq!(s.namespace_count, 0);
565        assert!(s.dimension.is_none());
566    }
567
568    #[test]
569    fn test_stats_counts_correctly() {
570        let mut store = VectorStore::new();
571        store
572            .upsert(entry("a", "ns1", vec![1.0, 2.0]))
573            .expect("should succeed");
574        store
575            .upsert(entry("b", "ns1", vec![3.0, 4.0]))
576            .expect("should succeed");
577        store
578            .upsert(entry("c", "ns2", vec![5.0, 6.0]))
579            .expect("should succeed");
580        let s = store.stats();
581        assert_eq!(s.total_vectors, 3);
582        assert_eq!(s.namespace_count, 2);
583        assert_eq!(s.dimension, Some(2));
584    }
585
586    #[test]
587    fn test_stats_with_configured_dimension() {
588        let store = VectorStore::with_dimension(128);
589        let s = store.stats();
590        assert_eq!(s.dimension, Some(128));
591    }
592
593    // ── additional coverage ───────────────────────────────────────────────────
594
595    #[test]
596    fn test_upsert_multiple_namespaces() {
597        let mut store = VectorStore::new();
598        store
599            .upsert(entry("a", "ns1", vec![1.0]))
600            .expect("should succeed");
601        store
602            .upsert(entry("b", "ns2", vec![2.0]))
603            .expect("should succeed");
604        assert!(store.contains("ns1", "a"));
605        assert!(store.contains("ns2", "b"));
606    }
607
608    #[test]
609    fn test_delete_from_wrong_namespace() {
610        let mut store = VectorStore::new();
611        store
612            .upsert(entry("a", "ns1", vec![1.0]))
613            .expect("should succeed");
614        assert!(!store.delete("ns2", "a")); // wrong namespace
615        assert!(store.contains("ns1", "a")); // still exists
616    }
617
618    #[test]
619    fn test_search_returns_correct_namespace() {
620        let mut store = VectorStore::new();
621        store
622            .upsert(entry("a", "ns1", vec![1.0_f32, 0.0]))
623            .expect("should succeed");
624        let results = store.search("ns1", &[1.0_f32, 0.0], 1);
625        assert_eq!(results[0].namespace, "ns1");
626    }
627
628    #[test]
629    fn test_search_all_namespaces_empty_store() {
630        let store = VectorStore::new();
631        assert!(store.search_all_namespaces(&[1.0_f32], 5).is_empty());
632    }
633
634    #[test]
635    fn test_metadata_stored_and_retrieved() {
636        let mut store = VectorStore::new();
637        let mut meta = HashMap::new();
638        meta.insert("source".into(), "test".into());
639        let mut e = entry("e1", "ns", vec![1.0]);
640        e.metadata = meta;
641        store.upsert(e).expect("should succeed");
642        let got = store.get("ns", "e1").expect("should succeed");
643        assert_eq!(got.metadata.get("source").map(|s| s.as_str()), Some("test"));
644    }
645
646    #[test]
647    fn test_created_at_stored() {
648        let mut store = VectorStore::new();
649        let mut e = entry("e1", "ns", vec![1.0]);
650        e.created_at = 12345678;
651        store.upsert(e).expect("should succeed");
652        assert_eq!(
653            store.get("ns", "e1").expect("should succeed").created_at,
654            12345678
655        );
656    }
657
658    #[test]
659    fn test_store_error_display_empty_vector() {
660        let e = StoreError::EmptyVector;
661        assert!(!e.to_string().is_empty());
662    }
663
664    #[test]
665    fn test_store_error_display_dimension_mismatch() {
666        let e = StoreError::DimensionMismatch {
667            expected: 3,
668            got: 5,
669        };
670        let s = e.to_string();
671        assert!(s.contains("3") && s.contains("5"));
672    }
673
674    #[test]
675    fn test_stats_namespace_count_after_delete() {
676        let mut store = VectorStore::new();
677        store
678            .upsert(entry("a", "ns1", vec![1.0]))
679            .expect("should succeed");
680        store
681            .upsert(entry("b", "ns2", vec![1.0]))
682            .expect("should succeed");
683        store.delete_namespace("ns1");
684        let s = store.stats();
685        assert_eq!(s.namespace_count, 1);
686    }
687
688    #[test]
689    fn test_search_scores_in_range() {
690        let mut store = VectorStore::new();
691        store
692            .upsert(entry("a", "ns", vec![1.0_f32, 0.0]))
693            .expect("should succeed");
694        store
695            .upsert(entry("b", "ns", vec![0.0_f32, 1.0]))
696            .expect("should succeed");
697        let results = store.search("ns", &[0.7_f32, 0.7], 2);
698        for r in &results {
699            assert!(r.score >= -1.0 && r.score <= 1.0);
700        }
701    }
702
703    #[test]
704    fn test_search_returns_metadata() {
705        let mut store = VectorStore::new();
706        let mut e = entry("e1", "ns", vec![1.0_f32, 0.0]);
707        e.metadata.insert("key".into(), "val".into());
708        store.upsert(e).expect("should succeed");
709        let results = store.search("ns", &[1.0_f32, 0.0], 1);
710        assert_eq!(
711            results[0].metadata.get("key").map(|s| s.as_str()),
712            Some("val")
713        );
714    }
715
716    #[test]
717    fn test_upsert_different_ids_same_namespace() {
718        let mut store = VectorStore::new();
719        store
720            .upsert(entry("a", "ns", vec![1.0]))
721            .expect("should succeed");
722        store
723            .upsert(entry("b", "ns", vec![2.0]))
724            .expect("should succeed");
725        store
726            .upsert(entry("c", "ns", vec![3.0]))
727            .expect("should succeed");
728        assert_eq!(store.list("ns").len(), 3);
729    }
730
731    #[test]
732    fn test_delete_reduces_list_count() {
733        let mut store = VectorStore::new();
734        store
735            .upsert(entry("a", "ns", vec![1.0]))
736            .expect("should succeed");
737        store
738            .upsert(entry("b", "ns", vec![2.0]))
739            .expect("should succeed");
740        store.delete("ns", "a");
741        assert_eq!(store.list("ns").len(), 1);
742    }
743
744    #[test]
745    fn test_search_all_namespaces_result_ids_correct() {
746        let mut store = VectorStore::new();
747        store
748            .upsert(entry("best", "ns1", vec![1.0_f32, 0.0]))
749            .expect("should succeed");
750        store
751            .upsert(entry("other", "ns2", vec![0.0_f32, 1.0]))
752            .expect("should succeed");
753        let results = store.search_all_namespaces(&[1.0_f32, 0.0], 1);
754        assert_eq!(results[0].id, "best");
755    }
756
757    #[test]
758    fn test_cosine_similarity_opposite_vectors() {
759        let mut store = VectorStore::new();
760        store
761            .upsert(entry("a", "ns", vec![1.0_f32, 0.0]))
762            .expect("should succeed");
763        store
764            .upsert(entry("b", "ns", vec![-1.0_f32, 0.0]))
765            .expect("should succeed");
766        let results = store.search("ns", &[1.0_f32, 0.0], 2);
767        // "a" should score higher than "b"
768        assert_eq!(results[0].id, "a");
769        assert!(results[0].score > results[1].score);
770    }
771
772    #[test]
773    fn test_with_dimension_rejects_extra_dims() {
774        let mut store = VectorStore::with_dimension(2);
775        let res = store.upsert(entry("e", "ns", vec![1.0, 2.0, 3.0]));
776        assert!(matches!(
777            res,
778            Err(StoreError::DimensionMismatch {
779                expected: 2,
780                got: 3
781            })
782        ));
783    }
784
785    #[test]
786    fn test_upsert_returns_new_flag_consistently() {
787        let mut store = VectorStore::new();
788        let r1 = store
789            .upsert(entry("e", "ns", vec![1.0]))
790            .expect("should succeed");
791        let r2 = store
792            .upsert(entry("e", "ns", vec![2.0]))
793            .expect("should succeed");
794        assert!(r1); // new
795        assert!(!r2); // update
796    }
797
798    #[test]
799    fn test_stats_total_includes_all_namespaces() {
800        let mut store = VectorStore::new();
801        for i in 0..5 {
802            store
803                .upsert(entry(&i.to_string(), &format!("ns{i}"), vec![i as f32]))
804                .expect("should succeed");
805        }
806        assert_eq!(store.stats().total_vectors, 5);
807    }
808
809    #[test]
810    fn test_get_after_update_returns_new_vector() {
811        let mut store = VectorStore::new();
812        store
813            .upsert(entry("e", "ns", vec![1.0, 0.0]))
814            .expect("should succeed");
815        store
816            .upsert(entry("e", "ns", vec![0.0, 1.0]))
817            .expect("should succeed");
818        let got = store.get("ns", "e").expect("should succeed");
819        assert_eq!(got.vector, vec![0.0_f32, 1.0]);
820    }
821
822    #[test]
823    fn test_search_zero_top_k() {
824        let mut store = VectorStore::new();
825        store
826            .upsert(entry("a", "ns", vec![1.0_f32]))
827            .expect("should succeed");
828        let results = store.search("ns", &[1.0_f32], 0);
829        assert!(results.is_empty());
830    }
831}
oxirs_embed/vector_store.rs

oxirs_embed/
vector_store.rs
