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        self.next
45            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |next| {
46                next.checked_add(1)
47            })
48            .map_err(|_| StorageBackendError::Other("DiskANN memory identity exhausted".into()))
49    }
50    fn generation(&self) -> StorageBackendResult<DiskANNGeneration> {
51        DiskANNGeneration::new(self.database.as_bytes(), 1, 1, self.allocate()?)
52    }
53    fn version(&self) -> StorageBackendResult<DiskANNVectorVersion> {
54        // This incarnation owns one mutable memory lineage. Every fork shares its non-reused revision allocator; the identity is not SQL commit order.
55        let writer = StorageTransactionId::new(self.database, 1)
56            .map_err(crate::mvcc::VersionError::into_storage_error)?;
57        DiskANNVectorVersion::new(writer, self.allocate()?)
58    }
59}
60
61/// 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.
62pub struct DiskANNMemoryIndex {
63    index: RetainedDiskANNIndex<Canonical>,
64    clock: Arc<Budgeted<Clock>>,
65    options: DiskANNMemoryOptions,
66    temporary: DiskANNTemporaryBudget,
67    control: StorageReadControl,
68    _memory: MemoryReservation,
69}
70
71impl DiskANNMemoryIndex {
72    pub fn new(
73        dimensions: u32,
74        options: DiskANNMemoryOptions,
75        temporary: &DiskANNTemporaryBudget,
76        control: &StorageReadControl,
77    ) -> StorageBackendResult<Self> {
78        control.check()?;
79        options.parameters.validate(dimensions)?;
80        let memory = control.memory().reserve(size_of::<Self>())?;
81        let database = DatabaseId::from_bytes(crate::catalog::new_nonzero_catalog_identity(
82            "DiskANN memory",
83            "incarnation",
84        )?);
85        let clock = Budgeted::new(
86            Clock {
87                database,
88                next: AtomicU64::new(1),
89            },
90            control.memory().empty_reservation(),
91        )
92        .into_shared()?;
93        let index = build::prepare(
94            Canonical::empty(dimensions, control),
95            clock.generation()?,
96            options,
97            temporary,
98            control,
99        )?;
100        Ok(Self {
101            index,
102            clock,
103            options,
104            temporary: temporary.clone(),
105            control: control.clone(),
106            _memory: memory,
107        })
108    }
109
110    pub fn manifest(&self) -> &DiskANNManifest {
111        self.index.manifest()
112    }
113
114    fn fork(&self) -> StorageBackendResult<Self> {
115        self.control.check()?;
116        let memory = self.control.memory().reserve(size_of::<Self>())?;
117        Ok(Self {
118            index: self.index.clone(),
119            clock: self.clock.clone(),
120            options: self.options,
121            temporary: self.temporary.clone(),
122            control: self.control.clone(),
123            _memory: memory,
124        })
125    }
126
127    fn replace(
128        &mut self,
129        document: DocId,
130        values: Vec<Vec<f32>>,
131        memory: MemoryReservation,
132    ) -> StorageBackendResult<()> {
133        self.control.check()?;
134        let canonical =
135            self.index
136                .canonical()
137                .replaced(document, values, self.clock.version()?, memory)?;
138        let candidate = self.index.with_canonical(canonical)?;
139        self.control.check()?;
140        self.index = candidate;
141        Ok(())
142    }
143
144    fn rebuild(&mut self, canonical: Canonical) -> StorageBackendResult<()> {
145        self.control.check()?;
146        let candidate = build::prepare(
147            canonical,
148            self.clock.generation()?,
149            self.options,
150            &self.temporary,
151            &self.control,
152        )?;
153        self.control.check()?;
154        self.index = candidate;
155        Ok(())
156    }
157}
158
159impl VectorIndex for DiskANNMemoryIndex {
160    fn dimensions(&self) -> u32 {
161        self.index.dimensions()
162    }
163    fn index_kind(&self) -> &'static str {
164        "diskann"
165    }
166    fn diskann_query_metadata(
167        &self,
168        control: &StorageReadControl,
169    ) -> StorageBackendResult<Option<super::DiskANNQueryMetadata>> {
170        self.index.diskann_query_metadata(control)
171    }
172    fn add(&mut self, document: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
173        // Admit the one-element tensor header before allocating this adapter-owned wrapper; coordinates remain caller-owned until replacement adopts them.
174        self.control.check()?;
175        let header = self.control.memory().reserve(size_of::<Vec<f32>>())?;
176        let values = vec![vector];
177        self.replace(document, values, header)
178    }
179    fn add_many(&mut self, document: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
180        self.replace(document, vectors, self.control.memory().empty_reservation())
181    }
182    fn delete(&mut self, document: DocId) -> StorageBackendResult<()> {
183        if !self.index.contains_document(document)? {
184            return Ok(());
185        }
186        self.replace(
187            document,
188            Vec::new(),
189            self.control.memory().empty_reservation(),
190        )
191    }
192    fn clear(&mut self) -> StorageBackendResult<()> {
193        self.rebuild(Canonical::empty(self.dimensions(), &self.control))
194    }
195    fn initialize(&mut self) -> StorageBackendResult<()> {
196        self.rebuild(self.index.canonical().clone())
197    }
198    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
199        self.index.search_knn(query, k)
200    }
201    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
202        self.index.search_threshold(query, threshold)
203    }
204    fn search_knn_with_control(
205        &self,
206        query: &[f32],
207        k: usize,
208        control: &StorageReadControl,
209    ) -> StorageBackendResult<PostingList> {
210        self.index.search_knn_with_control(query, k, control)
211    }
212    fn search_threshold_with_control(
213        &self,
214        query: &[f32],
215        threshold: f32,
216        control: &StorageReadControl,
217    ) -> StorageBackendResult<PostingList> {
218        self.index
219            .search_threshold_with_control(query, threshold, control)
220    }
221    fn search_knn_with_statistics(
222        &self,
223        query: &[f32],
224        k: usize,
225        control: Option<&StorageReadControl>,
226    ) -> StorageBackendResult<crate::vector_index::VectorQueryResult> {
227        self.index.search_knn_with_statistics(query, k, control)
228    }
229    fn search_threshold_with_statistics(
230        &self,
231        query: &[f32],
232        threshold: f32,
233        control: Option<&StorageReadControl>,
234    ) -> StorageBackendResult<crate::vector_index::VectorQueryResult> {
235        self.index
236            .search_threshold_with_statistics(query, threshold, control)
237    }
238    fn count(&self) -> StorageBackendResult<usize> {
239        self.index.count()
240    }
241    fn contains_document(&self, document: DocId) -> StorageBackendResult<bool> {
242        self.index.contains_document(document)
243    }
244    fn diskann_read_snapshot(
245        &self,
246        control: &StorageReadControl,
247    ) -> StorageBackendResult<Option<super::DiskANNReadSnapshot>> {
248        control.check()?;
249        self.index.diskann_read_snapshot(control)
250    }
251    fn snapshot_with_diskann_changes(
252        &self,
253        changes: &super::DiskANNReadChanges,
254        control: &StorageReadControl,
255    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
256        control.check()?;
257        self.index.snapshot_with_diskann_changes(changes, control)
258    }
259    fn snapshot_with_vector_read(
260        &self,
261        source: crate::vector_index::VectorReadSnapshot,
262        control: &StorageReadControl,
263    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
264        self.index.snapshot_with_vector_read(source, control)
265    }
266    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
267        self.index.snapshot()
268    }
269    fn snapshot_with_control(
270        &self,
271        control: &StorageReadControl,
272    ) -> StorageBackendResult<Arc<dyn VectorIndex>> {
273        self.index.snapshot_with_control(control)
274    }
275    fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn VectorIndex>> {
276        Ok(Box::new(self.fork()?))
277    }
278}
uqa_storage/diskann_index/memory.rs

uqa_storage/diskann_index/
memory.rs
