engine/

persistence.rs

//! Index persistence helpers for saving/loading indexes to storage
//!
//! This module provides utilities for persisting trained indexes to object storage
//! and loading them back, enabling:
//! - Cold start without retraining
//! - Index sharing across instances
//! - Checkpoint/restore workflows

use serde::{de::DeserializeOwned, Deserialize, Serialize};

use crate::hnsw::{HnswConfig, HnswIndex};
use crate::ivf::{IndexedVector, IvfConfig, IvfIndex};
use crate::pq::ProductQuantizer;
use crate::spfresh::{Cluster, SpFreshConfig, SpFreshIndex};
use common::{Vector, VectorId};
use std::collections::{HashMap, HashSet};

// Re-export for external use
pub use storage::IndexType;

/// Trait for indexes that can be persisted
pub trait Persistable: Sized {
    /// The snapshot type for this index
    type Snapshot: Serialize + DeserializeOwned;

    /// Create a snapshot of the current state
    fn to_snapshot(&self) -> Self::Snapshot;

    /// Restore from a snapshot
    fn from_snapshot(snapshot: Self::Snapshot) -> Result<Self, String>;

    /// Serialize to bytes (JSON format)
    fn to_bytes(&self) -> Result<Vec<u8>, String> {
        let snapshot = self.to_snapshot();
        serde_json::to_vec(&snapshot).map_err(|e| format!("Failed to serialize index: {}", e))
    }

    /// Deserialize from bytes (JSON format)
    fn from_bytes(data: &[u8]) -> Result<Self, String> {
        let snapshot: Self::Snapshot = serde_json::from_slice(data)
            .map_err(|e| format!("Failed to deserialize index: {}", e))?;
        Self::from_snapshot(snapshot)
    }
}

/// Serializable snapshot of a ProductQuantizer (the trained codebooks)
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct PQSnapshot {
    /// The trained quantizer with codebooks
    pub quantizer: ProductQuantizer,
}

impl Persistable for ProductQuantizer {
    type Snapshot = PQSnapshot;

    fn to_snapshot(&self) -> PQSnapshot {
        PQSnapshot {
            quantizer: self.clone(),
        }
    }

    fn from_snapshot(snapshot: PQSnapshot) -> Result<Self, String> {
        Ok(snapshot.quantizer)
    }
}

/// Serializable snapshot for IVF index training data
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct IvfTrainingSnapshot {
    pub config: IvfConfig,
    pub dimension: usize,
    pub centroids: Vec<Vec<f32>>,
}

/// Serializable snapshot for SPFresh index training data
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SpFreshTrainingSnapshot {
    pub config: SpFreshConfig,
    pub dimension: usize,
    pub centroids: Vec<Vec<f32>>,
}

/// Serializable snapshot for HNSW configuration only (lightweight)
/// Note: For full graph persistence, use HnswFullSnapshot instead
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct HnswConfigSnapshot {
    pub config: HnswConfig,
    pub dimension: usize,
}

/// Serializable representation of an HNSW node
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SerializableHnswNode {
    /// The vector ID
    pub id: String,
    /// The vector data
    pub vector: Vec<f32>,
    /// Connections at each layer (layer -> neighbor indices)
    pub connections: Vec<Vec<usize>>,
    /// Maximum layer this node exists in
    pub max_layer: usize,
}

/// Full HNSW graph snapshot for complete persistence
/// This allows loading the index without rebuilding the graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HnswFullSnapshot {
    /// Index configuration
    pub config: HnswConfig,
    /// Vector dimension
    pub dimension: usize,
    /// All nodes in the graph
    pub nodes: Vec<SerializableHnswNode>,
    /// Entry point node index
    pub entry_point: Option<usize>,
    /// Maximum level in the graph
    pub max_level: usize,
}

impl Persistable for HnswIndex {
    type Snapshot = HnswFullSnapshot;

