Skip to main content

uqa_storage/key_value/inverted_index/
view.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! One reader owns an occurrence operation; retained indexes own their selected bytes.
8
9use super::{AnalyzerBindings, Arc, DocId, FieldName, InvertedIndex, KeyValueInvertedIndex};
10use crate::key_value::{index_view, KeyValueBatch, KeyValueRead, OccurrenceStorage};
11use crate::StorageBackendResult;
12use uqa_core::memory::BudgetedVec;
13
14type Entry = (BudgetedVec<u8>, BudgetedVec<u8>);
15
16#[derive(Clone)]
17pub(super) enum OccurrenceSource {
18    Live(Arc<dyn OccurrenceStorage>),
19    Retained(Arc<dyn KeyValueRead + Send + Sync>),
20}
21
22pub(super) struct OccurrenceRead<'a> {
23    pub(super) store: &'a dyn KeyValueRead,
24    pub(super) table: &'a str,
25    pub(super) bindings: &'a AnalyzerBindings,
26    pub(super) format_requires_rebuild: std::cell::Cell<Option<bool>>,
27}
28
29impl OccurrenceRead<'_> {
30    pub(super) fn scan_prefix(&self, prefix: &[u8]) -> StorageBackendResult<BudgetedVec<Entry>> {
31        let mut rows = BudgetedVec::new(self.store.control().memory());
32        append_rows(self.store, prefix, &mut rows)?;
33        Ok(rows)
34    }
35}
36
37fn append_rows(
38    read: &dyn KeyValueRead,
39    prefix: &[u8],
40    rows: &mut BudgetedVec<Entry>,
41) -> StorageBackendResult<()> {
42    read.visit_prefix(prefix, &mut |key, value| {
43        read.control().check()?;
44        rows.reserve(1)?;
45        let mut owned_key = BudgetedVec::new(read.control().memory());
46        owned_key.extend_from_slice(key)?;
47        let mut owned_value = BudgetedVec::new(read.control().memory());
48        owned_value.extend_from_slice(value)?;
49        rows.push((owned_key, owned_value))?;
50        Ok(())
51    })
52}
53
54impl KeyValueInvertedIndex {
55    pub(super) fn read<T>(
56        &self,
57        operation: impl FnOnce(OccurrenceRead<'_>) -> StorageBackendResult<T>,
58    ) -> StorageBackendResult<T> {
59        let evaluate = |store: &dyn KeyValueRead| {
60            operation(OccurrenceRead {
61                store,
62                table: &self.table,
63                bindings: &self.bindings,
64                format_requires_rebuild: std::cell::Cell::new(None),
65            })
66        };
67        match &self.source {
68            OccurrenceSource::Live(store) => {
69                index_view::read_scope(|operation| store.read(operation), evaluate)
70            }
71            OccurrenceSource::Retained(read) => {
72                read.control().check()?;
73                evaluate(read.as_ref())
74            }
75        }
76    }
77
78    pub(super) fn ensure_writable(&self) -> StorageBackendResult<()> {
79        match self.source {
80            OccurrenceSource::Live(_) => Ok(()),
81            OccurrenceSource::Retained(_) => {
82                Err(super::other_error("inverted-index snapshots are read-only"))
83            }
84        }
85    }
86
87    pub(super) fn mutate(
88        &self,
89        operation: impl FnOnce(&OccurrenceRead<'_>, &mut dyn KeyValueBatch) -> StorageBackendResult<()>,
90    ) -> StorageBackendResult<()> {
91        self.ensure_writable()?;
92        let OccurrenceSource::Live(store) = &self.source else {
93            unreachable!()
94        };
95        index_view::mutation_scope(
96            |operation| store.mutate(operation),
97            |read, batch| {
98                operation(
99                    &OccurrenceRead {
100                        store: read,
101                        table: &self.table,
102                        bindings: &self.bindings,
103                        format_requires_rebuild: std::cell::Cell::new(None),
104                    },
105                    batch,
106                )
107            },
108        )
109    }
110
111    pub(super) fn add_documents(
112        &self,
113        documents: Vec<(DocId, std::collections::BTreeMap<FieldName, String>)>,
114    ) -> StorageBackendResult<()> {
115        self.mutate(|view, batch| view.add_documents(batch, documents, None))
116    }
117
118    pub(super) fn rebuild_documents(
119        &self,
120        documents: Vec<(DocId, std::collections::BTreeMap<FieldName, String>)>,
121    ) -> StorageBackendResult<()> {
122        self.rebuild_documents_inner(documents, None)
123    }
124
125    pub(super) fn rebuild_documents_inner(
126        &self,
127        documents: Vec<(DocId, std::collections::BTreeMap<FieldName, String>)>,
128        cancellation: Option<&uqa_core::CancellationToken>,
129    ) -> StorageBackendResult<()> {
130        self.mutate(|view, batch| view.rebuild_documents_inner(batch, documents, cancellation))
131    }
132
133    /// Capture a provider's selected occurrence view without first cloning its analyzer bindings or table metadata into an intermediate live index.
134    pub fn snapshot_from_read(
135        read: &dyn KeyValueRead,
136        table: &str,
137        bindings: &AnalyzerBindings,
138    ) -> StorageBackendResult<Arc<dyn InvertedIndex>> {
139        read.control().check()?;
140        let prefixes = super::format::retained_prefixes(table, read.control())?;
141        let borrowed: [&[u8]; 8] = std::array::from_fn(|slot| &prefixes[slot][..]);
142        let source = read.retain(&borrowed)?;
143        Self::snapshot_from_source(source, table, bindings)
144    }
145
146    pub(super) fn retained_snapshot(&self) -> StorageBackendResult<Arc<dyn InvertedIndex>> {
147        match &self.source {
148            OccurrenceSource::Retained(source) => {
149                Self::snapshot_from_source(Arc::clone(source), &self.table, &self.bindings)
150            }
151            OccurrenceSource::Live(_) => {
152                self.read(|view| Self::snapshot_from_read(view.store, view.table, view.bindings))
153            }
154        }
155    }
156
157    fn snapshot_from_source(
158        source: Arc<dyn KeyValueRead + Send + Sync>,
159        table: &str,
160        bindings: &AnalyzerBindings,
161    ) -> StorageBackendResult<Arc<dyn InvertedIndex>> {
162        use crate::ReadOnlySnapshot;
163        use uqa_core::memory::BudgetedString;
164
165        let control = source.control().clone();
166        control.check()?;
167        let bindings = bindings.retained(&control)?;
168        let mut name = BudgetedString::new(control.memory());
169        name.reserve(table.len())?;
170        for (offset, character) in table.chars().enumerate() {
171            if offset % 1024 == 0 {
172                control.check()?;
173            }
174            name.push(character)?;
175        }
176        let (table, mut memory) = name.into_parts();
177        memory.grow(size_of::<Self>())?;
178        memory.grow(size_of::<ReadOnlySnapshot<dyn InvertedIndex>>())?;
179        let snapshot: Arc<dyn InvertedIndex> = Arc::new(Self {
180            source: OccurrenceSource::Retained(source),
181            table,
182            bindings,
183        });
184        let snapshot = ReadOnlySnapshot::with_retention(snapshot, memory)?
185            .with_inverted_read_control(&control)?;
186        Ok(Arc::new(snapshot))
187    }
188}