use crate::{
hnsw::{HnswConfig, HnswIndex},
ivf::{IvfConfig, IvfIndex},
similarity::SimilarityMetric,
Vector, VectorIndex,
};
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs::File;
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::Path;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FaissIndexType {
IndexFlatL2,
IndexFlatIP,
IndexIVFFlat,
IndexIVFPQ,
IndexHNSWFlat,
IndexLSH,
IndexPCAFlat,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FaissIndexMetadata {
pub index_type: FaissIndexType,
pub dimension: usize,
pub num_vectors: usize,
pub metric_type: FaissMetricType,
pub parameters: HashMap<String, FaissParameter>,
pub version: String,
pub created_at: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FaissMetricType {
L2,
InnerProduct,
Cosine,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum FaissParameter {
Integer(i64),
Float(f64),
String(String),
Boolean(bool),
}
pub struct FaissCompatibility {
supported_formats: Vec<FaissIndexType>,
conversion_cache: HashMap<String, ConversionResult>,
}
#[derive(Debug, Clone)]
pub struct ConversionResult {
pub success: bool,
pub metadata: FaissIndexMetadata,
pub performance_metrics: ConversionMetrics,
pub warnings: Vec<String>,
}
#[derive(Debug, Clone, Default)]
pub struct ConversionMetrics {
pub conversion_time: std::time::Duration,
pub memory_used: usize,
pub vectors_processed: usize,
pub accuracy_preserved: f32, }
#[derive(Debug, Clone)]
pub struct FaissExportConfig {
pub target_format: FaissIndexType,
pub compression_level: CompressionLevel,
pub preserve_accuracy: bool,
pub include_metadata: bool,
pub chunk_size: usize,
}
#[derive(Debug, Clone)]
pub struct FaissImportConfig {
pub validate_format: bool,
pub preserve_performance: bool,
pub rebuild_if_incompatible: bool,
pub batch_size: usize,
}
#[derive(Debug, Clone, Copy)]
pub enum CompressionLevel {
None,
Low,
Medium,
High,
Maximum,
}
impl Default for FaissExportConfig {
fn default() -> Self {
Self {
target_format: FaissIndexType::IndexHNSWFlat,
compression_level: CompressionLevel::Medium,
preserve_accuracy: true,
include_metadata: true,
chunk_size: 10000,
}
}
}
impl Default for FaissImportConfig {
fn default() -> Self {
Self {
validate_format: true,
preserve_performance: true,
rebuild_if_incompatible: false,
batch_size: 5000,
}
}
}
impl FaissCompatibility {
pub fn new() -> Self {
Self {
supported_formats: vec![
FaissIndexType::IndexFlatL2,
FaissIndexType::IndexFlatIP,
FaissIndexType::IndexIVFFlat,
FaissIndexType::IndexIVFPQ,
FaissIndexType::IndexHNSWFlat,
FaissIndexType::IndexLSH,
],
conversion_cache: HashMap::new(),
}
}
pub fn export_to_faiss<T: VectorIndex>(
&mut self,
index: &T,
output_path: &Path,
config: &FaissExportConfig,
) -> Result<ConversionResult> {
let start_time = std::time::Instant::now();
let mut warnings = Vec::new();
let detected_format = self.detect_optimal_faiss_format(index)?;
let target_format = if detected_format != config.target_format {
warnings.push(format!(
"Requested format {:?} differs from optimal format {:?}",
config.target_format, detected_format
));
config.target_format
} else {
detected_format
};
let metadata = FaissIndexMetadata {
index_type: target_format,
dimension: self.get_index_dimension(index)?,
num_vectors: self.get_index_size(index)?,
metric_type: self.convert_similarity_metric(self.get_index_metric(index)?),
parameters: self.extract_index_parameters(index, target_format)?,
version: "oxirs-vec-1.0".to_string(),
created_at: chrono::Utc::now().to_rfc3339(),
};
self.write_faiss_index(index, output_path, &metadata, config)?;
let conversion_time = start_time.elapsed();
let performance_metrics = ConversionMetrics {
conversion_time,
memory_used: self.estimate_memory_usage(&metadata),
vectors_processed: metadata.num_vectors,
accuracy_preserved: self.estimate_accuracy_preservation(target_format, config),
};
let result = ConversionResult {
success: true,
metadata,
performance_metrics,
warnings,
};
let cache_key = format!("{:?}-{}", target_format, output_path.display());
self.conversion_cache.insert(cache_key, result.clone());
Ok(result)
}
pub fn import_from_faiss(
&mut self,
input_path: &Path,
config: &FaissImportConfig,
) -> Result<Box<dyn VectorIndex>> {
let metadata = self.read_faiss_metadata(input_path)?;
if config.validate_format && !self.is_format_supported(&metadata.index_type) {
return Err(anyhow!(
"Unsupported FAISS format: {:?}",
metadata.index_type
));
}
match metadata.index_type {
FaissIndexType::IndexHNSWFlat => self.import_hnsw_index(input_path, &metadata, config),
FaissIndexType::IndexIVFFlat | FaissIndexType::IndexIVFPQ => {
self.import_ivf_index(input_path, &metadata, config)
}
FaissIndexType::IndexFlatL2 | FaissIndexType::IndexFlatIP => {
self.import_flat_index(input_path, &metadata, config)
}
_ => Err(anyhow!(
"Import not yet implemented for {:?}",
metadata.index_type
)),
}
}
fn export_hnsw_to_faiss(
&self,
index: &HnswIndex,
output_path: &Path,
metadata: &FaissIndexMetadata,
config: &FaissExportConfig,
) -> Result<()> {
let file = File::create(output_path)?;
let mut writer = BufWriter::new(file);
self.write_faiss_header(&mut writer, metadata)?;
self.write_hnsw_data(&mut writer, index, config)?;
self.write_vectors_chunked(&mut writer, index, config.chunk_size)?;
writer.flush()?;
Ok(())
}
fn export_ivf_to_faiss(
&self,
index: &IvfIndex,
output_path: &Path,
metadata: &FaissIndexMetadata,
config: &FaissExportConfig,
) -> Result<()> {
let file = File::create(output_path)?;
let mut writer = BufWriter::new(file);
self.write_faiss_header(&mut writer, metadata)?;
self.write_ivf_data(&mut writer, index, config)?;
self.write_ivf_structure(&mut writer, index)?;
writer.flush()?;
Ok(())
}
fn import_hnsw_index(
&self,
input_path: &Path,
metadata: &FaissIndexMetadata,
config: &FaissImportConfig,
) -> Result<Box<dyn VectorIndex>> {
let file = File::open(input_path)?;
let mut reader = BufReader::new(file);
self.skip_faiss_header(&mut reader)?;
let hnsw_config = self.read_hnsw_config(&mut reader, metadata)?;
let mut index = HnswIndex::new(hnsw_config)?;
self.import_vectors_batched(&mut reader, &mut index, config.batch_size)?;
Ok(Box::new(index))
}
fn import_ivf_index(
&self,
input_path: &Path,
metadata: &FaissIndexMetadata,
config: &FaissImportConfig,
) -> Result<Box<dyn VectorIndex>> {
let file = File::open(input_path)?;
let mut reader = BufReader::new(file);
self.skip_faiss_header(&mut reader)?;
let ivf_config = self.read_ivf_config(&mut reader, metadata)?;
let mut index = IvfIndex::new(ivf_config)?;
self.read_ivf_structure(&mut reader, &mut index)?;
self.import_vectors_batched(&mut reader, &mut index, config.batch_size)?;
Ok(Box::new(index))
}
fn import_flat_index(
&self,
input_path: &Path,
metadata: &FaissIndexMetadata,
_config: &FaissImportConfig,
) -> Result<Box<dyn VectorIndex>> {
let file = File::open(input_path)?;
let mut reader = BufReader::new(file);
self.skip_faiss_header(&mut reader)?;
let mut vectors = Vec::new();
let mut uris = Vec::new();
for i in 0..metadata.num_vectors {
let vector = self.read_vector(&mut reader, metadata.dimension)?;
vectors.push(vector);
uris.push(format!("faiss_vector_{i}"));
}
Ok(Box::new(SimpleVectorIndex::new(vectors, uris)))
}
fn detect_optimal_faiss_format<T: VectorIndex>(&self, index: &T) -> Result<FaissIndexType> {
let size = self.get_index_size(index)?;
let dimension = self.get_index_dimension(index)?;
if size < 10000 {
Ok(FaissIndexType::IndexFlatL2)
} else if dimension > 1000 {
Ok(FaissIndexType::IndexIVFPQ)
} else if size > 100000 {
Ok(FaissIndexType::IndexHNSWFlat)
} else {
Ok(FaissIndexType::IndexIVFFlat)
}
}
fn is_format_supported(&self, format: &FaissIndexType) -> bool {
self.supported_formats.contains(format)
}
fn convert_similarity_metric(&self, metric: SimilarityMetric) -> FaissMetricType {
match metric {
SimilarityMetric::Cosine => FaissMetricType::Cosine,
SimilarityMetric::Euclidean => FaissMetricType::L2,
SimilarityMetric::DotProduct => FaissMetricType::InnerProduct,
SimilarityMetric::Manhattan => FaissMetricType::L2, _ => FaissMetricType::L2,
}
}
fn write_faiss_header(
&self,
writer: &mut BufWriter<File>,
metadata: &FaissIndexMetadata,
) -> Result<()> {
writer.write_all(b"FAISS")?;
writer.write_all(&1u32.to_le_bytes())?;
let type_id = self.faiss_type_to_id(metadata.index_type);
writer.write_all(&type_id.to_le_bytes())?;
writer.write_all(&(metadata.dimension as u32).to_le_bytes())?;
writer.write_all(&(metadata.num_vectors as u64).to_le_bytes())?;
let metric_id = self.faiss_metric_to_id(metadata.metric_type);
writer.write_all(&metric_id.to_le_bytes())?;
Ok(())
}
fn skip_faiss_header(&self, reader: &mut BufReader<File>) -> Result<()> {
let mut magic = [0u8; 5];
reader.read_exact(&mut magic)?;
if &magic != b"FAISS" {
return Err(anyhow!("Invalid FAISS file format"));
}
let mut buffer = [0u8; 21]; reader.read_exact(&mut buffer)?;
Ok(())
}
fn write_hnsw_data(
&self,
writer: &mut BufWriter<File>,
index: &HnswIndex,
_config: &FaissExportConfig,
) -> Result<()> {
let config = index.config();
writer.write_all(&(config.m as u32).to_le_bytes())?;
writer.write_all(&(config.m_l0 as u32).to_le_bytes())?;
writer.write_all(&(config.ef as u32).to_le_bytes())?;
writer.write_all(&config.ml.to_le_bytes())?;
Ok(())
}
fn write_ivf_data(
&self,
writer: &mut BufWriter<File>,
index: &IvfIndex,
_config: &FaissExportConfig,
) -> Result<()> {
let config = index.config();
writer.write_all(&(config.n_clusters as u32).to_le_bytes())?;
writer.write_all(&(config.n_probes as u32).to_le_bytes())?;
Ok(())
}
fn write_vectors_chunked<T: VectorIndex>(
&self,
writer: &mut BufWriter<File>,
index: &T,
chunk_size: usize,
) -> Result<()> {
let total_vectors = self.get_index_size(index)?;
for chunk_start in (0..total_vectors).step_by(chunk_size) {
let chunk_end = std::cmp::min(chunk_start + chunk_size, total_vectors);
for i in chunk_start..chunk_end {
if let Some(vector) = self.get_vector_at_index(index, i) {
self.write_vector(writer, &vector)?;
}
}
}
Ok(())
}
fn write_vector(&self, writer: &mut BufWriter<File>, vector: &Vector) -> Result<()> {
let data = vector.as_f32();
for &value in &data {
writer.write_all(&value.to_le_bytes())?;
}
Ok(())
}
fn read_vector(&self, reader: &mut BufReader<File>, dimension: usize) -> Result<Vector> {
let mut data = vec![0.0f32; dimension];
for value in &mut data {
let mut bytes = [0u8; 4];
reader.read_exact(&mut bytes)?;
*value = f32::from_le_bytes(bytes);
}
Ok(Vector::new(data))
}
fn get_index_dimension<T: VectorIndex>(&self, _index: &T) -> Result<usize> {
Ok(768) }
fn get_index_size<T: VectorIndex>(&self, _index: &T) -> Result<usize> {
Ok(0) }
fn get_index_metric<T: VectorIndex>(&self, _index: &T) -> Result<SimilarityMetric> {
Ok(SimilarityMetric::Cosine) }
fn get_vector_at_index<T: VectorIndex>(&self, _index: &T, _idx: usize) -> Option<Vector> {
None }
fn faiss_type_to_id(&self, faiss_type: FaissIndexType) -> u32 {
match faiss_type {
FaissIndexType::IndexFlatL2 => 0,
FaissIndexType::IndexFlatIP => 1,
FaissIndexType::IndexIVFFlat => 2,
FaissIndexType::IndexIVFPQ => 3,
FaissIndexType::IndexHNSWFlat => 4,
FaissIndexType::IndexLSH => 5,
FaissIndexType::IndexPCAFlat => 6,
}
}
fn faiss_metric_to_id(&self, metric: FaissMetricType) -> u8 {
match metric {
FaissMetricType::L2 => 0,
FaissMetricType::InnerProduct => 1,
FaissMetricType::Cosine => 2,
}
}
fn extract_index_parameters<T: VectorIndex>(
&self,
_index: &T,
_format: FaissIndexType,
) -> Result<HashMap<String, FaissParameter>> {
let mut params = HashMap::new();
params.insert(
"created_by".to_string(),
FaissParameter::String("oxirs-vec".to_string()),
);
Ok(params)
}
fn estimate_memory_usage(&self, metadata: &FaissIndexMetadata) -> usize {
metadata.num_vectors * metadata.dimension * 4 }
fn estimate_accuracy_preservation(
&self,
_format: FaissIndexType,
_config: &FaissExportConfig,
) -> f32 {
0.95 }
fn read_hnsw_config(
&self,
_reader: &mut BufReader<File>,
_metadata: &FaissIndexMetadata,
) -> Result<HnswConfig> {
Ok(HnswConfig::default()) }
fn read_ivf_config(
&self,
_reader: &mut BufReader<File>,
_metadata: &FaissIndexMetadata,
) -> Result<IvfConfig> {
Ok(IvfConfig::default()) }
fn write_ivf_structure(&self, _writer: &mut BufWriter<File>, _index: &IvfIndex) -> Result<()> {
Ok(()) }
fn read_ivf_structure(
&self,
_reader: &mut BufReader<File>,
_index: &mut IvfIndex,
) -> Result<()> {
Ok(()) }
fn import_vectors_batched<T: VectorIndex>(
&self,
_reader: &mut BufReader<File>,
_index: &mut T,
_batch_size: usize,
) -> Result<()> {
Ok(()) }
fn read_faiss_metadata(&self, _input_path: &Path) -> Result<FaissIndexMetadata> {
Ok(FaissIndexMetadata {
index_type: FaissIndexType::IndexHNSWFlat,
dimension: 768,
num_vectors: 0,
metric_type: FaissMetricType::Cosine,
parameters: HashMap::new(),
version: "1.0".to_string(),
created_at: chrono::Utc::now().to_rfc3339(),
}) }
fn write_faiss_index<T: VectorIndex>(
&self,
index: &T,
output_path: &Path,
metadata: &FaissIndexMetadata,
config: &FaissExportConfig,
) -> Result<()> {
match metadata.index_type {
FaissIndexType::IndexHNSWFlat => {
if let Some(hnsw_index) = self.try_cast_to_hnsw(index) {
self.export_hnsw_to_faiss(hnsw_index, output_path, metadata, config)
} else {
Err(anyhow!("Index is not an HNSW index"))
}
}
FaissIndexType::IndexIVFFlat | FaissIndexType::IndexIVFPQ => {
if let Some(ivf_index) = self.try_cast_to_ivf(index) {
self.export_ivf_to_faiss(ivf_index, output_path, metadata, config)
} else {
Err(anyhow!("Index is not an IVF index"))
}
}
_ => Err(anyhow!(
"Export format not yet implemented: {:?}",
metadata.index_type
)),
}
}
fn try_cast_to_hnsw<T: VectorIndex>(&self, _index: &T) -> Option<&HnswIndex> {
None }
fn try_cast_to_ivf<T: VectorIndex>(&self, _index: &T) -> Option<&IvfIndex> {
None }
}
impl Default for FaissCompatibility {
fn default() -> Self {
Self::new()
}
}
pub struct SimpleVectorIndex {
vectors: Vec<Vector>,
uris: Vec<String>,
}
impl SimpleVectorIndex {
pub fn new(vectors: Vec<Vector>, uris: Vec<String>) -> Self {
Self { vectors, uris }
}
}
impl VectorIndex for SimpleVectorIndex {
fn insert(&mut self, uri: String, vector: Vector) -> Result<()> {
self.uris.push(uri);
self.vectors.push(vector);
Ok(())
}
fn search_knn(&self, query: &Vector, k: usize) -> Result<Vec<(String, f32)>> {
let mut results = Vec::new();
for (i, vector) in self.vectors.iter().enumerate() {
let similarity = self.compute_similarity(query, vector);
results.push((self.uris[i].clone(), similarity));
}
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
results.truncate(k);
Ok(results)
}
fn search_threshold(&self, query: &Vector, threshold: f32) -> Result<Vec<(String, f32)>> {
let mut results = Vec::new();
for (i, vector) in self.vectors.iter().enumerate() {
let similarity = self.compute_similarity(query, vector);
if similarity >= threshold {
results.push((self.uris[i].clone(), similarity));
}
}
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
Ok(results)
}
fn get_vector(&self, uri: &str) -> Option<&Vector> {
self.uris
.iter()
.position(|u| u == uri)
.map(|i| &self.vectors[i])
}
}
impl SimpleVectorIndex {
fn compute_similarity(&self, v1: &Vector, v2: &Vector) -> f32 {
let v1_data = v1.as_f32();
let v2_data = v2.as_f32();
if v1_data.len() != v2_data.len() {
return 0.0;
}
let dot_product: f32 = v1_data.iter().zip(v2_data.iter()).map(|(a, b)| a * b).sum();
let magnitude1: f32 = v1_data.iter().map(|x| x * x).sum::<f32>().sqrt();
let magnitude2: f32 = v2_data.iter().map(|x| x * x).sum::<f32>().sqrt();
if magnitude1 == 0.0 || magnitude2 == 0.0 {
0.0
} else {
dot_product / (magnitude1 * magnitude2)
}
}
}
pub mod utils {
use super::*;
pub fn convert_metric(metric: SimilarityMetric) -> FaissMetricType {
match metric {
SimilarityMetric::Cosine => FaissMetricType::Cosine,
SimilarityMetric::Euclidean => FaissMetricType::L2,
SimilarityMetric::DotProduct => FaissMetricType::InnerProduct,
SimilarityMetric::Manhattan => FaissMetricType::L2,
_ => FaissMetricType::L2,
}
}
pub fn recommend_faiss_format(
num_vectors: usize,
dimension: usize,
memory_constraint: Option<usize>,
accuracy_requirement: f32,
) -> FaissIndexType {
if num_vectors < 1000 || accuracy_requirement > 0.99 {
FaissIndexType::IndexFlatL2
} else if dimension > 1000 || memory_constraint.is_some_and(|mem| mem < 1024 * 1024 * 1024)
{
FaissIndexType::IndexIVFPQ
} else if num_vectors > 100000 {
FaissIndexType::IndexHNSWFlat
} else {
FaissIndexType::IndexIVFFlat
}
}
pub fn estimate_memory_requirement(
format: FaissIndexType,
num_vectors: usize,
dimension: usize,
) -> usize {
let base_memory = num_vectors * dimension * 4;
match format {
FaissIndexType::IndexFlatL2 | FaissIndexType::IndexFlatIP => base_memory,
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, }
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_faiss_compatibility_creation() {
let faiss_compat = FaissCompatibility::new();
assert!(!faiss_compat.supported_formats.is_empty());
}
#[test]
fn test_metric_conversion() {
let faiss_compat = FaissCompatibility::new();
assert_eq!(
faiss_compat.convert_similarity_metric(SimilarityMetric::Cosine),
FaissMetricType::Cosine
);
assert_eq!(
faiss_compat.convert_similarity_metric(SimilarityMetric::Euclidean),
FaissMetricType::L2
);
assert_eq!(
faiss_compat.convert_similarity_metric(SimilarityMetric::DotProduct),
FaissMetricType::InnerProduct
);
}
#[test]
fn test_simple_vector_index() -> Result<()> {
let vectors = vec![
Vector::new(vec![1.0, 0.0, 0.0]),
Vector::new(vec![0.0, 1.0, 0.0]),
Vector::new(vec![0.0, 0.0, 1.0]),
];
let uris = vec!["v1".to_string(), "v2".to_string(), "v3".to_string()];
let index = SimpleVectorIndex::new(vectors, uris);
let query = Vector::new(vec![1.0, 0.0, 0.0]);
let results = index.search_knn(&query, 2)?;
assert_eq!(results.len(), 2);
assert_eq!(results[0].0, "v1"); Ok(())
}
#[test]
fn test_format_recommendation() {
use crate::faiss_compatibility::utils::recommend_faiss_format;
assert_eq!(
recommend_faiss_format(100, 128, None, 0.9),
FaissIndexType::IndexFlatL2
);
assert_eq!(
recommend_faiss_format(1000000, 128, None, 0.8),
FaissIndexType::IndexHNSWFlat
);
assert_eq!(
recommend_faiss_format(50000, 2048, Some(512 * 1024 * 1024), 0.8),
FaissIndexType::IndexIVFPQ
);
}
}