    fn to_snapshot(&self) -> HnswFullSnapshot {
        let node_snapshots = self.nodes_read();
        let serializable_nodes: Vec<SerializableHnswNode> = node_snapshots
            .into_iter()
            .map(|node| SerializableHnswNode {
                id: node.id,
                vector: node.vector,
                connections: node.connections,
                max_layer: node.max_layer,
            })
            .collect();

        HnswFullSnapshot {
            config: self.config().clone(),
            dimension: self.dimension().unwrap_or(0),
            nodes: serializable_nodes,
            entry_point: self.entry_point(),
            max_level: self.max_level(),
        }
    }

    fn from_snapshot(snapshot: HnswFullSnapshot) -> Result<Self, String> {
        HnswIndex::from_snapshot(snapshot)
    }
}

/// Serializable representation of an indexed vector in IVF
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SerializableIndexedVector {
    /// The vector ID
    pub id: String,
    /// The vector data
    pub values: Vec<f32>,
}

/// Full IVF index snapshot for complete persistence
/// This allows loading the trained index without retraining
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IvfFullSnapshot {
    /// Index configuration
    pub config: IvfConfig,
    /// Vector dimension
    pub dimension: Option<usize>,
    /// Trained cluster centroids
    pub centroids: Vec<Vec<f32>>,
    /// Inverted lists: cluster_id -> vectors in that cluster
    pub inverted_lists: HashMap<usize, Vec<IndexedVector>>,
    /// Total vector count
    pub vector_count: usize,
    /// Whether the index has been trained
    pub is_trained: bool,
}

impl Persistable for IvfIndex {
    type Snapshot = IvfFullSnapshot;

    fn to_snapshot(&self) -> IvfFullSnapshot {
        IvfFullSnapshot {
            config: self.config().clone(),
            dimension: self.dimension(),
            centroids: self.centroids_read(),
            inverted_lists: self.inverted_lists_read(),
            vector_count: self.len(),
            is_trained: self.is_trained(),
        }
    }

    fn from_snapshot(snapshot: IvfFullSnapshot) -> Result<Self, String> {
        IvfIndex::from_snapshot(snapshot)
    }
}

/// Full SPFresh index snapshot for complete persistence
/// This allows loading the trained index with all clusters and vectors
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpFreshFullSnapshot {
    /// Index configuration
    pub config: SpFreshConfig,
    /// All clusters with their vectors and tombstones
    pub clusters: Vec<Cluster>,
    /// Vector ID to cluster ID mapping
    pub vector_cluster_map: HashMap<VectorId, usize>,
    /// Global tombstones for vectors not yet assigned to clusters
    pub global_tombstones: HashSet<VectorId>,
    /// Pending vectors not yet assigned to clusters
    pub pending_vectors: Vec<Vector>,
    /// Whether the index has been trained
    pub trained: bool,
    /// Vector dimension
    pub dimension: Option<usize>,
}

impl Persistable for SpFreshIndex {
    type Snapshot = SpFreshFullSnapshot;

    fn to_snapshot(&self) -> SpFreshFullSnapshot {
        SpFreshFullSnapshot {
            config: self.config().clone(),
            clusters: self.clusters_read(),
            vector_cluster_map: self.vector_cluster_map_read(),
            global_tombstones: self.global_tombstones_read(),
            pending_vectors: self.pending_vectors_read(),
            trained: self.is_trained(),
            dimension: self.dimension(),
        }
    }

    fn from_snapshot(snapshot: SpFreshFullSnapshot) -> Result<Self, String> {
        SpFreshIndex::from_snapshot(snapshot)
    }
}

/// Index persistence manager for saving/loading indexes to object storage
pub struct IndexPersistenceManager<S> {
    storage: S,
}

impl<S> IndexPersistenceManager<S> {
    /// Create a new persistence manager with the given storage backend
    pub fn new(storage: S) -> Self {
        Self { storage }
    }
}

