sql_rs/vector/

store.rs

use super::{DistanceMetric, VectorCollection};
use crate::{Result, SqlRsError, Storage};
use std::collections::HashMap;
use std::time::{Duration, SystemTime};

#[derive(Debug, Clone)]
struct CacheEntry {
    collection: VectorCollection,
    last_accessed: SystemTime,
    dirty: bool,
}

/// Persistent vector store that saves collections to disk
pub struct VectorStore<S: Storage> {
    pub storage: S,
    collections_cache: HashMap<String, CacheEntry>,
    max_cache_entries: usize,
    cache_ttl: Duration,
}

impl<S: Storage> VectorStore<S> {
    pub fn new(storage: S) -> Self {
        Self {
            storage,
            collections_cache: HashMap::new(),
            max_cache_entries: 10,
            cache_ttl: Duration::from_secs(300), // 5 minutes
        }
    }

    /// Create a new VectorStore with custom cache settings
    pub fn with_cache_config(storage: S, max_entries: usize, ttl_secs: u64) -> Self {
        Self {
            storage,
            collections_cache: HashMap::new(),
            max_cache_entries: max_entries,
            cache_ttl: Duration::from_secs(ttl_secs),
        }
    }

    /// Create a new vector collection
    pub fn create_collection(
        &mut self,
        name: &str,
        dimension: usize,
        metric: DistanceMetric,
    ) -> Result<()> {
        let key = format!("__vector_collection__{}", name);

        // Check if collection already exists
        if self.storage.get(key.as_bytes())?.is_some() {
            return Err(SqlRsError::Vector(format!(
                "Collection '{}' already exists",
                name
            )));
        }

        let collection = VectorCollection::new(name, dimension, metric);
        let serialized = bincode::serialize(&collection)
            .map_err(|e| SqlRsError::Serialization(e.to_string()))?;

        self.storage.put(key.as_bytes(), &serialized)?;

        // Add to cache
        self.collections_cache.insert(
            name.to_string(),
            CacheEntry {
                collection,
                last_accessed: SystemTime::now(),
                dirty: false,
            },
        );

        Ok(())
    }

    /// Get a vector collection, loading from disk if not cached
    pub fn get_collection(&mut self, name: &str) -> Result<VectorCollection> {
        // Check cache first
        if let Some(entry) = self.collections_cache.get_mut(name) {
            entry.last_accessed = SystemTime::now();
            return Ok(entry.collection.clone());
        }

        // Load from storage
        let key = format!("__vector_collection__{}", name);
        let bytes = self
            .storage
            .get(key.as_bytes())?
            .ok_or_else(|| SqlRsError::Vector(format!("Collection '{}' not found", name)))?;

        let collection: VectorCollection =
            bincode::deserialize(&bytes).map_err(|e| SqlRsError::Serialization(e.to_string()))?;

        // Check cache size and evict if necessary
        self.evict_if_needed()?;

        // Add to cache
        self.collections_cache.insert(
            name.to_string(),
            CacheEntry {
                collection: collection.clone(),
                last_accessed: SystemTime::now(),
                dirty: false,
            },
        );

        Ok(collection)
    }

    /// Save a collection to disk and cache
    pub fn save_collection(&mut self, collection: VectorCollection) -> Result<()> {
        let name = collection.name();
        let key = format!("__vector_collection__{}", name);

        let serialized = bincode::serialize(&collection)
            .map_err(|e| SqlRsError::Serialization(e.to_string()))?;

        self.storage.put(key.as_bytes(), &serialized)?;

        // Update cache
        if let Some(entry) = self.collections_cache.get_mut(name) {
            entry.collection = collection;
            entry.last_accessed = SystemTime::now();
            entry.dirty = false;
        } else {
            // Add to cache if not present
            self.evict_if_needed()?;
            self.collections_cache.insert(
                name.to_string(),
                CacheEntry {
                    collection,
                    last_accessed: SystemTime::now(),
                    dirty: false,
                },
            );
        }

        Ok(())
    }

