uqa_storage/diskann_index/
memory.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Mutable memory roots share immutable pages and copy only changed canonical paths.
8
9mod build;
10mod source;
11#[cfg(test)]
12mod tests;
13
14use super::{
15    build::DiskANNTemporaryBudget,
16    format::{DiskANNGeneration, DiskANNManifest, DiskANNVectorVersion},
17    RetainedDiskANNIndex,
18};
19use crate::{
20    mvcc::{DatabaseId, StorageTransactionId},
21    read_control::StorageReadControl,
22    StorageBackendError, StorageBackendResult, VectorIndex,
23};
24use source::Canonical;
25use std::sync::{
26    atomic::{AtomicU64, Ordering},
27    Arc,
28};
29use uqa_core::{
30    memory::{Budgeted, MemoryReservation},
31    DocId, PostingList,
32};
33
34/// Memory construction uses the same host settings as persistent generation construction.
35pub type DiskANNMemoryOptions = super::DiskANNIndexOptions;
36
37struct Clock {
38    database: DatabaseId,
39    next: AtomicU64,
40}
41
42impl Clock {
43    fn allocate(&self) -> StorageBackendResult<u64> {
44        uqa_core::atomic::try_update_u64(&self.next, Ordering::Relaxed, Ordering::Relaxed, |next| {
45            next.checked_add(1)
46        })
47        .map_err(|_| StorageBackendError::Other("DiskANN memory identity exhausted".into()))
48    }
49    fn generation(&self) -> StorageBackendResult<DiskANNGeneration> {
50        DiskANNGeneration::new(self.database.as_bytes(), 1, 1, self.allocate()?)
51    }
52    fn version(&self) -> StorageBackendResult<DiskANNVectorVersion> {
53        // This incarnation owns one mutable memory lineage. Every fork shares its non-reused revision allocator; the identity is not SQL commit order.
54        let writer = StorageTransactionId::new(self.database, 1)
55            .map_err(crate::mvcc::VersionError::into_storage_error)?;
56        DiskANNVectorVersion::new(writer, self.allocate()?)
57    }
58}
59
60/// A writable memory index over the same sealed-page search as persistent providers. Ordinary mutations replace canonical tensors and current changes atomically; explicit initialization and clear construct replacement physical generations. Writable snapshots retain independent roots under the same original allowance.
61pub struct DiskANNMemoryIndex {
62    index: RetainedDiskANNIndex<Canonical>,
63    clock: Arc<Budgeted<Clock>>,
64    options: DiskANNMemoryOptions,
65    temporary: DiskANNTemporaryBudget,
66    control: StorageReadControl,
67    _memory: MemoryReservation,
68}
69
70impl DiskANNMemoryIndex {
71    pub fn new(
72        dimensions: u32,
73        options: DiskANNMemoryOptions,
74        temporary: &DiskANNTemporaryBudget,
75        control: &StorageReadControl,
76    ) -> StorageBackendResult<Self> {
77        control.check()?;
78        options.parameters.validate(dimensions)?;
79        let memory = control.memory().reserve(size_of::<Self>())?;
80        let database = DatabaseId::from_bytes(crate::catalog::new_nonzero_catalog_identity(
81            "DiskANN memory",
82            "incarnation",
83        )?);
84        let clock = Budgeted::new(
85            Clock {
86                database,
87                next: AtomicU64::new(1),
88            },
89            control.memory().empty_reservation(),
90        )
91        .into_shared()?;
92        let index = build::prepare(
93            Canonical::empty(dimensions, control),
94            clock.generation()?,
95            options,
96            temporary,
97            control,
98        )?;
99        Ok(Self {
100            index,
101            clock,
102            options,
103            temporary: temporary.clone(),
104            control: control.clone(),
105            _memory: memory,
106        })
107    }
108
109    pub fn manifest(&self) -> &DiskANNManifest {
110        self.index.manifest()
111    }
112
113    fn fork(&self) -> StorageBackendResult<Self> {
114        self.control.check()?;
115        let memory = self.control.memory().reserve(size_of::<Self>())?;
116        Ok(Self {
117            index: self.index.clone(),
118            clock: self.clock.clone(),
119            options: self.options,
120            temporary: self.temporary.clone(),
121            control: self.control.clone(),
122            _memory: memory,
123        })
124    }
125
126    fn replace(
127        &mut self,
128        document: DocId,
129        values: Vec<Vec<f32>>,
130        memory: MemoryReservation,
131    ) -> StorageBackendResult<()> {
132        self.control.check()?;
133        let canonical =
134            self.index
135                .canonical()
136                .replaced(document, values, self.clock.version()?, memory)?;
137        let candidate = self.index.with_canonical(canonical)?;
138        self.control.check()?;
139        self.index = candidate;
140        Ok(())
141    }
142
143    fn rebuild(&mut self, canonical: Canonical) -> StorageBackendResult<()> {
144        self.control.check()?;
145        let candidate = build::prepare(
146            canonical,
147            self.clock.generation()?,
148            self.options,
149            &self.temporary,
150            &self.control,
151        )?;
152        self.control.check()?;
153        self.index = candidate;
154        Ok(())
155    }
156}
157
158impl VectorIndex for DiskANNMemoryIndex {
159    fn dimensions(&self) -> u32 {
160        self.index.dimensions()
161    }
162    fn index_kind(&self) -> &'static str {
163        "diskann"
164    }
165    fn diskann_query_metadata(
166        &self,
167        control: &StorageReadControl,
168    ) -> StorageBackendResult<Option<super::DiskANNQueryMetadata>> {
169        self.index.diskann_query_metadata(control)
170    }
171    fn add(&mut self, document: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
172        // Admit the one-element tensor header before allocating this adapter-owned wrapper; coordinates remain caller-owned until replacement adopts them.
173        self.control.check()?;
174        let header = self.control.memory().reserve(size_of::<Vec<f32>>())?;
175        let values = vec![vector];
176        self.replace(document, values, header)
177    }
178    fn add_many(&mut self, document: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
179        self.replace(document, vectors, self.control.memory().empty_reservation())
180    }
181    fn delete(&mut self, document: DocId) -> StorageBackendResult<()> {
182        if !self.index.contains_document(document)? {
183            return Ok(());
184        }
185        self.replace(
186            document,
187            Vec::new(),
188            self.control.memory().empty_reservation(),
189        )
190    }
191    fn clear(&mut self) -> StorageBackendResult<()> {
192        self.rebuild(Canonical::empty(self.dimensions(), &self.control))
193    }
194    fn initialize(&mut self) -> StorageBackendResult<()> {
195        self.rebuild(self.index.canonical().clone())
196    }
197    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
198        self.index.search_knn(query, k)
199    }
200    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
201        self.index.search_threshold(query, threshold)
202    }
203    fn search_knn_with_control(
204        &self,
205        query: &[f32],
206        k: usize,
207        control: &StorageReadControl,
208    ) -> StorageBackendResult<PostingList> {
209        self.index.search_knn_with_control(query, k, control)
210    }
211    fn search_threshold_with_control(
212        &self,
213        query: &[f32],
214        threshold: f32,
215        control: &StorageReadControl,
216    ) -> StorageBackendResult<PostingList> {
217        self.index
218            .search_threshold_with_control(query, threshold, control)
219    }
220    fn search_knn_with_statistics(
221        &self,
222        query: &[f32],
223        k: usize,
224        control: Option<&StorageReadControl>,
225    ) -> StorageBackendResult<crate::vector_index::VectorQueryResult> {
226        self.index.search_knn_with_statistics(query, k, control)
227    }
228    fn search_threshold_with_statistics(
229        &self,
230        query: &[f32],
231        threshold: f32,
232        control: Option<&StorageReadControl>,
233    ) -> StorageBackendResult<crate::vector_index::VectorQueryResult> {
234        self.index
235            .search_threshold_with_statistics(query, threshold, control)
236    }
237    fn count(&self) -> StorageBackendResult<usize> {
238        self.index.count()
239    }
240    fn contains_document(&self, document: DocId) -> StorageBackendResult<bool> {
241        self.index.contains_document(document)
242    }
243    fn diskann_read_snapshot(
244        &self,
245        control: &StorageReadControl,
246    ) -> StorageBackendResult<Option<super::DiskANNReadSnapshot>> {
247        control.check()?;
248        self.index.diskann_read_snapshot(control)
249    }
250    fn snapshot_with_diskann_changes(
251        &self,
252        changes: &super::DiskANNReadChanges,
253        control: &StorageReadControl,
254    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
255        control.check()?;
256        self.index.snapshot_with_diskann_changes(changes, control)
257    }
258    fn snapshot_with_vector_read(
259        &self,
260        source: crate::vector_index::VectorReadSnapshot,
261        control: &StorageReadControl,
262    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
263        self.index.snapshot_with_vector_read(source, control)
264    }
265    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
266        self.index.snapshot()
267    }
268    fn snapshot_with_control(
269        &self,
270        control: &StorageReadControl,
271    ) -> StorageBackendResult<Arc<dyn VectorIndex>> {
272        self.index.snapshot_with_control(control)
273    }
274    fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn VectorIndex>> {
275        Ok(Box::new(self.fork()?))
276    }
277}
uqa_storage/diskann_index/memory.rs

uqa_storage/diskann_index/
memory.rs
