oxify_vector/

persistence.rs

//! Index Persistence
//!
//! Save and load vector search indexes to/from disk for faster startup
//! and sharing between processes.
//!
//! ## Features
//!
//! - **JSON Serialization**: Human-readable format for debugging
//! - **Binary Serialization (rkyv)**: Zero-copy deserialization for maximum performance
//! - **Memory-Mapped Files**: Lazy loading for large indexes with minimal memory overhead
//! - **File I/O**: Simple save/load operations
//! - **Type Safety**: Compile-time guarantees for index types
//!
//! ## Example
//!
//! ```rust,no_run
//! use oxify_vector::{HnswIndex, HnswConfig};
//! use oxify_vector::persistence::{save_index, load_index};
//! use std::collections::HashMap;
//!
//! # fn example() -> anyhow::Result<()> {
//! // Build an index
//! let mut embeddings = HashMap::new();
//! embeddings.insert("doc1".to_string(), vec![0.1, 0.2, 0.3]);
//! embeddings.insert("doc2".to_string(), vec![0.4, 0.5, 0.6]);
//!
//! let mut index = HnswIndex::new(HnswConfig::default());
//! index.build(&embeddings)?;
//!
//! // Save to disk
//! save_index(&index, "/tmp/my_index.json")?;
//!
//! // Load from disk
//! let loaded_index: HnswIndex = load_index("/tmp/my_index.json")?;
//!
//! // Use loaded index
//! let query = vec![0.2, 0.3, 0.4];
//! let results = loaded_index.search(&query, 5)?;
//! # Ok(())
//! # }
//! ```

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::fs::File;
use std::io::{BufReader, BufWriter};
use std::path::Path;
use tracing::info;

#[cfg(feature = "zerocopy")]
use std::io::Write;

#[cfg(feature = "mmap")]
use memmap2::Mmap;

#[cfg(feature = "zerocopy")]
use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};

/// Save an index to a JSON file
///
/// The index must implement `Serialize`. This works with all index types:
/// - `VectorSearchIndex`
/// - `HnswIndex`
/// - `IvfPqIndex`
/// - `HybridIndex`
/// - `ColbertIndex`
pub fn save_index<T: Serialize, P: AsRef<Path>>(index: &T, path: P) -> Result<()> {
    let path = path.as_ref();
    info!("Saving index to: {}", path.display());

    let file =
        File::create(path).with_context(|| format!("Failed to create file: {}", path.display()))?;

    let writer = BufWriter::new(file);
    serde_json::to_writer_pretty(writer, index)
        .with_context(|| format!("Failed to serialize index to: {}", path.display()))?;

    info!("Index saved successfully");
    Ok(())
}

/// Load an index from a JSON file
///
/// The index type must implement `Deserialize`. Specify the type explicitly:
/// ```rust,ignore
/// let index: HnswIndex = load_index("path/to/index.json")?;
/// ```
pub fn load_index<T: for<'de> Deserialize<'de>, P: AsRef<Path>>(path: P) -> Result<T> {
    let path = path.as_ref();
    info!("Loading index from: {}", path.display());

    let file =
        File::open(path).with_context(|| format!("Failed to open file: {}", path.display()))?;

    let reader = BufReader::new(file);
    let index = serde_json::from_reader(reader)
        .with_context(|| format!("Failed to deserialize index from: {}", path.display()))?;

    info!("Index loaded successfully");
    Ok(index)
}

/// Get the size of a serialized index without saving to disk
///
/// Useful for estimating storage requirements.
pub fn get_serialized_size<T: Serialize>(index: &T) -> Result<usize> {
    let json =
        serde_json::to_string(index).context("Failed to serialize index for size calculation")?;
    Ok(json.len())
}

/// Check if an index file exists and is readable
pub fn index_file_exists<P: AsRef<Path>>(path: P) -> bool {
    path.as_ref().exists() && path.as_ref().is_file()
}

// ============================================================================
// Zero-Copy Serialization (rkyv)
// ============================================================================

