uqa_storage/document_store/
retained_store.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Copied document snapshots retain their payload and entry capacity through the last reader.
8
9use std::{collections::BTreeMap, sync::Arc};
10
11use uqa_core::{
12    memory::{Budgeted, BudgetedVec, MemoryReservation},
13    ordering::sort_by_with_control,
14    DocId, Value,
15};
16
17use super::{Document, DocumentMetadata, DocumentStore, StoredDocument};
18use crate::{read_control::StorageReadControl, StorageBackendError, StorageBackendResult};
19
20type Documents = Vec<(DocId, StoredDocument)>;
21
22/// Adopt caller-owned rows without copying payloads. Finalization orders unique document identities and transfers the same allowance into immutable readers. Input decoding and mutable output allocations belong to their respective producers and callers.
23pub struct RetainedDocumentStoreBuilder {
24    entries: BudgetedVec<(DocId, StoredDocument)>,
25    payload: MemoryReservation,
26    control: StorageReadControl,
27}
28
29impl RetainedDocumentStoreBuilder {
30    pub fn new(control: &StorageReadControl) -> Self {
31        Self {
32            entries: BudgetedVec::new(control.memory()),
33            payload: control.memory().empty_reservation(),
34            control: control.clone(),
35        }
36    }
37
38    pub fn add_document(
39        &mut self,
40        id: DocId,
41        document: StoredDocument,
42    ) -> StorageBackendResult<()> {
43        self.control.check()?;
44        let (fields, metadata) = document.into_parts();
45        let value = Value::Map(fields);
46        let memory = value
47            .reserve_retained_payload(self.control.memory(), self.control.cancellation())
48            .map_err(super::retained::retention_error)?;
49        let Value::Map(fields) = value else {
50            unreachable!("document field map");
51        };
52        self.push_document(
53            id,
54            Budgeted::new(StoredDocument::with_metadata(fields, metadata), memory),
55        )
56    }
57
58    /// Adopt a controlled producer's complete tuple and existing payload reservation. The input lease must cover live map entries, field-name capacity and value payloads under this builder's allowance. Entry capacity is reserved while that lease remains live.
59    pub fn add_budgeted_document(
60        &mut self,
61        id: DocId,
62        document: Budgeted<StoredDocument>,
63    ) -> StorageBackendResult<()> {
64        self.control.check()?;
65        if !document.budget().shares_allowance(self.control.memory()) {
66            return Err(StorageBackendError::Other(
67                "retained document input belongs to a different allowance".into(),
68            ));
69        }
70        let (document, mut memory) = document.into_parts();
71        let (fields, metadata) = document.into_parts();
72        let value = Value::Map(fields);
73        let required = value
74            .retained_payload_bytes(self.control.memory(), self.control.cancellation())
75            .map_err(super::retained::retention_error)?;
76        if memory.bytes() < required {
77            return Err(StorageBackendError::Other(
78                "retained document input reservation does not cover its payload".into(),
79            ));
80        }
81        let surplus = memory.bytes() - required;
82        drop(memory.split(surplus));
83        let Value::Map(fields) = value else {
84            unreachable!("document field map");
85        };
86        self.push_document(
87            id,
88            Budgeted::new(StoredDocument::with_metadata(fields, metadata), memory),
89        )
90    }
91
92    /// Transfer a retained tuple into this corpus, copying only when another reader still owns its fields.
93    pub fn add_retained_document(
94        &mut self,
95        id: DocId,
96        document: super::RetainedStoredDocument,
97    ) -> StorageBackendResult<()> {
98        let document = document.into_budgeted(&self.control)?;
99        self.add_budgeted_document(id, document)
100    }
101
102    fn push_document(
103        &mut self,
104        id: DocId,
105        pending: Budgeted<StoredDocument>,
106    ) -> StorageBackendResult<()> {
107        // Drop pending fields before returning their payload reservation on failure.
108        self.entries.reserve(1)?;
109        self.control.check()?;
110        let (document, memory) = pending.into_parts();
111        self.entries.push((id, document))?;
112        self.payload.absorb(memory);
113        Ok(())
114    }
115
116    pub fn finish(mut self) -> StorageBackendResult<RetainedDocumentStore> {
117        sort_by_with_control(
118            &mut self.entries,
119            &mut || self.control.check(),
120            |left, right, _| Ok(left.0.cmp(&right.0)),
121        )?;
122        for pair in self.entries.windows(2) {
123            self.control.check()?;
124            if pair[0].0 == pair[1].0 {
125                return Err(StorageBackendError::Other(
126                    "retained document input repeats a document identity".into(),
127                ));
128            }
129        }
130        self.control.check()?;
131        let (entries, mut memory) = self.entries.into_parts();
132        memory.absorb(self.payload);
133        Ok(RetainedDocumentStore {
134            entries: Budgeted::new(entries, memory).into_shared()?,
135            control: self.control,
136        })
137    }
138}
139
140/// An immutable corpus whose nested snapshots share one entry/payload lease. Borrowed projections reserve only their reference buffer; field maps, tuple metadata and value buffers stay in the original corpus.
141#[derive(Clone)]
142pub struct RetainedDocumentStore {
143    entries: Arc<Budgeted<Documents>>,
144    control: StorageReadControl,
145}
146
147impl RetainedDocumentStore {
148    fn document(&self, id: DocId) -> Option<&StoredDocument> {
149        self.entries
150            .binary_search_by_key(&id, |entry| entry.0)
151            .ok()
152            .map(|index| &self.entries[index].1)
153    }
154
155    fn after(&self, after: Option<DocId>) -> &[(DocId, StoredDocument)] {
156        &self.entries[self
157            .entries
158            .partition_point(|(id, _)| after.is_some_and(|after| *id <= after))..]
159    }
160
161    fn visit<'a>(
162        &'a self,
163        rows: impl Iterator<Item = (DocId, Option<&'a StoredDocument>)>,
164        fields: &[&str],
165        visitor: &mut dyn FnMut(DocId, bool, &[&Value]) -> bool,
166    ) -> StorageBackendResult<()> {
167        self.control.check()?;
168        let mut projected = BudgetedVec::new(self.control.memory());
169        projected.reserve(fields.len())?;
170        for (id, document) in rows {
171            self.control.check()?;
172            projected.clear();
173            for field in fields {
174                self.control.check()?;
175                projected.push(
176                    document
177                        .and_then(|document| document.fields().get(*field))
178                        .unwrap_or(&Value::Null),
179                )?;
180            }
181            if !visitor(id, document.is_some(), &projected) {
182                break;
183            }
184        }
185        self.control.check()
186    }
187}
188
189fn read_only() -> StorageBackendError {
190    StorageBackendError::Other("cannot write a retained document snapshot".into())
191}
192
193impl DocumentStore for RetainedDocumentStore {
194    fn put(&mut self, _: DocId, _: Document) -> StorageBackendResult<()> {
195        Err(read_only())
196    }
197    fn put_stored(&mut self, _: DocId, _: StoredDocument) -> StorageBackendResult<()> {
198        Err(read_only())
199    }
200    fn patch_fields(&mut self, _: DocId, _: &Document) -> StorageBackendResult<bool> {
201        Err(read_only())
202    }
203    fn delete(&mut self, _: DocId) -> StorageBackendResult<()> {
204        Err(read_only())
205    }
206    fn clear(&mut self) -> StorageBackendResult<()> {
207        Err(read_only())
208    }
209
210    fn get_stored(&self, id: DocId) -> StorageBackendResult<Option<StoredDocument>> {
211        self.control.check()?;
212        Ok(self.document(id).cloned())
213    }
214
215    fn get_stored_many_controlled(
216        &self,
217        ids: &[DocId],
218        control: &StorageReadControl,
219    ) -> StorageBackendResult<super::RetainedDocumentPage> {
220        self.control.check()?;
221        control.check()?;
222        let mut page = BudgetedVec::new(control.memory());
223        page.reserve(ids.len())?;
224        for id in ids {
225            self.control.check()?;
226            control.check()?;
227            let row = self
228                .document(*id)
229                .map(|row| -> StorageBackendResult<_> {
230                    let fields = super::controlled_rows::copy_fields(
231                        row.fields()
232                            .iter()
233                            .map(|(name, value)| (name.as_str(), value)),
234                        control,
235                    )?;
236                    Ok(super::RetainedStoredDocument::with_metadata(
237                        fields,
238                        row.metadata(),
239                    ))
240                })
241                .transpose()?;
242            page.push(row)?;
243        }
244        self.control.check()?;
245        control.check()?;
246        Ok(page)
247    }
248    fn get_metadata(&self, id: DocId) -> StorageBackendResult<Option<DocumentMetadata>> {
249        self.control.check()?;
250        Ok(self.document(id).map(StoredDocument::metadata))
251    }
252    fn contains_doc_id(&self, id: DocId) -> StorageBackendResult<bool> {
253        self.control.check()?;
254        Ok(self.document(id).is_some())
255    }
256    fn get_field(&self, id: DocId, field: &str) -> StorageBackendResult<Option<Value>> {
257        self.control.check()?;
258        Ok(self
259            .document(id)
260            .and_then(|document| document.fields().get(field).cloned()))
261    }
262
263    fn get_fields_multi(
264        &self,
265        ids: &[DocId],
266        fields: &[&str],
267    ) -> StorageBackendResult<BTreeMap<DocId, Vec<Value>>> {
268        let mut rows = BTreeMap::new();
269        self.for_each_fields_multi_ref_with_presence(ids, fields, &mut |id, present, values| {
270            if present {
271                rows.insert(id, values.iter().map(|value| (*value).clone()).collect());
272            }
273            true
274        })?;
275        Ok(rows)
276    }
277    fn for_each_fields_multi(
278        &self,
279        ids: &[DocId],
280        fields: &[&str],
281        visitor: &mut dyn FnMut(DocId, Vec<Value>) -> bool,
282    ) -> StorageBackendResult<()> {
283        self.for_each_fields_multi_ref(ids, fields, &mut |id, values| {
284            visitor(id, values.iter().map(|value| (*value).clone()).collect())
285        })
286    }
287    fn for_each_fields_multi_ref(
288        &self,
289        ids: &[DocId],
290        fields: &[&str],
291        visitor: &mut dyn FnMut(DocId, &[&Value]) -> bool,
292    ) -> StorageBackendResult<()> {
293        self.for_each_fields_multi_ref_with_presence(ids, fields, &mut |id, _, values| {
294            visitor(id, values)
295        })
296    }
297    fn for_each_fields_multi_ref_with_presence(
298        &self,
299        ids: &[DocId],
300        fields: &[&str],
301        visitor: &mut dyn FnMut(DocId, bool, &[&Value]) -> bool,
302    ) -> StorageBackendResult<()> {
303        self.visit(
304            ids.iter().map(|id| (*id, self.document(*id))),
305            fields,
306            visitor,
307        )
308    }
309    fn for_each_next_fields(
310        &self,
311        after: Option<DocId>,
312        limit: usize,
313        fields: &[&str],
314        visitor: &mut dyn FnMut(DocId, &[&Value]) -> bool,
315    ) -> StorageBackendResult<Option<usize>> {
316        let mut visited = 0;
317        self.visit(
318            self.after(after)
319                .iter()
320                .take(limit)
321                .map(|(id, row)| (*id, Some(row))),
322            fields,
323            &mut |id, _, values| {
324                visited += 1;
325                visitor(id, values)
326            },
327        )?;
328        Ok(Some(visited))
329    }
330
331    fn with_field_ref_controlled(
332        &self,
333        id: DocId,
334        field: &str,
335        control: &StorageReadControl,
336        visitor: &mut dyn FnMut(Option<&Value>) -> StorageBackendResult<()>,
337    ) -> StorageBackendResult<()> {
338        self.control.check()?;
339        control.check()?;
340        visitor(self.document(id).and_then(|row| row.fields().get(field)))?;
341        self.control.check()?;
342        control.check()
343    }
344
345    fn field_presence_controlled(
346        &self,
347        ids: &[DocId],
348        fields: &[&str],
349        control: &StorageReadControl,
350    ) -> StorageBackendResult<BudgetedVec<bool>> {
351        self.control.check()?;
352        control.check()?;
353        let mut present = BudgetedVec::new(control.memory());
354        for id in ids {
355            let row = self.document(*id);
356            for field in fields {
357                self.control.check()?;
358                control.check()?;
359                present.push(row.is_some_and(|row| row.fields().contains_key(*field)))?;
360            }
361        }
362        self.control.check()?;
363        control.check()?;
364        Ok(present)
365    }
366
367    fn find_doc_id_by_field(
368        &self,
369        field: &str,
370        value: &Value,
371    ) -> StorageBackendResult<Option<DocId>> {
372        self.control.check()?;
373        for (id, document) in self.entries.iter() {
374            self.control.check()?;
375            if document.fields().get(field) == Some(value) {
376                return Ok(Some(*id));
377            }
378        }
379        Ok(None)
380    }
381    fn find_doc_id_by_fields(
382        &self,
383        fields: &[String],
384        values: &[Value],
385    ) -> StorageBackendResult<Option<DocId>> {
386        self.control.check()?;
387        if fields.is_empty() || fields.len() != values.len() {
388            return Ok(None);
389        }
390        for (id, document) in self.entries.iter() {
391            self.control.check()?;
392            let mut matched = true;
393            for (field, value) in fields.iter().zip(values) {
394                self.control.check()?;
395                if document.fields().get(field).unwrap_or(&Value::Null) != value {
396                    matched = false;
397                    break;
398                }
399            }
400            if matched {
401                return Ok(Some(*id));
402            }
403        }
404        Ok(None)
405    }
406    fn has_value(&self, field: &str, value: &Value) -> StorageBackendResult<bool> {
407        Ok(self.find_doc_id_by_field(field, value)?.is_some())
408    }
409
410    fn doc_ids(&self) -> StorageBackendResult<Vec<DocId>> {
411        self.next_doc_ids(None, self.entries.len())
412    }
413    fn next_doc_id(&self, after: Option<DocId>) -> StorageBackendResult<Option<DocId>> {
414        self.control.check()?;
415        Ok(self.after(after).first().map(|(id, _)| *id))
416    }
417    fn next_doc_ids(&self, after: Option<DocId>, limit: usize) -> StorageBackendResult<Vec<DocId>> {
418        let (ids, _memory) = self
419            .next_doc_ids_controlled(after, limit, &self.control)?
420            .into_parts();
421        Ok(ids)
422    }
423    fn next_doc_ids_controlled(
424        &self,
425        after: Option<DocId>,
426        limit: usize,
427        control: &StorageReadControl,
428    ) -> StorageBackendResult<BudgetedVec<DocId>> {
429        self.control.check()?;
430        control.check()?;
431        let mut ids = BudgetedVec::new(control.memory());
432        for (id, _) in self.after(after).iter().take(limit) {
433            self.control.check()?;
434            control.check()?;
435            ids.push(*id)?;
436        }
437        control.check()?;
438        Ok(ids)
439    }
440    fn max_doc_id(&self) -> StorageBackendResult<DocId> {
441        self.control.check()?;
442        Ok(self.entries.last().map_or(0, |(id, _)| *id))
443    }
444    fn len(&self) -> StorageBackendResult<usize> {
445        self.control.check()?;
446        Ok(self.entries.len())
447    }
448    fn iter_all(&self) -> StorageBackendResult<Box<dyn Iterator<Item = (DocId, Document)> + '_>> {
449        self.control.check()?;
450        Ok(Box::new(
451            self.entries
452                .iter()
453                .map(|(id, document)| (*id, document.fields().clone())),
454        ))
455    }
456    fn snapshot(&self) -> StorageBackendResult<Arc<dyn DocumentStore>> {
457        self.control.check()?;
458        Ok(Arc::new(self.clone()))
459    }
460
461    fn retained_snapshot(&self) -> StorageBackendResult<Option<Arc<dyn DocumentStore>>> {
462        self.snapshot().map(Some)
463    }
464}
465
466#[cfg(test)]
467mod tests;
uqa_storage/document_store/retained_store.rs

uqa_storage/document_store/
retained_store.rs
