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

uqa_storage/
document_store.rs