impl<S: storage::IndexStorage> IndexPersistenceManager<S> {
    /// Save an HNSW index to storage
    pub async fn save_hnsw(
        &self,
        namespace: &common::NamespaceId,
        index: &HnswIndex,
    ) -> common::Result<()> {
        let bytes = index.to_bytes().map_err(common::DakeraError::Storage)?;
        self.storage
            .save_index(namespace, storage::IndexType::Hnsw, bytes)
            .await
    }

    /// Load an HNSW index from storage
    pub async fn load_hnsw(
        &self,
        namespace: &common::NamespaceId,
    ) -> common::Result<Option<HnswIndex>> {
        match self
            .storage
            .load_index(namespace, storage::IndexType::Hnsw)
            .await?
        {
            Some(bytes) => {
                let index = HnswIndex::from_bytes(&bytes).map_err(common::DakeraError::Storage)?;
                Ok(Some(index))
            }
            None => Ok(None),
        }
    }

    /// Save a ProductQuantizer to storage
    pub async fn save_pq(
        &self,
        namespace: &common::NamespaceId,
        quantizer: &ProductQuantizer,
    ) -> common::Result<()> {
        let bytes = quantizer.to_bytes().map_err(common::DakeraError::Storage)?;
        self.storage
            .save_index(namespace, storage::IndexType::Pq, bytes)
            .await
    }

    /// Load a ProductQuantizer from storage
    pub async fn load_pq(
        &self,
        namespace: &common::NamespaceId,
    ) -> common::Result<Option<ProductQuantizer>> {
        match self
            .storage
            .load_index(namespace, storage::IndexType::Pq)
            .await?
        {
            Some(bytes) => {
                let pq =
                    ProductQuantizer::from_bytes(&bytes).map_err(common::DakeraError::Storage)?;
                Ok(Some(pq))
            }
            None => Ok(None),
        }
    }

    /// Save an IVF index to storage
    pub async fn save_ivf(
        &self,
        namespace: &common::NamespaceId,
        index: &IvfIndex,
    ) -> common::Result<()> {
        let bytes = index.to_bytes().map_err(common::DakeraError::Storage)?;
        self.storage
            .save_index(namespace, storage::IndexType::Ivf, bytes)
            .await
    }

    /// Load an IVF index from storage
    pub async fn load_ivf(
        &self,
        namespace: &common::NamespaceId,
    ) -> common::Result<Option<IvfIndex>> {
        match self
            .storage
            .load_index(namespace, storage::IndexType::Ivf)
            .await?
        {
            Some(bytes) => {
                let index = IvfIndex::from_bytes(&bytes).map_err(common::DakeraError::Storage)?;
                Ok(Some(index))
            }
            None => Ok(None),
        }
    }

    /// Save an SPFresh index to storage
    pub async fn save_spfresh(
        &self,
        namespace: &common::NamespaceId,
        index: &SpFreshIndex,
    ) -> common::Result<()> {
        let bytes = index.to_bytes().map_err(common::DakeraError::Storage)?;
        self.storage
            .save_index(namespace, storage::IndexType::SpFresh, bytes)
            .await
    }

    /// Load an SPFresh index from storage
    pub async fn load_spfresh(
        &self,
        namespace: &common::NamespaceId,
    ) -> common::Result<Option<SpFreshIndex>> {
        match self
            .storage
            .load_index(namespace, storage::IndexType::SpFresh)
            .await?
        {
            Some(bytes) => {
                let index =
                    SpFreshIndex::from_bytes(&bytes).map_err(common::DakeraError::Storage)?;
                Ok(Some(index))
            }
            None => Ok(None),
        }
    }

    /// Check if an index exists in storage
    pub async fn index_exists(
        &self,
        namespace: &common::NamespaceId,
        index_type: storage::IndexType,
    ) -> common::Result<bool> {
        self.storage.index_exists(namespace, index_type).await
    }

    /// Delete an index from storage
    pub async fn delete_index(
        &self,
        namespace: &common::NamespaceId,
        index_type: storage::IndexType,
    ) -> common::Result<bool> {
        self.storage.delete_index(namespace, index_type).await
    }

