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