1use crate::{
8 hnsw::{HnswConfig, HnswIndex},
9 ivf::{IvfConfig, IvfIndex},
10 similarity::SimilarityMetric,
11 Vector, VectorIndex,
12};
13use anyhow::{anyhow, Result};
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16use std::fs::File;
17use std::io::{BufReader, BufWriter, Read, Write};
18use std::path::Path;
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
22pub enum FaissIndexType {
23 IndexFlatL2,
25 IndexFlatIP,
26 IndexIVFFlat,
28 IndexIVFPQ,
30 IndexHNSWFlat,
32 IndexLSH,
34 IndexPCAFlat,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct FaissIndexMetadata {
41 pub index_type: FaissIndexType,
42 pub dimension: usize,
43 pub num_vectors: usize,
44 pub metric_type: FaissMetricType,
45 pub parameters: HashMap<String, FaissParameter>,
46 pub version: String,
47 pub created_at: String,
48}
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
52pub enum FaissMetricType {
53 L2,
55 InnerProduct,
57 Cosine,
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize)]
63pub enum FaissParameter {
64 Integer(i64),
65 Float(f64),
66 String(String),
67 Boolean(bool),
68}
69
70pub struct FaissCompatibility {
72 supported_formats: Vec<FaissIndexType>,
73 conversion_cache: HashMap<String, ConversionResult>,
74}
75
76#[derive(Debug, Clone)]
78pub struct ConversionResult {
79 pub success: bool,
80 pub metadata: FaissIndexMetadata,
81 pub performance_metrics: ConversionMetrics,
82 pub warnings: Vec<String>,
83}
84
85#[derive(Debug, Clone, Default)]
87pub struct ConversionMetrics {
88 pub conversion_time: std::time::Duration,
89 pub memory_used: usize,
90 pub vectors_processed: usize,
91 pub accuracy_preserved: f32, }
93
94#[derive(Debug, Clone)]
96pub struct FaissExportConfig {
97 pub target_format: FaissIndexType,
98 pub compression_level: CompressionLevel,
99 pub preserve_accuracy: bool,
100 pub include_metadata: bool,
101 pub chunk_size: usize,
102}
103
104#[derive(Debug, Clone)]
106pub struct FaissImportConfig {
107 pub validate_format: bool,
108 pub preserve_performance: bool,
109 pub rebuild_if_incompatible: bool,
110 pub batch_size: usize,
111}
112
113#[derive(Debug, Clone, Copy)]
115pub enum CompressionLevel {
116 None,
117 Low,
118 Medium,
119 High,
120 Maximum,
121}
122
123impl Default for FaissExportConfig {
124 fn default() -> Self {
125 Self {
126 target_format: FaissIndexType::IndexHNSWFlat,
127 compression_level: CompressionLevel::Medium,
128 preserve_accuracy: true,
129 include_metadata: true,
130 chunk_size: 10000,
131 }
132 }
133}
134
135impl Default for FaissImportConfig {
136 fn default() -> Self {
137 Self {
138 validate_format: true,
139 preserve_performance: true,
140 rebuild_if_incompatible: false,
141 batch_size: 5000,
142 }
143 }
144}
145
146impl FaissCompatibility {
147 pub fn new() -> Self {
149 Self {
150 supported_formats: vec![
151 FaissIndexType::IndexFlatL2,
152 FaissIndexType::IndexFlatIP,
153 FaissIndexType::IndexIVFFlat,
154 FaissIndexType::IndexIVFPQ,
155 FaissIndexType::IndexHNSWFlat,
156 FaissIndexType::IndexLSH,
157 ],
158 conversion_cache: HashMap::new(),
159 }
160 }
161
162 pub fn export_to_faiss<T: VectorIndex + 'static>(
164 &mut self,
165 index: &T,
166 output_path: &Path,
167 config: &FaissExportConfig,
168 ) -> Result<ConversionResult> {
169 let start_time = std::time::Instant::now();
170 let mut warnings = Vec::new();
171
172 let detected_format = self.detect_optimal_faiss_format(index)?;
174 let target_format = if detected_format != config.target_format {
175 warnings.push(format!(
176 "Requested format {:?} differs from optimal format {:?}",
177 config.target_format, detected_format
178 ));
179 config.target_format
180 } else {
181 detected_format
182 };
183
184 let metadata = FaissIndexMetadata {
186 index_type: target_format,
187 dimension: self.get_index_dimension(index)?,
188 num_vectors: self.get_index_size(index)?,
189 metric_type: self.convert_similarity_metric(self.get_index_metric(index)?),
190 parameters: self.extract_index_parameters(index, target_format)?,
191 version: "oxirs-vec-1.0".to_string(),
192 created_at: chrono::Utc::now().to_rfc3339(),
193 };
194
195 self.write_faiss_index(index, output_path, &metadata, config)?;
197
198 let conversion_time = start_time.elapsed();
199 let performance_metrics = ConversionMetrics {
200 conversion_time,
201 memory_used: self.estimate_memory_usage(&metadata),
202 vectors_processed: metadata.num_vectors,
203 accuracy_preserved: self.estimate_accuracy_preservation(target_format, config),
204 };
205
206 let result = ConversionResult {
207 success: true,
208 metadata,
209 performance_metrics,
210 warnings,
211 };
212
213 let cache_key = format!("{:?}-{}", target_format, output_path.display());
215 self.conversion_cache.insert(cache_key, result.clone());
216
217 Ok(result)
218 }
219
220 pub fn import_from_faiss(
222 &mut self,
223 input_path: &Path,
224 config: &FaissImportConfig,
225 ) -> Result<Box<dyn VectorIndex>> {
226 let metadata = self.read_faiss_metadata(input_path)?;
228
229 if config.validate_format && !self.is_format_supported(&metadata.index_type) {
230 return Err(anyhow!(
231 "Unsupported FAISS format: {:?}",
232 metadata.index_type
233 ));
234 }
235
236 match metadata.index_type {
238 FaissIndexType::IndexHNSWFlat => self.import_hnsw_index(input_path, &metadata, config),
239 FaissIndexType::IndexIVFFlat | FaissIndexType::IndexIVFPQ => {
240 self.import_ivf_index(input_path, &metadata, config)
241 }
242 FaissIndexType::IndexFlatL2 | FaissIndexType::IndexFlatIP => {
243 self.import_flat_index(input_path, &metadata, config)
244 }
245 _ => Err(anyhow!(
246 "Import not yet implemented for {:?}",
247 metadata.index_type
248 )),
249 }
250 }
251
252 fn export_hnsw_to_faiss(
254 &self,
255 index: &HnswIndex,
256 output_path: &Path,
257 metadata: &FaissIndexMetadata,
258 config: &FaissExportConfig,
259 ) -> Result<()> {
260 let file = File::create(output_path)?;
261 let mut writer = BufWriter::new(file);
262
263 self.write_faiss_header(&mut writer, metadata)?;
265
266 self.write_hnsw_data(&mut writer, index, config)?;
268
269 self.write_vectors_chunked(&mut writer, index, config.chunk_size)?;
271
272 writer.flush()?;
273 Ok(())
274 }
275
276 fn export_ivf_to_faiss(
278 &self,
279 index: &IvfIndex,
280 output_path: &Path,
281 metadata: &FaissIndexMetadata,
282 config: &FaissExportConfig,
283 ) -> Result<()> {
284 let file = File::create(output_path)?;
285 let mut writer = BufWriter::new(file);
286
287 self.write_faiss_header(&mut writer, metadata)?;
289
290 self.write_ivf_data(&mut writer, index, config)?;
292
293 self.write_ivf_structure(&mut writer, index)?;
295
296 writer.flush()?;
297 Ok(())
298 }
299
300 fn import_hnsw_index(
302 &self,
303 input_path: &Path,
304 metadata: &FaissIndexMetadata,
305 config: &FaissImportConfig,
306 ) -> Result<Box<dyn VectorIndex>> {
307 let file = File::open(input_path)?;
308 let mut reader = BufReader::new(file);
309
310 self.skip_faiss_header(&mut reader)?;
312
313 let hnsw_config = self.read_hnsw_config(&mut reader, metadata)?;
315
316 let mut index = HnswIndex::new(hnsw_config)?;
318
319 self.import_vectors_batched(&mut reader, &mut index, metadata, config.batch_size)?;
321
322 Ok(Box::new(index))
323 }
324
325 fn import_ivf_index(
327 &self,
328 input_path: &Path,
329 metadata: &FaissIndexMetadata,
330 config: &FaissImportConfig,
331 ) -> Result<Box<dyn VectorIndex>> {
332 let file = File::open(input_path)?;
333 let mut reader = BufReader::new(file);
334
335 self.skip_faiss_header(&mut reader)?;
337
338 let ivf_config = self.read_ivf_config(&mut reader, metadata)?;
340
341 let mut index = IvfIndex::new(ivf_config)?;
343
344 self.read_ivf_structure(&mut reader, &mut index)?;
346
347 self.import_vectors_batched(&mut reader, &mut index, metadata, config.batch_size)?;
352
353 Ok(Box::new(index))
354 }
355
356 fn import_flat_index(
358 &self,
359 input_path: &Path,
360 metadata: &FaissIndexMetadata,
361 _config: &FaissImportConfig,
362 ) -> Result<Box<dyn VectorIndex>> {
363 let file = File::open(input_path)?;
364 let mut reader = BufReader::new(file);
365
366 self.skip_faiss_header(&mut reader)?;
368
369 let mut vectors = Vec::new();
371 let mut uris = Vec::new();
372
373 for i in 0..metadata.num_vectors {
375 let vector = self.read_vector(&mut reader, metadata.dimension)?;
376 vectors.push(vector);
377 uris.push(format!("faiss_vector_{i}"));
378 }
379
380 Ok(Box::new(SimpleVectorIndex::new(vectors, uris)))
382 }
383
384 fn detect_optimal_faiss_format<T: VectorIndex>(&self, index: &T) -> Result<FaissIndexType> {
386 let size = self.get_index_size(index)?;
387 let dimension = self.get_index_dimension(index)?;
388
389 if size < 10000 {
391 Ok(FaissIndexType::IndexFlatL2)
393 } else if dimension > 1000 {
394 Ok(FaissIndexType::IndexIVFPQ)
396 } else if size > 100000 {
397 Ok(FaissIndexType::IndexHNSWFlat)
399 } else {
400 Ok(FaissIndexType::IndexIVFFlat)
402 }
403 }
404
405 fn is_format_supported(&self, format: &FaissIndexType) -> bool {
407 self.supported_formats.contains(format)
408 }
409
410 fn convert_similarity_metric(&self, metric: SimilarityMetric) -> FaissMetricType {
412 match metric {
413 SimilarityMetric::Cosine => FaissMetricType::Cosine,
414 SimilarityMetric::Euclidean => FaissMetricType::L2,
415 SimilarityMetric::DotProduct => FaissMetricType::InnerProduct,
416 SimilarityMetric::Manhattan => FaissMetricType::L2, _ => FaissMetricType::L2,
419 }
420 }
421
422 fn write_faiss_header(
424 &self,
425 writer: &mut BufWriter<File>,
426 metadata: &FaissIndexMetadata,
427 ) -> Result<()> {
428 writer.write_all(b"FAISS")?;
430
431 writer.write_all(&1u32.to_le_bytes())?;
433
434 let type_id = self.faiss_type_to_id(metadata.index_type);
436 writer.write_all(&type_id.to_le_bytes())?;
437
438 writer.write_all(&(metadata.dimension as u32).to_le_bytes())?;
440
441 writer.write_all(&(metadata.num_vectors as u64).to_le_bytes())?;
443
444 let metric_id = self.faiss_metric_to_id(metadata.metric_type);
446 writer.write_all(&metric_id.to_le_bytes())?;
447
448 Ok(())
449 }
450
451 fn skip_faiss_header(&self, reader: &mut BufReader<File>) -> Result<()> {
453 let mut magic = [0u8; 5];
454 reader.read_exact(&mut magic)?;
455
456 if &magic != b"FAISS" {
457 return Err(anyhow!("Invalid FAISS file format"));
458 }
459
460 let mut buffer = [0u8; 21]; reader.read_exact(&mut buffer)?;
463
464 Ok(())
465 }
466
467 fn write_hnsw_data(
469 &self,
470 writer: &mut BufWriter<File>,
471 index: &HnswIndex,
472 _config: &FaissExportConfig,
473 ) -> Result<()> {
474 let config = index.config();
476 writer.write_all(&(config.m as u32).to_le_bytes())?;
477 writer.write_all(&(config.m_l0 as u32).to_le_bytes())?;
478 writer.write_all(&(config.ef as u32).to_le_bytes())?;
479 writer.write_all(&config.ml.to_le_bytes())?;
480
481 Ok(())
482 }
483
484 fn write_ivf_data(
486 &self,
487 writer: &mut BufWriter<File>,
488 index: &IvfIndex,
489 _config: &FaissExportConfig,
490 ) -> Result<()> {
491 let config = index.config();
493 writer.write_all(&(config.n_clusters as u32).to_le_bytes())?;
494 writer.write_all(&(config.n_probes as u32).to_le_bytes())?;
495
496 Ok(())
497 }
498
499 fn write_vectors_chunked<T: VectorIndex>(
504 &self,
505 writer: &mut BufWriter<File>,
506 index: &T,
507 chunk_size: usize,
508 ) -> Result<()> {
509 let vectors = index.iter_vectors();
510 let chunk_size = chunk_size.max(1);
511
512 for chunk in vectors.chunks(chunk_size) {
513 for (_uri, vector) in chunk {
514 self.write_vector(writer, vector)?;
515 }
516 }
517
518 Ok(())
519 }
520
521 fn write_vector(&self, writer: &mut BufWriter<File>, vector: &Vector) -> Result<()> {
523 let data = vector.as_f32();
524 for &value in &data {
525 writer.write_all(&value.to_le_bytes())?;
526 }
527 Ok(())
528 }
529
530 fn read_vector(&self, reader: &mut BufReader<File>, dimension: usize) -> Result<Vector> {
532 let mut data = vec![0.0f32; dimension];
533 for value in &mut data {
534 let mut bytes = [0u8; 4];
535 reader.read_exact(&mut bytes)?;
536 *value = f32::from_le_bytes(bytes);
537 }
538
539 Ok(Vector::new(data))
540 }
541
542 fn get_index_dimension<T: VectorIndex>(&self, index: &T) -> Result<usize> {
549 Ok(index
550 .iter_vectors()
551 .first()
552 .map(|(_, v)| v.dimensions)
553 .unwrap_or(768))
554 }
555
556 fn get_index_size<T: VectorIndex>(&self, index: &T) -> Result<usize> {
560 Ok(index.iter_vectors().len())
561 }
562
563 fn get_index_metric<T: VectorIndex>(&self, _index: &T) -> Result<SimilarityMetric> {
567 Ok(SimilarityMetric::Cosine) }
569
570 fn faiss_type_to_id(&self, faiss_type: FaissIndexType) -> u32 {
572 match faiss_type {
573 FaissIndexType::IndexFlatL2 => 0,
574 FaissIndexType::IndexFlatIP => 1,
575 FaissIndexType::IndexIVFFlat => 2,
576 FaissIndexType::IndexIVFPQ => 3,
577 FaissIndexType::IndexHNSWFlat => 4,
578 FaissIndexType::IndexLSH => 5,
579 FaissIndexType::IndexPCAFlat => 6,
580 }
581 }
582
583 fn faiss_metric_to_id(&self, metric: FaissMetricType) -> u8 {
584 match metric {
585 FaissMetricType::L2 => 0,
586 FaissMetricType::InnerProduct => 1,
587 FaissMetricType::Cosine => 2,
588 }
589 }
590
591 fn extract_index_parameters<T: VectorIndex>(
592 &self,
593 _index: &T,
594 _format: FaissIndexType,
595 ) -> Result<HashMap<String, FaissParameter>> {
596 let mut params = HashMap::new();
598 params.insert(
599 "created_by".to_string(),
600 FaissParameter::String("oxirs-vec".to_string()),
601 );
602 Ok(params)
603 }
604
605 fn estimate_memory_usage(&self, metadata: &FaissIndexMetadata) -> usize {
606 metadata.num_vectors * metadata.dimension * 4 }
609
610 fn estimate_accuracy_preservation(
611 &self,
612 _format: FaissIndexType,
613 _config: &FaissExportConfig,
614 ) -> f32 {
615 0.95 }
618
619 fn read_hnsw_config(
634 &self,
635 reader: &mut BufReader<File>,
636 _metadata: &FaissIndexMetadata,
637 ) -> Result<HnswConfig> {
638 let mut m_bytes = [0u8; 4];
639 reader.read_exact(&mut m_bytes)?;
640 let m = u32::from_le_bytes(m_bytes) as usize;
641
642 let mut m_l0_bytes = [0u8; 4];
643 reader.read_exact(&mut m_l0_bytes)?;
644 let m_l0 = u32::from_le_bytes(m_l0_bytes) as usize;
645
646 let mut ef_bytes = [0u8; 4];
647 reader.read_exact(&mut ef_bytes)?;
648 let ef = u32::from_le_bytes(ef_bytes) as usize;
649
650 let mut ml_bytes = [0u8; 8];
651 reader.read_exact(&mut ml_bytes)?;
652 let ml = f64::from_le_bytes(ml_bytes);
653
654 Ok(HnswConfig {
655 m,
656 m_l0,
657 ef,
658 ml,
659 ..HnswConfig::default()
660 })
661 }
662
663 fn read_ivf_config(
669 &self,
670 reader: &mut BufReader<File>,
671 _metadata: &FaissIndexMetadata,
672 ) -> Result<IvfConfig> {
673 let mut n_clusters_bytes = [0u8; 4];
674 reader.read_exact(&mut n_clusters_bytes)?;
675 let n_clusters = u32::from_le_bytes(n_clusters_bytes) as usize;
676
677 let mut n_probes_bytes = [0u8; 4];
678 reader.read_exact(&mut n_probes_bytes)?;
679 let n_probes = u32::from_le_bytes(n_probes_bytes) as usize;
680
681 Ok(IvfConfig {
682 n_clusters,
683 n_probes,
684 ..IvfConfig::default()
685 })
686 }
687
688 fn write_ivf_structure(&self, _writer: &mut BufWriter<File>, _index: &IvfIndex) -> Result<()> {
689 Ok(()) }
692
693 fn read_ivf_structure(
694 &self,
695 _reader: &mut BufReader<File>,
696 _index: &mut IvfIndex,
697 ) -> Result<()> {
698 Ok(()) }
701
702 fn import_vectors_batched<T: VectorIndex>(
708 &self,
709 reader: &mut BufReader<File>,
710 index: &mut T,
711 metadata: &FaissIndexMetadata,
712 batch_size: usize,
713 ) -> Result<()> {
714 let batch_size = batch_size.max(1);
715 let mut imported = 0usize;
716
717 while imported < metadata.num_vectors {
718 let batch_end = (imported + batch_size).min(metadata.num_vectors);
719 for i in imported..batch_end {
720 let vector = self.read_vector(reader, metadata.dimension)?;
721 index.insert(format!("faiss_vector_{i}"), vector)?;
722 }
723 imported = batch_end;
724 }
725
726 Ok(())
727 }
728
729 fn read_faiss_metadata(&self, input_path: &Path) -> Result<FaissIndexMetadata> {
733 let file = File::open(input_path)?;
734 let mut reader = BufReader::new(file);
735
736 let mut magic = [0u8; 5];
737 reader.read_exact(&mut magic)?;
738 if &magic != b"FAISS" {
739 return Err(anyhow!("Invalid FAISS file format: bad magic bytes"));
740 }
741
742 let mut version_bytes = [0u8; 4];
743 reader.read_exact(&mut version_bytes)?;
744 let _version = u32::from_le_bytes(version_bytes);
745
746 let mut type_bytes = [0u8; 4];
747 reader.read_exact(&mut type_bytes)?;
748 let index_type = self.faiss_id_to_type(u32::from_le_bytes(type_bytes))?;
749
750 let mut dim_bytes = [0u8; 4];
751 reader.read_exact(&mut dim_bytes)?;
752 let dimension = u32::from_le_bytes(dim_bytes) as usize;
753
754 let mut count_bytes = [0u8; 8];
755 reader.read_exact(&mut count_bytes)?;
756 let num_vectors = u64::from_le_bytes(count_bytes) as usize;
757
758 let mut metric_byte = [0u8; 1];
759 reader.read_exact(&mut metric_byte)?;
760 let metric_type = self.faiss_id_to_metric(metric_byte[0])?;
761
762 Ok(FaissIndexMetadata {
763 index_type,
764 dimension,
765 num_vectors,
766 metric_type,
767 parameters: HashMap::new(),
768 version: "1.0".to_string(),
769 created_at: chrono::Utc::now().to_rfc3339(),
770 })
771 }
772
773 fn write_faiss_index<T: VectorIndex + 'static>(
774 &self,
775 index: &T,
776 output_path: &Path,
777 metadata: &FaissIndexMetadata,
778 config: &FaissExportConfig,
779 ) -> Result<()> {
780 match metadata.index_type {
781 FaissIndexType::IndexHNSWFlat => {
782 if let Some(hnsw_index) = self.try_cast_to_hnsw(index) {
785 self.export_hnsw_to_faiss(hnsw_index, output_path, metadata, config)
786 } else {
787 Err(anyhow!("Index is not an HNSW index"))
788 }
789 }
790 FaissIndexType::IndexIVFFlat | FaissIndexType::IndexIVFPQ => {
791 if let Some(ivf_index) = self.try_cast_to_ivf(index) {
792 self.export_ivf_to_faiss(ivf_index, output_path, metadata, config)
793 } else {
794 Err(anyhow!("Index is not an IVF index"))
795 }
796 }
797 _ => Err(anyhow!(
798 "Export format not yet implemented: {:?}",
799 metadata.index_type
800 )),
801 }
802 }
803
804 fn try_cast_to_hnsw<'a, T: VectorIndex + 'static>(
810 &self,
811 index: &'a T,
812 ) -> Option<&'a HnswIndex> {
813 (index as &dyn std::any::Any).downcast_ref::<HnswIndex>()
814 }
815
816 fn try_cast_to_ivf<'a, T: VectorIndex + 'static>(&self, index: &'a T) -> Option<&'a IvfIndex> {
817 (index as &dyn std::any::Any).downcast_ref::<IvfIndex>()
818 }
819
820 fn faiss_id_to_type(&self, id: u32) -> Result<FaissIndexType> {
822 match id {
823 0 => Ok(FaissIndexType::IndexFlatL2),
824 1 => Ok(FaissIndexType::IndexFlatIP),
825 2 => Ok(FaissIndexType::IndexIVFFlat),
826 3 => Ok(FaissIndexType::IndexIVFPQ),
827 4 => Ok(FaissIndexType::IndexHNSWFlat),
828 5 => Ok(FaissIndexType::IndexLSH),
829 6 => Ok(FaissIndexType::IndexPCAFlat),
830 other => Err(anyhow!("Unknown FAISS index type id: {}", other)),
831 }
832 }
833
834 fn faiss_id_to_metric(&self, id: u8) -> Result<FaissMetricType> {
836 match id {
837 0 => Ok(FaissMetricType::L2),
838 1 => Ok(FaissMetricType::InnerProduct),
839 2 => Ok(FaissMetricType::Cosine),
840 other => Err(anyhow!("Unknown FAISS metric type id: {}", other)),
841 }
842 }
843}
844
845impl Default for FaissCompatibility {
846 fn default() -> Self {
847 Self::new()
848 }
849}
850
851pub struct SimpleVectorIndex {
853 vectors: Vec<Vector>,
854 uris: Vec<String>,
855}
856
857impl SimpleVectorIndex {
858 pub fn new(vectors: Vec<Vector>, uris: Vec<String>) -> Self {
859 Self { vectors, uris }
860 }
861}
862
863impl VectorIndex for SimpleVectorIndex {
864 fn insert(&mut self, uri: String, vector: Vector) -> Result<()> {
865 self.uris.push(uri);
866 self.vectors.push(vector);
867 Ok(())
868 }
869
870 fn search_knn(&self, query: &Vector, k: usize) -> Result<Vec<(String, f32)>> {
871 let mut results = Vec::new();
872
873 for (i, vector) in self.vectors.iter().enumerate() {
874 let similarity = self.compute_similarity(query, vector);
875 results.push((self.uris[i].clone(), similarity));
876 }
877
878 results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
880 results.truncate(k);
881
882 Ok(results)
883 }
884
885 fn search_threshold(&self, query: &Vector, threshold: f32) -> Result<Vec<(String, f32)>> {
886 let mut results = Vec::new();
887
888 for (i, vector) in self.vectors.iter().enumerate() {
889 let similarity = self.compute_similarity(query, vector);
890 if similarity >= threshold {
891 results.push((self.uris[i].clone(), similarity));
892 }
893 }
894
895 results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
897
898 Ok(results)
899 }
900
901 fn get_vector(&self, uri: &str) -> Option<&Vector> {
902 self.uris
903 .iter()
904 .position(|u| u == uri)
905 .map(|i| &self.vectors[i])
906 }
907
908 fn iter_vectors(&self) -> Vec<(String, Vector)> {
909 self.uris
910 .iter()
911 .cloned()
912 .zip(self.vectors.iter().cloned())
913 .collect()
914 }
915
916 fn supports_enumeration(&self) -> bool {
917 true
918 }
919}
920
921impl SimpleVectorIndex {
922 fn compute_similarity(&self, v1: &Vector, v2: &Vector) -> f32 {
923 let v1_data = v1.as_f32();
925 let v2_data = v2.as_f32();
926
927 if v1_data.len() != v2_data.len() {
928 return 0.0;
929 }
930
931 let dot_product: f32 = v1_data.iter().zip(v2_data.iter()).map(|(a, b)| a * b).sum();
932 let magnitude1: f32 = v1_data.iter().map(|x| x * x).sum::<f32>().sqrt();
933 let magnitude2: f32 = v2_data.iter().map(|x| x * x).sum::<f32>().sqrt();
934
935 if magnitude1 == 0.0 || magnitude2 == 0.0 {
936 0.0
937 } else {
938 dot_product / (magnitude1 * magnitude2)
939 }
940 }
941}
942
943pub mod utils {
945 use super::*;
946
947 pub fn convert_metric(metric: SimilarityMetric) -> FaissMetricType {
949 match metric {
950 SimilarityMetric::Cosine => FaissMetricType::Cosine,
951 SimilarityMetric::Euclidean => FaissMetricType::L2,
952 SimilarityMetric::DotProduct => FaissMetricType::InnerProduct,
953 SimilarityMetric::Manhattan => FaissMetricType::L2,
954 _ => FaissMetricType::L2,
956 }
957 }
958
959 pub fn recommend_faiss_format(
961 num_vectors: usize,
962 dimension: usize,
963 memory_constraint: Option<usize>,
964 accuracy_requirement: f32,
965 ) -> FaissIndexType {
966 if num_vectors < 1000 || accuracy_requirement > 0.99 {
967 FaissIndexType::IndexFlatL2
968 } else if dimension > 1000 || memory_constraint.is_some_and(|mem| mem < 1024 * 1024 * 1024)
969 {
970 FaissIndexType::IndexIVFPQ
971 } else if num_vectors > 100000 {
972 FaissIndexType::IndexHNSWFlat
973 } else {
974 FaissIndexType::IndexIVFFlat
975 }
976 }
977
978 pub fn estimate_memory_requirement(
980 format: FaissIndexType,
981 num_vectors: usize,
982 dimension: usize,
983 ) -> usize {
984 let base_memory = num_vectors * dimension * 4; match format {
987 FaissIndexType::IndexFlatL2 | FaissIndexType::IndexFlatIP => base_memory,
988 FaissIndexType::IndexIVFFlat => base_memory + (num_vectors / 100) * dimension * 4, FaissIndexType::IndexIVFPQ => base_memory / 8 + (num_vectors / 100) * dimension * 4, FaissIndexType::IndexHNSWFlat => base_memory * 2, FaissIndexType::IndexLSH => base_memory / 2, FaissIndexType::IndexPCAFlat => base_memory, }
994 }
995}
996
997#[cfg(test)]
998mod tests {
999 use super::*;
1000
1001 #[test]
1002 fn test_faiss_compatibility_creation() {
1003 let faiss_compat = FaissCompatibility::new();
1004 assert!(!faiss_compat.supported_formats.is_empty());
1005 }
1006
1007 #[test]
1008 fn test_metric_conversion() {
1009 let faiss_compat = FaissCompatibility::new();
1010
1011 assert_eq!(
1012 faiss_compat.convert_similarity_metric(SimilarityMetric::Cosine),
1013 FaissMetricType::Cosine
1014 );
1015 assert_eq!(
1016 faiss_compat.convert_similarity_metric(SimilarityMetric::Euclidean),
1017 FaissMetricType::L2
1018 );
1019 assert_eq!(
1020 faiss_compat.convert_similarity_metric(SimilarityMetric::DotProduct),
1021 FaissMetricType::InnerProduct
1022 );
1023 }
1024
1025 #[test]
1026 fn test_simple_vector_index() -> Result<()> {
1027 let vectors = vec![
1028 Vector::new(vec![1.0, 0.0, 0.0]),
1029 Vector::new(vec![0.0, 1.0, 0.0]),
1030 Vector::new(vec![0.0, 0.0, 1.0]),
1031 ];
1032 let uris = vec!["v1".to_string(), "v2".to_string(), "v3".to_string()];
1033
1034 let index = SimpleVectorIndex::new(vectors, uris);
1035
1036 let query = Vector::new(vec![1.0, 0.0, 0.0]);
1037 let results = index.search_knn(&query, 2)?;
1038
1039 assert_eq!(results.len(), 2);
1040 assert_eq!(results[0].0, "v1"); Ok(())
1042 }
1043
1044 #[test]
1045 fn test_format_recommendation() {
1046 use crate::faiss_compatibility::utils::recommend_faiss_format;
1047
1048 assert_eq!(
1050 recommend_faiss_format(100, 128, None, 0.9),
1051 FaissIndexType::IndexFlatL2
1052 );
1053
1054 assert_eq!(
1056 recommend_faiss_format(1000000, 128, None, 0.8),
1057 FaissIndexType::IndexHNSWFlat
1058 );
1059
1060 assert_eq!(
1062 recommend_faiss_format(50000, 2048, Some(512 * 1024 * 1024), 0.8),
1063 FaissIndexType::IndexIVFPQ
1064 );
1065 }
1066
1067 #[test]
1072 fn test_try_cast_to_hnsw_succeeds_for_real_hnsw_index() -> Result<()> {
1073 use crate::hnsw::{HnswConfig, HnswIndex};
1074
1075 let faiss_compat = FaissCompatibility::new();
1076 let index = HnswIndex::new(HnswConfig::default())?;
1077
1078 assert!(
1079 faiss_compat.try_cast_to_hnsw(&index).is_some(),
1080 "downcasting a real HnswIndex must succeed, not return None"
1081 );
1082
1083 let flat = SimpleVectorIndex::new(Vec::new(), Vec::new());
1085 assert!(faiss_compat.try_cast_to_ivf(&index).is_none());
1086 assert!(faiss_compat.try_cast_to_hnsw(&flat).is_none());
1087 Ok(())
1088 }
1089
1090 #[test]
1096 fn test_faiss_export_import_round_trip_hnsw() -> Result<()> {
1097 use crate::hnsw::{HnswConfig, HnswIndex};
1098
1099 let mut faiss_compat = FaissCompatibility::new();
1100 let mut index = HnswIndex::new(HnswConfig::default())?;
1101 index.insert("doc_a".to_string(), Vector::new(vec![1.0, 0.0, 0.0, 0.0]))?;
1102 index.insert("doc_b".to_string(), Vector::new(vec![0.0, 1.0, 0.0, 0.0]))?;
1103 index.insert("doc_c".to_string(), Vector::new(vec![0.0, 0.0, 1.0, 0.0]))?;
1104
1105 let dir = std::env::temp_dir().join(format!(
1106 "oxirs_vec_faiss_roundtrip_{}",
1107 uuid::Uuid::new_v4()
1108 ));
1109 std::fs::create_dir_all(&dir)?;
1110 let output_path = dir.join("index.faiss");
1111
1112 let export_config = FaissExportConfig::default(); let export_result = faiss_compat.export_to_faiss(&index, &output_path, &export_config)?;
1114 assert!(export_result.success);
1115 assert_eq!(export_result.metadata.num_vectors, 3);
1116 assert_eq!(export_result.metadata.dimension, 4);
1117
1118 let import_config = FaissImportConfig::default();
1119 let imported = faiss_compat.import_from_faiss(&output_path, &import_config)?;
1120
1121 let recovered = imported.iter_vectors();
1127 assert_eq!(recovered.len(), 3);
1128 let mut recovered_values: Vec<Vec<f32>> =
1129 recovered.iter().map(|(_, v)| v.as_f32().to_vec()).collect();
1130 recovered_values.sort_by(|a, b| a.partial_cmp(b).expect("no NaNs in test vectors"));
1131
1132 let mut expected_values = vec![
1133 vec![1.0f32, 0.0, 0.0, 0.0],
1134 vec![0.0, 1.0, 0.0, 0.0],
1135 vec![0.0, 0.0, 1.0, 0.0],
1136 ];
1137 expected_values.sort_by(|a, b| a.partial_cmp(b).expect("no NaNs in test vectors"));
1138
1139 assert_eq!(recovered_values, expected_values);
1140
1141 std::fs::remove_dir_all(&dir).ok();
1142 Ok(())
1143 }
1144}