weavatrix_search_vector/mutable/

mutations.rs

use super::support::{
    current_vector, key_exists, prepare_records, prepare_vector, replace_metadata,
};
use super::{MutableVectorIndex, MutationOutcome, VectorRecord};
use crate::error::SearchError;
use crate::metadata::Metadata;

impl MutableVectorIndex {
    /// Inserts a new key.
    ///
    /// # Errors
    ///
    /// Returns [`SearchError::DuplicateKey`] when the live key already exists,
    /// or a typed vector validation error.
    pub fn insert(&self, key: u64, vector: &[f32], metadata: Metadata) -> Result<(), SearchError> {
        let normalized = prepare_vector(self.config.dimensions, self.config.metric, vector)?;
        let mut state = self
            .state
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if key_exists(&state, key) {
            return Err(SearchError::DuplicateKey(key));
        }
        state.deleted.remove(&key);
        state.pending.insert(key, normalized);
        replace_metadata(&mut state.metadata, key, metadata);
        state.generation = state.generation.wrapping_add(1);
        Ok(())
    }

    /// Inserts or replaces one key without rebuilding the immutable base.
    ///
    /// # Errors
    ///
    /// Returns a typed vector validation error.
    pub fn upsert(
        &self,
        key: u64,
        vector: &[f32],
        metadata: Metadata,
    ) -> Result<MutationOutcome, SearchError> {
        let normalized = prepare_vector(self.config.dimensions, self.config.metric, vector)?;
        let mut state = self
            .state
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let outcome = if key_exists(&state, key) {
            MutationOutcome::Updated
        } else {
            MutationOutcome::Inserted
        };
        state.deleted.remove(&key);
        state.pending.insert(key, normalized);
        replace_metadata(&mut state.metadata, key, metadata);
        state.generation = state.generation.wrapping_add(1);
        Ok(outcome)
    }

    /// Atomically inserts a batch after validating every record.
    ///
    /// # Errors
    ///
    /// Returns without applying any record when the batch contains duplicate
    /// keys, an existing live key, or an invalid vector.
    pub fn insert_batch(&self, records: &[VectorRecord]) -> Result<(), SearchError> {
        let prepared = prepare_records(&self.config, records)?;
        let mut state = self
            .state
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        for (key, _, _) in &prepared {
            if key_exists(&state, *key) {
                return Err(SearchError::DuplicateKey(*key));
            }
        }
        for (key, vector, metadata) in prepared {
            state.deleted.remove(&key);
            state.pending.insert(key, vector);
            replace_metadata(&mut state.metadata, key, metadata);
        }
        state.generation = state.generation.wrapping_add(1);
        Ok(())
    }

    /// Atomically inserts or replaces a batch after validating every record.
    ///
    /// Results follow input order.
    ///
    /// # Errors
    ///
    /// Returns without applying any record when the batch contains duplicate
    /// keys or an invalid vector.
    pub fn upsert_batch(
        &self,
        records: &[VectorRecord],
    ) -> Result<Vec<MutationOutcome>, SearchError> {
        let prepared = prepare_records(&self.config, records)?;
        let mut state = self
            .state
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let outcomes = prepared
            .iter()
            .map(|(key, _, _)| {
                if key_exists(&state, *key) {
                    MutationOutcome::Updated
                } else {
                    MutationOutcome::Inserted
                }
            })
            .collect::<Vec<_>>();
        for (key, vector, metadata) in prepared {
            state.deleted.remove(&key);
            state.pending.insert(key, vector);
            replace_metadata(&mut state.metadata, key, metadata);
        }
        state.generation = state.generation.wrapping_add(1);
        Ok(outcomes)
    }

    /// Tombstones a live key. Returns whether a vector was removed.
    pub fn delete(&self, key: u64) -> bool {
        let mut state = self
            .state
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if !key_exists(&state, key) {
            return false;
        }
        state.pending.remove(&key);
        if state.base.vector(key).is_some()
            || state
                .sealed
                .as_ref()
                .is_some_and(|index| index.vector(key).is_some())
        {
            state.deleted.insert(key);
        }
        state.metadata.remove(key);
        state.generation = state.generation.wrapping_add(1);
        true
    }

    /// Tombstones every live key in `keys` and returns the number removed.
    pub fn delete_batch(&self, keys: &[u64]) -> usize {
        let mut state = self
            .state
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let mut removed = 0;
        for key in keys {
            if !key_exists(&state, *key) {
                continue;
            }
            state.pending.remove(key);
            if state.base.vector(*key).is_some()
                || state
                    .sealed
                    .as_ref()
                    .is_some_and(|index| index.vector(*key).is_some())
            {
                state.deleted.insert(*key);
            }
            state.metadata.remove(*key);
            removed += 1;
        }
        if removed != 0 {
            state.generation = state.generation.wrapping_add(1);
        }
        removed
    }

    /// Renames one live vector key without rebuilding either graph.
    ///
    /// # Errors
    ///
    /// Returns [`SearchError::DuplicateKey`] when `to` is already live.
    pub fn rename(&self, from: u64, to: u64) -> Result<bool, SearchError> {
        if from == to {
            let state = self
                .state
                .read()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            return Ok(key_exists(&state, from));
        }
        let mut state = self
            .state
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if !key_exists(&state, from) {
            return Ok(false);
        }
        if key_exists(&state, to) {
            return Err(SearchError::DuplicateKey(to));
        }
        let vector = current_vector(&state, from)
            .ok_or(SearchError::MissingKey(from))?
            .to_vec();
        let metadata = state.metadata.remove(from);
        state.pending.remove(&from);
        if state.base.vector(from).is_some()
            || state
                .sealed
                .as_ref()
                .is_some_and(|index| index.vector(from).is_some())
        {
            state.deleted.insert(from);
        }
        state.deleted.remove(&to);
        state.pending.insert(to, vector);
        if let Some(metadata) = metadata {
            state.metadata.insert(to, metadata);
        }
        state.generation = state.generation.wrapping_add(1);
        Ok(true)
    }

    /// Replaces metadata without changing the vector.
    ///
    /// # Errors
    ///
    /// Returns [`SearchError::MissingKey`] when `key` is not live.
    pub fn set_metadata(&self, key: u64, metadata: Metadata) -> Result<(), SearchError> {
        let mut state = self
            .state
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if !key_exists(&state, key) {
            return Err(SearchError::MissingKey(key));
        }
        replace_metadata(&mut state.metadata, key, metadata);
        state.generation = state.generation.wrapping_add(1);
        Ok(())
    }

    #[must_use]
    pub fn metadata(&self, key: u64) -> Option<Metadata> {
        let state = self
            .state
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        state.metadata.get(key).cloned()
    }
}