uqa_storage/inverted_index/source/
store.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! The text of a document store's indexed fields, read as a source a page of documents at a time.
8
9use std::collections::{BTreeMap, VecDeque};
10use std::sync::Arc;
11
12use uqa_core::{CancellationToken, DocId, FieldName, Value};
13
14use super::TextIndexSource;
15use crate::backend::StorageBackendResult;
16use crate::DocumentStore;
17
18/// The documents a page holds at most.
19const PAGE_DOCUMENTS: usize = 1024;
20/// The text a page holds, after which it ends at the next document.
21const PAGE_TEXT_BYTES: usize = 4 << 20;
22
23/// The text of the indexed fields of a document store's documents, read a page at a time in identity order. A document holding no text in those fields is not read.
24pub struct DocumentTextSource {
25    store: Arc<dyn DocumentStore>,
26    fields: Vec<FieldName>,
27    /// The last document a page read.
28    after: Option<DocId>,
29    page: VecDeque<(DocId, BTreeMap<FieldName, String>)>,
30    done: bool,
31    cancellation: Option<CancellationToken>,
32}
33
34impl DocumentTextSource {
35    /// The text of `fields` in the documents of `store`, which should be an immutable view of them, as a snapshot is.
36    pub fn new(
37        store: Arc<dyn DocumentStore>,
38        fields: Vec<FieldName>,
39        cancellation: Option<CancellationToken>,
40    ) -> Self {
41        Self {
42            store,
43            fields,
44            after: None,
45            page: VecDeque::new(),
46            done: false,
47            cancellation,
48        }
49    }
50
51    fn check(&self) -> StorageBackendResult<()> {
52        if let Some(cancellation) = &self.cancellation {
53            cancellation.check()?;
54        }
55        Ok(())
56    }
57
58    fn read_page(&mut self) -> StorageBackendResult<()> {
59        self.check()?;
60        let ids = self.store.next_doc_ids(self.after, PAGE_DOCUMENTS)?;
61        let Some(&last) = ids.last() else {
62            self.done = true;
63            return Ok(());
64        };
65        let Self {
66            store,
67            fields,
68            page,
69            cancellation,
70            ..
71        } = self;
72        let names = fields.iter().map(String::as_str).collect::<Vec<_>>();
73        let mut bytes = 0_usize;
74        let mut ended = None;
75        store.for_each_fields_multi_ref(&ids, &names, &mut |doc_id, values| {
76            if cancellation
77                .as_ref()
78                .is_some_and(CancellationToken::is_cancelled)
79            {
80                ended = Some(doc_id);
81                return false;
82            }
83            let mut text = BTreeMap::new();
84            for (field, value) in fields.iter().zip(values) {
85                if let Value::Str(value) = value {
86                    bytes = bytes.saturating_add(value.len());
87                    text.insert(field.clone(), value.clone());
88                }
89            }
90            if !text.is_empty() {
91                page.push_back((doc_id, text));
92            }
93            if bytes >= PAGE_TEXT_BYTES {
94                ended = Some(doc_id);
95                return false;
96            }
97            true
98        })?;
99        self.check()?;
100        self.after = Some(ended.unwrap_or(last));
101        Ok(())
102    }
103}
104
105impl TextIndexSource for DocumentTextSource {
106    fn next_document(
107        &mut self,
108    ) -> StorageBackendResult<Option<(DocId, BTreeMap<FieldName, String>)>> {
109        while self.page.is_empty() && !self.done {
110            self.read_page()?;
111        }
112        Ok(self.page.pop_front())
113    }
114}
uqa_storage/inverted_index/source/store.rs

uqa_storage/inverted_index/source/
store.rs
