1use std::collections::{BTreeMap, BTreeSet};
15use std::sync::Arc;
16
17use parking_lot::Mutex;
18use serde::{Deserialize, Serialize};
19use uqa_analysis::Analyzer;
20use uqa_core::{DocId, FieldName, IndexStats, Payload, PostingEntry, PostingList, Value};
21
22use crate::backend::{PersistentStorageBackend, PersistentStorageIdentity};
23use crate::document_store::{Document, DocumentMetadata, DocumentStore, StoredDocument};
24use crate::inverted_index::{AnalyzerPhase, InvertedIndex};
25use crate::vector_index::{
26 cosine_similarity, validate_vector_values, VectorIndex, VectorIndexOpenMode, VectorIndexSpec,
27};
28use crate::{StorageBackendError, StorageBackendResult};
29
30mod catalog;
31pub use catalog::KeyValueCatalog;
32
33const TAG_METADATA: u8 = b'm';
34const TAG_TABLE: u8 = b't';
35const TAG_MODEL: u8 = b'M';
36const TAG_SCORING_PARAMS: u8 = b'S';
37const TAG_NAMED_GRAPH: u8 = b'g';
38const TAG_VERTEX: u8 = b'V';
39const TAG_EDGE: u8 = b'E';
40const TAG_GRAPH_MEMBERSHIP: u8 = b'G';
41const TAG_GRAPH_LOOKUP: u8 = b'J';
42const TAG_ANALYZER: u8 = b'a';
43const TAG_TABLE_FIELD_ANALYZER: u8 = b'A';
44const TAG_FOREIGN_SERVER: u8 = b'F';
45const TAG_FOREIGN_TABLE: u8 = b'T';
46const TAG_CATALOG_INDEX: u8 = b'C';
47const TAG_PATH_INDEX: u8 = b'P';
48const TAG_PATH_INDEX_DATA: u8 = b'Q';
49const TAG_COLUMN_STATS: u8 = b'c';
50const TAG_SCHEMA: u8 = b's';
51const TAG_SEQUENCE: u8 = b'q';
52const TAG_RELATION: u8 = b'R';
53const TAG_VIEW: u8 = b'w';
54const TAG_DOCUMENT: u8 = b'd';
55const TAG_POSTING: u8 = b'p';
56const TAG_POSTING_CLUSTER_SCORE: u8 = b'k';
57const TAG_POSTING_CLUSTER_POSITIONS: u8 = b'o';
58const TAG_POSTING_DOCUMENT: u8 = b'x';
59const TAG_DOC_LENGTH: u8 = b'l';
60const TAG_FIELD_STATS: u8 = b'f';
61const TAG_REVERSE_POSTING: u8 = b'r';
62const TAG_VECTOR: u8 = b'v';
63const TAG_BTREE_INDEX: u8 = b'B';
64const TAG_BTREE_ENTRY: u8 = b'b';
65const TAG_NAMED_BTREE_INDEX: u8 = b'N';
66const TAG_NAMED_BTREE_ENTRY: u8 = b'n';
67const TAG_IVF_METADATA: u8 = b'I';
68const TAG_IVF_CENTROID: u8 = b'i';
69const TAG_IVF_ASSIGNMENT: u8 = b'j';
70const TAG_HNSW_METADATA: u8 = b'H';
71const TAG_HNSW_NODE: u8 = b'h';
72
73const DOCUMENT_VALUE_V1_PREFIX: &[u8] = b"\0uqa-document-json-v1\0";
78const DOCUMENT_VALUE_V2_PREFIX: &[u8] = b"\0uqa-document-record-v2\0";
79
80#[derive(Debug, Clone)]
81enum KeyValueBatchOperation {
82 Put(Vec<u8>, Vec<u8>),
83 Delete(Vec<u8>),
84 DeletePrefix(Vec<u8>),
85}
86
87pub trait KeyValueBatch {
89 fn put(&mut self, key: &[u8], value: &[u8]) -> StorageBackendResult<()>;
90 fn delete(&mut self, key: &[u8]) -> StorageBackendResult<()>;
91 fn delete_prefix(&mut self, prefix: &[u8]) -> StorageBackendResult<()>;
92 fn commit(self: Box<Self>) -> StorageBackendResult<()>;
93}
94
95pub trait KeyValueStore: Send + Sync {
97 fn storage_identity(&self) -> StorageBackendResult<Option<PersistentStorageIdentity>> {
98 Ok(None)
99 }
100
101 fn open_session(&self) -> StorageBackendResult<Arc<dyn KeyValueStore>> {
103 Err(StorageBackendError::Other(
104 "independent sessions are not implemented for this KeyValue store".into(),
105 ))
106 }
107
108 fn get(&self, key: &[u8]) -> StorageBackendResult<Option<Vec<u8>>>;
109 fn contains_key(&self, key: &[u8]) -> StorageBackendResult<bool> {
110 self.get(key).map(|value| value.is_some())
111 }
112 fn put(&self, key: &[u8], value: &[u8]) -> StorageBackendResult<()>;
113 fn delete(&self, key: &[u8]) -> StorageBackendResult<()>;
114 fn scan_prefix(&self, prefix: &[u8]) -> StorageBackendResult<Vec<(Vec<u8>, Vec<u8>)>>;
115 fn scan_prefix_after(
119 &self,
120 prefix: &[u8],
121 after: Option<&[u8]>,
122 limit: usize,
123 ) -> StorageBackendResult<Vec<(Vec<u8>, Vec<u8>)>> {
124 if limit == 0 {
125 return Ok(Vec::new());
126 }
127 Ok(self
128 .scan_prefix(prefix)?
129 .into_iter()
130 .filter(|(key, _)| after.is_none_or(|after| key.as_slice() > after))
131 .take(limit)
132 .collect())
133 }
134 fn scan_prefix_keys_after(
139 &self,
140 prefix: &[u8],
141 after: Option<&[u8]>,
142 limit: usize,
143 ) -> StorageBackendResult<Vec<Vec<u8>>> {
144 if limit == 0 {
145 return Ok(Vec::new());
146 }
147 Ok(self
148 .scan_prefix(prefix)?
149 .into_iter()
150 .filter(|(key, _)| after.is_none_or(|after| key.as_slice() > after))
151 .take(limit)
152 .map(|(key, _)| key)
153 .collect())
154 }
155 fn first_prefix_after(
156 &self,
157 prefix: &[u8],
158 after: Option<&[u8]>,
159 ) -> StorageBackendResult<Option<(Vec<u8>, Vec<u8>)>> {
160 Ok(self
161 .scan_prefix(prefix)?
162 .into_iter()
163 .find(|(key, _)| after.is_none_or(|after| key.as_slice() > after)))
164 }
165 fn delete_prefix(&self, prefix: &[u8]) -> StorageBackendResult<usize>;
166 fn batch(&self) -> Box<dyn KeyValueBatch + '_>;
167
168 fn begin_transaction(&self) -> StorageBackendResult<()> {
169 Err(StorageBackendError::Other(
170 "KeyValue transaction begin is not implemented for this store".into(),
171 ))
172 }
173
174 fn begin_read_transaction(&self) -> StorageBackendResult<()> {
175 self.begin_transaction()
176 }
177
178 fn begin_upgradeable_transaction(&self) -> StorageBackendResult<()> {
179 self.begin_transaction()
180 }
181
182 fn in_transaction(&self) -> bool;
183
184 fn transaction_has_written(&self) -> StorageBackendResult<bool>;
185
186 fn change_version(&self) -> StorageBackendResult<Option<u64>> {
187 Ok(None)
188 }
189
190 fn change_version_monitor_is_nonblocking(&self) -> StorageBackendResult<bool> {
191 Ok(true)
192 }
193
194 fn pin_transaction_snapshot(&self) -> StorageBackendResult<()> {
195 Ok(())
196 }
197
198 fn commit_transaction(&self) -> StorageBackendResult<()> {
199 Err(StorageBackendError::Other(
200 "KeyValue transaction commit is not implemented for this store".into(),
201 ))
202 }
203
204 fn rollback_transaction(&self) -> StorageBackendResult<()> {
205 Err(StorageBackendError::Other(
206 "KeyValue transaction rollback is not implemented for this store".into(),
207 ))
208 }
209
210 fn savepoint(&self, _name: &str) -> StorageBackendResult<()> {
211 Err(StorageBackendError::Other(
212 "KeyValue savepoints are not implemented for this store".into(),
213 ))
214 }
215
216 fn release_savepoint(&self, _name: &str) -> StorageBackendResult<()> {
217 Err(StorageBackendError::Other(
218 "KeyValue savepoint release is not implemented for this store".into(),
219 ))
220 }
221
222 fn rollback_to_savepoint(&self, _name: &str) -> StorageBackendResult<()> {
223 Err(StorageBackendError::Other(
224 "KeyValue savepoint rollback is not implemented for this store".into(),
225 ))
226 }
227}
228
229mod btree_index;
230mod codec;
231pub mod conformance;
232mod document_store;
233mod hnsw_index;
234mod hnsw_persistence;
235mod index_keys;
236mod inverted_index;
237mod ivf_index;
238mod ivf_persistence;
239mod memory_store;
240mod storage_backend;
241mod vector_index;
242
243pub use codec::prefix_upper_bound;
244pub use document_store::KeyValueDocumentStore;
245pub use hnsw_index::KeyValueHNSWIndex;
246pub use inverted_index::KeyValueInvertedIndex;
247pub use ivf_index::KeyValueIVFIndex;
248pub use memory_store::MemoryKeyValueStore;
249pub use storage_backend::KeyValueStorageBackend;
250pub use vector_index::KeyValueVectorIndex;
251
252#[cfg(test)]
253mod tests;