1use crate::error::VectorError;
18use crate::{PgVectorStore, SearchResult, VectorMetric, VectorRecord};
19use std::collections::HashMap;
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum AnnIndexType {
33 Ivfflat,
35 Hnsw,
37}
38
39impl AnnIndexType {
40 pub fn as_sql_method(&self) -> &'static str {
42 match self {
43 AnnIndexType::Ivfflat => "ivfflat",
44 AnnIndexType::Hnsw => "hnsw",
45 }
46 }
47
48 pub fn as_str(&self) -> &'static str {
50 match self {
51 AnnIndexType::Ivfflat => "IVFFlat",
52 AnnIndexType::Hnsw => "HNSW",
53 }
54 }
55}
56
57#[derive(Debug, Clone)]
63pub struct IvfflatParams {
64 pub lists: usize,
66 pub probes: usize,
68}
69
70impl IvfflatParams {
71 pub fn new(lists: usize, probes: usize) -> Self {
73 Self { lists, probes }
74 }
75
76 pub fn recommended_for_rows(rows: usize) -> Self {
78 let lists = if rows < 1_000_000 {
79 (rows as f64).sqrt() as usize
80 } else {
81 rows / 1000
82 };
83 let lists = lists.max(1);
84 Self {
85 lists,
86 probes: (lists / 10).max(1),
87 }
88 }
89
90 pub fn to_sql_options(&self) -> String {
92 format!("(lists = {})", self.lists)
93 }
94}
95
96impl Default for IvfflatParams {
97 fn default() -> Self {
98 Self::new(100, 10)
99 }
100}
101
102#[derive(Debug, Clone)]
107pub struct HnswParams {
108 pub m: usize,
110 pub ef_construction: usize,
112}
113
114impl HnswParams {
115 pub fn new(m: usize, ef_construction: usize) -> Self {
117 Self { m, ef_construction }
118 }
119
120 pub fn to_sql_options(&self) -> String {
122 format!(
123 "(m = {}, ef_construction = {})",
124 self.m, self.ef_construction
125 )
126 }
127}
128
129impl Default for HnswParams {
130 fn default() -> Self {
131 Self::new(16, 64)
132 }
133}
134
135#[derive(Debug, Clone)]
137pub struct AnnIndexDef {
138 pub index_name: String,
140 pub collection: String,
142 pub index_type: AnnIndexType,
144 pub metric: VectorMetric,
146 pub ivfflat_params: Option<IvfflatParams>,
148 pub hnsw_params: Option<HnswParams>,
150 pub concurrently: bool,
152}
153
154impl AnnIndexDef {
155 pub fn new_ivfflat(collection: &str, metric: VectorMetric, params: IvfflatParams) -> Self {
157 Self {
158 index_name: format!("idx_{}_ivfflat", collection),
159 collection: collection.to_string(),
160 index_type: AnnIndexType::Ivfflat,
161 metric,
162 ivfflat_params: Some(params),
163 hnsw_params: None,
164 concurrently: false,
165 }
166 }
167
168 pub fn new_hnsw(collection: &str, metric: VectorMetric, params: HnswParams) -> Self {
170 Self {
171 index_name: format!("idx_{}_hnsw", collection),
172 collection: collection.to_string(),
173 index_type: AnnIndexType::Hnsw,
174 metric,
175 ivfflat_params: None,
176 hnsw_params: Some(params),
177 concurrently: false,
178 }
179 }
180
181 pub fn with_concurrently(mut self) -> Self {
183 self.concurrently = true;
184 self
185 }
186
187 pub fn to_create_sql(&self) -> String {
189 let concurrently_str = if self.concurrently {
190 "CONCURRENTLY "
191 } else {
192 ""
193 };
194 let op_class = match self.metric {
195 VectorMetric::Cosine => "vector_cosine_ops",
196 VectorMetric::Euclidean => "vector_l2_ops",
197 VectorMetric::DotProduct => "vector_ip_ops",
198 };
199 let options = match self.index_type {
200 AnnIndexType::Ivfflat => self
201 .ivfflat_params
202 .as_ref()
203 .map(|p| p.to_sql_options())
204 .unwrap_or_default(),
205 AnnIndexType::Hnsw => self
206 .hnsw_params
207 .as_ref()
208 .map(|p| p.to_sql_options())
209 .unwrap_or_default(),
210 };
211 let options_clause = if options.is_empty() {
212 String::new()
213 } else {
214 format!(" WITH {}", options)
215 };
216 format!(
217 "CREATE INDEX {}{} ON vectors_{} USING {} (embedding {}){}",
218 concurrently_str,
219 self.index_name,
220 self.collection,
221 self.index_type.as_sql_method(),
222 op_class,
223 options_clause
224 )
225 }
226
227 pub fn to_drop_sql(&self) -> String {
229 let concurrently_str = if self.concurrently {
230 "CONCURRENTLY "
231 } else {
232 ""
233 };
234 format!("DROP INDEX {}{}", concurrently_str, self.index_name)
235 }
236
237 pub fn to_reindex_sql(&self) -> String {
239 format!("REINDEX INDEX {}", self.index_name)
240 }
241}
242
243#[derive(Debug, Clone, Default)]
245pub struct AnnIndexRegistry {
246 indexes: HashMap<String, AnnIndexDef>,
248}
249
250impl AnnIndexRegistry {
251 pub fn new() -> Self {
253 Self::default()
254 }
255
256 pub fn register(&mut self, def: AnnIndexDef) -> Result<(), VectorError> {
258 if self.indexes.contains_key(&def.index_name) {
259 return Err(VectorError::Query(format!(
260 "ANN index already exists: {}",
261 def.index_name
262 )));
263 }
264 self.indexes.insert(def.index_name.clone(), def);
265 Ok(())
266 }
267
268 pub fn unregister(&mut self, index_name: &str) -> Result<AnnIndexDef, VectorError> {
270 self.indexes
271 .remove(index_name)
272 .ok_or_else(|| VectorError::Query(format!("ANN index not found: {}", index_name)))
273 }
274
275 pub fn exists(&self, index_name: &str) -> bool {
277 self.indexes.contains_key(index_name)
278 }
279
280 pub fn list_for_collection(&self, collection: &str) -> Vec<&AnnIndexDef> {
282 self.indexes
283 .values()
284 .filter(|def| def.collection == collection)
285 .collect()
286 }
287
288 pub fn list_all(&self) -> Vec<&AnnIndexDef> {
290 self.indexes.values().collect()
291 }
292
293 pub fn reindex_all_sql(&self) -> Vec<String> {
295 self.indexes
296 .values()
297 .map(|def| def.to_reindex_sql())
298 .collect()
299 }
300}
301
302pub struct SimilarityAlgorithms;
310
311impl SimilarityAlgorithms {
312 pub fn l2_distance(a: &[f32], b: &[f32]) -> f32 {
316 if a.len() != b.len() || a.is_empty() {
317 return f32::MAX;
318 }
319 a.iter()
320 .zip(b.iter())
321 .map(|(x, y)| (x - y) * (x - y))
322 .sum::<f32>()
323 .sqrt()
324 }
325
326 pub fn inner_product(a: &[f32], b: &[f32]) -> f32 {
330 if a.len() != b.len() || a.is_empty() {
331 return 0.0;
332 }
333 a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
334 }
335
336 pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
340 if a.len() != b.len() || a.is_empty() {
341 return 0.0;
342 }
343 let dot = Self::inner_product(a, b);
344 let na = Self::l2_norm(a);
345 let nb = Self::l2_norm(b);
346 if na == 0.0 || nb == 0.0 {
347 return 0.0;
348 }
349 dot / (na * nb)
350 }
351
352 pub fn manhattan_distance(a: &[f32], b: &[f32]) -> f32 {
356 if a.len() != b.len() || a.is_empty() {
357 return f32::MAX;
358 }
359 a.iter().zip(b.iter()).map(|(x, y)| (x - y).abs()).sum()
360 }
361
362 pub fn l2_norm(v: &[f32]) -> f32 {
366 v.iter().map(|x| x * x).sum::<f32>().sqrt()
367 }
368
369 pub fn distance_to_similarity(metric: VectorMetric, distance: f32) -> f32 {
375 match metric {
376 VectorMetric::Cosine => 1.0 - distance,
377 VectorMetric::Euclidean => 1.0 / (1.0 + distance),
378 VectorMetric::DotProduct => distance,
379 }
380 }
381
382 pub fn similarity(metric: VectorMetric, a: &[f32], b: &[f32]) -> f32 {
384 match metric {
385 VectorMetric::Cosine => Self::cosine_similarity(a, b),
386 VectorMetric::Euclidean => {
387 let dist = Self::l2_distance(a, b);
388 Self::distance_to_similarity(metric, dist)
389 }
390 VectorMetric::DotProduct => Self::inner_product(a, b),
391 }
392 }
393
394 pub fn batch_similarity(
398 metric: VectorMetric,
399 query: &[f32],
400 candidates: &[Vec<f32>],
401 ) -> Vec<(usize, f32)> {
402 candidates
403 .iter()
404 .enumerate()
405 .map(|(i, v)| (i, Self::similarity(metric, query, v)))
406 .collect()
407 }
408
409 pub fn batch_top_k(
411 metric: VectorMetric,
412 query: &[f32],
413 candidates: &[Vec<f32>],
414 top_k: usize,
415 ) -> Vec<(usize, f32)> {
416 let mut scored = Self::batch_similarity(metric, query, candidates);
417 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
419 scored.truncate(top_k);
420 scored
421 }
422}
423
424pub struct DimensionValidator;
430
431pub const MAX_VECTOR_DIMENSION: usize = 16000;
435pub const MIN_VECTOR_DIMENSION: usize = 1;
437
438impl DimensionValidator {
439 pub fn validate_dimension(dim: usize) -> Result<(), VectorError> {
441 if !(MIN_VECTOR_DIMENSION..=MAX_VECTOR_DIMENSION).contains(&dim) {
442 return Err(VectorError::InvalidConfig(format!(
443 "dimension must be between {} and {}, got {}",
444 MIN_VECTOR_DIMENSION, MAX_VECTOR_DIMENSION, dim
445 )));
446 }
447 Ok(())
448 }
449
450 pub fn validate_vector(vector: &[f32], expected_dim: usize) -> Result<(), VectorError> {
452 Self::validate_dimension(expected_dim)?;
453 if vector.len() != expected_dim {
454 return Err(VectorError::DimensionMismatch {
455 expected: expected_dim,
456 actual: vector.len(),
457 });
458 }
459 Ok(())
460 }
461
462 pub fn validate_batch(vectors: &[Vec<f32>], expected_dim: usize) -> Result<(), VectorError> {
466 Self::validate_dimension(expected_dim)?;
467 for (i, v) in vectors.iter().enumerate() {
468 if v.len() != expected_dim {
469 return Err(VectorError::DimensionMismatch {
470 expected: expected_dim,
471 actual: v.len(),
472 });
473 }
474 if v.iter().any(|x| x.is_nan() || x.is_infinite()) {
476 return Err(VectorError::InvalidConfig(format!(
477 "vector at index {} contains NaN or Inf values",
478 i
479 )));
480 }
481 }
482 Ok(())
483 }
484
485 pub fn validate_collection_compatibility(
487 dim_a: usize,
488 dim_b: usize,
489 ) -> Result<(), VectorError> {
490 Self::validate_dimension(dim_a)?;
491 Self::validate_dimension(dim_b)?;
492 if dim_a != dim_b {
493 return Err(VectorError::DimensionMismatch {
494 expected: dim_a,
495 actual: dim_b,
496 });
497 }
498 Ok(())
499 }
500
501 pub fn validate_query(query: &[f32], collection_dim: usize) -> Result<(), VectorError> {
503 Self::validate_vector(query, collection_dim)
504 }
505}
506
507pub struct VectorNormalizer;
513
514impl VectorNormalizer {
515 pub fn l2_normalize(v: &[f32]) -> Vec<f32> {
519 let norm = SimilarityAlgorithms::l2_norm(v);
520 if norm == 0.0 {
521 return v.to_vec();
522 }
523 v.iter().map(|x| x / norm).collect()
524 }
525
526 pub fn l2_normalize_in_place(v: &mut [f32]) {
528 let norm = SimilarityAlgorithms::l2_norm(v);
529 if norm != 0.0 {
530 for x in v.iter_mut() {
531 *x /= norm;
532 }
533 }
534 }
535
536 pub fn is_normalized(v: &[f32], tolerance: f32) -> bool {
538 let norm = SimilarityAlgorithms::l2_norm(v);
539 (norm - 1.0).abs() < tolerance
540 }
541
542 pub fn batch_l2_normalize(vectors: &[Vec<f32>]) -> Vec<Vec<f32>> {
544 vectors.iter().map(|v| Self::l2_normalize(v)).collect()
545 }
546}
547
548#[async_trait::async_trait]
556pub trait BatchOpsExt: PgVectorStore {
557 async fn batch_search(
561 &self,
562 collection: &str,
563 queries: &[Vec<f32>],
564 top_k: usize,
565 ) -> Result<Vec<Vec<SearchResult>>, VectorError>;
566
567 async fn batch_get(
569 &self,
570 collection: &str,
571 ids: &[String],
572 ) -> Result<Vec<Option<VectorRecord>>, VectorError>;
573
574 async fn search_with_filter(
578 &self,
579 collection: &str,
580 query: &[f32],
581 top_k: usize,
582 filter_key: &str,
583 filter_value: &serde_json::Value,
584 ) -> Result<Vec<SearchResult>, VectorError>;
585}
586
587pub struct MemoryBatchOps {
596 store: crate::InMemoryVectorStore,
597}
598
599impl MemoryBatchOps {
600 pub fn new(store: crate::InMemoryVectorStore) -> Self {
601 Self { store }
602 }
603
604 pub fn from_new() -> Self {
605 Self::new(crate::InMemoryVectorStore::new())
606 }
607
608 pub fn inner(&self) -> &crate::InMemoryVectorStore {
610 &self.store
611 }
612}
613
614#[async_trait::async_trait]
616impl PgVectorStore for MemoryBatchOps {
617 async fn create_collection(
618 &self,
619 name: &str,
620 dimension: usize,
621 metric: Option<VectorMetric>,
622 ) -> Result<(), VectorError> {
623 self.store.create_collection(name, dimension, metric).await
624 }
625
626 async fn delete_collection(&self, name: &str) -> Result<(), VectorError> {
627 self.store.delete_collection(name).await
628 }
629
630 async fn insert(
631 &self,
632 collection: &str,
633 records: Vec<VectorRecord>,
634 ) -> Result<(), VectorError> {
635 self.store.insert(collection, records).await
636 }
637
638 async fn search(
639 &self,
640 collection: &str,
641 query: &[f32],
642 top_k: usize,
643 ) -> Result<Vec<SearchResult>, VectorError> {
644 self.store.search(collection, query, top_k).await
645 }
646
647 async fn get(&self, collection: &str, id: &str) -> Result<Option<VectorRecord>, VectorError> {
648 self.store.get(collection, id).await
649 }
650
651 async fn delete(&self, collection: &str, ids: Vec<String>) -> Result<u64, VectorError> {
652 self.store.delete(collection, ids).await
653 }
654
655 async fn count(&self, collection: &str) -> Result<usize, VectorError> {
656 self.store.count(collection).await
657 }
658}
659
660#[async_trait::async_trait]
661impl BatchOpsExt for MemoryBatchOps {
662 async fn batch_search(
663 &self,
664 collection: &str,
665 queries: &[Vec<f32>],
666 top_k: usize,
667 ) -> Result<Vec<Vec<SearchResult>>, VectorError> {
668 let mut all_results = Vec::with_capacity(queries.len());
669 for query in queries {
670 let results = self.store.search(collection, query, top_k).await?;
671 all_results.push(results);
672 }
673 Ok(all_results)
674 }
675
676 async fn batch_get(
677 &self,
678 collection: &str,
679 ids: &[String],
680 ) -> Result<Vec<Option<VectorRecord>>, VectorError> {
681 let mut results = Vec::with_capacity(ids.len());
682 for id in ids {
683 let record = self.store.get(collection, id).await?;
684 results.push(record);
685 }
686 Ok(results)
687 }
688
689 async fn search_with_filter(
690 &self,
691 collection: &str,
692 query: &[f32],
693 top_k: usize,
694 filter_key: &str,
695 filter_value: &serde_json::Value,
696 ) -> Result<Vec<SearchResult>, VectorError> {
697 let expanded_k = top_k.saturating_mul(4).max(top_k);
699 let candidates = self.store.search(collection, query, expanded_k).await?;
700
701 let filtered: Vec<SearchResult> = candidates
703 .into_iter()
704 .filter(|r| {
705 r.metadata
706 .as_ref()
707 .and_then(|m| m.get(filter_key))
708 .map(|v| v == filter_value)
709 .unwrap_or(false)
710 })
711 .take(top_k)
712 .collect();
713
714 Ok(filtered)
715 }
716}
717
718#[cfg(test)]
719mod tests {
720 use super::*;
721 use serde_json::json;
722
723 #[test]
726 fn test_ann_index_type_sql_method() {
727 assert_eq!(AnnIndexType::Ivfflat.as_sql_method(), "ivfflat");
728 assert_eq!(AnnIndexType::Hnsw.as_sql_method(), "hnsw");
729 }
730
731 #[test]
732 fn test_ann_index_type_as_str() {
733 assert_eq!(AnnIndexType::Ivfflat.as_str(), "IVFFlat");
734 assert_eq!(AnnIndexType::Hnsw.as_str(), "HNSW");
735 }
736
737 #[test]
740 fn test_ivfflat_params_new() {
741 let p = IvfflatParams::new(100, 10);
742 assert_eq!(p.lists, 100);
743 assert_eq!(p.probes, 10);
744 }
745
746 #[test]
747 fn test_ivfflat_params_default() {
748 let p = IvfflatParams::default();
749 assert_eq!(p.lists, 100);
750 assert_eq!(p.probes, 10);
751 }
752
753 #[test]
754 fn test_ivfflat_params_recommended_for_small_dataset() {
755 let p = IvfflatParams::recommended_for_rows(10_000);
757 assert_eq!(p.lists, 100);
758 assert_eq!(p.probes, 10);
759 }
760
761 #[test]
762 fn test_ivfflat_params_recommended_for_large_dataset() {
763 let p = IvfflatParams::recommended_for_rows(2_000_000);
765 assert_eq!(p.lists, 2000);
766 }
767
768 #[test]
769 fn test_ivfflat_params_recommended_for_empty() {
770 let p = IvfflatParams::recommended_for_rows(0);
771 assert_eq!(p.lists, 1); }
773
774 #[test]
775 fn test_ivfflat_params_to_sql_options() {
776 let p = IvfflatParams::new(50, 5);
777 assert_eq!(p.to_sql_options(), "(lists = 50)");
778 }
779
780 #[test]
783 fn test_hnsw_params_new() {
784 let p = HnswParams::new(32, 128);
785 assert_eq!(p.m, 32);
786 assert_eq!(p.ef_construction, 128);
787 }
788
789 #[test]
790 fn test_hnsw_params_default() {
791 let p = HnswParams::default();
792 assert_eq!(p.m, 16);
793 assert_eq!(p.ef_construction, 64);
794 }
795
796 #[test]
797 fn test_hnsw_params_to_sql_options() {
798 let p = HnswParams::new(32, 128);
799 assert_eq!(p.to_sql_options(), "(m = 32, ef_construction = 128)");
800 }
801
802 #[test]
805 fn test_ann_index_def_new_ivfflat() {
806 let def =
807 AnnIndexDef::new_ivfflat("docs", VectorMetric::Cosine, IvfflatParams::new(100, 10));
808 assert_eq!(def.index_name, "idx_docs_ivfflat");
809 assert_eq!(def.index_type, AnnIndexType::Ivfflat);
810 assert_eq!(def.metric, VectorMetric::Cosine);
811 assert!(def.ivfflat_params.is_some());
812 assert!(def.hnsw_params.is_none());
813 assert!(!def.concurrently);
814 }
815
816 #[test]
817 fn test_ann_index_def_new_hnsw() {
818 let def = AnnIndexDef::new_hnsw("docs", VectorMetric::Euclidean, HnswParams::default());
819 assert_eq!(def.index_name, "idx_docs_hnsw");
820 assert_eq!(def.index_type, AnnIndexType::Hnsw);
821 assert_eq!(def.metric, VectorMetric::Euclidean);
822 assert!(def.ivfflat_params.is_none());
823 assert!(def.hnsw_params.is_some());
824 }
825
826 #[test]
827 fn test_ann_index_def_with_concurrently() {
828 let def = AnnIndexDef::new_hnsw("docs", VectorMetric::Cosine, HnswParams::default())
829 .with_concurrently();
830 assert!(def.concurrently);
831 }
832
833 #[test]
834 fn test_ann_index_def_to_create_sql_ivfflat() {
835 let def =
836 AnnIndexDef::new_ivfflat("docs", VectorMetric::Cosine, IvfflatParams::new(100, 10));
837 let sql = def.to_create_sql();
838 assert!(sql.contains("CREATE INDEX"));
839 assert!(sql.contains("idx_docs_ivfflat"));
840 assert!(sql.contains("USING ivfflat"));
841 assert!(sql.contains("vector_cosine_ops"));
842 assert!(sql.contains("WITH (lists = 100)"));
843 }
844
845 #[test]
846 fn test_ann_index_def_to_create_sql_hnsw() {
847 let def = AnnIndexDef::new_hnsw("docs", VectorMetric::Euclidean, HnswParams::new(16, 64));
848 let sql = def.to_create_sql();
849 assert!(sql.contains("USING hnsw"));
850 assert!(sql.contains("vector_l2_ops"));
851 assert!(sql.contains("WITH (m = 16, ef_construction = 64)"));
852 }
853
854 #[test]
855 fn test_ann_index_def_to_drop_sql() {
856 let def =
857 AnnIndexDef::new_ivfflat("docs", VectorMetric::DotProduct, IvfflatParams::default());
858 let sql = def.to_drop_sql();
859 assert_eq!(sql, "DROP INDEX idx_docs_ivfflat");
860 }
861
862 #[test]
863 fn test_ann_index_def_to_reindex_sql() {
864 let def = AnnIndexDef::new_hnsw("docs", VectorMetric::Cosine, HnswParams::default());
865 let sql = def.to_reindex_sql();
866 assert_eq!(sql, "REINDEX INDEX idx_docs_hnsw");
867 }
868
869 #[test]
872 fn test_ann_index_registry_register_and_exists() {
873 let mut reg = AnnIndexRegistry::new();
874 let def = AnnIndexDef::new_ivfflat("docs", VectorMetric::Cosine, IvfflatParams::default());
875 reg.register(def).unwrap();
876 assert!(reg.exists("idx_docs_ivfflat"));
877 assert!(!reg.exists("nonexistent"));
878 }
879
880 #[test]
881 fn test_ann_index_registry_duplicate_register_fails() {
882 let mut reg = AnnIndexRegistry::new();
883 let def = AnnIndexDef::new_ivfflat("docs", VectorMetric::Cosine, IvfflatParams::default());
884 reg.register(def).unwrap();
885 let def2 = AnnIndexDef::new_ivfflat("docs", VectorMetric::Cosine, IvfflatParams::default());
887 assert!(reg.register(def2).is_err());
888 }
889
890 #[test]
891 fn test_ann_index_registry_unregister() {
892 let mut reg = AnnIndexRegistry::new();
893 let def = AnnIndexDef::new_ivfflat("docs", VectorMetric::Cosine, IvfflatParams::default());
894 reg.register(def).unwrap();
895 let removed = reg.unregister("idx_docs_ivfflat").unwrap();
896 assert_eq!(removed.collection, "docs");
897 assert!(!reg.exists("idx_docs_ivfflat"));
898 }
899
900 #[test]
901 fn test_ann_index_registry_unregister_nonexistent_fails() {
902 let mut reg = AnnIndexRegistry::new();
903 assert!(reg.unregister("nonexistent").is_err());
904 }
905
906 #[test]
907 fn test_ann_index_registry_list_for_collection() {
908 let mut reg = AnnIndexRegistry::new();
909 reg.register(AnnIndexDef::new_ivfflat(
910 "docs",
911 VectorMetric::Cosine,
912 IvfflatParams::default(),
913 ))
914 .unwrap();
915 reg.register(AnnIndexDef::new_hnsw(
916 "docs",
917 VectorMetric::Euclidean,
918 HnswParams::default(),
919 ))
920 .unwrap();
921 reg.register(AnnIndexDef::new_ivfflat(
922 "images",
923 VectorMetric::Cosine,
924 IvfflatParams::default(),
925 ))
926 .unwrap();
927
928 let docs_indexes = reg.list_for_collection("docs");
929 assert_eq!(docs_indexes.len(), 2);
930 let images_indexes = reg.list_for_collection("images");
931 assert_eq!(images_indexes.len(), 1);
932 }
933
934 #[test]
935 fn test_ann_index_registry_list_all() {
936 let mut reg = AnnIndexRegistry::new();
937 reg.register(AnnIndexDef::new_ivfflat(
938 "docs",
939 VectorMetric::Cosine,
940 IvfflatParams::default(),
941 ))
942 .unwrap();
943 reg.register(AnnIndexDef::new_hnsw(
944 "images",
945 VectorMetric::Euclidean,
946 HnswParams::default(),
947 ))
948 .unwrap();
949 assert_eq!(reg.list_all().len(), 2);
950 }
951
952 #[test]
953 fn test_ann_index_registry_reindex_all_sql() {
954 let mut reg = AnnIndexRegistry::new();
955 reg.register(AnnIndexDef::new_ivfflat(
956 "docs",
957 VectorMetric::Cosine,
958 IvfflatParams::default(),
959 ))
960 .unwrap();
961 reg.register(AnnIndexDef::new_hnsw(
962 "images",
963 VectorMetric::Euclidean,
964 HnswParams::default(),
965 ))
966 .unwrap();
967 let sqls = reg.reindex_all_sql();
968 assert_eq!(sqls.len(), 2);
969 assert!(sqls[0].contains("REINDEX INDEX"));
970 }
971
972 #[test]
975 fn test_l2_distance_identical() {
976 let a = vec![1.0, 2.0, 3.0];
977 assert!((SimilarityAlgorithms::l2_distance(&a, &a)).abs() < 1e-10);
978 }
979
980 #[test]
981 fn test_l2_distance_known() {
982 let a = vec![0.0, 0.0];
983 let b = vec![3.0, 4.0];
984 assert!((SimilarityAlgorithms::l2_distance(&a, &b) - 5.0).abs() < 1e-6);
986 }
987
988 #[test]
989 fn test_l2_distance_mismatched_dims() {
990 assert_eq!(
991 SimilarityAlgorithms::l2_distance(&[1.0], &[1.0, 2.0]),
992 f32::MAX
993 );
994 }
995
996 #[test]
997 fn test_inner_product() {
998 let a = vec![1.0, 2.0, 3.0];
999 let b = vec![4.0, 5.0, 6.0];
1000 assert!((SimilarityAlgorithms::inner_product(&a, &b) - 32.0).abs() < 1e-6);
1002 }
1003
1004 #[test]
1005 fn test_cosine_similarity_identical() {
1006 let a = vec![1.0, 2.0, 3.0];
1007 assert!((SimilarityAlgorithms::cosine_similarity(&a, &a) - 1.0).abs() < 1e-6);
1008 }
1009
1010 #[test]
1011 fn test_cosine_similarity_orthogonal() {
1012 let a = vec![1.0, 0.0];
1013 let b = vec![0.0, 1.0];
1014 assert!(SimilarityAlgorithms::cosine_similarity(&a, &b).abs() < 1e-6);
1015 }
1016
1017 #[test]
1018 fn test_cosine_similarity_opposite() {
1019 let a = vec![1.0, 0.0];
1020 let b = vec![-1.0, 0.0];
1021 assert!((SimilarityAlgorithms::cosine_similarity(&a, &b) + 1.0).abs() < 1e-6);
1022 }
1023
1024 #[test]
1025 fn test_cosine_similarity_zero_vector() {
1026 let a = vec![0.0, 0.0];
1027 let b = vec![1.0, 0.0];
1028 assert_eq!(SimilarityAlgorithms::cosine_similarity(&a, &b), 0.0);
1029 }
1030
1031 #[test]
1032 fn test_manhattan_distance() {
1033 let a = vec![0.0, 0.0];
1034 let b = vec![3.0, 4.0];
1035 assert!((SimilarityAlgorithms::manhattan_distance(&a, &b) - 7.0).abs() < 1e-6);
1036 }
1037
1038 #[test]
1039 fn test_l2_norm() {
1040 let v = vec![3.0, 4.0];
1041 assert!((SimilarityAlgorithms::l2_norm(&v) - 5.0).abs() < 1e-6);
1042 }
1043
1044 #[test]
1045 fn test_l2_norm_zero() {
1046 assert_eq!(SimilarityAlgorithms::l2_norm(&[0.0, 0.0]), 0.0);
1047 }
1048
1049 #[test]
1050 fn test_distance_to_similarity_cosine() {
1051 assert!(
1053 (SimilarityAlgorithms::distance_to_similarity(VectorMetric::Cosine, 0.0) - 1.0).abs()
1054 < 1e-6
1055 );
1056 assert!(
1058 (SimilarityAlgorithms::distance_to_similarity(VectorMetric::Cosine, 1.0)).abs() < 1e-6
1059 );
1060 }
1061
1062 #[test]
1063 fn test_distance_to_similarity_euclidean() {
1064 assert!(
1066 (SimilarityAlgorithms::distance_to_similarity(VectorMetric::Euclidean, 0.0) - 1.0)
1067 .abs()
1068 < 1e-6
1069 );
1070 assert!(
1072 (SimilarityAlgorithms::distance_to_similarity(VectorMetric::Euclidean, 1.0) - 0.5)
1073 .abs()
1074 < 1e-6
1075 );
1076 }
1077
1078 #[test]
1079 fn test_similarity_by_metric_cosine() {
1080 let a = vec![1.0, 0.0];
1081 let b = vec![1.0, 0.0];
1082 assert!(
1083 (SimilarityAlgorithms::similarity(VectorMetric::Cosine, &a, &b) - 1.0).abs() < 1e-6
1084 );
1085 }
1086
1087 #[test]
1088 fn test_similarity_by_metric_dot_product() {
1089 let a = vec![2.0, 3.0];
1090 let b = vec![1.0, 1.0];
1091 assert!(
1093 (SimilarityAlgorithms::similarity(VectorMetric::DotProduct, &a, &b) - 5.0).abs() < 1e-6
1094 );
1095 }
1096
1097 #[test]
1098 fn test_batch_similarity() {
1099 let query = vec![1.0, 0.0];
1100 let candidates = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]];
1101 let scored =
1102 SimilarityAlgorithms::batch_similarity(VectorMetric::Cosine, &query, &candidates);
1103 assert_eq!(scored.len(), 3);
1104 assert_eq!(scored[0].0, 0);
1106 assert!((scored[0].1 - 1.0).abs() < 1e-6);
1107 }
1108
1109 #[test]
1110 fn test_batch_top_k() {
1111 let query = vec![1.0, 0.0];
1112 let candidates = vec![
1113 vec![0.0, 1.0],
1114 vec![1.0, 0.0],
1115 vec![1.0, 1.0],
1116 vec![0.9, 0.1],
1117 ];
1118 let top = SimilarityAlgorithms::batch_top_k(VectorMetric::Cosine, &query, &candidates, 2);
1119 assert_eq!(top.len(), 2);
1120 assert_eq!(top[0].0, 1);
1122 }
1123
1124 #[test]
1127 fn test_validate_dimension_valid() {
1128 assert!(DimensionValidator::validate_dimension(1).is_ok());
1129 assert!(DimensionValidator::validate_dimension(128).is_ok());
1130 assert!(DimensionValidator::validate_dimension(1536).is_ok());
1131 assert!(DimensionValidator::validate_dimension(MAX_VECTOR_DIMENSION).is_ok());
1132 }
1133
1134 #[test]
1135 fn test_validate_dimension_invalid() {
1136 assert!(DimensionValidator::validate_dimension(0).is_err());
1137 assert!(DimensionValidator::validate_dimension(MAX_VECTOR_DIMENSION + 1).is_err());
1138 }
1139
1140 #[test]
1141 fn test_validate_vector_match() {
1142 let v = vec![1.0, 2.0, 3.0];
1143 assert!(DimensionValidator::validate_vector(&v, 3).is_ok());
1144 }
1145
1146 #[test]
1147 fn test_validate_vector_mismatch() {
1148 let v = vec![1.0, 2.0];
1149 let result = DimensionValidator::validate_vector(&v, 3);
1150 assert!(matches!(result, Err(VectorError::DimensionMismatch { .. })));
1151 }
1152
1153 #[test]
1154 fn test_validate_batch_all_match() {
1155 let vectors = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]];
1156 assert!(DimensionValidator::validate_batch(&vectors, 2).is_ok());
1157 }
1158
1159 #[test]
1160 fn test_validate_batch_one_mismatch() {
1161 let vectors = vec![vec![1.0, 0.0], vec![0.0, 1.0, 0.0]];
1162 let result = DimensionValidator::validate_batch(&vectors, 2);
1163 assert!(matches!(result, Err(VectorError::DimensionMismatch { .. })));
1164 }
1165
1166 #[test]
1167 fn test_validate_batch_with_nan() {
1168 let vectors = vec![vec![1.0, f32::NAN]];
1169 let result = DimensionValidator::validate_batch(&vectors, 2);
1170 assert!(matches!(result, Err(VectorError::InvalidConfig(_))));
1171 }
1172
1173 #[test]
1174 fn test_validate_batch_with_inf() {
1175 let vectors = vec![vec![1.0, f32::INFINITY]];
1176 let result = DimensionValidator::validate_batch(&vectors, 2);
1177 assert!(matches!(result, Err(VectorError::InvalidConfig(_))));
1178 }
1179
1180 #[test]
1181 fn test_validate_collection_compatibility_match() {
1182 assert!(DimensionValidator::validate_collection_compatibility(128, 128).is_ok());
1183 }
1184
1185 #[test]
1186 fn test_validate_collection_compatibility_mismatch() {
1187 let result = DimensionValidator::validate_collection_compatibility(128, 256);
1188 assert!(matches!(result, Err(VectorError::DimensionMismatch { .. })));
1189 }
1190
1191 #[test]
1192 fn test_validate_query_match() {
1193 let query = vec![1.0, 2.0, 3.0];
1194 assert!(DimensionValidator::validate_query(&query, 3).is_ok());
1195 }
1196
1197 #[test]
1198 fn test_validate_query_mismatch() {
1199 let query = vec![1.0, 2.0];
1200 assert!(DimensionValidator::validate_query(&query, 3).is_err());
1201 }
1202
1203 #[test]
1206 fn test_l2_normalize_unit_vector() {
1207 let v = vec![1.0, 0.0];
1208 let n = VectorNormalizer::l2_normalize(&v);
1209 assert!((n[0] - 1.0).abs() < 1e-6);
1210 assert!(n[1].abs() < 1e-6);
1211 }
1212
1213 #[test]
1214 fn test_l2_normalize_non_unit() {
1215 let v = vec![3.0, 4.0];
1216 let n = VectorNormalizer::l2_normalize(&v);
1217 assert!((n[0] - 0.6).abs() < 1e-6);
1219 assert!((n[1] - 0.8).abs() < 1e-6);
1220 }
1221
1222 #[test]
1223 fn test_l2_normalize_zero_vector() {
1224 let v = vec![0.0, 0.0];
1225 let n = VectorNormalizer::l2_normalize(&v);
1226 assert_eq!(n, vec![0.0, 0.0]);
1228 }
1229
1230 #[test]
1231 fn test_l2_normalize_in_place() {
1232 let mut v = vec![3.0, 4.0];
1233 VectorNormalizer::l2_normalize_in_place(&mut v);
1234 assert!((v[0] - 0.6).abs() < 1e-6);
1235 assert!((v[1] - 0.8).abs() < 1e-6);
1236 }
1237
1238 #[test]
1239 fn test_is_normalized_true() {
1240 let v = vec![1.0, 0.0];
1241 assert!(VectorNormalizer::is_normalized(&v, 1e-5));
1242 }
1243
1244 #[test]
1245 fn test_is_normalized_false() {
1246 let v = vec![3.0, 4.0];
1247 assert!(!VectorNormalizer::is_normalized(&v, 1e-5));
1248 }
1249
1250 #[test]
1251 fn test_batch_l2_normalize() {
1252 let vectors = vec![vec![3.0, 4.0], vec![1.0, 0.0]];
1253 let normalized = VectorNormalizer::batch_l2_normalize(&vectors);
1254 assert!((normalized[0][0] - 0.6).abs() < 1e-6);
1255 assert!((normalized[0][1] - 0.8).abs() < 1e-6);
1256 assert!((normalized[1][0] - 1.0).abs() < 1e-6);
1257 }
1258
1259 #[tokio::test]
1262 async fn test_memory_batch_ops_batch_search() {
1263 let batch_ops = MemoryBatchOps::from_new();
1264 batch_ops
1265 .store
1266 .create_collection("docs", 3, Some(VectorMetric::Cosine))
1267 .await
1268 .unwrap();
1269 batch_ops
1270 .store
1271 .insert(
1272 "docs",
1273 vec![
1274 VectorRecord::new("a", vec![1.0, 0.0, 0.0]),
1275 VectorRecord::new("b", vec![0.0, 1.0, 0.0]),
1276 VectorRecord::new("c", vec![0.0, 0.0, 1.0]),
1277 ],
1278 )
1279 .await
1280 .unwrap();
1281
1282 let queries = vec![vec![1.0, 0.0, 0.0], vec![0.0, 0.0, 1.0]];
1283 let results = batch_ops.batch_search("docs", &queries, 2).await.unwrap();
1284 assert_eq!(results.len(), 2);
1285 assert_eq!(results[0][0].id, "a");
1286 assert_eq!(results[1][0].id, "c");
1287 }
1288
1289 #[tokio::test]
1290 async fn test_memory_batch_ops_batch_get() {
1291 let batch_ops = MemoryBatchOps::from_new();
1292 batch_ops
1293 .store
1294 .create_collection("docs", 2, None)
1295 .await
1296 .unwrap();
1297 batch_ops
1298 .store
1299 .insert(
1300 "docs",
1301 vec![
1302 VectorRecord::new("a", vec![1.0, 0.0]),
1303 VectorRecord::new("b", vec![0.0, 1.0]),
1304 ],
1305 )
1306 .await
1307 .unwrap();
1308
1309 let ids = vec!["a".to_string(), "missing".to_string(), "b".to_string()];
1310 let results = batch_ops.batch_get("docs", &ids).await.unwrap();
1311 assert_eq!(results.len(), 3);
1312 assert!(results[0].is_some());
1313 assert!(results[1].is_none());
1314 assert!(results[2].is_some());
1315 }
1316
1317 #[tokio::test]
1318 async fn test_memory_batch_ops_search_with_filter() {
1319 let batch_ops = MemoryBatchOps::from_new();
1320 batch_ops
1321 .store
1322 .create_collection("docs", 2, Some(VectorMetric::Cosine))
1323 .await
1324 .unwrap();
1325
1326 let mut meta_a = HashMap::new();
1327 meta_a.insert("category".to_string(), json!("science"));
1328 let mut meta_b = HashMap::new();
1329 meta_b.insert("category".to_string(), json!("sports"));
1330 let mut meta_c = HashMap::new();
1331 meta_c.insert("category".to_string(), json!("science"));
1332
1333 batch_ops
1334 .store
1335 .insert(
1336 "docs",
1337 vec![
1338 VectorRecord::new("a", vec![1.0, 0.0]).with_metadata(meta_a),
1339 VectorRecord::new("b", vec![0.0, 1.0]).with_metadata(meta_b),
1340 VectorRecord::new("c", vec![0.9, 0.1]).with_metadata(meta_c),
1341 ],
1342 )
1343 .await
1344 .unwrap();
1345
1346 let results = batch_ops
1347 .search_with_filter("docs", &[1.0, 0.0], 5, "category", &json!("science"))
1348 .await
1349 .unwrap();
1350
1351 assert_eq!(results.len(), 2);
1353 assert!(results.iter().all(|r| r.id == "a" || r.id == "c"));
1354 }
1355
1356 #[tokio::test]
1357 async fn test_memory_batch_ops_filter_no_match() {
1358 let batch_ops = MemoryBatchOps::from_new();
1359 batch_ops
1360 .store
1361 .create_collection("docs", 2, None)
1362 .await
1363 .unwrap();
1364
1365 let mut meta = HashMap::new();
1366 meta.insert("category".to_string(), json!("science"));
1367 batch_ops
1368 .store
1369 .insert(
1370 "docs",
1371 vec![VectorRecord::new("a", vec![1.0, 0.0]).with_metadata(meta)],
1372 )
1373 .await
1374 .unwrap();
1375
1376 let results = batch_ops
1377 .search_with_filter("docs", &[1.0, 0.0], 5, "category", &json!("nonexistent"))
1378 .await
1379 .unwrap();
1380 assert!(results.is_empty());
1381 }
1382}