uqa_storage/key_value/inverted_index/
view.rs1use 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 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}