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 for_each_fields_multi_borrowed(
284        &self,
285        ids: &[DocId],
286        fields: &[&str],
287        visitor: &mut dyn FnMut(DocId, bool, &[&Value]) -> bool,
288    ) -> StorageBackendResult<Option<usize>> {
289        let mut visited = 0;
290        self.visit_fields_multi_ref_with_presence(ids, fields, &mut |id, present, values| {
291            visited += 1;
292            visitor(id, present, values)
293        });
294        Ok(Some(visited))
295    }
296
297    fn with_field_ref_controlled(
298        &self,
299        id: DocId,
300        field: &str,
301        control: &StorageReadControl,
302        visitor: &mut dyn FnMut(Option<&Value>) -> StorageBackendResult<()>,
303    ) -> StorageBackendResult<()> {
304        control.check()?;
305        let value = self
306            .state
307            .documents
308            .get(&id)
309            .and_then(|row| self.field(row, field));
310        visitor(value)?;
311        control.check()
312    }
313
314    fn field_presence_controlled(
315        &self,
316        ids: &[DocId],
317        fields: &[&str],
318        control: &StorageReadControl,
319    ) -> StorageBackendResult<BudgetedVec<bool>> {
320        control.check()?;
321        let mut present = BudgetedVec::new(control.memory());
322        for id in ids {
323            let row = self.state.documents.get(id);
324            for field in fields {
325                control.check()?;
326                present.push(row.is_some_and(|row| self.field(row, field).is_some()))?;
327            }
328        }
329        control.check()?;
330        Ok(present)
331    }
332
333    fn get_shared_fields(
334        &self,
335        doc_ids: &[DocId],
336        fields: &[&str],
337    ) -> StorageBackendResult<Option<Vec<Option<SharedDocumentRow>>>> {
338        Ok(Some(self.shared_fields(doc_ids, fields)))
339    }
340
341    fn find_doc_id_by_field(
342        &self,
343        field: &str,
344        value: &Value,
345    ) -> StorageBackendResult<Option<DocId>> {
346        Ok(self.state.documents.iter().find_map(|(doc_id, stored)| {
347            (self.field(stored, field) == Some(value)).then_some(*doc_id)
348        }))
349    }
350
351    fn find_doc_id_by_fields(
352        &self,
353        fields: &[String],
354        values: &[Value],
355    ) -> StorageBackendResult<Option<DocId>> {
356        if fields.is_empty() || fields.len() != values.len() {
357            return Ok(None);
358        }
359        Ok(self.state.documents.iter().find_map(|(doc_id, stored)| {
360            fields
361                .iter()
362                .zip(values.iter())
363                .all(|(field, value)| self.field(stored, field).unwrap_or(&Value::Null) == value)
364                .then_some(*doc_id)
365        }))
366    }
367
368    fn patch_fields(
369        &mut self,
370        doc_id: DocId,
371        updates: &BTreeMap<String, Value>,
372    ) -> StorageBackendResult<bool> {
373        let Some(mut document) = self
374            .state
375            .documents
376            .get(&doc_id)
377            .map(|stored| self.materialize_stored_document(stored))
378        else {
379            return Ok(false);
380        };
381        for (field, value) in updates {
382            if matches!(value, Value::Null) {
383                document.fields_mut().remove(field);
384            } else {
385                document.fields_mut().insert(field.clone(), value.clone());
386            }
387        }
388        self.put_stored(doc_id, document)?;
389        Ok(true)
390    }
391
392    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
393        if self.state.documents.contains_key(&doc_id) {
394            Arc::make_mut(&mut self.state).documents.remove(&doc_id);
395        }
396        Ok(())
397    }
398
399    fn clear(&mut self) -> StorageBackendResult<()> {
400        if let Some(state) = Arc::get_mut(&mut self.state) {
401            *state = MemoryDocumentState::default();
402        } else {
403            self.state = Arc::default();
404        }
405        Ok(())
406    }
407
408    fn doc_ids(&self) -> StorageBackendResult<Vec<DocId>> {
409        Ok(self.state.documents.keys().copied().collect())
410    }
411
412    fn next_doc_id(&self, after: Option<DocId>) -> StorageBackendResult<Option<DocId>> {
413        use std::ops::Bound::{Excluded, Unbounded};
414
415        Ok(match after {
416            Some(after) => self
417                .state
418                .documents
419                .range((Excluded(after), Unbounded))
420                .next()
421                .map(|(doc_id, _)| *doc_id),
422            None => self.state.documents.keys().next().copied(),
423        })
424    }
425
426    fn next_doc_ids(&self, after: Option<DocId>, limit: usize) -> StorageBackendResult<Vec<DocId>> {
427        use std::ops::Bound::{Excluded, Unbounded};
428
429        if limit == 0 {
430            return Ok(Vec::new());
431        }
432        Ok(match after {
433            Some(after) => self
434                .state
435                .documents
436                .range((Excluded(after), Unbounded))
437                .take(limit)
438                .map(|(doc_id, _)| *doc_id)
439                .collect(),
440            None => self.state.documents.keys().take(limit).copied().collect(),
441        })
442    }
443
444    fn next_doc_ids_controlled(
445        &self,
446        after: Option<DocId>,
447        limit: usize,
448        control: &StorageReadControl,
449    ) -> StorageBackendResult<BudgetedVec<DocId>> {
450        use std::ops::Bound::{Excluded, Unbounded};
451
452        control.check()?;
453        let mut ids = BudgetedVec::new(control.memory());
454        let after = after.map_or(Unbounded, Excluded);
455        for (&id, _) in self.state.documents.range((after, Unbounded)).take(limit) {
456            control.check()?;
457            ids.push(id)?;
458        }
459        control.check()?;
460        Ok(ids)
461    }
462
463    fn next_shared_fields(
464        &self,
465        after: Option<DocId>,
466        limit: usize,
467        fields: &[&str],
468    ) -> StorageBackendResult<Option<Vec<(DocId, SharedDocumentRow)>>> {
469        Ok(Some(self.next_shared_rows(after, limit, fields)))
470    }
471
472    fn for_each_next_fields(
473        &self,
474        after: Option<DocId>,
475        limit: usize,
476        fields: &[&str],
477        visitor: &mut dyn FnMut(DocId, &[&Value]) -> bool,
478    ) -> StorageBackendResult<Option<usize>> {
479        Ok(Some(self.visit_next_rows(after, limit, fields, visitor)))
480    }
481
482    fn len(&self) -> StorageBackendResult<usize> {
483        Ok(self.state.documents.len())
484    }
485
486    fn snapshot(&self) -> StorageBackendResult<Arc<dyn DocumentStore>> {
487        Ok(Arc::new(self.shared_snapshot()))
488    }
489
490    fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn DocumentStore>> {
491        Ok(Box::new(self.shared_snapshot()))
492    }
493}
494
495#[cfg(test)]
496mod tests;
uqa_storage/document_store/memory.rs

uqa_storage/document_store/
memory.rs
