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}
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    /// Capture a provider's selected occurrence view without first cloning its analyzer bindings or table metadata into an intermediate live index.
131    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}