/// Save an index to a binary file using rkyv for zero-copy deserialization
///
/// This is significantly faster than JSON for large indexes. The binary format
/// allows for instant loading without parsing overhead.
///
/// **Note**: Requires the `zerocopy` feature to be enabled.
#[cfg(feature = "zerocopy")]
pub fn save_index_binary<T, P>(index: &T, path: P) -> Result<()>
where
    T: for<'a> RkyvSerialize<
        rkyv::rancor::Strategy<
            rkyv::ser::Serializer<
                rkyv::util::AlignedVec,
                rkyv::ser::allocator::ArenaHandle<'a>,
                rkyv::ser::sharing::Share,
            >,
            rkyv::rancor::Error,
        >,
    >,
    P: AsRef<Path>,
{
    let path = path.as_ref();
    info!("Saving index (binary) to: {}", path.display());

    // Serialize to bytes using rkyv 0.8 API
    let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(index)
        .map_err(|e| anyhow::anyhow!("Failed to serialize index: {}", e))?;

    // Write to file
    let mut file =
        File::create(path).with_context(|| format!("Failed to create file: {}", path.display()))?;
    file.write_all(&bytes)
        .with_context(|| format!("Failed to write to file: {}", path.display()))?;

    info!("Index saved successfully ({} bytes)", bytes.len());
    Ok(())
}

/// Load an index from a binary file using zero-copy deserialization
///
/// This is much faster than JSON loading, especially for large indexes.
/// The data is deserialized in-place without allocating intermediate structures.
///
/// **Note**: Requires the `zerocopy` feature to be enabled.
///
/// # Safety
///
/// This uses unchecked access for performance. Only use with trusted data.
#[cfg(feature = "zerocopy")]
pub fn load_index_binary<T, P>(path: P) -> Result<T>
where
    T: Archive,
    T::Archived: RkyvDeserialize<T, rkyv::rancor::Strategy<rkyv::de::Pool, rkyv::rancor::Error>>,
    P: AsRef<Path>,
{
    let path = path.as_ref();
    info!("Loading index (binary) from: {}", path.display());

    // Read file into memory
    let bytes =
        std::fs::read(path).with_context(|| format!("Failed to read file: {}", path.display()))?;

    // Deserialize using unchecked access (safe for trusted data)
    // SAFETY: We control the serialization format and only deserialize our own data
    let archived = unsafe { rkyv::access_unchecked::<T::Archived>(&bytes) };

    let mut deserializer = rkyv::de::Pool::new();
    let index: T = archived
        .deserialize(rkyv::rancor::Strategy::wrap(&mut deserializer))
        .map_err(|e| anyhow::anyhow!("Failed to deserialize archived data: {}", e))?;

    info!("Index loaded successfully");
    Ok(index)
}

// ============================================================================
// Memory-Mapped File Support
// ============================================================================

/// Memory-mapped index for zero-copy lazy loading
///
/// This struct holds a memory-mapped view of an index file. The OS handles
/// paging data in/out of memory as needed, reducing memory footprint for
/// large indexes.
///
/// **Note**: Requires the `mmap` feature to be enabled.
#[cfg(feature = "mmap")]
pub struct MappedIndex {
    _mmap: Mmap,
    data: Vec<u8>,
}

#[cfg(feature = "mmap")]
impl MappedIndex {
    /// Create a memory-mapped view of an index file
    ///
    /// The file is memory-mapped, allowing the OS to lazily load pages
    /// as they're accessed. This is ideal for very large indexes.
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
        let path = path.as_ref();
        info!("Memory-mapping index from: {}", path.display());

        let file =
            File::open(path).with_context(|| format!("Failed to open file: {}", path.display()))?;

        // SAFETY: We're opening the file in read-only mode and not modifying it
        let mmap = unsafe {
            Mmap::map(&file)
                .with_context(|| format!("Failed to memory-map file: {}", path.display()))?
        };

        // For this implementation, we'll copy the data to allow safe access
        // A more advanced implementation could use rkyv with the mmap directly
        let data = mmap.to_vec();

