uqa_storage/diskann_index/
persistent.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Runtime vector operations share the provider's original catalog and transaction owner.
8
9use super::{build::DiskANNTemporaryBudget, pages::DiskANNReadLimits, DiskANNIndexOptions};
10use crate::{
11    read_control::StorageReadControl, vector_index::DiskANNIndexParams, StorageBackendError,
12    StorageBackendResult, VectorIndex,
13};
14use std::sync::Arc;
15use uqa_core::{memory::MemoryReservation, DocId, PostingList};
16
17/// Actual SQL catalog address and the caller's retained allowance. Providers resolve immutable ownership from the stored definition through the supplied SQL owner.
18pub struct DiskANNIndexBinding<'a> {
19    pub table: &'a str,
20    pub field: &'a str,
21    pub dimensions: u32,
22    pub index: &'a crate::RelationIdentity,
23    pub resolver: Arc<dyn super::catalog::DiskANNIndexResolver + Send + Sync>,
24    pub control: &'a StorageReadControl,
25}
26
27/// Provider boundary for a bound mutable index. Structural operations preserve the active caller transaction and restore their own private effects on failure.
28pub trait DiskANNPersistentOwner: Send + Sync {
29    type Snapshot: VectorIndex + 'static;
30    fn dimensions(&self) -> u32;
31    fn parameters(&self) -> DiskANNIndexParams;
32    fn read_limits(&self) -> DiskANNReadLimits;
33    fn control(&self) -> &StorageReadControl;
34    /// Read the selected manifest and canonical identity on one provider view without preparing graph, PQ or origin readers.
35    fn query_metadata(&self) -> StorageBackendResult<super::DiskANNQueryMetadata>;
36    fn snapshot(&self) -> StorageBackendResult<Self::Snapshot>;
37    fn count(&self) -> StorageBackendResult<usize>;
38    fn contains_document(&self, document: DocId) -> StorageBackendResult<bool>;
39    fn replace(&self, document: DocId, vectors: &[Vec<f32>]) -> StorageBackendResult<()>;
40    fn rebuild(
41        &self,
42        options: DiskANNIndexOptions,
43        temporary: &DiskANNTemporaryBudget,
44        clear: bool,
45    ) -> StorageBackendResult<()>;
46}
47
48/// A mutable persistent `VectorIndex`: ordinary writes preserve the base; explicit initialize/clear build and publish replacement generations through the same bound provider session.
49pub struct PersistentDiskANNIndex<P> {
50    owner: P,
51    options: DiskANNIndexOptions,
52    temporary: DiskANNTemporaryBudget,
53    _memory: MemoryReservation,
54}
55
56impl<P: DiskANNPersistentOwner> PersistentDiskANNIndex<P> {
57    pub fn new(
58        owner: P,
59        options: DiskANNIndexOptions,
60        temporary: &DiskANNTemporaryBudget,
61    ) -> StorageBackendResult<Self> {
62        owner.control().check()?;
63        if options.parameters != owner.parameters() || options.read != owner.read_limits() {
64            return Err(StorageBackendError::Other(
65                "DiskANN runtime options differ from its bound catalog or read limits".into(),
66            ));
67        }
68        options.parameters.validate(owner.dimensions())?;
69        let memory = owner.control().memory().reserve(size_of::<Self>())?;
70        // Restoration validates actual resident data under the original allowance before exposing a live runtime handle. It never rebuilds missing or corrupt state.
71        drop(owner.snapshot()?);
72        Ok(Self {
73            owner,
74            options,
75            temporary: temporary.clone(),
76            _memory: memory,
77        })
78    }
79}
80
81impl<P: DiskANNPersistentOwner> VectorIndex for PersistentDiskANNIndex<P> {
82    fn dimensions(&self) -> u32 {
83        self.owner.dimensions()
84    }
85    fn index_kind(&self) -> &'static str {
86        "diskann"
87    }
88    fn diskann_query_metadata(
89        &self,
90        control: &StorageReadControl,
91    ) -> StorageBackendResult<Option<super::DiskANNQueryMetadata>> {
92        control.check()?;
93        let metadata = self.owner.query_metadata()?;
94        control.check()?;
95        Ok(Some(metadata))
96    }
97    fn add(&mut self, document: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
98        self.owner.replace(document, std::slice::from_ref(&vector))
99    }
100    fn add_many(&mut self, document: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
101        self.owner.replace(document, &vectors)
102    }
103    fn delete(&mut self, document: DocId) -> StorageBackendResult<()> {
104        self.owner.replace(document, &[])
105    }
106    fn clear(&mut self) -> StorageBackendResult<()> {
107        self.owner.rebuild(self.options, &self.temporary, true)
108    }
109    fn initialize(&mut self) -> StorageBackendResult<()> {
110        self.owner.rebuild(self.options, &self.temporary, false)
111    }
112    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
113        self.owner.snapshot()?.search_knn(query, k)
114    }
115    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
116        self.owner.snapshot()?.search_threshold(query, threshold)
117    }
118    fn search_knn_with_control(
119        &self,
120        query: &[f32],
121        k: usize,
122        control: &StorageReadControl,
123    ) -> StorageBackendResult<PostingList> {
124        control.check()?;
125        self.owner
126            .snapshot()?
127            .search_knn_with_control(query, k, control)
128    }
129    fn search_threshold_with_control(
130        &self,
131        query: &[f32],
132        threshold: f32,
133        control: &StorageReadControl,
134    ) -> StorageBackendResult<PostingList> {
135        control.check()?;
136        self.owner
137            .snapshot()?
138            .search_threshold_with_control(query, threshold, control)
139    }
140    fn search_knn_with_statistics(
141        &self,
142        query: &[f32],
143        k: usize,
144        control: Option<&StorageReadControl>,
145    ) -> StorageBackendResult<crate::vector_index::VectorQueryResult> {
146        if let Some(control) = control {
147            control.check()?;
148        }
149        self.owner
150            .snapshot()?
151            .search_knn_with_statistics(query, k, control)
152    }
153    fn search_threshold_with_statistics(
154        &self,
155        query: &[f32],
156        threshold: f32,
157        control: Option<&StorageReadControl>,
158    ) -> StorageBackendResult<crate::vector_index::VectorQueryResult> {
159        if let Some(control) = control {
160            control.check()?;
161        }
162        self.owner
163            .snapshot()?
164            .search_threshold_with_statistics(query, threshold, control)
165    }
166    fn count(&self) -> StorageBackendResult<usize> {
167        self.owner.count()
168    }
169    fn contains_document(&self, document: DocId) -> StorageBackendResult<bool> {
170        self.owner.contains_document(document)
171    }
172    fn diskann_read_snapshot(
173        &self,
174        control: &StorageReadControl,
175    ) -> StorageBackendResult<Option<super::DiskANNReadSnapshot>> {
176        control.check()?;
177        self.owner.snapshot()?.diskann_read_snapshot(control)
178    }
179    fn snapshot_with_diskann_changes(
180        &self,
181        changes: &super::DiskANNReadChanges,
182        control: &StorageReadControl,
183    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
184        control.check()?;
185        self.owner
186            .snapshot()?
187            .snapshot_with_diskann_changes(changes, control)
188    }
189    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
190        self.owner.snapshot()?.snapshot()
191    }
192    fn snapshot_with_vector_read(
193        &self,
194        source: crate::vector_index::VectorReadSnapshot,
195        control: &StorageReadControl,
196    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
197        control.check()?;
198        self.owner
199            .snapshot()?
200            .snapshot_with_vector_read(source, control)
201    }
202    fn snapshot_with_control(
203        &self,
204        control: &StorageReadControl,
205    ) -> StorageBackendResult<Arc<dyn VectorIndex>> {
206        control.check()?;
207        self.owner.snapshot()?.snapshot_with_control(control)
208    }
209}
uqa_storage/diskann_index/persistent.rs

uqa_storage/diskann_index/
persistent.rs
