uqa_storage/inverted_index/
source_rebuild.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! A source rebuild of a text index staged before it is written. The field records of its documents and its postings are private records ordered as the index stores them, which spill to encrypted temporary files once the allowance is under pressure, as a transaction's own changes do. Writing the index then needs memory for one cluster at a time instead of for every posting.
8
9use std::collections::BTreeMap;
10
11use uqa_core::memory::BudgetedVec;
12use uqa_core::{DocId, FieldName, TokenOccurrence};
13
14use super::{IndexedFieldMetadata, IndexedFieldRevision};
15use crate::backend::{StorageBackendError, StorageBackendResult};
16use crate::clustered_postings::{cluster_id, encode_posting, OccurrenceClusterBuilder};
17use crate::mvcc::{PrivateRecordChanges, RecordWrite, VersionError};
18use crate::read_control::StorageReadControl;
19use crate::TokenTermKey;
20
21/// The records of documents' fields, keyed by document and field.
22const DOCUMENT: u8 = 0;
23/// The records of postings, keyed by field, term and document.
24const POSTING: u8 = 1;
25/// The records a page of a visit keeps at most.
26const PAGE_RECORDS: usize = 1024;
27/// The bytes after which a page of a visit ends: a sixteenth of the allowance, within these bounds.
28const PAGE_BYTES: std::ops::RangeInclusive<usize> = (64 << 10)..=(1 << 20);
29
30/// The indexed documents of one field and their total length.
31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
32pub struct FieldTotals {
33    pub revision: IndexedFieldRevision,
34    pub doc_count: u64,
35    pub total_length: u64,
36}
37
38/// The analyzed fields of a document a rebuild stages: each field's name, metadata and the occurrences of each of its terms.
39pub type StagedFields<'a> = (
40    &'a str,
41    &'a IndexedFieldMetadata,
42    &'a BTreeMap<TokenTermKey, Vec<TokenOccurrence>>,
43);
44
45/// A staged field of a document: its metadata and the encoded keys of its terms.
46pub struct StagedFieldRecord<'a> {
47    pub doc_id: DocId,
48    pub field: &'a str,
49    pub metadata: IndexedFieldMetadata,
50    pub terms: &'a [u8],
51}
52
53/// A staged cluster with its encoded score and positions values.
54pub struct StagedCluster<'a> {
55    pub field: &'a str,
56    pub term: &'a TokenTermKey,
57    pub cluster: u64,
58    pub score: &'a [u8],
59    pub positions: &'a [u8],
60}
61
62/// A visit of staged records, each a key without its kind and a value.
63type RecordVisit<'a> = &'a mut dyn FnMut(&[u8], &[u8]) -> StorageBackendResult<()>;
64
65/// A source rebuild of a text index staged before it is written: documents are staged in ascending identity order, then visited by field record and by cluster in the order the index stores them.
66pub struct SourceRebuild {
67    records: PrivateRecordChanges,
68    totals: BTreeMap<FieldName, FieldTotals>,
69    last: Option<DocId>,
70    control: StorageReadControl,
71}
72
73fn invalid(message: &str) -> StorageBackendError {
74    StorageBackendError::Other(message.into())
75}
76
77/// Append `bytes` so that encoded keys order as their byte strings do: a zero byte is followed by 0xff, and the component ends with a zero byte followed by one.
78fn push_component(key: &mut BudgetedVec<u8>, bytes: &[u8]) -> StorageBackendResult<()> {
79    key.reserve(bytes.len().saturating_add(2))?;
80    for &byte in bytes {
81        key.push(byte)?;
82        if byte == 0 {
83            key.push(0xff)?;
84        }
85    }
86    key.extend_from_slice(&[0, 1])?;
87    Ok(())
88}
89
90/// Read a component `push_component` appended, returning it and the rest of the key.
91fn read_component(key: &[u8]) -> StorageBackendResult<(Vec<u8>, &[u8])> {
92    let mut bytes = Vec::new();
93    let mut index = 0;
94    while index < key.len() {
95        match (key[index], key.get(index + 1)) {
96            (0, Some(0xff)) => bytes.push(0),
97            (0, Some(1)) => return Ok((bytes, &key[index + 2..])),
98            (0, _) => return Err(invalid("a staged rebuild key component is malformed")),
99            (byte, _) => {
100                bytes.push(byte);
101                index += 1;
102                continue;
103            }
104        }
105        index += 2;
106    }
107    Err(invalid("a staged rebuild key component is unterminated"))
108}
109
110/// Count a staged field of a document in `totals_by_field`.
111fn count(
112    totals_by_field: &mut BTreeMap<FieldName, FieldTotals>,
113    field: &str,
114    metadata: &IndexedFieldMetadata,
115) -> StorageBackendResult<()> {
116    let revision = metadata.revision();
117    if !totals_by_field.contains_key(field) {
118        totals_by_field.insert(
119            field.to_owned(),
120            FieldTotals {
121                revision,
122                doc_count: 0,
123                total_length: 0,
124            },
125        );
126    }
127    let totals = totals_by_field.get_mut(field).expect("a counted field");
128    if totals.revision != revision {
129        return Err(invalid(
130            "indexed field revisions disagree during replacement",
131        ));
132    }
133    totals.doc_count = totals
134        .doc_count
135        .checked_add(1)
136        .ok_or_else(|| invalid("field document count overflow"))?;
137    totals.total_length = totals
138        .total_length
139        .checked_add(metadata.length)
140        .ok_or_else(|| invalid("total field length overflow"))?;
141    Ok(())
142}
143
144fn read_doc_id(bytes: &[u8]) -> StorageBackendResult<DocId> {
145    bytes
146        .try_into()
147        .map(DocId::from_be_bytes)
148        .map_err(|_| invalid("a staged rebuild key identity is malformed"))
149}
150
151impl SourceRebuild {
152    pub fn new(control: &StorageReadControl) -> Self {
153        Self {
154            records: PrivateRecordChanges::new(control.memory()),
155            totals: BTreeMap::new(),
156            last: None,
157            control: control.clone(),
158        }
159    }
160
161    /// Stage document `doc_id` with its analyzed fields. Documents are staged in strictly ascending identity order.
162    pub fn stage<'a>(
163        &mut self,
164        doc_id: DocId,
165        fields: impl IntoIterator<Item = StagedFields<'a>>,
166    ) -> StorageBackendResult<()> {
167        self.control.check()?;
168        if self.last.is_some_and(|last| last >= doc_id) {
169            return Err(invalid(
170                "a source rebuild reads documents in strictly ascending identity order",
171            ));
172        }
173        self.last = Some(doc_id);
174        let memory = self.control.memory();
175        let mut records = BudgetedVec::<(BudgetedVec<u8>, BudgetedVec<u8>)>::new(memory);
176        for (field, metadata, terms) in fields {
177            count(&mut self.totals, field, metadata)?;
178            let mut key = BudgetedVec::new(memory);
179            key.push(DOCUMENT)?;
180            key.extend_from_slice(&doc_id.to_be_bytes())?;
181            push_component(&mut key, field.as_bytes())?;
182            let mut value = BudgetedVec::new(memory);
183            let encoded = metadata.to_bytes()?;
184            value.push(
185                u8::try_from(encoded.len())
186                    .map_err(|_| invalid("indexed field metadata exceeds a staged record"))?,
187            )?;
188            value.extend_from_slice(&encoded)?;
189            value.extend_from_slice(&crate::clustered_postings::encode_term_key_refs(
190                terms.keys(),
191            )?)?;
192            records.push((key, value))?;
193            for (term, occurrences) in terms {
194                self.control.check()?;
195                let mut key = BudgetedVec::new(memory);
196                key.push(POSTING)?;
197                push_component(&mut key, field.as_bytes())?;
198                push_component(&mut key, term.as_bytes())?;
199                key.extend_from_slice(&doc_id.to_be_bytes())?;
200                let mut value = BudgetedVec::new(memory);
201                encode_posting(metadata.length, occurrences, &mut value, &self.control)?;
202                records.push((key, value))?;
203            }
204        }
205        let mut writes = BudgetedVec::new(memory);
206        writes.reserve(records.len())?;
207        for (key, value) in records.iter() {
208            writes.push(RecordWrite {
209                key,
210                expected: None,
211                value: Some(value),
212            })?;
213        }
214        self.records
215            .apply(&writes, &self.control)
216            .map_err(VersionError::into_storage_error)
217    }
218
219    /// The indexed documents and total length of each staged field.
220    pub fn totals(&self) -> &BTreeMap<FieldName, FieldTotals> {
221        &self.totals
222    }
223
224    /// Visit the staged field records in order of document and field, each with its metadata and the encoded keys of its terms.
225    pub fn visit_documents(
226        &self,
227        visit: &mut dyn FnMut(StagedFieldRecord<'_>) -> StorageBackendResult<()>,
228    ) -> StorageBackendResult<()> {
229        self.visit_records(DOCUMENT, &mut |key, value| {
230            let (id, field) = key
231                .split_at_checked(size_of::<DocId>())
232                .ok_or_else(|| invalid("a staged field record key is malformed"))?;
233            let (field, trailing) = read_component(field)?;
234            let (metadata, terms) = value
235                .split_first()
236                .and_then(|(&length, encoded)| encoded.split_at_checked(usize::from(length)))
237                .filter(|_| trailing.is_empty())
238                .ok_or_else(|| invalid("a staged field record is malformed"))?;
239            let field = String::from_utf8(field)
240                .map_err(|_| invalid("a staged field name is not UTF-8"))?;
241            visit(StagedFieldRecord {
242                doc_id: read_doc_id(id)?,
243                field: &field,
244                metadata: IndexedFieldMetadata::from_bytes(metadata)?,
245                terms,
246            })
247        })
248    }
249
250    /// Visit the staged postings as clusters in order of field, term and cluster, each with its encoded score and positions values.
251    pub fn visit_clusters(
252        &self,
253        visit: &mut dyn FnMut(StagedCluster<'_>) -> StorageBackendResult<()>,
254    ) -> StorageBackendResult<()> {
255        let mut current: Option<(String, TokenTermKey, u64, OccurrenceClusterBuilder)> = None;
256        let mut finish = |cluster: Option<(
257            String,
258            TokenTermKey,
259            u64,
260            OccurrenceClusterBuilder,
261        )>| {
262            match cluster {
263                Some((field, term, cluster, builder)) => {
264                    let (score, positions) = builder.finish(&self.control)?;
265                    visit(StagedCluster {
266                        field: &field,
267                        term: &term,
268                        cluster,
269                        score: &score,
270                        positions: &positions,
271                    })
272                }
273                None => Ok(()),
274            }
275        };
276        self.visit_records(POSTING, &mut |key, value| {
277            let (field, rest) = read_component(key)?;
278            let (term, id) = read_component(rest)?;
279            let doc_id = read_doc_id(id)?;
280            let cluster = cluster_id(doc_id);
281            let same = current
282                .as_ref()
283                .is_some_and(|(current_field, current_term, id, _)| {
284                    *id == cluster
285                        && current_field.as_bytes() == field.as_slice()
286                        && current_term.as_bytes() == term.as_slice()
287                });
288            if !same {
289                finish(current.take())?;
290                let field = String::from_utf8(field)
291                    .map_err(|_| invalid("a staged field name is not UTF-8"))?;
292                current = Some((
293                    field,
294                    TokenTermKey::from_bytes(term)?,
295                    cluster,
296                    OccurrenceClusterBuilder::new(&self.control)?,
297                ));
298            }
299            let (_, _, _, builder) = current.as_mut().expect("a current cluster");
300            builder.push_encoded(doc_id, value)
301        })?;
302        finish(current.take())
303    }
304
305    /// Visit the staged records of `kind` in key order, without its kind byte, a bounded page at a time so that no visit runs while the records are read.
306    fn visit_records(&self, kind: u8, visit: RecordVisit<'_>) -> StorageBackendResult<()> {
307        let snapshot = self
308            .records
309            .snapshot()
310            .map_err(VersionError::into_storage_error)?;
311        let memory = self.control.memory();
312        let page_bytes = (memory.limit() / 16).clamp(*PAGE_BYTES.start(), *PAGE_BYTES.end());
313        let mut after: Option<BudgetedVec<u8>> = None;
314        loop {
315            let mut page = BudgetedVec::<(BudgetedVec<u8>, BudgetedVec<u8>)>::new(memory);
316            let mut bytes = 0_usize;
317            let mut full = false;
318            snapshot
319                .visit(&[kind], after.as_deref(), &self.control, &mut |write| {
320                    let mut key = BudgetedVec::new(memory);
321                    key.extend_from_slice(write.key())?;
322                    let mut value = BudgetedVec::new(memory);
323                    value.extend_from_slice(write.value().ok_or_else(|| {
324                        VersionError::from(invalid("a staged rebuild record is deleted"))
325                    })?)?;
326                    bytes = bytes.saturating_add(key.len() + value.len());
327                    page.push((key, value))?;
328                    full = page.len() == PAGE_RECORDS || bytes >= page_bytes;
329                    Ok(!full)
330                })
331                .map_err(VersionError::into_storage_error)?;
332            for (key, value) in page.iter() {
333                self.control.check()?;
334                visit(&key[1..], value)?;
335            }
336            let Some((last, _)) = page.last().filter(|_| full) else {
337                return Ok(());
338            };
339            let mut next = BudgetedVec::new(memory);
340            next.extend_from_slice(last)?;
341            after = Some(next);
342        }
343    }
344}
345
346#[cfg(test)]
347mod tests;
uqa_storage/inverted_index/source_rebuild.rs

uqa_storage/inverted_index/
source_rebuild.rs
