Skip to main content

uqa_storage/key_value/diskann/source/
selection.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Select physical pages from the canonical query boundary, including private publication resources.
8
9use super::{invalid, DiskANNStageStatus, KeyValueDiskANNSource, ROOT};
10use crate::diskann_index::{
11    catalog::DiskANNIndexScope,
12    format::DiskANNManifest,
13    pages::{DiskANNPageSource, DiskANNRecordKey},
14};
15use crate::key_value::{diskann::publication, KeyValueRead};
16use crate::{
17    read_control::StorageReadControl, vector_index::DiskANNIndexParams, StorageBackendResult,
18};
19use std::sync::Arc;
20
21impl KeyValueDiskANNSource {
22    /// Provider boundary: select the head, configuration and actual physical source on one retained catalog/canonical view. A committed head retains that same view; a private head requires the source attached by its original publication. Both original and invoking controls remain binding. This reads fixed metadata only, leaving graph/PQ/origin batches lazy.
23    pub fn select(
24        scope: &DiskANNIndexScope,
25        dimensions: u32,
26        parameters: DiskANNIndexParams,
27        read: &dyn KeyValueRead,
28        original: &StorageReadControl,
29        control: &StorageReadControl,
30    ) -> StorageBackendResult<Option<Arc<Self>>> {
31        original.check()?;
32        let Some(generation) = publication::selected_generation(scope, read, control)? else {
33            original.check()?;
34            return Ok(None);
35        };
36        let key = publication::head_key(scope);
37        let revision = read
38            .record_revision(&key)?
39            .ok_or_else(|| invalid("selected head has no record identity"))?;
40        let (physical, status) = if revision.has_private_changes() {
41            (
42                read.retained_source(&key)?
43                    .ok_or_else(|| invalid("private head has no retained physical source"))?,
44                DiskANNStageStatus::Sealed,
45            )
46        } else {
47            (read.retain(&[ROOT])?, DiskANNStageStatus::Published)
48        };
49        let source = Self::from_read(physical, generation, status, Some(original), control)?;
50        super::super::identity::validate_mapping(scope, generation, &*source.read, control)?;
51        source.read_record(
52            DiskANNRecordKey::Manifest,
53            DiskANNManifest::MAX_ENCODED_BYTES,
54            control,
55            &mut |bytes| {
56                let manifest = DiskANNManifest::decode(generation, bytes, control)?;
57                if manifest.origins().is_none()
58                    || manifest.input().dimensions != dimensions
59                    || manifest.input().parameters != parameters
60                {
61                    return Err(invalid(
62                        "selected generation differs from the captured index configuration",
63                    ));
64                }
65                Ok(())
66            },
67        )?;
68        source.check(control)?;
69        Ok(Some(source))
70    }
71}