        info!("Index memory-mapped successfully ({} bytes)", data.len());
        Ok(Self { _mmap: mmap, data })
    }

    /// Get the raw bytes of the memory-mapped index
    pub fn as_bytes(&self) -> &[u8] {
        &self.data
    }

    /// Deserialize the memory-mapped index
    ///
    /// This works with JSON-serialized indexes.
    pub fn deserialize<T: for<'de> Deserialize<'de>>(&self) -> Result<T> {
        serde_json::from_slice(&self.data).context("Failed to deserialize memory-mapped index")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::hnsw::{HnswConfig, HnswIndex};
    use crate::ivf::{IvfPqConfig, IvfPqIndex};
    use crate::search::VectorSearchIndex;
    use crate::types::SearchConfig;
    use std::collections::HashMap;
    use tempfile::TempDir;

    fn create_test_embeddings() -> HashMap<String, Vec<f32>> {
        let mut embeddings = HashMap::new();
        embeddings.insert("doc1".to_string(), vec![0.1, 0.2, 0.3]);
        embeddings.insert("doc2".to_string(), vec![0.4, 0.5, 0.6]);
        embeddings.insert("doc3".to_string(), vec![0.7, 0.8, 0.9]);
        embeddings
    }

    #[test]
    fn test_save_and_load_hnsw() {
        let temp_dir = TempDir::new().unwrap();
        let index_path = temp_dir.path().join("hnsw_index.json");

        // Build and save index
        let embeddings = create_test_embeddings();
        let mut index = HnswIndex::new(HnswConfig::default());
        index.build(&embeddings).unwrap();

        save_index(&index, &index_path).unwrap();
        assert!(index_file_exists(&index_path));

        // Load index
        let loaded_index: HnswIndex = load_index(&index_path).unwrap();

        // Verify search works
        let query = vec![0.2, 0.3, 0.4];
        let results = loaded_index.search(&query, 2).unwrap();
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn test_save_and_load_exact_search() {
        let temp_dir = TempDir::new().unwrap();
        let index_path = temp_dir.path().join("exact_index.json");

        // Build and save index
        let embeddings = create_test_embeddings();
        let mut index = VectorSearchIndex::new(SearchConfig::default());
        index.build(&embeddings).unwrap();

        save_index(&index, &index_path).unwrap();

        // Load index
        let loaded_index: VectorSearchIndex = load_index(&index_path).unwrap();

        // Verify search works
        let query = vec![0.5, 0.6, 0.7];
        let results = loaded_index.search(&query, 2).unwrap();
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn test_save_and_load_ivf_pq() {
        // Optimized test with reduced parameters for fast execution
        let temp_dir = TempDir::new().unwrap();
        let index_path = temp_dir.path().join("ivf_index.json");

        let mut embeddings = HashMap::new();
        for i in 0..500 {
            let vec = vec![
                i as f32 * 0.001,
                (i + 1) as f32 * 0.001,
                (i + 2) as f32 * 0.001,
                (i + 3) as f32 * 0.001,
            ];
            embeddings.insert(format!("doc{}", i), vec);
        }

        let config = IvfPqConfig {
            nclusters: 8, // Reduced from 16
            nsubvectors: 4,
            nbits: 4,                  // Explicitly set to 4 (vs default 8 = 256 centroids!)
            nprobe: 2,                 // Reduced from 4
            max_kmeans_iterations: 20, // Reduced from 100
            ..IvfPqConfig::default()
        };
        let mut index = IvfPqIndex::new(config);
        index.build(&embeddings).unwrap();

        save_index(&index, &index_path).unwrap();

        // Load index
        let loaded_index: IvfPqIndex = load_index(&index_path).unwrap();

        // Verify search works
        let query = vec![0.5, 0.6, 0.7, 0.8];
        let results = loaded_index.search(&query, 5).unwrap();
        assert!(!results.is_empty());
    }

    #[test]
    #[ignore]
    fn test_save_and_load_ivf_pq_full() {
        // Slow comprehensive test with default parameters (51s+)
        // Run with: cargo test test_save_and_load_ivf_pq_full -- --ignored
        let temp_dir = TempDir::new().unwrap();
        let index_path = temp_dir.path().join("ivf_index_full.json");

        // Build and save index (need many vectors for IVF with default clusters)
        let mut embeddings = HashMap::new();
        for i in 0..1000 {
            let vec = vec![
                i as f32 * 0.001,
                (i + 1) as f32 * 0.001,
                (i + 2) as f32 * 0.001,
                (i + 3) as f32 * 0.001,
            ];
            embeddings.insert(format!("doc{}", i), vec);
        }

        let config = IvfPqConfig {
            nclusters: 16, // Use fewer clusters for test
            nsubvectors: 4,
            nprobe: 4,
            ..IvfPqConfig::default()
        };
        let mut index = IvfPqIndex::new(config);
        index.build(&embeddings).unwrap();

        save_index(&index, &index_path).unwrap();

        // Load index
        let loaded_index: IvfPqIndex = load_index(&index_path).unwrap();

        // Verify search works
        let query = vec![0.5, 0.6, 0.7, 0.8];
        let results = loaded_index.search(&query, 5).unwrap();
        assert!(!results.is_empty());
    }

    #[test]
    fn test_get_serialized_size() {
        let embeddings = create_test_embeddings();
        let mut index = HnswIndex::new(HnswConfig::default());
        index.build(&embeddings).unwrap();

        let size = get_serialized_size(&index).unwrap();
        assert!(size > 0);
        assert!(size < 100000); // Should be reasonably small for 3 vectors
    }

    #[test]
    fn test_index_file_exists() {
        let temp_dir = TempDir::new().unwrap();
        let index_path = temp_dir.path().join("test_index.json");

        assert!(!index_file_exists(&index_path));

        // Create file
        let embeddings = create_test_embeddings();
        let mut index = HnswIndex::new(HnswConfig::default());
        index.build(&embeddings).unwrap();
        save_index(&index, &index_path).unwrap();

        assert!(index_file_exists(&index_path));
    }

    #[test]
    fn test_load_nonexistent_file() {
        let result: Result<HnswIndex> = load_index("/nonexistent/path/index.json");
        assert!(result.is_err());
    }

    #[test]
    fn test_save_to_invalid_path() {
        let embeddings = create_test_embeddings();
        let mut index = HnswIndex::new(HnswConfig::default());
        index.build(&embeddings).unwrap();

        let result = save_index(&index, "/invalid/nonexistent/path/index.json");
        assert!(result.is_err());
    }

    // ========================================================================
    // Memory-Mapped File Tests
    // ========================================================================

    #[test]
    #[cfg(feature = "mmap")]
    fn test_mmap_index_creation() {
        let temp_dir = TempDir::new().unwrap();
        let index_path = temp_dir.path().join("mmap_index.json");

        // Create and save an index
        let embeddings = create_test_embeddings();
        let mut index = HnswIndex::new(HnswConfig::default());
        index.build(&embeddings).unwrap();
        save_index(&index, &index_path).unwrap();

        // Memory-map the index
        let mapped = MappedIndex::new(&index_path).unwrap();
        assert!(!mapped.as_bytes().is_empty());

        // Deserialize from memory-mapped data
        let loaded_index: HnswIndex = mapped.deserialize().unwrap();
        let query = vec![0.2, 0.3, 0.4];
        let results = loaded_index.search(&query, 2).unwrap();
        assert_eq!(results.len(), 2);
    }

    #[test]
    #[cfg(feature = "mmap")]
    fn test_mmap_nonexistent_file() {
        let result = MappedIndex::new("/nonexistent/file.json");
        assert!(result.is_err());
    }

    #[test]
    #[cfg(feature = "mmap")]
    fn test_mmap_large_index() {
        let temp_dir = TempDir::new().unwrap();
        let index_path = temp_dir.path().join("mmap_large_index.json");

        // Create a larger index
        let mut embeddings = HashMap::new();
        for i in 0..1000 {
            embeddings.insert(
                format!("doc{}", i),
                vec![
                    i as f32 * 0.001,
                    (i + 1) as f32 * 0.001,
                    (i + 2) as f32 * 0.001,
                ],
            );
        }

        let mut index = HnswIndex::new(HnswConfig::default());
        index.build(&embeddings).unwrap();
        save_index(&index, &index_path).unwrap();

        // Memory-map and verify
        let mapped = MappedIndex::new(&index_path).unwrap();
        let loaded_index: HnswIndex = mapped.deserialize().unwrap();

        let query = vec![0.5, 0.6, 0.7];
        let results = loaded_index.search(&query, 10).unwrap();
        assert_eq!(results.len(), 10);
    }
}