uqa-storage 0.5.2

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Runtime vector operations share the provider's original catalog and transaction owner.

use super::{build::DiskANNTemporaryBudget, pages::DiskANNReadLimits, DiskANNIndexOptions};
use crate::{
    read_control::StorageReadControl, vector_index::DiskANNIndexParams, StorageBackendError,
    StorageBackendResult, VectorIndex,
};
use std::sync::Arc;
use uqa_core::{memory::MemoryReservation, DocId, PostingList};

/// Actual SQL catalog address and the caller's retained allowance. Providers resolve immutable ownership from the stored definition through the supplied SQL owner.
pub struct DiskANNIndexBinding<'a> {
    pub table: &'a str,
    pub field: &'a str,
    pub dimensions: u32,
    pub index: &'a crate::RelationIdentity,
    pub resolver: Arc<dyn super::catalog::DiskANNIndexResolver + Send + Sync>,
    pub control: &'a StorageReadControl,
}

/// Provider boundary for a bound mutable index. Structural operations preserve the active caller transaction and restore their own private effects on failure.
pub trait DiskANNPersistentOwner: Send + Sync {
    type Snapshot: VectorIndex + 'static;
    fn dimensions(&self) -> u32;
    fn parameters(&self) -> DiskANNIndexParams;
    fn read_limits(&self) -> DiskANNReadLimits;
    fn control(&self) -> &StorageReadControl;
    /// Read the selected manifest and canonical identity on one provider view without preparing graph, PQ or origin readers.
    fn query_metadata(&self) -> StorageBackendResult<super::DiskANNQueryMetadata>;
    fn snapshot(&self) -> StorageBackendResult<Self::Snapshot>;
    fn count(&self) -> StorageBackendResult<usize>;
    fn contains_document(&self, document: DocId) -> StorageBackendResult<bool>;
    fn replace(&self, document: DocId, vectors: &[Vec<f32>]) -> StorageBackendResult<()>;
    fn rebuild(
        &self,
        options: DiskANNIndexOptions,
        temporary: &DiskANNTemporaryBudget,
        clear: bool,
    ) -> StorageBackendResult<()>;
}

/// A mutable persistent `VectorIndex`: ordinary writes preserve the base; explicit initialize/clear build and publish replacement generations through the same bound provider session.
pub struct PersistentDiskANNIndex<P> {
    owner: P,
    options: DiskANNIndexOptions,
    temporary: DiskANNTemporaryBudget,
    _memory: MemoryReservation,
}

impl<P: DiskANNPersistentOwner> PersistentDiskANNIndex<P> {
    pub fn new(
        owner: P,
        options: DiskANNIndexOptions,
        temporary: &DiskANNTemporaryBudget,
    ) -> StorageBackendResult<Self> {
        owner.control().check()?;
        if options.parameters != owner.parameters() || options.read != owner.read_limits() {
            return Err(StorageBackendError::Other(
                "DiskANN runtime options differ from its bound catalog or read limits".into(),
            ));
        }
        options.parameters.validate(owner.dimensions())?;
        let memory = owner.control().memory().reserve(size_of::<Self>())?;
        // Restoration validates actual resident data under the original allowance before exposing a live runtime handle. It never rebuilds missing or corrupt state.
        drop(owner.snapshot()?);
        Ok(Self {
            owner,
            options,
            temporary: temporary.clone(),
            _memory: memory,
        })
    }
}

impl<P: DiskANNPersistentOwner> VectorIndex for PersistentDiskANNIndex<P> {
    fn dimensions(&self) -> u32 {
        self.owner.dimensions()
    }
    fn index_kind(&self) -> &'static str {
        "diskann"
    }
    fn diskann_query_metadata(
        &self,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<super::DiskANNQueryMetadata>> {
        control.check()?;
        let metadata = self.owner.query_metadata()?;
        control.check()?;
        Ok(Some(metadata))
    }
    fn add(&mut self, document: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
        self.owner.replace(document, std::slice::from_ref(&vector))
    }
    fn add_many(&mut self, document: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
        self.owner.replace(document, &vectors)
    }
    fn delete(&mut self, document: DocId) -> StorageBackendResult<()> {
        self.owner.replace(document, &[])
    }
    fn clear(&mut self) -> StorageBackendResult<()> {
        self.owner.rebuild(self.options, &self.temporary, true)
    }
    fn initialize(&mut self) -> StorageBackendResult<()> {
        self.owner.rebuild(self.options, &self.temporary, false)
    }
    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
        self.owner.snapshot()?.search_knn(query, k)
    }
    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
        self.owner.snapshot()?.search_threshold(query, threshold)
    }
    fn search_knn_with_control(
        &self,
        query: &[f32],
        k: usize,
        control: &StorageReadControl,
    ) -> StorageBackendResult<PostingList> {
        control.check()?;
        self.owner
            .snapshot()?
            .search_knn_with_control(query, k, control)
    }
    fn search_threshold_with_control(
        &self,
        query: &[f32],
        threshold: f32,
        control: &StorageReadControl,
    ) -> StorageBackendResult<PostingList> {
        control.check()?;
        self.owner
            .snapshot()?
            .search_threshold_with_control(query, threshold, control)
    }
    fn search_knn_with_statistics(
        &self,
        query: &[f32],
        k: usize,
        control: Option<&StorageReadControl>,
    ) -> StorageBackendResult<crate::vector_index::VectorQueryResult> {
        if let Some(control) = control {
            control.check()?;
        }
        self.owner
            .snapshot()?
            .search_knn_with_statistics(query, k, control)
    }
    fn search_threshold_with_statistics(
        &self,
        query: &[f32],
        threshold: f32,
        control: Option<&StorageReadControl>,
    ) -> StorageBackendResult<crate::vector_index::VectorQueryResult> {
        if let Some(control) = control {
            control.check()?;
        }
        self.owner
            .snapshot()?
            .search_threshold_with_statistics(query, threshold, control)
    }
    fn count(&self) -> StorageBackendResult<usize> {
        self.owner.count()
    }
    fn contains_document(&self, document: DocId) -> StorageBackendResult<bool> {
        self.owner.contains_document(document)
    }
    fn diskann_read_snapshot(
        &self,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<super::DiskANNReadSnapshot>> {
        control.check()?;
        self.owner.snapshot()?.diskann_read_snapshot(control)
    }
    fn snapshot_with_diskann_changes(
        &self,
        changes: &super::DiskANNReadChanges,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
        control.check()?;
        self.owner
            .snapshot()?
            .snapshot_with_diskann_changes(changes, control)
    }
    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
        self.owner.snapshot()?.snapshot()
    }
    fn snapshot_with_vector_read(
        &self,
        source: crate::vector_index::VectorReadSnapshot,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
        control.check()?;
        self.owner
            .snapshot()?
            .snapshot_with_vector_read(source, control)
    }
    fn snapshot_with_control(
        &self,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Arc<dyn VectorIndex>> {
        control.check()?;
        self.owner.snapshot()?.snapshot_with_control(control)
    }
}