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, DocumentStore};
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_ANALYZER: u8 = b'a';
42const TAG_TABLE_FIELD_ANALYZER: u8 = b'A';
43const TAG_FOREIGN_SERVER: u8 = b'F';
44const TAG_FOREIGN_TABLE: u8 = b'T';
45const TAG_CATALOG_INDEX: u8 = b'C';
46const TAG_PATH_INDEX: u8 = b'P';
47const TAG_COLUMN_STATS: u8 = b'c';
48const TAG_SCHEMA: u8 = b's';
49const TAG_SEQUENCE: u8 = b'q';
50const TAG_RELATION: u8 = b'R';
51const TAG_VIEW: u8 = b'w';
52const TAG_DOCUMENT: u8 = b'd';
53const TAG_POSTING: u8 = b'p';
54const TAG_POSTING_CLUSTER_SCORE: u8 = b'k';
55const TAG_POSTING_CLUSTER_POSITIONS: u8 = b'o';
56const TAG_POSTING_DOCUMENT: u8 = b'x';
57const TAG_DOC_LENGTH: u8 = b'l';
58const TAG_FIELD_STATS: u8 = b'f';
59const TAG_REVERSE_POSTING: u8 = b'r';
60const TAG_VECTOR: u8 = b'v';
61const TAG_BTREE_INDEX: u8 = b'B';
62const TAG_BTREE_ENTRY: u8 = b'b';
63const TAG_IVF_METADATA: u8 = b'I';
64const TAG_IVF_CENTROID: u8 = b'i';
65const TAG_IVF_ASSIGNMENT: u8 = b'j';
66const TAG_HNSW_METADATA: u8 = b'H';
67const TAG_HNSW_NODE: u8 = b'h';
68
69const DOCUMENT_VALUE_V1_PREFIX: &[u8] = b"\0uqa-document-json-v1\0";
74
75#[derive(Debug, Clone)]
76enum KeyValueBatchOperation {
77 Put(Vec<u8>, Vec<u8>),
78 Delete(Vec<u8>),
79 DeletePrefix(Vec<u8>),
80}
81
82pub trait KeyValueBatch {
84 fn put(&mut self, key: &[u8], value: &[u8]) -> StorageBackendResult<()>;
85 fn delete(&mut self, key: &[u8]) -> StorageBackendResult<()>;
86 fn delete_prefix(&mut self, prefix: &[u8]) -> StorageBackendResult<()>;
87 fn commit(self: Box<Self>) -> StorageBackendResult<()>;
88}
89
90pub trait KeyValueStore: Send + Sync {
92 fn storage_identity(&self) -> StorageBackendResult<Option<PersistentStorageIdentity>> {
93 Ok(None)
94 }
95
96 fn open_session(&self) -> StorageBackendResult<Arc<dyn KeyValueStore>> {
98 Err(StorageBackendError::Other(
99 "independent sessions are not implemented for this KeyValue store".into(),
100 ))
101 }
102
103 fn get(&self, key: &[u8]) -> StorageBackendResult<Option<Vec<u8>>>;
104 fn contains_key(&self, key: &[u8]) -> StorageBackendResult<bool> {
105 self.get(key).map(|value| value.is_some())
106 }
107 fn put(&self, key: &[u8], value: &[u8]) -> StorageBackendResult<()>;
108 fn delete(&self, key: &[u8]) -> StorageBackendResult<()>;
109 fn scan_prefix(&self, prefix: &[u8]) -> StorageBackendResult<Vec<(Vec<u8>, Vec<u8>)>>;
110 fn scan_prefix_after(
114 &self,
115 prefix: &[u8],
116 after: Option<&[u8]>,
117 limit: usize,
118 ) -> StorageBackendResult<Vec<(Vec<u8>, Vec<u8>)>> {
119 if limit == 0 {
120 return Ok(Vec::new());
121 }
122 Ok(self
123 .scan_prefix(prefix)?
124 .into_iter()
125 .filter(|(key, _)| after.is_none_or(|after| key.as_slice() > after))
126 .take(limit)
127 .collect())
128 }
129 fn scan_prefix_keys_after(
134 &self,
135 prefix: &[u8],
136 after: Option<&[u8]>,
137 limit: usize,
138 ) -> StorageBackendResult<Vec<Vec<u8>>> {
139 if limit == 0 {
140 return Ok(Vec::new());
141 }
142 Ok(self
143 .scan_prefix(prefix)?
144 .into_iter()
145 .filter(|(key, _)| after.is_none_or(|after| key.as_slice() > after))
146 .take(limit)
147 .map(|(key, _)| key)
148 .collect())
149 }
150 fn first_prefix_after(
151 &self,
152 prefix: &[u8],
153 after: Option<&[u8]>,
154 ) -> StorageBackendResult<Option<(Vec<u8>, Vec<u8>)>> {
155 Ok(self
156 .scan_prefix(prefix)?
157 .into_iter()
158 .find(|(key, _)| after.is_none_or(|after| key.as_slice() > after)))
159 }
160 fn delete_prefix(&self, prefix: &[u8]) -> StorageBackendResult<usize>;
161 fn batch(&self) -> Box<dyn KeyValueBatch + '_>;
162
163 fn begin_transaction(&self) -> StorageBackendResult<()> {
164 Err(StorageBackendError::Other(
165 "KeyValue transaction begin is not implemented for this store".into(),
166 ))
167 }
168
169 fn begin_read_transaction(&self) -> StorageBackendResult<()> {
170 self.begin_transaction()
171 }
172
173 fn in_transaction(&self) -> bool;
174
175 fn transaction_has_written(&self) -> StorageBackendResult<bool>;
176
177 fn change_version(&self) -> StorageBackendResult<Option<u64>> {
178 Ok(None)
179 }
180
181 fn change_version_monitor_is_nonblocking(&self) -> StorageBackendResult<bool> {
182 Ok(true)
183 }
184
185 fn pin_transaction_snapshot(&self) -> StorageBackendResult<()> {
186 Ok(())
187 }
188
189 fn commit_transaction(&self) -> StorageBackendResult<()> {
190 Err(StorageBackendError::Other(
191 "KeyValue transaction commit is not implemented for this store".into(),
192 ))
193 }
194
195 fn rollback_transaction(&self) -> StorageBackendResult<()> {
196 Err(StorageBackendError::Other(
197 "KeyValue transaction rollback is not implemented for this store".into(),
198 ))
199 }
200
201 fn savepoint(&self, _name: &str) -> StorageBackendResult<()> {
202 Err(StorageBackendError::Other(
203 "KeyValue savepoints are not implemented for this store".into(),
204 ))
205 }
206
207 fn release_savepoint(&self, _name: &str) -> StorageBackendResult<()> {
208 Err(StorageBackendError::Other(
209 "KeyValue savepoint release is not implemented for this store".into(),
210 ))
211 }
212
213 fn rollback_to_savepoint(&self, _name: &str) -> StorageBackendResult<()> {
214 Err(StorageBackendError::Other(
215 "KeyValue savepoint rollback is not implemented for this store".into(),
216 ))
217 }
218}
219
220mod btree_index;
221mod codec;
222pub mod conformance;
223mod document_store;
224mod hnsw_index;
225mod hnsw_persistence;
226mod index_keys;
227mod inverted_index;
228mod ivf_index;
229mod ivf_persistence;
230mod memory_store;
231mod storage_backend;
232mod vector_index;
233
234pub use codec::prefix_upper_bound;
235pub use document_store::KeyValueDocumentStore;
236pub use hnsw_index::KeyValueHNSWIndex;
237pub use inverted_index::KeyValueInvertedIndex;
238pub use ivf_index::KeyValueIVFIndex;
239pub use memory_store::MemoryKeyValueStore;
240pub use storage_backend::KeyValueStorageBackend;
241pub use vector_index::KeyValueVectorIndex;
242
243#[cfg(test)]
244mod tests;