    /// Evict old or dirty entries from cache
    fn evict_if_needed(&mut self) -> Result<()> {
        // First, save any dirty entries
        let dirty_entries: Vec<_> = self
            .collections_cache
            .iter()
            .filter(|(_, entry)| entry.dirty)
            .map(|(name, _)| name.clone())
            .collect();

        for name in dirty_entries {
            if let Some(entry) = self.collections_cache.remove(&name) {
                let key = format!("__vector_collection__{}", name);
                let serialized = bincode::serialize(&entry.collection)
                    .map_err(|e| SqlRsError::Serialization(e.to_string()))?;
                self.storage.put(key.as_bytes(), &serialized)?;

                // Re-add as clean
                self.collections_cache.insert(
                    name,
                    CacheEntry {
                        collection: entry.collection,
                        last_accessed: SystemTime::now(),
                        dirty: false,
                    },
                );
            }
        }

        // If cache is still full, remove oldest entries
        if self.collections_cache.len() >= self.max_cache_entries {
            let mut entries_by_age: Vec<_> = self
                .collections_cache
                .iter()
                .map(|(name, entry)| (name.clone(), entry.last_accessed))
                .collect();

            entries_by_age.sort_by_key(|(_, time)| *time);

            let to_remove = entries_by_age.len() - self.max_cache_entries + 1;
            for (name, _) in entries_by_age.into_iter().take(to_remove) {
                self.collections_cache.remove(&name);
            }
        }

        Ok(())
    }

    /// Flush all cached collections to disk
    pub fn flush_all(&mut self) -> Result<()> {
        let collections: Vec<_> = self.collections_cache.keys().cloned().collect();

        for name in collections {
            if let Some(entry) = self.collections_cache.remove(&name) {
                let key = format!("__vector_collection__{}", name);
                let serialized = bincode::serialize(&entry.collection)
                    .map_err(|e| SqlRsError::Serialization(e.to_string()))?;
                self.storage.put(key.as_bytes(), &serialized)?;
            }
        }

        Ok(())
    }

    /// Add a vector to a collection
    pub fn add_vector(&mut self, collection_name: &str, vector: super::Vector) -> Result<()> {
        let mut collection = self.get_collection(collection_name)?;
        collection.add(vector)?;
        self.save_collection(collection)?;
        Ok(())
    }

    /// Add multiple vectors to a collection in batch
    pub fn add_vectors_batch(
        &mut self,
        collection_name: &str,
        vectors: Vec<super::Vector>,
    ) -> Result<usize> {
        let mut collection = self.get_collection(collection_name)?;
        let added_count = collection.add_batch(vectors)?;
        self.save_collection(collection)?;
        Ok(added_count)
    }

    /// Search for similar vectors in a collection
    pub fn search(
        &mut self,
        collection_name: &str,
        query: &[f32],
        k: usize,
    ) -> Result<Vec<super::SearchResult>> {
        let collection = self.get_collection(collection_name)?;
        collection.search(query, k)
    }

    /// List all collection names
    pub fn list_collections(&self) -> Result<Vec<String>> {
        let prefix = b"__vector_collection__";
        let scan_results = self.storage.scan(prefix, &[0xFF; 256])?;

        let mut names = Vec::new();
        for (key, _) in scan_results {
            if let Ok(key_str) = std::str::from_utf8(&key) {
                if let Some(name) = key_str.strip_prefix("__vector_collection__") {
                    names.push(name.to_string());
                }
            }
        }

        Ok(names)
    }

    /// Delete a collection
    pub fn delete_collection(&mut self, name: &str) -> Result<()> {
        let key = format!("__vector_collection__{}", name);
        self.storage.delete(key.as_bytes())?;
        self.collections_cache.remove(name);
        Ok(())
    }

    /// Get statistics about all vector collections
    pub fn get_collection_stats(&mut self) -> Result<Vec<(String, usize, usize, DistanceMetric)>> {
        let collection_names = self.list_collections()?;
        let mut stats = Vec::new();

        for name in collection_names {
            let collection = self.get_collection(&name)?;
            stats.push((
                name,
                collection.len(),
                collection.dimension(),
                collection.metric().clone(),
            ));
        }

        Ok(stats)
    }
}