uqa_storage/document_store/
memory.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! In-memory document storage and borrowed projection scans.
8
9use std::collections::BTreeMap;
10use std::sync::Arc;
11
12use uqa_core::{memory::BudgetedVec, DocId, Value};
13
14use crate::backend::StorageBackendResult;
15use crate::read_control::StorageReadControl;
16
17use super::{Document, DocumentMetadata, DocumentStore, SharedDocumentRow, StoredDocument};
18
19mod projection;
20
21#[derive(Debug, Default)]
22pub struct MemoryDocumentStore {
23    /// Rows and their layout IDs advance together when a shared snapshot is changed.
24    state: Arc<MemoryDocumentState>,
25}
26
27/// Cloning retains independent row maps and layouts; snapshot APIs share them until a write.
28impl Clone for MemoryDocumentStore {
29    fn clone(&self) -> Self {
30        Self {
31            state: Arc::new((*self.state).clone()),
32        }
33    }
34}
35
36#[derive(Debug, Default, Clone)]
37struct MemoryDocumentState {
38    documents: BTreeMap<DocId, MemoryDocumentRow>,
39    layouts: Vec<Vec<String>>,
40}
41
42#[derive(Debug, Clone)]
43pub(super) struct MemoryDocumentRow {
44    layout_id: usize,
45    values: Arc<Vec<Value>>,
46    metadata: DocumentMetadata,
47}
48
49impl MemoryDocumentStore {
50    pub fn new() -> Self {
51        Self::default()
52    }
53
54    fn shared_snapshot(&self) -> Self {
55        Self {
56            state: Arc::clone(&self.state),
57        }
58    }
59
60    pub fn iter(&self) -> impl Iterator<Item = (DocId, Document)> + '_ {
61        self.state
62            .documents
63            .iter()
64            .map(|(doc_id, stored)| (*doc_id, self.materialize_document(stored)))
65    }
66
67    fn materialize_document(&self, stored: &MemoryDocumentRow) -> Document {
68        self.state.layouts[stored.layout_id]
69            .iter()
70            .cloned()
71            .zip(stored.values.iter().cloned())
72            .collect()
73    }
74
75    fn materialize_stored_document(&self, stored: &MemoryDocumentRow) -> StoredDocument {
76        StoredDocument::with_metadata(self.materialize_document(stored), stored.metadata)
77    }
78
79    fn field<'a>(&'a self, stored: &'a MemoryDocumentRow, field: &str) -> Option<&'a Value> {
80        let slot = self.state.layouts[stored.layout_id]
81            .binary_search_by(|stored| stored.as_str().cmp(field))
82            .ok()?;
83        stored.values.get(slot)
84    }
85
86    fn put_stored_inner(&mut self, doc_id: DocId, document: StoredDocument) {
87        let state = Arc::make_mut(&mut self.state);
88        let (document, metadata) = document.into_parts();
89        let (layout_id, values) =
90            if let Some(layout_id) = document_layout_id(&state.layouts, &document) {
91                (layout_id, document.into_values().collect())
92            } else {
93                let (layout, values): (Vec<_>, Vec<_>) = document.into_iter().unzip();
94                let layout_id = state.layouts.len();
95                state.layouts.push(layout);
96                (layout_id, values)
97            };
98        state.documents.insert(
99            doc_id,
100            MemoryDocumentRow {
101                layout_id,
102                values: Arc::new(values),
103                metadata,
104            },
105        );
106    }
107}
108
109fn document_layout_id(layouts: &[Vec<String>], document: &Document) -> Option<usize> {
110    layouts.iter().position(|layout| {
111        layout.len() == document.len()
112            && layout
113                .iter()
114                .map(String::as_str)
115                .eq(document.keys().map(String::as_str))
116    })
117}
118
119impl DocumentStore for MemoryDocumentStore {
120    fn put(&mut self, doc_id: DocId, document: Document) -> StorageBackendResult<()> {
121        let metadata = self
122            .state
123            .documents
124            .get(&doc_id)
125            .map_or_else(DocumentMetadata::default, |stored| stored.metadata);
126        self.put_stored_inner(doc_id, StoredDocument::with_metadata(document, metadata));
127        Ok(())
128    }
129
130    fn get(&self, doc_id: DocId) -> StorageBackendResult<Option<Document>> {
131        Ok(self
132            .state
133            .documents
134            .get(&doc_id)
135            .map(|stored| self.materialize_document(stored)))
136    }
137
138    fn put_stored(&mut self, doc_id: DocId, document: StoredDocument) -> StorageBackendResult<()> {
139        self.put_stored_inner(doc_id, document);
140        Ok(())
141    }
142
143    fn get_stored(&self, doc_id: DocId) -> StorageBackendResult<Option<StoredDocument>> {
144        Ok(self
145            .state
146            .documents
147            .get(&doc_id)
148            .map(|stored| self.materialize_stored_document(stored)))
149    }
150
151    fn get_stored_many_controlled(
152        &self,
153        ids: &[DocId],
154        control: &StorageReadControl,
155    ) -> StorageBackendResult<super::RetainedDocumentPage> {
156        control.check()?;
157        let mut page = BudgetedVec::new(control.memory());
158        page.reserve(ids.len())?;
159        for id in ids {
160            control.check()?;
161            let row = self
162                .state
163                .documents
164                .get(id)
165                .map(|row| -> StorageBackendResult<_> {
166                    let fields = super::controlled_rows::copy_fields(
167                        self.state.layouts[row.layout_id]
168                            .iter()
169                            .map(String::as_str)
170                            .zip(row.values.iter()),
171                        control,
172                    )?;
173                    Ok(super::RetainedStoredDocument::with_metadata(
174                        fields,
175                        row.metadata,
176                    ))
177                })
178                .transpose()?;
179            page.push(row)?;
180        }
181        control.check()?;
182        Ok(page)
183    }
184
185    fn get_stored_many(
186        &self,
187        doc_ids: &[DocId],
188    ) -> StorageBackendResult<BTreeMap<DocId, StoredDocument>> {
189        Ok(doc_ids
190            .iter()
191            .filter_map(|doc_id| {
192                self.state
193                    .documents
194                    .get(doc_id)
195                    .map(|stored| (*doc_id, self.materialize_stored_document(stored)))
196            })
197            .collect())
198    }
199
200    fn get_metadata(&self, doc_id: DocId) -> StorageBackendResult<Option<DocumentMetadata>> {
201        Ok(self
202            .state
203            .documents
204            .get(&doc_id)
205            .map(|stored| stored.metadata))
206    }
207
208    fn contains_doc_id(&self, doc_id: DocId) -> StorageBackendResult<bool> {
209        Ok(self.state.documents.contains_key(&doc_id))
210    }
211
212    fn get_field(&self, doc_id: DocId, field: &str) -> StorageBackendResult<Option<Value>> {
213        Ok(self
214            .state
215            .documents
216            .get(&doc_id)
217            .and_then(|stored| self.field(stored, field).cloned()))
218    }
219
220    fn get_fields_multi(
221        &self,
222        doc_ids: &[DocId],
223        fields: &[&str],
224    ) -> StorageBackendResult<BTreeMap<DocId, Vec<Value>>> {
225        Ok(doc_ids
226            .iter()
227            .filter_map(|doc_id| {
228                let stored = self.state.documents.get(doc_id)?;
229                let values = fields
230                    .iter()
231                    .map(|field| self.field(stored, field).cloned().unwrap_or(Value::Null))
232                    .collect();
233                Some((*doc_id, values))
234            })
235            .collect())
236    }
237
238    fn for_each_fields_multi(
239        &self,
240        doc_ids: &[DocId],
241        fields: &[&str],
242        visitor: &mut dyn FnMut(DocId, Vec<Value>) -> bool,
243    ) -> StorageBackendResult<()> {
244        for doc_id in doc_ids {
245            let values = self.state.documents.get(doc_id).map_or_else(
246                || vec![Value::Null; fields.len()],
247                |stored| {
248                    fields
249                        .iter()
250                        .map(|field| self.field(stored, field).cloned().unwrap_or(Value::Null))
251                        .collect()
252                },
253            );
254            if !visitor(*doc_id, values) {
255                break;
256            }
257        }
258        Ok(())
259    }
260
261    fn for_each_fields_multi_ref(
262        &self,
263        doc_ids: &[DocId],
264        fields: &[&str],
265        visitor: &mut dyn FnMut(DocId, &[&Value]) -> bool,
266    ) -> StorageBackendResult<()> {
267        self.visit_fields_multi_ref_with_presence(doc_ids, fields, &mut |doc_id, _, values| {
268            visitor(doc_id, values)
269        });
270        Ok(())
271    }
272
273    fn for_each_fields_multi_ref_with_presence(
274        &self,
275        doc_ids: &[DocId],
276        fields: &[&str],
277        visitor: &mut dyn FnMut(DocId, bool, &[&Value]) -> bool,
278    ) -> StorageBackendResult<()> {
279        self.visit_fields_multi_ref_with_presence(doc_ids, fields, visitor);
280        Ok(())
281    }
282
283    fn with_field_ref_controlled(
284        &self,
285        id: DocId,
286        field: &str,
287        control: &StorageReadControl,
288        visitor: &mut dyn FnMut(Option<&Value>) -> StorageBackendResult<()>,
289    ) -> StorageBackendResult<()> {
290        control.check()?;
291        let value = self
292            .state
293            .documents
294            .get(&id)
295            .and_then(|row| self.field(row, field));
296        visitor(value)?;
297        control.check()
298    }
299
300    fn field_presence_controlled(
301        &self,
302        ids: &[DocId],
303        fields: &[&str],
304        control: &StorageReadControl,
305    ) -> StorageBackendResult<BudgetedVec<bool>> {
306        control.check()?;
307        let mut present = BudgetedVec::new(control.memory());
308        for id in ids {
309            let row = self.state.documents.get(id);
310            for field in fields {
311                control.check()?;
312                present.push(row.is_some_and(|row| self.field(row, field).is_some()))?;
313            }
314        }
315        control.check()?;
316        Ok(present)
317    }
318
319    fn get_shared_fields(
320        &self,
321        doc_ids: &[DocId],
322        fields: &[&str],
323    ) -> StorageBackendResult<Option<Vec<Option<SharedDocumentRow>>>> {
324        Ok(Some(self.shared_fields(doc_ids, fields)))
325    }
326
327    fn find_doc_id_by_field(
328        &self,
329        field: &str,
330        value: &Value,
331    ) -> StorageBackendResult<Option<DocId>> {
332        Ok(self.state.documents.iter().find_map(|(doc_id, stored)| {
333            (self.field(stored, field) == Some(value)).then_some(*doc_id)
334        }))
335    }
336
337    fn find_doc_id_by_fields(
338        &self,
339        fields: &[String],
340        values: &[Value],
341    ) -> StorageBackendResult<Option<DocId>> {
342        if fields.is_empty() || fields.len() != values.len() {
343            return Ok(None);
344        }
345        Ok(self.state.documents.iter().find_map(|(doc_id, stored)| {
346            fields
347                .iter()
348                .zip(values.iter())
349                .all(|(field, value)| self.field(stored, field).unwrap_or(&Value::Null) == value)
350                .then_some(*doc_id)
351        }))
352    }
353
354    fn patch_fields(
355        &mut self,
356        doc_id: DocId,
357        updates: &BTreeMap<String, Value>,
358    ) -> StorageBackendResult<bool> {
359        let Some(mut document) = self
360            .state
361            .documents
362            .get(&doc_id)
363            .map(|stored| self.materialize_stored_document(stored))
364        else {
365            return Ok(false);
366        };
367        for (field, value) in updates {
368            if matches!(value, Value::Null) {
369                document.fields_mut().remove(field);
370            } else {
371                document.fields_mut().insert(field.clone(), value.clone());
372            }
373        }
374        self.put_stored(doc_id, document)?;
375        Ok(true)
376    }
377
378    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
379        if self.state.documents.contains_key(&doc_id) {
380            Arc::make_mut(&mut self.state).documents.remove(&doc_id);
381        }
382        Ok(())
383    }
384
385    fn clear(&mut self) -> StorageBackendResult<()> {
386        if let Some(state) = Arc::get_mut(&mut self.state) {
387            *state = MemoryDocumentState::default();
388        } else {
389            self.state = Arc::default();
390        }
391        Ok(())
392    }
393
394    fn doc_ids(&self) -> StorageBackendResult<Vec<DocId>> {
395        Ok(self.state.documents.keys().copied().collect())
396    }
397
398    fn next_doc_id(&self, after: Option<DocId>) -> StorageBackendResult<Option<DocId>> {
399        use std::ops::Bound::{Excluded, Unbounded};
400
401        Ok(match after {
402            Some(after) => self
403                .state
404                .documents
405                .range((Excluded(after), Unbounded))
406                .next()
407                .map(|(doc_id, _)| *doc_id),
408            None => self.state.documents.keys().next().copied(),
409        })
410    }
411
412    fn next_doc_ids(&self, after: Option<DocId>, limit: usize) -> StorageBackendResult<Vec<DocId>> {
413        use std::ops::Bound::{Excluded, Unbounded};
414
415        if limit == 0 {
416            return Ok(Vec::new());
417        }
418        Ok(match after {
419            Some(after) => self
420                .state
421                .documents
422                .range((Excluded(after), Unbounded))
423                .take(limit)
424                .map(|(doc_id, _)| *doc_id)
425                .collect(),
426            None => self.state.documents.keys().take(limit).copied().collect(),
427        })
428    }
429
430    fn next_doc_ids_controlled(
431        &self,
432        after: Option<DocId>,
433        limit: usize,
434        control: &StorageReadControl,
435    ) -> StorageBackendResult<BudgetedVec<DocId>> {
436        use std::ops::Bound::{Excluded, Unbounded};
437
438        control.check()?;
439        let mut ids = BudgetedVec::new(control.memory());
440        let after = after.map_or(Unbounded, Excluded);
441        for (&id, _) in self.state.documents.range((after, Unbounded)).take(limit) {
442            control.check()?;
443            ids.push(id)?;
444        }
445        control.check()?;
446        Ok(ids)
447    }
448
449    fn next_shared_fields(
450        &self,
451        after: Option<DocId>,
452        limit: usize,
453        fields: &[&str],
454    ) -> StorageBackendResult<Option<Vec<(DocId, SharedDocumentRow)>>> {
455        Ok(Some(self.next_shared_rows(after, limit, fields)))
456    }
457
458    fn for_each_next_fields(
459        &self,
460        after: Option<DocId>,
461        limit: usize,
462        fields: &[&str],
463        visitor: &mut dyn FnMut(DocId, &[&Value]) -> bool,
464    ) -> StorageBackendResult<Option<usize>> {
465        Ok(Some(self.visit_next_rows(after, limit, fields, visitor)))
466    }
467
468    fn len(&self) -> StorageBackendResult<usize> {
469        Ok(self.state.documents.len())
470    }
471
472    fn snapshot(&self) -> StorageBackendResult<Arc<dyn DocumentStore>> {
473        Ok(Arc::new(self.shared_snapshot()))
474    }
475
476    fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn DocumentStore>> {
477        Ok(Box::new(self.shared_snapshot()))
478    }
479}
480
481#[cfg(test)]
482mod tests;
uqa_storage/document_store/memory.rs

uqa_storage/document_store/
memory.rs