    /// List all indexes for a namespace
    pub async fn list_indexes(
        &self,
        namespace: &common::NamespaceId,
    ) -> common::Result<Vec<storage::IndexType>> {
        self.storage.list_indexes(namespace).await
    }
}

/// Helper to serialize any Serialize type to bytes
pub fn serialize_to_bytes<T: Serialize>(value: &T) -> Result<Vec<u8>, String> {
    serde_json::to_vec(value).map_err(|e| format!("Serialization failed: {}", e))
}

/// Helper to deserialize from bytes to any DeserializeOwned type
pub fn deserialize_from_bytes<T: DeserializeOwned>(data: &[u8]) -> Result<T, String> {
    serde_json::from_slice(data).map_err(|e| format!("Deserialization failed: {}", e))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::pq::PQConfig;
    use common::DistanceMetric;

    #[test]
    fn test_pq_quantizer_persistence() {
        use common::Vector;

        let config = PQConfig {
            num_subquantizers: 4,
            num_centroids: 16,
            kmeans_iterations: 10,
            distance_metric: DistanceMetric::Euclidean,
        };

        let mut pq = ProductQuantizer::new(config, 32).unwrap();

        // Create training vectors
        let vectors: Vec<Vector> = (0..100)
            .map(|i| Vector {
                id: format!("v{}", i),
                values: (0..32).map(|j| ((i + j) as f32 * 0.1).sin()).collect(),
                metadata: None,
                ttl_seconds: None,
                expires_at: None,
            })
            .collect();

        // Train the quantizer
        pq.train(&vectors).unwrap();
        assert!(pq.is_trained());

        // Serialize
        let bytes = pq.to_bytes().unwrap();
        assert!(!bytes.is_empty());

        // Deserialize
        let restored = ProductQuantizer::from_bytes(&bytes).unwrap();

        // Verify
        assert!(restored.is_trained());
        assert_eq!(restored.dimension, 32);
        assert_eq!(restored.config.num_subquantizers, 4);
        assert_eq!(restored.codebooks.len(), 4);
    }

    #[test]
    fn test_ivf_training_snapshot() {
        let snapshot = IvfTrainingSnapshot {
            config: IvfConfig {
                n_clusters: 8,
                n_probe: 2,
                metric: DistanceMetric::Euclidean,
                ..Default::default()
            },
            dimension: 64,
            centroids: vec![vec![0.0; 64]; 8],
        };

        // Serialize and deserialize
        let bytes = serialize_to_bytes(&snapshot).unwrap();
        let restored: IvfTrainingSnapshot = deserialize_from_bytes(&bytes).unwrap();

        assert_eq!(restored.config.n_clusters, 8);
        assert_eq!(restored.dimension, 64);
        assert_eq!(restored.centroids.len(), 8);
    }

    #[test]
    fn test_hnsw_config_snapshot() {
        let snapshot = HnswConfigSnapshot {
            config: HnswConfig::default(),
            dimension: 128,
        };

        // Serialize and deserialize
        let bytes = serialize_to_bytes(&snapshot).unwrap();
        let restored: HnswConfigSnapshot = deserialize_from_bytes(&bytes).unwrap();

        assert_eq!(restored.config.m, 16);
        assert_eq!(restored.dimension, 128);
    }

    #[test]
    fn test_spfresh_training_snapshot() {
        let snapshot = SpFreshTrainingSnapshot {
            config: SpFreshConfig::default(),
            dimension: 32,
            centroids: vec![vec![1.0; 32]; 16],
        };

        let bytes = serialize_to_bytes(&snapshot).unwrap();
        let restored: SpFreshTrainingSnapshot = deserialize_from_bytes(&bytes).unwrap();

        assert_eq!(restored.dimension, 32);
        assert_eq!(restored.centroids.len(), 16);
    }

    #[test]
    fn test_hnsw_full_persistence() {
        use crate::hnsw::HnswIndex;

        // Create and populate an HNSW index
        let index = HnswIndex::new();

        // Insert some vectors
        for i in 0..50 {
            let vector: Vec<f32> = (0..64).map(|j| ((i + j) as f32 * 0.1).sin()).collect();
            index.insert(format!("vec_{}", i), vector);
        }

        assert_eq!(index.len(), 50);

        // Serialize to bytes
        let bytes = index.to_bytes().unwrap();
        assert!(!bytes.is_empty());

        // Deserialize
        let restored = HnswIndex::from_bytes(&bytes).unwrap();

        // Verify structure
        assert_eq!(restored.len(), 50);
        assert_eq!(restored.dimension(), index.dimension());
        assert_eq!(restored.max_level(), index.max_level());

        // Verify search still works
        let query: Vec<f32> = (0..64).map(|j| (j as f32 * 0.1).sin()).collect();
        let original_results = index.search(&query, 5);
        let restored_results = restored.search(&query, 5);

        // Results should be identical
        assert_eq!(original_results.len(), restored_results.len());
        for (orig, rest) in original_results.iter().zip(restored_results.iter()) {
            assert_eq!(orig.0, rest.0); // Same IDs
            assert!((orig.1 - rest.1).abs() < 1e-6); // Same distances
        }
    }

    #[test]
    fn test_hnsw_empty_persistence() {
        use crate::hnsw::HnswIndex;

        let index = HnswIndex::new();

        // Serialize empty index
        let bytes = index.to_bytes().unwrap();

        // Deserialize
        let restored = HnswIndex::from_bytes(&bytes).unwrap();

        assert_eq!(restored.len(), 0);
        assert!(restored.is_empty());
    }

    #[test]
    fn test_ivf_full_persistence() {
        use crate::ivf::{IvfConfig, IvfIndex};

        // Create and train an IVF index
        let training_vectors: Vec<Vec<f32>> = (0..100)
            .map(|i| (0..32).map(|j| ((i + j) as f32 * 0.1).sin()).collect())
            .collect();

        let mut index = IvfIndex::new(IvfConfig {
            n_clusters: 10,
            n_probe: 3,
            ..Default::default()
        });

        index.train(&training_vectors).unwrap();
        assert!(index.is_trained());

        // Add vectors to the index
        for (i, v) in training_vectors.iter().enumerate() {
            index.add(format!("vec_{}", i), v.clone()).unwrap();
        }

        assert_eq!(index.len(), 100);

        // Serialize to bytes
        let bytes = index.to_bytes().unwrap();
        assert!(!bytes.is_empty());

        // Deserialize
        let restored = IvfIndex::from_bytes(&bytes).unwrap();

        // Verify structure
        assert_eq!(restored.len(), 100);
        assert!(restored.is_trained());
        assert_eq!(restored.n_clusters(), 10);

        // Verify search still works
        let query = &training_vectors[0];
        let original_results = index.search(query, 5).unwrap();
        let restored_results = restored.search(query, 5).unwrap();

        // First result should be the same (the query vector itself)
        assert_eq!(original_results[0].id, restored_results[0].id);
        assert_eq!(original_results[0].id, "vec_0");
    }

    #[test]
    fn test_ivf_empty_persistence() {
        use crate::ivf::{IvfConfig, IvfIndex};

        // Create untrained index
        let index = IvfIndex::new(IvfConfig::default());

        // Serialize
        let bytes = index.to_bytes().unwrap();

        // Deserialize
        let restored = IvfIndex::from_bytes(&bytes).unwrap();

        assert_eq!(restored.len(), 0);
        assert!(!restored.is_trained());
    }

    #[test]
    fn test_spfresh_full_persistence() {
        use crate::spfresh::{SpFreshConfig, SpFreshIndex};
        use common::Vector;

        // Create training vectors
        let training_vectors: Vec<Vector> = (0..100)
            .map(|i| Vector {
                id: format!("vec_{}", i),
                values: (0..32).map(|j| ((i + j) as f32 * 0.1).sin()).collect(),
                metadata: None,
                ttl_seconds: None,
                expires_at: None,
            })
            .collect();

        // Create and train an SPFresh index
        let index = SpFreshIndex::new(SpFreshConfig {
            num_clusters: 4,
            n_probe: 2,
            ..Default::default()
        });

        index.train(&training_vectors).unwrap();
        assert!(index.is_trained());

        let stats = index.stats();
        assert_eq!(stats.total_vectors, 100);
        assert_eq!(stats.num_clusters, 4);

        // Serialize to bytes
        let bytes = index.to_bytes().unwrap();
        assert!(!bytes.is_empty());

        // Deserialize
        let restored = SpFreshIndex::from_bytes(&bytes).unwrap();

        // Verify structure
        let restored_stats = restored.stats();
        assert!(restored.is_trained());
        assert_eq!(restored_stats.total_vectors, 100);
        assert_eq!(restored_stats.num_clusters, 4);
        assert_eq!(restored_stats.dimension, Some(32));

        // Verify search still works
        let query = &training_vectors[50].values;
        let original_results = index.search(query, 10).unwrap();
        let restored_results = restored.search(query, 10).unwrap();

        // Results should have the same length
        assert_eq!(original_results.len(), restored_results.len());

        // Top results should be consistent between original and restored
        // (SPFresh is cluster-based, so we check that results overlap significantly)
        let original_ids: std::collections::HashSet<_> =
            original_results.iter().map(|r| &r.id).collect();
        let restored_ids: std::collections::HashSet<_> =
            restored_results.iter().map(|r| &r.id).collect();
        let overlap = original_ids.intersection(&restored_ids).count();
        assert!(
            overlap >= 8,
            "Expected at least 80% overlap in top-10 results, got {}/10",
            overlap
        );
    }

    #[test]
    fn test_spfresh_empty_persistence() {
        use crate::spfresh::{SpFreshConfig, SpFreshIndex};

        // Create untrained index
        let index = SpFreshIndex::new(SpFreshConfig::default());

        // Serialize
        let bytes = index.to_bytes().unwrap();

        // Deserialize
        let restored = SpFreshIndex::from_bytes(&bytes).unwrap();

        let stats = restored.stats();
        assert_eq!(stats.total_vectors, 0);
        assert!(!restored.is_trained());
    }

    #[test]
    fn test_spfresh_persistence_with_tombstones() {
        use crate::spfresh::{SpFreshConfig, SpFreshIndex};
        use common::Vector;

        // Create training vectors
        let training_vectors: Vec<Vector> = (0..50)
            .map(|i| Vector {
                id: format!("vec_{}", i),
                values: (0..16).map(|j| ((i + j) as f32 * 0.1).cos()).collect(),
                metadata: None,
                ttl_seconds: None,
                expires_at: None,
            })
            .collect();

        // Create and train index
        let index = SpFreshIndex::new(SpFreshConfig {
            num_clusters: 2,
            ..Default::default()
        });

        index.train(&training_vectors).unwrap();

        // Remove some vectors (creates tombstones)
        let ids_to_remove: Vec<String> = (0..10).map(|i| format!("vec_{}", i)).collect();
        let removed = index.remove(&ids_to_remove);
        assert_eq!(removed, 10);

        let stats = index.stats();
        assert_eq!(stats.total_vectors, 40);
        assert_eq!(stats.total_tombstones, 10);

        // Serialize
        let bytes = index.to_bytes().unwrap();

        // Deserialize
        let restored = SpFreshIndex::from_bytes(&bytes).unwrap();

        // Verify tombstones are preserved
        let restored_stats = restored.stats();
        assert_eq!(restored_stats.total_vectors, 40);
        assert_eq!(restored_stats.total_tombstones, 10);

        // Verify removed vectors don't appear in search results
        let results = restored.search(&training_vectors[0].values, 50).unwrap();
        for result in &results {
            // vec_0 through vec_9 should not appear
            let id_num: usize = result.id.strip_prefix("vec_").unwrap().parse().unwrap();
            assert!(
                id_num >= 10,
                "Tombstoned vector {} appeared in results",
                result.id
            );
        }
    }
}