1use std::collections::BTreeMap;
14use std::sync::Arc;
15
16use uqa_core::{memory::BudgetedVec, DocId, FieldName, PathSegment, Value};
17
18use crate::backend::{StorageBackendError, StorageBackendResult};
19use crate::read_control::StorageReadControl;
20
21mod controlled_ids;
22pub use controlled_ids::read_document_ids;
23mod controlled_field;
24pub use controlled_field::read_selected_field;
25mod controlled_rows;
26pub use controlled_rows::{read_stored_documents, RetainedDocumentPage};
27mod controlled_presence;
28pub use controlled_presence::read_field_presence;
29pub mod decoding;
30pub mod identifiers;
31mod retained;
32pub use retained::{RetainedDocumentFields, RetainedStoredDocument};
33mod retained_store;
34pub use retained_store::{RetainedDocumentStore, RetainedDocumentStoreBuilder};
35
36pub type Document = BTreeMap<FieldName, Value>;
38
39#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
41pub struct DocumentMetadata {
42 tuple_xmin: Option<u32>,
43}
44
45impl DocumentMetadata {
46 #[must_use]
47 pub const fn with_tuple_xmin(tuple_xmin: u32) -> Self {
48 Self {
49 tuple_xmin: Some(tuple_xmin),
50 }
51 }
52
53 #[must_use]
54 pub const fn tuple_xmin(self) -> Option<u32> {
55 self.tuple_xmin
56 }
57}
58
59#[derive(Debug, Clone, Default, PartialEq)]
61pub struct StoredDocument {
62 fields: Document,
63 metadata: DocumentMetadata,
64}
65
66impl StoredDocument {
67 #[must_use]
68 pub fn new(fields: Document) -> Self {
69 Self {
70 fields,
71 metadata: DocumentMetadata::default(),
72 }
73 }
74
75 #[must_use]
76 pub fn with_metadata(fields: Document, metadata: DocumentMetadata) -> Self {
77 Self { fields, metadata }
78 }
79
80 #[must_use]
81 pub fn fields(&self) -> &Document {
82 &self.fields
83 }
84
85 #[must_use]
86 pub fn fields_mut(&mut self) -> &mut Document {
87 &mut self.fields
88 }
89
90 #[must_use]
91 pub fn metadata(&self) -> DocumentMetadata {
92 self.metadata
93 }
94
95 #[must_use]
96 pub fn into_fields(self) -> Document {
97 self.fields
98 }
99
100 #[must_use]
101 pub fn into_parts(self) -> (Document, DocumentMetadata) {
102 (self.fields, self.metadata)
103 }
104}
105
106const MISSING_SHARED_SLOT: usize = usize::MAX;
107static SHARED_NULL_VALUE: Value = Value::Null;
108
109#[derive(Debug, Clone, PartialEq)]
116pub struct SharedDocumentRow {
117 values: Arc<Vec<Value>>,
118 projection: Arc<[usize]>,
119}
120
121impl SharedDocumentRow {
122 pub(crate) fn new(values: Arc<Vec<Value>>, projection: Arc<[usize]>) -> Self {
123 debug_assert!(projection
124 .iter()
125 .all(|slot| *slot == MISSING_SHARED_SLOT || *slot < values.len()));
126 Self { values, projection }
127 }
128
129 pub fn project<'a>(&'a self, output: &mut Vec<&'a Value>) {
131 output.clear();
132 output.extend(self.projection.iter().map(|slot| {
133 if *slot == MISSING_SHARED_SLOT {
134 &SHARED_NULL_VALUE
135 } else {
136 &self.values[*slot]
137 }
138 }));
139 }
140
141 pub fn with_projected<R>(&self, visitor: impl FnOnce(&[&Value]) -> R) -> R {
145 const INLINE_FIELDS: usize = 32;
146 if self.projection.len() <= INLINE_FIELDS {
147 let mut projected = [&SHARED_NULL_VALUE; INLINE_FIELDS];
148 for (output, slot) in projected.iter_mut().zip(self.projection.iter()) {
149 if *slot != MISSING_SHARED_SLOT {
150 *output = &self.values[*slot];
151 }
152 }
153 visitor(&projected[..self.projection.len()])
154 } else {
155 let projected = self
156 .projection
157 .iter()
158 .map(|slot| {
159 if *slot == MISSING_SHARED_SLOT {
160 &SHARED_NULL_VALUE
161 } else {
162 &self.values[*slot]
163 }
164 })
165 .collect::<Vec<_>>();
166 visitor(&projected)
167 }
168 }
169
170 pub fn indexed_values(&self) -> (&[Value], &[usize]) {
172 (&self.values, &self.projection)
173 }
174
175 pub fn into_parts(self) -> (Arc<Vec<Value>>, Arc<[usize]>) {
178 (self.values, self.projection)
179 }
180}
181
182pub trait DocumentStore: Send + Sync {
184 fn put_stored(&mut self, doc_id: DocId, document: StoredDocument) -> StorageBackendResult<()>;
186
187 fn get_stored(&self, doc_id: DocId) -> StorageBackendResult<Option<StoredDocument>>;
189
190 fn get_stored_many_controlled(
192 &self,
193 doc_ids: &[DocId],
194 control: &StorageReadControl,
195 ) -> StorageBackendResult<RetainedDocumentPage> {
196 control.check()?;
197 if doc_ids.is_empty() {
198 return Ok(BudgetedVec::new(control.memory()));
199 }
200 Err(StorageBackendError::Other(
201 "controlled whole-document reads are not supported by this store".into(),
202 ))
203 }
204
205 fn field_presence_controlled(
207 &self,
208 doc_ids: &[DocId],
209 fields: &[&str],
210 control: &StorageReadControl,
211 ) -> StorageBackendResult<BudgetedVec<bool>> {
212 controlled_presence::from_controlled_rows(self, doc_ids, fields, control)
213 }
214
215 fn with_field_ref_controlled(
217 &self,
218 doc_id: DocId,
219 field: &str,
220 control: &StorageReadControl,
221 visitor: &mut dyn FnMut(Option<&Value>) -> StorageBackendResult<()>,
222 ) -> StorageBackendResult<()> {
223 controlled_field::visit_selected_field(self, doc_id, field, control, visitor)
224 }
225
226 fn put(&mut self, doc_id: DocId, document: Document) -> StorageBackendResult<()> {
228 let metadata = self.get_metadata(doc_id)?.unwrap_or_default();
229 self.put_stored(doc_id, StoredDocument::with_metadata(document, metadata))
230 }
231
232 fn get(&self, doc_id: DocId) -> StorageBackendResult<Option<Document>> {
233 self.get_stored(doc_id)
234 .map(|document| document.map(StoredDocument::into_fields))
235 }
236
237 fn get_stored_many(
239 &self,
240 doc_ids: &[DocId],
241 ) -> StorageBackendResult<BTreeMap<DocId, StoredDocument>> {
242 let mut out = BTreeMap::new();
243 for doc_id in doc_ids {
244 if let Some(document) = self.get_stored(*doc_id)? {
245 out.insert(*doc_id, document);
246 }
247 }
248 Ok(out)
249 }
250
251 fn get_metadata(&self, doc_id: DocId) -> StorageBackendResult<Option<DocumentMetadata>> {
253 self.get_stored(doc_id)
254 .map(|document| document.map(|document| document.metadata()))
255 }
256 fn contains_doc_id(&self, doc_id: DocId) -> StorageBackendResult<bool> {
257 Ok(self.get(doc_id)?.is_some())
258 }
259 fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()>;
260 fn clear(&mut self) -> StorageBackendResult<()>;
261
262 fn get_field(&self, doc_id: DocId, field: &str) -> StorageBackendResult<Option<Value>> {
266 Ok(self
267 .get(doc_id)?
268 .and_then(|document| document.get(field).cloned()))
269 }
270
271 fn find_doc_id_by_field(
275 &self,
276 field: &str,
277 value: &Value,
278 ) -> StorageBackendResult<Option<DocId>> {
279 for doc_id in self.doc_ids()? {
280 if self.get_field(doc_id, field)?.as_ref() == Some(value) {
281 return Ok(Some(doc_id));
282 }
283 }
284 Ok(None)
285 }
286
287 fn patch_fields(
292 &mut self,
293 doc_id: DocId,
294 updates: &BTreeMap<String, Value>,
295 ) -> StorageBackendResult<bool> {
296 let Some(mut document) = self.get_stored(doc_id)? else {
297 return Ok(false);
298 };
299 for (field, value) in updates {
300 if matches!(value, Value::Null) {
301 document.fields_mut().remove(field);
302 } else {
303 document.fields_mut().insert(field.clone(), value.clone());
304 }
305 }
306 self.put_stored(doc_id, document)?;
307 Ok(true)
308 }
309
310 fn get_many(&self, doc_ids: &[DocId]) -> StorageBackendResult<BTreeMap<DocId, Document>> {
315 let mut out = BTreeMap::new();
316 for doc_id in doc_ids {
317 if let Some(document) = self.get(*doc_id)? {
318 out.insert(*doc_id, document);
319 }
320 }
321 Ok(out)
322 }
323
324 fn get_fields_multi(
330 &self,
331 doc_ids: &[DocId],
332 fields: &[&str],
333 ) -> StorageBackendResult<BTreeMap<DocId, Vec<Value>>> {
334 let mut out = BTreeMap::new();
335 for doc_id in doc_ids {
336 let Some(document) = self.get(*doc_id)? else {
337 continue;
338 };
339 let values = fields
340 .iter()
341 .map(|field| document.get(*field).cloned().unwrap_or(Value::Null))
342 .collect();
343 out.insert(*doc_id, values);
344 }
345 Ok(out)
346 }
347
348 fn for_each_fields_multi(
354 &self,
355 doc_ids: &[DocId],
356 fields: &[&str],
357 visitor: &mut dyn FnMut(DocId, Vec<Value>) -> bool,
358 ) -> StorageBackendResult<()> {
359 let mut projected = self.get_fields_multi(doc_ids, fields)?;
360 for doc_id in doc_ids {
361 let values = projected
362 .remove(doc_id)
363 .unwrap_or_else(|| vec![Value::Null; fields.len()]);
364 if !visitor(*doc_id, values) {
365 break;
366 }
367 }
368 Ok(())
369 }
370
371 fn for_each_fields_multi_ref(
376 &self,
377 doc_ids: &[DocId],
378 fields: &[&str],
379 visitor: &mut dyn FnMut(DocId, &[&Value]) -> bool,
380 ) -> StorageBackendResult<()> {
381 self.for_each_fields_multi(doc_ids, fields, &mut |doc_id, values| {
382 let references: Vec<&Value> = values.iter().collect();
383 visitor(doc_id, &references)
384 })
385 }
386
387 fn for_each_fields_multi_ref_with_presence(
392 &self,
393 doc_ids: &[DocId],
394 fields: &[&str],
395 visitor: &mut dyn FnMut(DocId, bool, &[&Value]) -> bool,
396 ) -> StorageBackendResult<()> {
397 if fields.is_empty() {
398 for doc_id in doc_ids {
399 if !visitor(*doc_id, self.contains_doc_id(*doc_id)?, &[]) {
400 break;
401 }
402 }
403 return Ok(());
404 }
405
406 let projected = self.get_fields_multi(doc_ids, fields)?;
407 let null = Value::Null;
408 let missing = vec![&null; fields.len()];
409 for doc_id in doc_ids {
410 let Some(values) = projected.get(doc_id) else {
411 if !visitor(*doc_id, false, &missing) {
412 break;
413 }
414 continue;
415 };
416 let references = values.iter().collect::<Vec<_>>();
417 if !visitor(*doc_id, true, &references) {
418 break;
419 }
420 }
421 Ok(())
422 }
423
424 fn get_shared_fields(
429 &self,
430 _doc_ids: &[DocId],
431 _fields: &[&str],
432 ) -> StorageBackendResult<Option<Vec<Option<SharedDocumentRow>>>> {
433 Ok(None)
434 }
435
436 fn get_fields_bulk(
440 &self,
441 doc_ids: &[DocId],
442 field: &str,
443 ) -> StorageBackendResult<BTreeMap<DocId, Value>> {
444 let mut out = BTreeMap::new();
445 for doc_id in doc_ids {
446 out.insert(
447 *doc_id,
448 self.get_field(*doc_id, field)?.unwrap_or(Value::Null),
449 );
450 }
451 Ok(out)
452 }
453
454 fn has_value(&self, field: &str, value: &Value) -> StorageBackendResult<bool> {
456 for doc_id in self.doc_ids()? {
457 if self.get_field(doc_id, field)?.as_ref() == Some(value) {
458 return Ok(true);
459 }
460 }
461 Ok(false)
462 }
463
464 fn find_doc_id_by_fields(
467 &self,
468 fields: &[String],
469 values: &[Value],
470 ) -> StorageBackendResult<Option<DocId>> {
471 if fields.is_empty() || fields.len() != values.len() {
472 return Ok(None);
473 }
474 for doc_id in self.doc_ids()? {
475 let mut matches = true;
476 for (field, value) in fields.iter().zip(values) {
477 if self.get_field(doc_id, field)?.unwrap_or(Value::Null) != *value {
478 matches = false;
479 break;
480 }
481 }
482 if matches {
483 return Ok(Some(doc_id));
484 }
485 }
486 Ok(None)
487 }
488
489 fn eval_path(
491 &self,
492 doc_id: DocId,
493 path: &[PathSegment],
494 ) -> StorageBackendResult<Option<Value>> {
495 let Some(document) = self.get(doc_id)? else {
496 return Ok(None);
497 };
498 Ok(eval_path_in_document(&document, path))
499 }
500
501 fn doc_ids(&self) -> StorageBackendResult<Vec<DocId>>;
502
503 fn next_doc_id(&self, after: Option<DocId>) -> StorageBackendResult<Option<DocId>> {
508 Ok(self
509 .doc_ids()?
510 .into_iter()
511 .filter(|doc_id| after.is_none_or(|after| *doc_id > after))
512 .min())
513 }
514
515 fn next_doc_ids(&self, after: Option<DocId>, limit: usize) -> StorageBackendResult<Vec<DocId>> {
519 if limit == 0 {
520 return Ok(Vec::new());
521 }
522 let mut doc_ids = self.doc_ids()?;
523 doc_ids.sort_unstable();
524 Ok(doc_ids
525 .into_iter()
526 .filter(|doc_id| after.is_none_or(|after| *doc_id > after))
527 .take(limit)
528 .collect())
529 }
530
531 fn next_doc_ids_controlled(
533 &self,
534 _after: Option<DocId>,
535 limit: usize,
536 control: &StorageReadControl,
537 ) -> StorageBackendResult<BudgetedVec<DocId>> {
538 control.check()?;
539 if limit == 0 {
540 return Ok(BudgetedVec::new(control.memory()));
541 }
542 Err(StorageBackendError::Other(
543 "controlled document identity reads are not supported".into(),
544 ))
545 }
546
547 fn next_shared_fields(
551 &self,
552 _after: Option<DocId>,
553 _limit: usize,
554 _fields: &[&str],
555 ) -> StorageBackendResult<Option<Vec<(DocId, SharedDocumentRow)>>> {
556 Ok(None)
557 }
558
559 fn for_each_next_fields(
561 &self,
562 _after: Option<DocId>,
563 _limit: usize,
564 _fields: &[&str],
565 _visitor: &mut dyn FnMut(DocId, &[&Value]) -> bool,
566 ) -> StorageBackendResult<Option<usize>> {
567 Ok(None)
568 }
569
570 fn max_doc_id(&self) -> StorageBackendResult<DocId> {
571 Ok(self.doc_ids()?.into_iter().max().unwrap_or(0))
572 }
573
574 fn len(&self) -> StorageBackendResult<usize>;
575
576 fn is_empty(&self) -> StorageBackendResult<bool> {
577 Ok(self.len()? == 0)
578 }
579
580 fn iter_all(&self) -> StorageBackendResult<Box<dyn Iterator<Item = (DocId, Document)> + '_>> {
584 let mut ids = self.doc_ids()?;
585 ids.sort_unstable();
586 let snapshot = self.snapshot()?;
587 let mut rows = Vec::with_capacity(ids.len());
588 for doc_id in ids {
589 if let Some(document) = snapshot.get(doc_id)? {
590 rows.push((doc_id, document));
591 }
592 }
593 Ok(Box::new(rows.into_iter()))
594 }
595
596 fn snapshot(&self) -> StorageBackendResult<Arc<dyn DocumentStore>>;
600
601 fn retained_snapshot(&self) -> StorageBackendResult<Option<Arc<dyn DocumentStore>>> {
603 Ok(None)
604 }
605
606 fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn DocumentStore>> {
610 Err(StorageBackendError::Other(
611 "writable document-store snapshots are not supported by this backend".into(),
612 ))
613 }
614}
615
616pub fn eval_path_in_document(doc: &Document, path: &[PathSegment]) -> Option<Value> {
620 let mut current: Value = match path.first()? {
621 PathSegment::Key(k) => doc.get(k)?.clone(),
622 PathSegment::Index(_) => return None,
623 };
624 for seg in path.iter().skip(1) {
625 current = match (current, seg) {
626 (Value::Map(m), PathSegment::Key(k)) => m.get(k)?.clone(),
627 (Value::List(items), PathSegment::Index(i)) => items.get(*i)?.clone(),
628 (Value::List(items), PathSegment::Key(k)) => {
629 let collected: Vec<Value> = items
630 .into_iter()
631 .filter_map(|v| match v {
632 Value::Map(m) => m.get(k).cloned(),
633 _ => None,
634 })
635 .collect();
636 Value::List(collected)
637 }
638 _ => return None,
639 };
640 }
641 Some(current)
642}
643
644mod memory;
645
646pub use memory::MemoryDocumentStore;