uqa_storage/diskann_index/pages/
memory.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::{Budgeted, BudgetedMap, BudgetedVec, MemoryBudget};
10
11use super::{
12    copy, invalid, DiskANNArtifactSeal, DiskANNArtifactSealer, DiskANNPageSource,
13    DiskANNPageVisitor, DiskANNReadCapabilities, DiskANNRecordKey, DiskANNRecordVisitor,
14};
15use crate::diskann_index::format::{DiskANNGeneration, DiskANNManifest, PAGE_BYTES};
16use crate::{read_control::StorageReadControl, StorageBackendResult};
17
18struct Records {
19    records: BudgetedMap<DiskANNRecordKey, BudgetedVec<u8>>,
20    graph: BudgetedMap<u64, BudgetedVec<u8>>,
21}
22
23pub struct DiskANNMemoryBuilder {
24    generation: DiskANNGeneration,
25    records: Records,
26    budget: MemoryBudget,
27}
28
29#[derive(Clone)]
30pub struct DiskANNMemorySource {
31    seal: DiskANNArtifactSeal,
32    records: Arc<Budgeted<Records>>,
33}
34
35impl DiskANNMemoryBuilder {
36    pub fn generation(&self) -> DiskANNGeneration {
37        self.generation
38    }
39
40    pub fn new(generation: DiskANNGeneration, budget: &MemoryBudget) -> Self {
41        Self {
42            generation,
43            records: Records {
44                records: BudgetedMap::new(budget),
45                graph: BudgetedMap::new(budget),
46            },
47            budget: budget.clone(),
48        }
49    }
50
51    pub fn write_record(
52        &mut self,
53        key: DiskANNRecordKey,
54        bytes: &[u8],
55        control: &StorageReadControl,
56    ) -> StorageBackendResult<()> {
57        control.check()?;
58        if key == DiskANNRecordKey::Manifest || self.records.records.contains_key(&key) {
59            return Err(invalid(
60                "manifest is written at sealing; staging records cannot be replaced",
61            ));
62        }
63        let owner = StorageReadControl::new(&self.budget, control.cancellation());
64        let bytes = copy(bytes, usize::MAX, &owner)?;
65        self.records.records.insert(key, bytes)?;
66        Ok(())
67    }
68
69    pub fn write_graph_page(
70        &mut self,
71        id: u64,
72        bytes: &[u8],
73        control: &StorageReadControl,
74    ) -> StorageBackendResult<()> {
75        control.check()?;
76        if bytes.len() != PAGE_BYTES || self.records.graph.contains_key(&id) {
77            return Err(invalid("staged page length or identity differs"));
78        }
79        let owner = StorageReadControl::new(&self.budget, control.cancellation());
80        let bytes = copy(bytes, PAGE_BYTES, &owner)?;
81        self.records.graph.insert(id, bytes)?;
82        Ok(())
83    }
84
85    pub fn finish(
86        mut self,
87        manifest: DiskANNManifest,
88        control: &StorageReadControl,
89    ) -> StorageBackendResult<DiskANNMemorySource> {
90        control.check()?;
91        if manifest.input().generation != self.generation {
92            return Err(invalid("manifest does not belong to staging owner"));
93        }
94        let mut sealer = DiskANNArtifactSealer::new(manifest, control)?;
95        for (key, bytes) in &self.records.records {
96            match key {
97                DiskANNRecordKey::Codebook => sealer.codebook(bytes)?,
98                DiskANNRecordKey::Codes(first) => sealer.code_batch(*first, bytes)?,
99                DiskANNRecordKey::Side(first) => sealer.side_batch(*first, bytes)?,
100                DiskANNRecordKey::Origins(first) => sealer.origin_batch(*first, bytes)?,
101                DiskANNRecordKey::Manifest => {
102                    return Err(invalid("staged manifest before sealing"))
103                }
104            }
105        }
106        for (id, bytes) in &self.records.graph {
107            sealer.graph_page(*id, bytes)?;
108        }
109        let seal = sealer.finish()?;
110        let owner = StorageReadControl::new(&self.budget, control.cancellation());
111        self.records
112            .records
113            .insert(DiskANNRecordKey::Manifest, manifest.encode(&owner)?)?;
114        let records = Budgeted::new(self.records, self.budget.empty_reservation()).into_shared()?;
115        control.check()?;
116        Ok(DiskANNMemorySource { seal, records })
117    }
118}
119
120impl DiskANNPageSource for DiskANNMemorySource {
121    fn generation(&self) -> DiskANNGeneration {
122        self.seal.manifest().input().generation
123    }
124    fn capabilities(&self) -> DiskANNReadCapabilities {
125        DiskANNReadCapabilities {
126            max_batch_pages: 32,
127            read_concurrency: 1,
128        }
129    }
130    fn read_record(
131        &self,
132        key: DiskANNRecordKey,
133        max_bytes: usize,
134        control: &StorageReadControl,
135        visit: &mut DiskANNRecordVisitor<'_>,
136    ) -> StorageBackendResult<()> {
137        control.check()?;
138        let bytes = self
139            .records
140            .records
141            .get(&key)
142            .ok_or_else(|| invalid("missing memory record"))?;
143        if bytes.len() > max_bytes {
144            return Err(uqa_core::memory::MemoryError::Limit {
145                required: bytes.len(),
146                limit: max_bytes,
147            }
148            .into());
149        }
150        visit(bytes)?;
151        control.check()
152    }
153    fn read_graph_pages(
154        &self,
155        pages: &[u64],
156        control: &StorageReadControl,
157        visit: &mut DiskANNPageVisitor<'_>,
158    ) -> StorageBackendResult<()> {
159        if pages.len() > self.capabilities().max_batch_pages() {
160            return Err(invalid("page request exceeds source batch limit"));
161        }
162        for &id in pages {
163            control.check()?;
164            let bytes = self
165                .records
166                .graph
167                .get(&id)
168                .ok_or_else(|| invalid("missing memory page"))?;
169            visit(id, bytes)?;
170        }
171        control.check()
172    }
173}
uqa_storage/diskann_index/pages/memory.rs

uqa_storage/diskann_index/pages/
memory.rs
