uqa_storage/
document_store.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Document storage abstraction.
8//!
9//! A `DocumentStore` maps [`DocId`] keys to field maps and supports
10//! field-level access. Implementations include in-memory, provider-owned `SQLite`,
11//! and Key/Value-backed implementations behind the same trait.
12
13use std::collections::BTreeMap;
14use std::sync::Arc;
15
16use uqa_core::{DocId, FieldName, PathSegment, Value};
17
18use crate::backend::{StorageBackendError, StorageBackendResult};
19
20/// Document field map. Keys are field names; values are dynamic.
21pub type Document = BTreeMap<FieldName, Value>;
22
23/// Storage-owned tuple metadata that must never share the user field namespace.
24#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
25pub struct DocumentMetadata {
26    tuple_xmin: Option<u32>,
27}
28
29impl DocumentMetadata {
30    #[must_use]
31    pub const fn with_tuple_xmin(tuple_xmin: u32) -> Self {
32        Self {
33            tuple_xmin: Some(tuple_xmin),
34        }
35    }
36
37    #[must_use]
38    pub const fn tuple_xmin(self) -> Option<u32> {
39        self.tuple_xmin
40    }
41}
42
43/// One persisted tuple, split into its public fields and storage-owned metadata.
44#[derive(Debug, Clone, Default, PartialEq)]
45pub struct StoredDocument {
46    fields: Document,
47    metadata: DocumentMetadata,
48}
49
50impl StoredDocument {
51    #[must_use]
52    pub fn new(fields: Document) -> Self {
53        Self {
54            fields,
55            metadata: DocumentMetadata::default(),
56        }
57    }
58
59    #[must_use]
60    pub fn with_metadata(fields: Document, metadata: DocumentMetadata) -> Self {
61        Self { fields, metadata }
62    }
63
64    #[must_use]
65    pub fn fields(&self) -> &Document {
66        &self.fields
67    }
68
69    #[must_use]
70    pub fn fields_mut(&mut self) -> &mut Document {
71        &mut self.fields
72    }
73
74    #[must_use]
75    pub fn metadata(&self) -> DocumentMetadata {
76        self.metadata
77    }
78
79    #[must_use]
80    pub fn into_fields(self) -> Document {
81        self.fields
82    }
83
84    #[must_use]
85    pub fn into_parts(self) -> (Document, DocumentMetadata) {
86        (self.fields, self.metadata)
87    }
88}
89
90const MISSING_SHARED_SLOT: usize = usize::MAX;
91static SHARED_NULL_VALUE: Value = Value::Null;
92
93/// A positional projection that shares an in-memory document's stored values.
94///
95/// Persistent backends still decode owned rows through the ordinary bulk
96/// methods. The memory backend exposes this optional representation so a
97/// physical scan can carry storage-owned values through joins without
98/// cloning them.
99#[derive(Debug, Clone, PartialEq)]
100pub struct SharedDocumentRow {
101    values: Arc<Vec<Value>>,
102    projection: Arc<[usize]>,
103}
104
105impl SharedDocumentRow {
106    pub(crate) fn new(values: Arc<Vec<Value>>, projection: Arc<[usize]>) -> Self {
107        debug_assert!(projection
108            .iter()
109            .all(|slot| *slot == MISSING_SHARED_SLOT || *slot < values.len()));
110        Self { values, projection }
111    }
112
113    /// Populate a reusable borrowed projection for predicate evaluation.
114    pub fn project<'a>(&'a self, output: &mut Vec<&'a Value>) {
115        output.clear();
116        output.extend(self.projection.iter().map(|slot| {
117            if *slot == MISSING_SHARED_SLOT {
118                &SHARED_NULL_VALUE
119            } else {
120                &self.values[*slot]
121            }
122        }));
123    }
124
125    /// Borrow the projection through an inline reference array. Relational
126    /// tables normally have far fewer than 32 requested fields, so predicate
127    /// evaluation needs no scratch allocation.
128    pub fn with_projected<R>(&self, visitor: impl FnOnce(&[&Value]) -> R) -> R {
129        const INLINE_FIELDS: usize = 32;
130        if self.projection.len() <= INLINE_FIELDS {
131            let mut projected = [&SHARED_NULL_VALUE; INLINE_FIELDS];
132            for (output, slot) in projected.iter_mut().zip(self.projection.iter()) {
133                if *slot != MISSING_SHARED_SLOT {
134                    *output = &self.values[*slot];
135                }
136            }
137            visitor(&projected[..self.projection.len()])
138        } else {
139            let projected = self
140                .projection
141                .iter()
142                .map(|slot| {
143                    if *slot == MISSING_SHARED_SLOT {
144                        &SHARED_NULL_VALUE
145                    } else {
146                        &self.values[*slot]
147                    }
148                })
149                .collect::<Vec<_>>();
150            visitor(&projected)
151        }
152    }
153
154    /// Borrow the storage-owned values and the projection from requested field positions to value slots. A `usize::MAX` projection slot represents a missing field and therefore SQL NULL.
155    pub fn indexed_values(&self) -> (&[Value], &[usize]) {
156        (&self.values, &self.projection)
157    }
158
159    /// Transfer the shared vector and its fragment-local projection into a
160    /// physical row without cloning either allocation.
161    pub fn into_parts(self) -> (Arc<Vec<Value>>, Arc<[usize]>) {
162        (self.values, self.projection)
163    }
164}
165
166/// Mutating methods are fallible: persistent backends surface their write failures so callers (engine DML, upserts, referential rewrites) can abort the enclosing transaction instead of silently committing a partially-applied statement. A rewrite that deletes a row and then fails to re-insert it must never look like success.
167pub trait DocumentStore: Send + Sync {
168    /// Persist one typed storage record. Every backend owns the physical representation of tuple metadata and must keep it outside the public field map.
169    fn put_stored(&mut self, doc_id: DocId, document: StoredDocument) -> StorageBackendResult<()>;
170
171    /// Read one typed storage record without projecting metadata into user fields.
172    fn get_stored(&self, doc_id: DocId) -> StorageBackendResult<Option<StoredDocument>>;
173
174    /// Replace public fields while preserving metadata already owned by the stored tuple. Engine code that creates a new tuple version must call [`DocumentStore::put_stored`] with the new metadata explicitly.
175    fn put(&mut self, doc_id: DocId, document: Document) -> StorageBackendResult<()> {
176        let metadata = self.get_metadata(doc_id)?.unwrap_or_default();
177        self.put_stored(doc_id, StoredDocument::with_metadata(document, metadata))
178    }
179
180    fn get(&self, doc_id: DocId) -> StorageBackendResult<Option<Document>> {
181        self.get_stored(doc_id)
182            .map(|document| document.map(StoredDocument::into_fields))
183    }
184
185    /// Bulk variant of [`DocumentStore::get_stored`].
186    fn get_stored_many(
187        &self,
188        doc_ids: &[DocId],
189    ) -> StorageBackendResult<BTreeMap<DocId, StoredDocument>> {
190        let mut out = BTreeMap::new();
191        for doc_id in doc_ids {
192            if let Some(document) = self.get_stored(*doc_id)? {
193                out.insert(*doc_id, document);
194            }
195        }
196        Ok(out)
197    }
198
199    /// Read one tuple's storage metadata without exposing it as a field.
200    fn get_metadata(&self, doc_id: DocId) -> StorageBackendResult<Option<DocumentMetadata>> {
201        self.get_stored(doc_id)
202            .map(|document| document.map(|document| document.metadata()))
203    }
204    fn contains_doc_id(&self, doc_id: DocId) -> StorageBackendResult<bool> {
205        Ok(self.get(doc_id)?.is_some())
206    }
207    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()>;
208    fn clear(&mut self) -> StorageBackendResult<()>;
209
210    /// Read a single field. Returns an owned [`Value`] so persistent
211    /// backends (`SQLite`, ...) can decode on demand without reaching
212    /// for a reference into a transient row.
213    fn get_field(&self, doc_id: DocId, field: &str) -> StorageBackendResult<Option<Value>> {
214        Ok(self
215            .get(doc_id)?
216            .and_then(|document| document.get(field).cloned()))
217    }
218
219    /// Find the first document whose top-level field equals `value`.
220    /// Persistent stores can override this with an indexed or JSON-path
221    /// lookup so point updates do not have to materialise every row.
222    fn find_doc_id_by_field(
223        &self,
224        field: &str,
225        value: &Value,
226    ) -> StorageBackendResult<Option<DocId>> {
227        for doc_id in self.doc_ids()? {
228            if self.get_field(doc_id, field)?.as_ref() == Some(value) {
229                return Ok(Some(doc_id));
230            }
231        }
232        Ok(None)
233    }
234
235    /// Apply top-level field updates without requiring callers to
236    /// materialise the whole document. `Value::Null` matches `put` by
237    /// removing the stored field. `Ok(false)` means the document does
238    /// not exist; write failures surface as `Err`.
239    fn patch_fields(
240        &mut self,
241        doc_id: DocId,
242        updates: &BTreeMap<String, Value>,
243    ) -> StorageBackendResult<bool> {
244        let Some(mut document) = self.get_stored(doc_id)? else {
245            return Ok(false);
246        };
247        for (field, value) in updates {
248            if matches!(value, Value::Null) {
249                document.fields_mut().remove(field);
250            } else {
251                document.fields_mut().insert(field.clone(), value.clone());
252            }
253        }
254        self.put_stored(doc_id, document)?;
255        Ok(true)
256    }
257
258    /// Bulk variant of [`DocumentStore::get`]. Ids without a stored
259    /// document are absent from the result. The default implementation
260    /// walks each id one at a time; persistent backends should
261    /// override to batch the reads into few queries.
262    fn get_many(&self, doc_ids: &[DocId]) -> StorageBackendResult<BTreeMap<DocId, Document>> {
263        let mut out = BTreeMap::new();
264        for doc_id in doc_ids {
265            if let Some(document) = self.get(*doc_id)? {
266                out.insert(*doc_id, document);
267            }
268        }
269        Ok(out)
270    }
271
272    /// Fetch several top-level fields for many documents. The result
273    /// vector is aligned with `fields`; missing fields come back as
274    /// [`Value::Null`], ids without a document are absent. Persistent
275    /// backends override this to extract all fields in one scan
276    /// instead of materialising whole documents.
277    fn get_fields_multi(
278        &self,
279        doc_ids: &[DocId],
280        fields: &[&str],
281    ) -> StorageBackendResult<BTreeMap<DocId, Vec<Value>>> {
282        let mut out = BTreeMap::new();
283        for doc_id in doc_ids {
284            let Some(document) = self.get(*doc_id)? else {
285                continue;
286            };
287            let values = fields
288                .iter()
289                .map(|field| document.get(*field).cloned().unwrap_or(Value::Null))
290                .collect();
291            out.insert(*doc_id, values);
292        }
293        Ok(out)
294    }
295
296    /// Visit a column projection in the caller's document-id order.
297    /// The callback receives one owned row at a time, allowing scan and
298    /// aggregate pipelines to avoid materialising a second doc-id map.
299    /// Returning `false` stops the visit early. Missing documents yield
300    /// a row of NULLs, matching row-evaluator semantics.
301    fn for_each_fields_multi(
302        &self,
303        doc_ids: &[DocId],
304        fields: &[&str],
305        visitor: &mut dyn FnMut(DocId, Vec<Value>) -> bool,
306    ) -> StorageBackendResult<()> {
307        let mut projected = self.get_fields_multi(doc_ids, fields)?;
308        for doc_id in doc_ids {
309            let values = projected
310                .remove(doc_id)
311                .unwrap_or_else(|| vec![Value::Null; fields.len()]);
312            if !visitor(*doc_id, values) {
313                break;
314            }
315        }
316        Ok(())
317    }
318
319    /// Visit a column projection by reference when the backend can keep
320    /// decoded values alive for the duration of the callback. The default
321    /// adapter preserves the backend's owned/batched projection path;
322    /// in-memory stores override it to avoid cloning every projected value.
323    fn for_each_fields_multi_ref(
324        &self,
325        doc_ids: &[DocId],
326        fields: &[&str],
327        visitor: &mut dyn FnMut(DocId, &[&Value]) -> bool,
328    ) -> StorageBackendResult<()> {
329        self.for_each_fields_multi(doc_ids, fields, &mut |doc_id, values| {
330            let references: Vec<&Value> = values.iter().collect();
331            visitor(doc_id, &references)
332        })
333    }
334
335    /// Visit a projection together with whether each requested document
336    /// actually exists. This avoids a separate `contains_doc_id` probe when a
337    /// caller must distinguish a missing document from an existing document
338    /// whose requested fields are all NULL.
339    fn for_each_fields_multi_ref_with_presence(
340        &self,
341        doc_ids: &[DocId],
342        fields: &[&str],
343        visitor: &mut dyn FnMut(DocId, bool, &[&Value]) -> bool,
344    ) -> StorageBackendResult<()> {
345        if fields.is_empty() {
346            for doc_id in doc_ids {
347                if !visitor(*doc_id, self.contains_doc_id(*doc_id)?, &[]) {
348                    break;
349                }
350            }
351            return Ok(());
352        }
353
354        let projected = self.get_fields_multi(doc_ids, fields)?;
355        let null = Value::Null;
356        let missing = vec![&null; fields.len()];
357        for doc_id in doc_ids {
358            let Some(values) = projected.get(doc_id) else {
359                if !visitor(*doc_id, false, &missing) {
360                    break;
361                }
362                continue;
363            };
364            let references = values.iter().collect::<Vec<_>>();
365            if !visitor(*doc_id, true, &references) {
366                break;
367            }
368        }
369        Ok(())
370    }
371
372    /// Return rows aligned with `doc_ids` as shared positional projections
373    /// when the backend owns stable decoded value vectors. `None` means the
374    /// backend does not support zero-copy projection; entries inside the
375    /// returned vector are `None` only for missing document ids.
376    fn get_shared_fields(
377        &self,
378        _doc_ids: &[DocId],
379        _fields: &[&str],
380    ) -> StorageBackendResult<Option<Vec<Option<SharedDocumentRow>>>> {
381        Ok(None)
382    }
383
384    /// Bulk variant of [`DocumentStore::get_field`]. The default
385    /// implementation walks each id one at a time; persistent backends
386    /// should override to run a single batched query.
387    fn get_fields_bulk(
388        &self,
389        doc_ids: &[DocId],
390        field: &str,
391    ) -> StorageBackendResult<BTreeMap<DocId, Value>> {
392        let mut out = BTreeMap::new();
393        for doc_id in doc_ids {
394            out.insert(
395                *doc_id,
396                self.get_field(*doc_id, field)?.unwrap_or(Value::Null),
397            );
398        }
399        Ok(out)
400    }
401
402    /// Return `true` if any document has `field == value`.
403    fn has_value(&self, field: &str, value: &Value) -> StorageBackendResult<bool> {
404        for doc_id in self.doc_ids()? {
405            if self.get_field(doc_id, field)?.as_ref() == Some(value) {
406                return Ok(true);
407            }
408        }
409        Ok(false)
410    }
411
412    /// Find the first document whose top-level fields match every
413    /// requested value.
414    fn find_doc_id_by_fields(
415        &self,
416        fields: &[String],
417        values: &[Value],
418    ) -> StorageBackendResult<Option<DocId>> {
419        if fields.is_empty() || fields.len() != values.len() {
420            return Ok(None);
421        }
422        for doc_id in self.doc_ids()? {
423            let mut matches = true;
424            for (field, value) in fields.iter().zip(values) {
425                if self.get_field(doc_id, field)?.unwrap_or(Value::Null) != *value {
426                    matches = false;
427                    break;
428                }
429            }
430            if matches {
431                return Ok(Some(doc_id));
432            }
433        }
434        Ok(None)
435    }
436
437    /// Evaluate a hierarchical path expression against a document.
438    fn eval_path(
439        &self,
440        doc_id: DocId,
441        path: &[PathSegment],
442    ) -> StorageBackendResult<Option<Value>> {
443        let Some(document) = self.get(doc_id)? else {
444            return Ok(None);
445        };
446        Ok(eval_path_in_document(&document, path))
447    }
448
449    fn doc_ids(&self) -> StorageBackendResult<Vec<DocId>>;
450
451    /// Return the first stored document id strictly greater than `after`, or
452    /// the first id when `after` is `None`. Scan operators use this cursor API
453    /// so a full table scan does not need a cardinality-sized id vector before
454    /// it can yield its first row.
455    fn next_doc_id(&self, after: Option<DocId>) -> StorageBackendResult<Option<DocId>> {
456        Ok(self
457            .doc_ids()?
458            .into_iter()
459            .filter(|doc_id| after.is_none_or(|after| *doc_id > after))
460            .min())
461    }
462
463    /// Return up to `limit` document ids strictly greater than `after`, in
464    /// ascending order. Scan operators use this bounded cursor instead of
465    /// reacquiring their store lock and issuing one backend lookup per row.
466    fn next_doc_ids(&self, after: Option<DocId>, limit: usize) -> StorageBackendResult<Vec<DocId>> {
467        if limit == 0 {
468            return Ok(Vec::new());
469        }
470        let mut doc_ids = self.doc_ids()?;
471        doc_ids.sort_unstable();
472        Ok(doc_ids
473            .into_iter()
474            .filter(|doc_id| after.is_none_or(|after| *doc_id > after))
475            .take(limit)
476            .collect())
477    }
478
479    /// Return the next bounded id range and its shared positional projections
480    /// in one storage traversal when stable decoded rows are available.
481    /// `None` lets persistent backends use the ordinary id + projection path.
482    fn next_shared_fields(
483        &self,
484        _after: Option<DocId>,
485        _limit: usize,
486        _fields: &[&str],
487    ) -> StorageBackendResult<Option<Vec<(DocId, SharedDocumentRow)>>> {
488        Ok(None)
489    }
490
491    /// Visit the next bounded id range through a reusable borrowed projection of each backend-owned row. Missing fields are exposed as SQL NULL. `Some(count)` means the backend supports this borrowed cursor and reports how many rows it visited; `None` selects the ordinary cursor path without invoking `visitor`.
492    fn for_each_next_fields(
493        &self,
494        _after: Option<DocId>,
495        _limit: usize,
496        _fields: &[&str],
497        _visitor: &mut dyn FnMut(DocId, &[&Value]) -> bool,
498    ) -> StorageBackendResult<Option<usize>> {
499        Ok(None)
500    }
501
502    fn max_doc_id(&self) -> StorageBackendResult<DocId> {
503        Ok(self.doc_ids()?.into_iter().max().unwrap_or(0))
504    }
505
506    fn len(&self) -> StorageBackendResult<usize>;
507
508    fn is_empty(&self) -> StorageBackendResult<bool> {
509        Ok(self.len()? == 0)
510    }
511
512    /// Iterate over `(doc_id, document)` pairs in id order. The default
513    /// implementation fetches each document individually; SQLite-backed
514    /// stores override with a single query.
515    fn iter_all(&self) -> StorageBackendResult<Box<dyn Iterator<Item = (DocId, Document)> + '_>> {
516        let mut ids = self.doc_ids()?;
517        ids.sort_unstable();
518        let snapshot = self.snapshot()?;
519        let mut rows = Vec::with_capacity(ids.len());
520        for doc_id in ids {
521            if let Some(document) = snapshot.get(doc_id)? {
522                rows.push((doc_id, document));
523            }
524        }
525        Ok(Box::new(rows.into_iter()))
526    }
527
528    /// Read-only handle suitable for an `ExecutionContext`. Persistent
529    /// backends share their connection; memory backends deep-clone so the
530    /// snapshot is isolated from later mutations.
531    fn snapshot(&self) -> StorageBackendResult<Arc<dyn DocumentStore>>;
532
533    /// Independent writable copy used by the in-memory engine transaction
534    /// rollback path. Persistent engines restore through their backend
535    /// transaction and need not implement this operation.
536    fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn DocumentStore>> {
537        Err(StorageBackendError::Other(
538            "writable document-store snapshots are not supported by this backend".into(),
539        ))
540    }
541}
542
543/// Walk a document along a [`PathSegment`] sequence: strings descend into
544/// maps, integers descend into lists, and the implicit array-wildcard rule
545/// applies a string component over every map element of an array.
546pub fn eval_path_in_document(doc: &Document, path: &[PathSegment]) -> Option<Value> {
547    let mut current: Value = match path.first()? {
548        PathSegment::Key(k) => doc.get(k)?.clone(),
549        PathSegment::Index(_) => return None,
550    };
551    for seg in path.iter().skip(1) {
552        current = match (current, seg) {
553            (Value::Map(m), PathSegment::Key(k)) => m.get(k)?.clone(),
554            (Value::List(items), PathSegment::Index(i)) => items.get(*i)?.clone(),
555            (Value::List(items), PathSegment::Key(k)) => {
556                let collected: Vec<Value> = items
557                    .into_iter()
558                    .filter_map(|v| match v {
559                        Value::Map(m) => m.get(k).cloned(),
560                        _ => None,
561                    })
562                    .collect();
563                Value::List(collected)
564            }
565            _ => return None,
566        };
567    }
568    Some(current)
569}
570
571mod memory;
572
573pub use memory::MemoryDocumentStore;
uqa_storage/document_store.rs

uqa_storage/
document_store.rs
