uqa_storage/key_value/diskann/
source.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use std::sync::Arc;
8
9use uqa_core::memory::{BudgetedVec, MemoryReservation};
10
11use crate::diskann_index::format::{DiskANNGeneration, PAGE_BYTES};
12use crate::diskann_index::pages::{
13    DiskANNPageSource, DiskANNPageVisitor, DiskANNReadCapabilities, DiskANNRecordKey,
14    DiskANNRecordVisitor,
15};
16use crate::key_value::KeyValueRead;
17use crate::read_control::StorageReadControl;
18use crate::StorageBackendResult;
19
20use super::keys::{Key, Keys, Kind, ROOT};
21use super::state::{fixed, State, STATE_BYTES};
22use super::{
23    invalid, read_data_identity, validate_session, DiskANNStageStatus, KeyValueDiskANNStore,
24};
25
26mod owned;
27mod selection;
28
29pub(super) const KEY_PAGE_LIMIT: usize = 64;
30
31/// A physical generation pinned to an existing MVCC lease. Opening it reads only fixed metadata; pages and record batches remain lazy.
32pub struct KeyValueDiskANNSource {
33    pub(in crate::key_value::diskann) read: Arc<dyn KeyValueRead + Send + Sync>,
34    generation: DiskANNGeneration,
35    query_control: Option<StorageReadControl>,
36    _memory: MemoryReservation,
37}
38
39impl KeyValueDiskANNSource {
40    /// Retain and fully validate this exact generation's origin stream without reopening a newer provider view or dropping its build lease.
41    pub fn origin_reader(
42        &self,
43        maximum: usize,
44        control: &StorageReadControl,
45    ) -> StorageBackendResult<crate::diskann_index::pages::DiskANNOriginReader> {
46        self.check(control)?;
47        let source = Self::from_read(
48            self.read.clone(),
49            self.generation,
50            DiskANNStageStatus::Sealed,
51            self.query_control.as_ref(),
52            control,
53        )?;
54        crate::diskann_index::pages::DiskANNOriginReader::open(source, maximum, control)
55    }
56
57    pub(super) fn capture(
58        repository: &KeyValueDiskANNStore,
59        generation: DiskANNGeneration,
60        status: DiskANNStageStatus,
61        control: &StorageReadControl,
62    ) -> StorageBackendResult<Arc<Self>> {
63        control.check()?;
64        let mut state = None;
65        {
66            let _writer = repository.owner.writer.lock();
67            repository.idle(control)?;
68            repository.owner.store.with_read_view(&mut |read| {
69                state = super::staging::load_state(read, generation, control)?;
70                Ok(())
71            })?;
72        }
73        let state = state.ok_or_else(|| invalid("generation state is missing"))?;
74        if state.status != status
75            && !(status == DiskANNStageStatus::Sealed && state.status.is_complete())
76        {
77            return Err(invalid(
78                "generation has not reached the requested physical state",
79            ));
80        }
81        let lease = if matches!(
82            state.status,
83            DiskANNStageStatus::Published
84                | DiskANNStageStatus::Retired
85                | DiskANNStageStatus::Discarding
86        ) {
87            None
88        } else if matches!(state.owner, super::state::StageOwner::Legacy(_)) {
89            repository.resume_stage(generation, control)?.lease
90        } else {
91            Some(
92                repository
93                    .acquire_owner(
94                        state.owner,
95                        crate::mvcc::ResourceLeaseRequest::Share,
96                        control,
97                    )?
98                    .ok_or_else(|| invalid("generation recovery is still in progress"))?,
99            )
100        };
101        let store = {
102            let _writer = repository.owner.writer.lock();
103            repository.idle(control)?;
104            let retained = repository
105                .owner
106                .store
107                .open_retained_read_session(control.cancellation())?;
108            validate_session(&*repository.owner.store, &*retained)?;
109            retained
110        };
111        let mut read = None;
112        store.with_read_view(&mut |view| {
113            read = Some(view.retain(&[ROOT])?);
114            Ok(())
115        })?;
116        let mut read = read.ok_or_else(|| invalid("source did not expose a read view"))?;
117        if let Some(lease) = lease {
118            read = owned::OwnedRead::retain(read, lease, control)?;
119        }
120        Self::from_read(read, generation, status, None, control)
121    }
122
123    fn from_read(
124        read: Arc<dyn KeyValueRead + Send + Sync>,
125        generation: DiskANNGeneration,
126        status: DiskANNStageStatus,
127        original: Option<&StorageReadControl>,
128        control: &StorageReadControl,
129    ) -> StorageBackendResult<Arc<Self>> {
130        if let Some(original) = original {
131            original.check()?;
132        }
133        read.control().check()?;
134        control.check()?;
135        let memory = control.memory().reserve(std::mem::size_of::<Self>())?;
136        if read_data_identity(&*read, control)? != Some(generation.database()) {
137            return Err(invalid("source belongs to another data identity"));
138        }
139        let key = Keys::new(generation).key(Kind::State);
140        let state = fixed(control, |visit| {
141            read.visit_value_bounded(key.as_ref(), STATE_BYTES, control, visit)
142        })?
143        .map(State::decode)
144        .transpose()?
145        .ok_or_else(|| invalid("generation state is missing"))?;
146        if state.status != status
147            && !(status == DiskANNStageStatus::Sealed && state.status.is_complete())
148        {
149            return Err(invalid(
150                "generation has not reached the requested physical state",
151            ));
152        }
153        control.check()?;
154        Ok(Arc::new(Self {
155            read,
156            generation,
157            query_control: original.cloned(),
158            _memory: memory,
159        }))
160    }
161
162    pub(super) fn keys_after(
163        &self,
164        after: Option<&[u8]>,
165        control: &StorageReadControl,
166    ) -> StorageBackendResult<BudgetedVec<(Key, Kind)>> {
167        self.check(control)?;
168        let result = key_page(
169            &*self.read,
170            Keys::new(self.generation),
171            after,
172            KEY_PAGE_LIMIT,
173            control,
174        )?;
175        self.check(control)?;
176        Ok(result)
177    }
178
179    pub(super) fn value(
180        &self,
181        kind: Kind,
182        maximum: usize,
183        control: &StorageReadControl,
184        visit: &mut DiskANNRecordVisitor<'_>,
185    ) -> StorageBackendResult<()> {
186        self.check(control)?;
187        let key = Keys::new(self.generation).key(kind);
188        let mut seen = false;
189        let mut failure = None;
190        let source = self
191            .read
192            .visit_value_bounded(key.as_ref(), maximum, control, &mut |value| {
193                if failure.is_some() {
194                    return Err(invalid("record visitor already failed"));
195                }
196                let outcome = (|| {
197                    if seen {
198                        return Err(invalid("record was returned more than once"));
199                    }
200                    seen = true;
201                    let bytes = value.ok_or_else(|| invalid("generation record is missing"))?;
202                    control.check_value_size(bytes.len(), maximum)?;
203                    visit(bytes)
204                })();
205                if let Err(error) = outcome {
206                    failure = Some(error);
207                    return Err(invalid("record visitor rejected data"));
208                }
209                Ok(())
210            });
211        if let Some(error) = failure {
212            return Err(error);
213        }
214        source?;
215        if !seen {
216            return Err(invalid("record was not returned"));
217        }
218        self.check(control)
219    }
220
221    pub(in crate::key_value::diskann) fn check(
222        &self,
223        control: &StorageReadControl,
224    ) -> StorageBackendResult<()> {
225        if let Some(original) = &self.query_control {
226            original.check()?;
227        }
228        self.read.control().check()?;
229        control.check()
230    }
231}
232
233impl DiskANNPageSource for KeyValueDiskANNSource {
234    fn generation(&self) -> DiskANNGeneration {
235        self.generation
236    }
237    fn capabilities(&self) -> DiskANNReadCapabilities {
238        DiskANNReadCapabilities::new(32, 1).expect("fixed sequential read capability")
239    }
240    fn read_record(
241        &self,
242        key: DiskANNRecordKey,
243        max_bytes: usize,
244        control: &StorageReadControl,
245        visit: &mut DiskANNRecordVisitor<'_>,
246    ) -> StorageBackendResult<()> {
247        self.value(Kind::Record(key), max_bytes, control, visit)
248    }
249    fn read_graph_pages(
250        &self,
251        pages: &[u64],
252        control: &StorageReadControl,
253        visit: &mut DiskANNPageVisitor<'_>,
254    ) -> StorageBackendResult<()> {
255        self.check(control)?;
256        if pages.len() > self.capabilities().max_batch_pages()
257            || pages
258                .iter()
259                .enumerate()
260                .any(|(index, id)| pages[..index].contains(id))
261        {
262            return Err(invalid(
263                "page request exceeds its batch bound or repeats a page",
264            ));
265        }
266        for &id in pages {
267            self.value(Kind::Graph(id), PAGE_BYTES, control, &mut |bytes| {
268                if bytes.len() != PAGE_BYTES {
269                    return Err(invalid("graph page length differs"));
270                }
271                visit(id, bytes)
272            })?;
273        }
274        self.check(control)
275    }
276}
277
278/// Copy only a bounded page of fixed-size keys, releasing provider guards before any value decoding or further reads.
279pub(super) fn key_page(
280    read: &dyn KeyValueRead,
281    keys: Keys,
282    after: Option<&[u8]>,
283    limit: usize,
284    control: &StorageReadControl,
285) -> StorageBackendResult<BudgetedVec<(Key, Kind)>> {
286    control.check()?;
287    let mut page = BudgetedVec::<(Key, Kind)>::new(control.memory());
288    let mut failure = None;
289    let source = read.visit_keys_after(keys.prefix(), after, limit, control, &mut |bytes| {
290        if failure.is_some() {
291            return Err(invalid("key visitor already failed"));
292        }
293        let outcome = (|| {
294            control.check()?;
295            let previous = page.last().map(|(key, _)| key.as_ref()).or(after);
296            if page.len() >= limit || previous.is_some_and(|previous| previous >= bytes) {
297                return Err(invalid("key page exceeds its bound or ordering"));
298            }
299            let kind = keys.decode(bytes)?;
300            page.push((keys.key(kind), kind))?;
301            Ok(())
302        })();
303        if let Err(error) = outcome {
304            failure = Some(error);
305            return Err(invalid("key visitor rejected data"));
306        }
307        Ok(())
308    });
309    if let Some(error) = failure {
310        return Err(error);
311    }
312    source?;
313    control.check()?;
314    Ok(page)
315}
uqa_storage/key_value/diskann/source.rs

uqa_storage/key_value/diskann/
source.rs
