vectlite 0.1.3
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/LICENSE +21 -0
- package/README.md +75 -0
- package/index.d.ts +182 -0
- package/index.js +284 -0
- package/native/Cargo.toml +23 -0
- package/native/build.rs +3 -0
- package/native/src/lib.rs +1214 -0
- package/native/vectlite-core/Cargo.toml +17 -0
- package/native/vectlite-core/src/lib.rs +3545 -0
- package/package.json +43 -0
- package/scripts/build-addon.mjs +52 -0
- package/scripts/clean-package.mjs +14 -0
- package/scripts/prepare-package.mjs +52 -0
|
@@ -0,0 +1,3545 @@
|
|
|
1
|
+
use std::collections::{BTreeMap, BTreeSet};
|
|
2
|
+
use std::error::Error as StdError;
|
|
3
|
+
use std::fmt;
|
|
4
|
+
use std::fs::{self, File, OpenOptions};
|
|
5
|
+
use std::io::{self, BufReader, BufWriter, ErrorKind, Read, Write};
|
|
6
|
+
use std::path::{Path, PathBuf};
|
|
7
|
+
use std::time::Instant;
|
|
8
|
+
|
|
9
|
+
use fs2::FileExt;
|
|
10
|
+
use hnsw_rs::prelude::*;
|
|
11
|
+
|
|
12
|
+
const MAGIC: &[u8; 4] = b"VDB1";
|
|
13
|
+
const VERSION: u16 = 4;
|
|
14
|
+
const WAL_MAGIC: &[u8; 4] = b"VWL1";
|
|
15
|
+
const TYPE_STRING: u8 = 1;
|
|
16
|
+
const TYPE_INTEGER: u8 = 2;
|
|
17
|
+
const TYPE_FLOAT: u8 = 3;
|
|
18
|
+
const TYPE_BOOLEAN: u8 = 4;
|
|
19
|
+
const TYPE_NULL: u8 = 5;
|
|
20
|
+
const TYPE_LIST: u8 = 6;
|
|
21
|
+
const TYPE_MAP: u8 = 7;
|
|
22
|
+
const DEFAULT_NAMESPACE: &str = "";
|
|
23
|
+
const DEFAULT_VECTOR_NAME: &str = "";
|
|
24
|
+
const ANN_MIN_POINTS: usize = 32;
|
|
25
|
+
const ANN_SEARCH_MIN_POINTS: usize = 128;
|
|
26
|
+
const ANN_OVERSAMPLE: usize = 8;
|
|
27
|
+
const ANN_MIN_CANDIDATES: usize = 64;
|
|
28
|
+
const ANN_M: usize = 16;
|
|
29
|
+
const ANN_EF_CONSTRUCTION: usize = 200;
|
|
30
|
+
const BM25_K1: f32 = 1.2;
|
|
31
|
+
const BM25_B: f32 = 0.75;
|
|
32
|
+
|
|
33
|
+
pub type Result<T> = std::result::Result<T, VectLiteError>;
|
|
34
|
+
pub type Metadata = BTreeMap<String, MetadataValue>;
|
|
35
|
+
pub type SparseVector = BTreeMap<String, f32>;
|
|
36
|
+
pub type NamedVectors = BTreeMap<String, Vec<f32>>;
|
|
37
|
+
type RecordKey = (String, String);
|
|
38
|
+
|
|
39
|
+
#[derive(Clone, Debug)]
|
|
40
|
+
enum WalOp {
|
|
41
|
+
Upsert(Record),
|
|
42
|
+
Delete { namespace: String, id: String },
|
|
43
|
+
}
|
|
44
|
+
|
|
45
|
+
#[derive(Clone, Debug)]
|
|
46
|
+
pub enum WriteOperation {
|
|
47
|
+
Insert(Record),
|
|
48
|
+
Upsert(Record),
|
|
49
|
+
Delete { namespace: String, id: String },
|
|
50
|
+
}
|
|
51
|
+
|
|
52
|
+
#[derive(Default)]
|
|
53
|
+
struct SparseIndex {
|
|
54
|
+
postings: BTreeMap<String, Vec<SparsePosting>>,
|
|
55
|
+
doc_lengths: BTreeMap<RecordKey, f32>,
|
|
56
|
+
avg_doc_len: f32,
|
|
57
|
+
doc_count: usize,
|
|
58
|
+
}
|
|
59
|
+
|
|
60
|
+
#[derive(Clone, Debug)]
|
|
61
|
+
struct SparsePosting {
|
|
62
|
+
key: RecordKey,
|
|
63
|
+
term_weight: f32,
|
|
64
|
+
}
|
|
65
|
+
|
|
66
|
+
#[derive(Debug)]
|
|
67
|
+
pub enum VectLiteError {
|
|
68
|
+
Io(io::Error),
|
|
69
|
+
InvalidFormat(String),
|
|
70
|
+
DimensionMismatch { expected: usize, found: usize },
|
|
71
|
+
DuplicateId { namespace: String, id: String },
|
|
72
|
+
ReadOnly,
|
|
73
|
+
LockContention(String),
|
|
74
|
+
}
|
|
75
|
+
|
|
76
|
+
impl fmt::Display for VectLiteError {
|
|
77
|
+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
78
|
+
match self {
|
|
79
|
+
Self::Io(err) => write!(f, "io error: {err}"),
|
|
80
|
+
Self::InvalidFormat(message) => write!(f, "invalid .vdb format: {message}"),
|
|
81
|
+
Self::DimensionMismatch { expected, found } => {
|
|
82
|
+
write!(
|
|
83
|
+
f,
|
|
84
|
+
"vector dimension mismatch: expected {expected}, found {found}"
|
|
85
|
+
)
|
|
86
|
+
}
|
|
87
|
+
Self::DuplicateId { namespace, id } => {
|
|
88
|
+
if namespace.is_empty() {
|
|
89
|
+
write!(f, "a record with id '{id}' already exists")
|
|
90
|
+
} else {
|
|
91
|
+
write!(
|
|
92
|
+
f,
|
|
93
|
+
"a record with id '{id}' already exists in namespace '{namespace}'"
|
|
94
|
+
)
|
|
95
|
+
}
|
|
96
|
+
}
|
|
97
|
+
Self::ReadOnly => write!(f, "database is opened in read-only mode"),
|
|
98
|
+
Self::LockContention(msg) => write!(f, "lock contention: {msg}"),
|
|
99
|
+
}
|
|
100
|
+
}
|
|
101
|
+
}
|
|
102
|
+
|
|
103
|
+
impl StdError for VectLiteError {
|
|
104
|
+
fn source(&self) -> Option<&(dyn StdError + 'static)> {
|
|
105
|
+
match self {
|
|
106
|
+
Self::Io(err) => Some(err),
|
|
107
|
+
Self::InvalidFormat(_)
|
|
108
|
+
| Self::DimensionMismatch { .. }
|
|
109
|
+
| Self::DuplicateId { .. }
|
|
110
|
+
| Self::ReadOnly
|
|
111
|
+
| Self::LockContention(_) => None,
|
|
112
|
+
}
|
|
113
|
+
}
|
|
114
|
+
}
|
|
115
|
+
|
|
116
|
+
impl From<io::Error> for VectLiteError {
|
|
117
|
+
fn from(value: io::Error) -> Self {
|
|
118
|
+
Self::Io(value)
|
|
119
|
+
}
|
|
120
|
+
}
|
|
121
|
+
|
|
122
|
+
#[derive(Clone, Debug, PartialEq)]
|
|
123
|
+
pub enum MetadataValue {
|
|
124
|
+
String(String),
|
|
125
|
+
Integer(i64),
|
|
126
|
+
Float(f64),
|
|
127
|
+
Boolean(bool),
|
|
128
|
+
Null,
|
|
129
|
+
List(Vec<MetadataValue>),
|
|
130
|
+
Map(BTreeMap<String, MetadataValue>),
|
|
131
|
+
}
|
|
132
|
+
|
|
133
|
+
impl MetadataValue {
|
|
134
|
+
fn type_tag(&self) -> u8 {
|
|
135
|
+
match self {
|
|
136
|
+
Self::String(_) => TYPE_STRING,
|
|
137
|
+
Self::Integer(_) => TYPE_INTEGER,
|
|
138
|
+
Self::Float(_) => TYPE_FLOAT,
|
|
139
|
+
Self::Boolean(_) => TYPE_BOOLEAN,
|
|
140
|
+
Self::Null => TYPE_NULL,
|
|
141
|
+
Self::List(_) => TYPE_LIST,
|
|
142
|
+
Self::Map(_) => TYPE_MAP,
|
|
143
|
+
}
|
|
144
|
+
}
|
|
145
|
+
|
|
146
|
+
fn as_number(&self) -> Option<f64> {
|
|
147
|
+
match self {
|
|
148
|
+
Self::Integer(value) => Some(*value as f64),
|
|
149
|
+
Self::Float(value) => Some(*value),
|
|
150
|
+
Self::String(_) | Self::Boolean(_) | Self::Null | Self::List(_) | Self::Map(_) => None,
|
|
151
|
+
}
|
|
152
|
+
}
|
|
153
|
+
}
|
|
154
|
+
|
|
155
|
+
impl fmt::Display for MetadataValue {
|
|
156
|
+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
157
|
+
match self {
|
|
158
|
+
Self::String(value) => write!(f, "{value}"),
|
|
159
|
+
Self::Integer(value) => write!(f, "{value}"),
|
|
160
|
+
Self::Float(value) => write!(f, "{value}"),
|
|
161
|
+
Self::Boolean(value) => write!(f, "{value}"),
|
|
162
|
+
Self::Null => write!(f, "null"),
|
|
163
|
+
Self::List(values) => {
|
|
164
|
+
write!(f, "[")?;
|
|
165
|
+
for (i, value) in values.iter().enumerate() {
|
|
166
|
+
if i > 0 {
|
|
167
|
+
write!(f, ", ")?;
|
|
168
|
+
}
|
|
169
|
+
write!(f, "{value}")?;
|
|
170
|
+
}
|
|
171
|
+
write!(f, "]")
|
|
172
|
+
}
|
|
173
|
+
Self::Map(entries) => {
|
|
174
|
+
write!(f, "{{")?;
|
|
175
|
+
for (i, (key, value)) in entries.iter().enumerate() {
|
|
176
|
+
if i > 0 {
|
|
177
|
+
write!(f, ", ")?;
|
|
178
|
+
}
|
|
179
|
+
write!(f, "{key}: {value}")?;
|
|
180
|
+
}
|
|
181
|
+
write!(f, "}}")
|
|
182
|
+
}
|
|
183
|
+
}
|
|
184
|
+
}
|
|
185
|
+
}
|
|
186
|
+
|
|
187
|
+
impl From<String> for MetadataValue {
|
|
188
|
+
fn from(value: String) -> Self {
|
|
189
|
+
Self::String(value)
|
|
190
|
+
}
|
|
191
|
+
}
|
|
192
|
+
|
|
193
|
+
impl From<&str> for MetadataValue {
|
|
194
|
+
fn from(value: &str) -> Self {
|
|
195
|
+
Self::String(value.to_owned())
|
|
196
|
+
}
|
|
197
|
+
}
|
|
198
|
+
|
|
199
|
+
impl From<bool> for MetadataValue {
|
|
200
|
+
fn from(value: bool) -> Self {
|
|
201
|
+
Self::Boolean(value)
|
|
202
|
+
}
|
|
203
|
+
}
|
|
204
|
+
|
|
205
|
+
impl From<i64> for MetadataValue {
|
|
206
|
+
fn from(value: i64) -> Self {
|
|
207
|
+
Self::Integer(value)
|
|
208
|
+
}
|
|
209
|
+
}
|
|
210
|
+
|
|
211
|
+
impl From<i32> for MetadataValue {
|
|
212
|
+
fn from(value: i32) -> Self {
|
|
213
|
+
Self::Integer(value.into())
|
|
214
|
+
}
|
|
215
|
+
}
|
|
216
|
+
|
|
217
|
+
impl From<f64> for MetadataValue {
|
|
218
|
+
fn from(value: f64) -> Self {
|
|
219
|
+
Self::Float(value)
|
|
220
|
+
}
|
|
221
|
+
}
|
|
222
|
+
|
|
223
|
+
impl From<f32> for MetadataValue {
|
|
224
|
+
fn from(value: f32) -> Self {
|
|
225
|
+
Self::Float(value.into())
|
|
226
|
+
}
|
|
227
|
+
}
|
|
228
|
+
|
|
229
|
+
#[derive(Clone, Debug, PartialEq)]
|
|
230
|
+
pub struct Record {
|
|
231
|
+
pub namespace: String,
|
|
232
|
+
pub id: String,
|
|
233
|
+
pub vector: Vec<f32>,
|
|
234
|
+
pub vectors: NamedVectors,
|
|
235
|
+
pub sparse: SparseVector,
|
|
236
|
+
pub metadata: Metadata,
|
|
237
|
+
}
|
|
238
|
+
|
|
239
|
+
impl Record {
|
|
240
|
+
fn vector_for(&self, vector_name: Option<&str>) -> Option<&[f32]> {
|
|
241
|
+
match vector_name {
|
|
242
|
+
Some(vector_name) if !vector_name.is_empty() => {
|
|
243
|
+
self.vectors.get(vector_name).map(Vec::as_slice)
|
|
244
|
+
}
|
|
245
|
+
Some(_) | None => Some(self.vector.as_slice()),
|
|
246
|
+
}
|
|
247
|
+
}
|
|
248
|
+
|
|
249
|
+
fn dense_vectors(&self) -> impl Iterator<Item = (&str, &Vec<f32>)> {
|
|
250
|
+
std::iter::once((DEFAULT_VECTOR_NAME, &self.vector)).chain(
|
|
251
|
+
self.vectors
|
|
252
|
+
.iter()
|
|
253
|
+
.map(|(name, vector)| (name.as_str(), vector)),
|
|
254
|
+
)
|
|
255
|
+
}
|
|
256
|
+
}
|
|
257
|
+
|
|
258
|
+
#[derive(Clone, Debug, PartialEq)]
|
|
259
|
+
pub enum MetadataFilter {
|
|
260
|
+
Eq {
|
|
261
|
+
key: String,
|
|
262
|
+
value: MetadataValue,
|
|
263
|
+
},
|
|
264
|
+
NotEq {
|
|
265
|
+
key: String,
|
|
266
|
+
value: MetadataValue,
|
|
267
|
+
},
|
|
268
|
+
In {
|
|
269
|
+
key: String,
|
|
270
|
+
values: Vec<MetadataValue>,
|
|
271
|
+
},
|
|
272
|
+
NotIn {
|
|
273
|
+
key: String,
|
|
274
|
+
values: Vec<MetadataValue>,
|
|
275
|
+
},
|
|
276
|
+
Contains {
|
|
277
|
+
key: String,
|
|
278
|
+
needle: String,
|
|
279
|
+
},
|
|
280
|
+
GreaterThan {
|
|
281
|
+
key: String,
|
|
282
|
+
value: f64,
|
|
283
|
+
},
|
|
284
|
+
GreaterThanOrEqual {
|
|
285
|
+
key: String,
|
|
286
|
+
value: f64,
|
|
287
|
+
},
|
|
288
|
+
LessThan {
|
|
289
|
+
key: String,
|
|
290
|
+
value: f64,
|
|
291
|
+
},
|
|
292
|
+
LessThanOrEqual {
|
|
293
|
+
key: String,
|
|
294
|
+
value: f64,
|
|
295
|
+
},
|
|
296
|
+
Exists {
|
|
297
|
+
key: String,
|
|
298
|
+
},
|
|
299
|
+
/// Matches if a list-typed field has at least one element satisfying `filter`.
|
|
300
|
+
ElemMatch {
|
|
301
|
+
key: String,
|
|
302
|
+
filter: Box<MetadataFilter>,
|
|
303
|
+
},
|
|
304
|
+
/// Matches if a list-typed field has exactly `size` elements.
|
|
305
|
+
Size {
|
|
306
|
+
key: String,
|
|
307
|
+
size: usize,
|
|
308
|
+
},
|
|
309
|
+
Not(Box<MetadataFilter>),
|
|
310
|
+
And(Vec<MetadataFilter>),
|
|
311
|
+
Or(Vec<MetadataFilter>),
|
|
312
|
+
}
|
|
313
|
+
|
|
314
|
+
impl MetadataFilter {
|
|
315
|
+
pub fn eq(key: impl Into<String>, value: impl Into<MetadataValue>) -> Self {
|
|
316
|
+
Self::Eq {
|
|
317
|
+
key: key.into(),
|
|
318
|
+
value: value.into(),
|
|
319
|
+
}
|
|
320
|
+
}
|
|
321
|
+
|
|
322
|
+
pub fn ne(key: impl Into<String>, value: impl Into<MetadataValue>) -> Self {
|
|
323
|
+
Self::NotEq {
|
|
324
|
+
key: key.into(),
|
|
325
|
+
value: value.into(),
|
|
326
|
+
}
|
|
327
|
+
}
|
|
328
|
+
|
|
329
|
+
pub fn r#in(key: impl Into<String>, values: Vec<MetadataValue>) -> Self {
|
|
330
|
+
Self::In {
|
|
331
|
+
key: key.into(),
|
|
332
|
+
values,
|
|
333
|
+
}
|
|
334
|
+
}
|
|
335
|
+
|
|
336
|
+
pub fn nin(key: impl Into<String>, values: Vec<MetadataValue>) -> Self {
|
|
337
|
+
Self::NotIn {
|
|
338
|
+
key: key.into(),
|
|
339
|
+
values,
|
|
340
|
+
}
|
|
341
|
+
}
|
|
342
|
+
|
|
343
|
+
pub fn contains(key: impl Into<String>, needle: impl Into<String>) -> Self {
|
|
344
|
+
Self::Contains {
|
|
345
|
+
key: key.into(),
|
|
346
|
+
needle: needle.into(),
|
|
347
|
+
}
|
|
348
|
+
}
|
|
349
|
+
|
|
350
|
+
pub fn gt(key: impl Into<String>, value: f64) -> Self {
|
|
351
|
+
Self::GreaterThan {
|
|
352
|
+
key: key.into(),
|
|
353
|
+
value,
|
|
354
|
+
}
|
|
355
|
+
}
|
|
356
|
+
|
|
357
|
+
pub fn gte(key: impl Into<String>, value: f64) -> Self {
|
|
358
|
+
Self::GreaterThanOrEqual {
|
|
359
|
+
key: key.into(),
|
|
360
|
+
value,
|
|
361
|
+
}
|
|
362
|
+
}
|
|
363
|
+
|
|
364
|
+
pub fn lt(key: impl Into<String>, value: f64) -> Self {
|
|
365
|
+
Self::LessThan {
|
|
366
|
+
key: key.into(),
|
|
367
|
+
value,
|
|
368
|
+
}
|
|
369
|
+
}
|
|
370
|
+
|
|
371
|
+
pub fn lte(key: impl Into<String>, value: f64) -> Self {
|
|
372
|
+
Self::LessThanOrEqual {
|
|
373
|
+
key: key.into(),
|
|
374
|
+
value,
|
|
375
|
+
}
|
|
376
|
+
}
|
|
377
|
+
|
|
378
|
+
pub fn exists(key: impl Into<String>) -> Self {
|
|
379
|
+
Self::Exists { key: key.into() }
|
|
380
|
+
}
|
|
381
|
+
|
|
382
|
+
pub fn elem_match(key: impl Into<String>, filter: MetadataFilter) -> Self {
|
|
383
|
+
Self::ElemMatch {
|
|
384
|
+
key: key.into(),
|
|
385
|
+
filter: Box::new(filter),
|
|
386
|
+
}
|
|
387
|
+
}
|
|
388
|
+
|
|
389
|
+
pub fn size(key: impl Into<String>, size: usize) -> Self {
|
|
390
|
+
Self::Size {
|
|
391
|
+
key: key.into(),
|
|
392
|
+
size,
|
|
393
|
+
}
|
|
394
|
+
}
|
|
395
|
+
|
|
396
|
+
pub fn not(filter: MetadataFilter) -> Self {
|
|
397
|
+
Self::Not(Box::new(filter))
|
|
398
|
+
}
|
|
399
|
+
|
|
400
|
+
pub fn and(filters: Vec<MetadataFilter>) -> Self {
|
|
401
|
+
Self::And(filters)
|
|
402
|
+
}
|
|
403
|
+
|
|
404
|
+
pub fn or(filters: Vec<MetadataFilter>) -> Self {
|
|
405
|
+
Self::Or(filters)
|
|
406
|
+
}
|
|
407
|
+
|
|
408
|
+
fn matches(&self, metadata: &Metadata) -> bool {
|
|
409
|
+
match self {
|
|
410
|
+
Self::Eq { key, value } => resolve_dot_path(metadata, key) == Some(value),
|
|
411
|
+
Self::NotEq { key, value } => {
|
|
412
|
+
resolve_dot_path(metadata, key).is_some_and(|candidate| candidate != value)
|
|
413
|
+
}
|
|
414
|
+
Self::In { key, values } => {
|
|
415
|
+
resolve_dot_path(metadata, key).is_some_and(|candidate| values.contains(candidate))
|
|
416
|
+
}
|
|
417
|
+
Self::NotIn { key, values } => {
|
|
418
|
+
resolve_dot_path(metadata, key).is_some_and(|candidate| !values.contains(candidate))
|
|
419
|
+
}
|
|
420
|
+
Self::Contains { key, needle } => resolve_dot_path(metadata, key)
|
|
421
|
+
.and_then(|value| match value {
|
|
422
|
+
MetadataValue::String(value) => Some(value.contains(needle)),
|
|
423
|
+
MetadataValue::Integer(_)
|
|
424
|
+
| MetadataValue::Float(_)
|
|
425
|
+
| MetadataValue::Boolean(_)
|
|
426
|
+
| MetadataValue::Null
|
|
427
|
+
| MetadataValue::List(_)
|
|
428
|
+
| MetadataValue::Map(_) => None,
|
|
429
|
+
})
|
|
430
|
+
.unwrap_or(false),
|
|
431
|
+
Self::GreaterThan { key, value } => resolve_dot_path(metadata, key)
|
|
432
|
+
.and_then(MetadataValue::as_number)
|
|
433
|
+
.map(|candidate| candidate > *value)
|
|
434
|
+
.unwrap_or(false),
|
|
435
|
+
Self::GreaterThanOrEqual { key, value } => resolve_dot_path(metadata, key)
|
|
436
|
+
.and_then(MetadataValue::as_number)
|
|
437
|
+
.map(|candidate| candidate >= *value)
|
|
438
|
+
.unwrap_or(false),
|
|
439
|
+
Self::LessThan { key, value } => resolve_dot_path(metadata, key)
|
|
440
|
+
.and_then(MetadataValue::as_number)
|
|
441
|
+
.map(|candidate| candidate < *value)
|
|
442
|
+
.unwrap_or(false),
|
|
443
|
+
Self::LessThanOrEqual { key, value } => resolve_dot_path(metadata, key)
|
|
444
|
+
.and_then(MetadataValue::as_number)
|
|
445
|
+
.map(|candidate| candidate <= *value)
|
|
446
|
+
.unwrap_or(false),
|
|
447
|
+
Self::Exists { key } => resolve_dot_path(metadata, key).is_some(),
|
|
448
|
+
Self::ElemMatch { key, filter } => {
|
|
449
|
+
resolve_dot_path(metadata, key)
|
|
450
|
+
.and_then(|value| match value {
|
|
451
|
+
MetadataValue::List(items) => Some(items.iter().any(|item| {
|
|
452
|
+
// Wrap the single element in a temporary metadata map
|
|
453
|
+
// so the sub-filter can match against it as a virtual record.
|
|
454
|
+
let mut elem_meta = Metadata::new();
|
|
455
|
+
// Flatten: if the item is a Map, use it directly.
|
|
456
|
+
// Otherwise, create a pseudo-map with the element
|
|
457
|
+
// as value under every key the filter references.
|
|
458
|
+
match item {
|
|
459
|
+
MetadataValue::Map(map) => {
|
|
460
|
+
for (k, v) in map {
|
|
461
|
+
elem_meta.insert(k.clone(), v.clone());
|
|
462
|
+
}
|
|
463
|
+
}
|
|
464
|
+
_ => {
|
|
465
|
+
// Put the scalar value as "_" so simple
|
|
466
|
+
// equality filters like {"$eq": 42} work.
|
|
467
|
+
elem_meta.insert("_".to_owned(), item.clone());
|
|
468
|
+
}
|
|
469
|
+
}
|
|
470
|
+
filter.matches(&elem_meta)
|
|
471
|
+
})),
|
|
472
|
+
_ => None,
|
|
473
|
+
})
|
|
474
|
+
.unwrap_or(false)
|
|
475
|
+
}
|
|
476
|
+
Self::Size { key, size } => resolve_dot_path(metadata, key)
|
|
477
|
+
.and_then(|value| match value {
|
|
478
|
+
MetadataValue::List(items) => Some(items.len() == *size),
|
|
479
|
+
_ => None,
|
|
480
|
+
})
|
|
481
|
+
.unwrap_or(false),
|
|
482
|
+
Self::Not(filter) => !filter.matches(metadata),
|
|
483
|
+
Self::And(filters) => filters.iter().all(|filter| filter.matches(metadata)),
|
|
484
|
+
Self::Or(filters) => filters.iter().any(|filter| filter.matches(metadata)),
|
|
485
|
+
}
|
|
486
|
+
}
|
|
487
|
+
}
|
|
488
|
+
|
|
489
|
+
/// Resolves a dot-separated key path against a metadata map.
|
|
490
|
+
/// E.g. "extra.nested_key" first looks up "extra" then "nested_key" inside it.
|
|
491
|
+
fn resolve_dot_path<'a>(metadata: &'a Metadata, key: &str) -> Option<&'a MetadataValue> {
|
|
492
|
+
// Fast path: no dots → direct lookup (also handles keys that literally contain dots
|
|
493
|
+
// which were stored before this feature).
|
|
494
|
+
if !key.contains('.') || metadata.contains_key(key) {
|
|
495
|
+
return metadata.get(key);
|
|
496
|
+
}
|
|
497
|
+
|
|
498
|
+
let mut parts = key.split('.');
|
|
499
|
+
let first = parts.next()?;
|
|
500
|
+
let mut current = metadata.get(first)?;
|
|
501
|
+
|
|
502
|
+
for part in parts {
|
|
503
|
+
match current {
|
|
504
|
+
MetadataValue::Map(map) => {
|
|
505
|
+
current = map.get(part)?;
|
|
506
|
+
}
|
|
507
|
+
_ => return None,
|
|
508
|
+
}
|
|
509
|
+
}
|
|
510
|
+
|
|
511
|
+
Some(current)
|
|
512
|
+
}
|
|
513
|
+
|
|
514
|
+
#[derive(Clone, Debug)]
|
|
515
|
+
pub struct SearchOptions {
|
|
516
|
+
pub top_k: usize,
|
|
517
|
+
pub filter: Option<MetadataFilter>,
|
|
518
|
+
}
|
|
519
|
+
|
|
520
|
+
impl Default for SearchOptions {
|
|
521
|
+
fn default() -> Self {
|
|
522
|
+
Self {
|
|
523
|
+
top_k: 10,
|
|
524
|
+
filter: None,
|
|
525
|
+
}
|
|
526
|
+
}
|
|
527
|
+
}
|
|
528
|
+
|
|
529
|
+
#[derive(Clone, Debug)]
|
|
530
|
+
pub struct HybridSearchOptions {
|
|
531
|
+
pub top_k: usize,
|
|
532
|
+
pub filter: Option<MetadataFilter>,
|
|
533
|
+
pub dense_weight: f32,
|
|
534
|
+
pub sparse_weight: f32,
|
|
535
|
+
pub fetch_k: usize,
|
|
536
|
+
pub mmr_lambda: Option<f32>,
|
|
537
|
+
pub vector_name: Option<String>,
|
|
538
|
+
pub fusion: FusionStrategy,
|
|
539
|
+
/// Multi-vector search: provide per-vector-name queries and their weights.
|
|
540
|
+
/// When non-empty, the dense score is the weighted sum of cosine
|
|
541
|
+
/// similarities over the given vector spaces, and `vector_name` is ignored.
|
|
542
|
+
pub multi_vector_queries: BTreeMap<String, (Vec<f32>, f32)>,
|
|
543
|
+
}
|
|
544
|
+
|
|
545
|
+
impl Default for HybridSearchOptions {
|
|
546
|
+
fn default() -> Self {
|
|
547
|
+
Self {
|
|
548
|
+
top_k: 10,
|
|
549
|
+
filter: None,
|
|
550
|
+
dense_weight: 1.0,
|
|
551
|
+
sparse_weight: 1.0,
|
|
552
|
+
fetch_k: 0,
|
|
553
|
+
mmr_lambda: None,
|
|
554
|
+
vector_name: None,
|
|
555
|
+
fusion: FusionStrategy::Linear,
|
|
556
|
+
multi_vector_queries: BTreeMap::new(),
|
|
557
|
+
}
|
|
558
|
+
}
|
|
559
|
+
}
|
|
560
|
+
|
|
561
|
+
#[derive(Clone, Debug, PartialEq)]
|
|
562
|
+
pub enum FusionStrategy {
|
|
563
|
+
Linear,
|
|
564
|
+
Rrf { rank_constant: usize },
|
|
565
|
+
}
|
|
566
|
+
|
|
567
|
+
#[derive(Clone, Debug, PartialEq)]
|
|
568
|
+
pub struct SearchResult {
|
|
569
|
+
pub namespace: String,
|
|
570
|
+
pub id: String,
|
|
571
|
+
pub score: f32,
|
|
572
|
+
pub dense_score: f32,
|
|
573
|
+
pub sparse_score: f32,
|
|
574
|
+
pub vector_name: Option<String>,
|
|
575
|
+
pub matched_terms: Vec<String>,
|
|
576
|
+
pub dense_rank: Option<usize>,
|
|
577
|
+
pub sparse_rank: Option<usize>,
|
|
578
|
+
pub metadata: Metadata,
|
|
579
|
+
/// Per-term BM25 contribution for this result (only when explain requested).
|
|
580
|
+
pub bm25_term_scores: BTreeMap<String, f32>,
|
|
581
|
+
}
|
|
582
|
+
|
|
583
|
+
#[derive(Clone, Debug, PartialEq, Default)]
|
|
584
|
+
pub struct SearchStats {
|
|
585
|
+
pub used_ann: bool,
|
|
586
|
+
pub ann_candidate_count: usize,
|
|
587
|
+
pub exact_fallback: bool,
|
|
588
|
+
pub considered_count: usize,
|
|
589
|
+
pub fetch_k: usize,
|
|
590
|
+
pub mmr_applied: bool,
|
|
591
|
+
pub sparse_candidate_count: usize,
|
|
592
|
+
pub ann_loaded_from_disk: bool,
|
|
593
|
+
pub wal_entries_replayed: usize,
|
|
594
|
+
pub fusion: String,
|
|
595
|
+
/// Timing breakdown in microseconds.
|
|
596
|
+
pub timings: SearchTimings,
|
|
597
|
+
}
|
|
598
|
+
|
|
599
|
+
#[derive(Clone, Debug, Default, PartialEq)]
|
|
600
|
+
pub struct SearchTimings {
|
|
601
|
+
/// Microseconds spent on dense (ANN or brute-force) scoring.
|
|
602
|
+
pub dense_us: u64,
|
|
603
|
+
/// Microseconds spent on sparse (BM25) scoring.
|
|
604
|
+
pub sparse_us: u64,
|
|
605
|
+
/// Microseconds spent on fusion (combining dense + sparse).
|
|
606
|
+
pub fusion_us: u64,
|
|
607
|
+
/// Total end-to-end microseconds.
|
|
608
|
+
pub total_us: u64,
|
|
609
|
+
}
|
|
610
|
+
|
|
611
|
+
#[derive(Clone, Debug, PartialEq)]
|
|
612
|
+
pub struct SearchOutcome {
|
|
613
|
+
pub results: Vec<SearchResult>,
|
|
614
|
+
pub stats: SearchStats,
|
|
615
|
+
}
|
|
616
|
+
|
|
617
|
+
/// A store manages a directory of independent physical collections, each
|
|
618
|
+
/// backed by its own `.vdb` file with its own dimension, WAL, and ANN index.
|
|
619
|
+
pub struct Store {
|
|
620
|
+
root: PathBuf,
|
|
621
|
+
}
|
|
622
|
+
|
|
623
|
+
impl Store {
|
|
624
|
+
/// Open (or create) a store rooted at the given directory.
|
|
625
|
+
pub fn open(root: impl AsRef<Path>) -> Result<Self> {
|
|
626
|
+
let root = root.as_ref().to_path_buf();
|
|
627
|
+
if !root.exists() {
|
|
628
|
+
fs::create_dir_all(&root)?;
|
|
629
|
+
}
|
|
630
|
+
Ok(Self { root })
|
|
631
|
+
}
|
|
632
|
+
|
|
633
|
+
/// Return the root directory.
|
|
634
|
+
pub fn root(&self) -> &Path {
|
|
635
|
+
&self.root
|
|
636
|
+
}
|
|
637
|
+
|
|
638
|
+
/// Create a new collection. Returns an error if it already exists.
|
|
639
|
+
pub fn create_collection(&self, name: &str, dimension: usize) -> Result<Database> {
|
|
640
|
+
let path = self.collection_path(name);
|
|
641
|
+
if path.exists() {
|
|
642
|
+
return Err(VectLiteError::InvalidFormat(format!(
|
|
643
|
+
"collection '{name}' already exists"
|
|
644
|
+
)));
|
|
645
|
+
}
|
|
646
|
+
Database::create(path, dimension)
|
|
647
|
+
}
|
|
648
|
+
|
|
649
|
+
/// Open an existing collection.
|
|
650
|
+
pub fn open_collection(&self, name: &str) -> Result<Database> {
|
|
651
|
+
let path = self.collection_path(name);
|
|
652
|
+
if !path.exists() {
|
|
653
|
+
return Err(VectLiteError::InvalidFormat(format!(
|
|
654
|
+
"collection '{name}' does not exist"
|
|
655
|
+
)));
|
|
656
|
+
}
|
|
657
|
+
Database::open(path)
|
|
658
|
+
}
|
|
659
|
+
|
|
660
|
+
/// Open an existing collection in read-only mode.
|
|
661
|
+
pub fn open_collection_read_only(&self, name: &str) -> Result<Database> {
|
|
662
|
+
let path = self.collection_path(name);
|
|
663
|
+
if !path.exists() {
|
|
664
|
+
return Err(VectLiteError::InvalidFormat(format!(
|
|
665
|
+
"collection '{name}' does not exist"
|
|
666
|
+
)));
|
|
667
|
+
}
|
|
668
|
+
Database::open_read_only(path)
|
|
669
|
+
}
|
|
670
|
+
|
|
671
|
+
/// Open an existing collection or create it with the given dimension.
|
|
672
|
+
pub fn open_or_create_collection(&self, name: &str, dimension: usize) -> Result<Database> {
|
|
673
|
+
Database::open_or_create(self.collection_path(name), dimension)
|
|
674
|
+
}
|
|
675
|
+
|
|
676
|
+
/// Drop a collection, deleting all its files.
|
|
677
|
+
pub fn drop_collection(&self, name: &str) -> Result<bool> {
|
|
678
|
+
let path = self.collection_path(name);
|
|
679
|
+
if !path.exists() {
|
|
680
|
+
return Ok(false);
|
|
681
|
+
}
|
|
682
|
+
// Remove main file, WAL, ANN sidecar files
|
|
683
|
+
let _ = fs::remove_file(&path);
|
|
684
|
+
let wal = wal_path(&path);
|
|
685
|
+
let _ = fs::remove_file(&wal);
|
|
686
|
+
let manifest = ann_manifest_path(&path);
|
|
687
|
+
let _ = fs::remove_file(&manifest);
|
|
688
|
+
// Remove any .hnsw.* sidecar files
|
|
689
|
+
if let Some(parent) = path.parent() {
|
|
690
|
+
if let Some(stem) = path.file_name().and_then(|n| n.to_str()) {
|
|
691
|
+
if let Ok(entries) = fs::read_dir(parent) {
|
|
692
|
+
for entry in entries.flatten() {
|
|
693
|
+
if let Some(fname) = entry.file_name().to_str() {
|
|
694
|
+
if fname.starts_with(&format!("{stem}.ann.")) {
|
|
695
|
+
let _ = fs::remove_file(entry.path());
|
|
696
|
+
}
|
|
697
|
+
}
|
|
698
|
+
}
|
|
699
|
+
}
|
|
700
|
+
}
|
|
701
|
+
}
|
|
702
|
+
Ok(true)
|
|
703
|
+
}
|
|
704
|
+
|
|
705
|
+
/// List all collection names in this store.
|
|
706
|
+
pub fn collections(&self) -> Result<Vec<String>> {
|
|
707
|
+
let mut names = Vec::new();
|
|
708
|
+
for entry in fs::read_dir(&self.root)? {
|
|
709
|
+
let entry = entry?;
|
|
710
|
+
let path = entry.path();
|
|
711
|
+
if path.extension().and_then(|ext| ext.to_str()) == Some("vdb") {
|
|
712
|
+
if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) {
|
|
713
|
+
names.push(stem.to_owned());
|
|
714
|
+
}
|
|
715
|
+
}
|
|
716
|
+
}
|
|
717
|
+
names.sort();
|
|
718
|
+
Ok(names)
|
|
719
|
+
}
|
|
720
|
+
|
|
721
|
+
fn collection_path(&self, name: &str) -> PathBuf {
|
|
722
|
+
self.root.join(format!("{name}.vdb"))
|
|
723
|
+
}
|
|
724
|
+
}
|
|
725
|
+
|
|
726
|
+
pub struct Database {
|
|
727
|
+
path: PathBuf,
|
|
728
|
+
wal_path: PathBuf,
|
|
729
|
+
dimension: usize,
|
|
730
|
+
records: BTreeMap<(String, String), Record>,
|
|
731
|
+
ann: AnnCatalog,
|
|
732
|
+
sparse_index: SparseIndex,
|
|
733
|
+
wal_entries_replayed: usize,
|
|
734
|
+
ann_loaded_from_disk: bool,
|
|
735
|
+
read_only: bool,
|
|
736
|
+
/// Holds the lock file open for the lifetime of the database.
|
|
737
|
+
/// Dropping this releases the advisory lock.
|
|
738
|
+
_lock_file: Option<File>,
|
|
739
|
+
}
|
|
740
|
+
|
|
741
|
+
#[derive(Default)]
|
|
742
|
+
struct AnnCatalog {
|
|
743
|
+
global: BTreeMap<String, AnnIndex>,
|
|
744
|
+
namespaces: BTreeMap<String, BTreeMap<String, AnnIndex>>,
|
|
745
|
+
}
|
|
746
|
+
|
|
747
|
+
struct AnnIndex {
|
|
748
|
+
hnsw: Hnsw<'static, f32, DistCosine>,
|
|
749
|
+
keys: Vec<RecordKey>,
|
|
750
|
+
}
|
|
751
|
+
|
|
752
|
+
struct AnnManifestEntry {
|
|
753
|
+
namespace: Option<String>,
|
|
754
|
+
vector_name: String,
|
|
755
|
+
record_count: usize,
|
|
756
|
+
key_signature: u64,
|
|
757
|
+
keys: Vec<RecordKey>,
|
|
758
|
+
}
|
|
759
|
+
|
|
760
|
+
#[derive(Clone)]
|
|
761
|
+
struct ScoredRecord<'a> {
|
|
762
|
+
record: &'a Record,
|
|
763
|
+
score: f32,
|
|
764
|
+
dense_score: f32,
|
|
765
|
+
sparse_score: f32,
|
|
766
|
+
vector_name: Option<String>,
|
|
767
|
+
matched_terms: Vec<String>,
|
|
768
|
+
dense_rank: Option<usize>,
|
|
769
|
+
sparse_rank: Option<usize>,
|
|
770
|
+
bm25_term_scores: BTreeMap<String, f32>,
|
|
771
|
+
}
|
|
772
|
+
|
|
773
|
+
impl Database {
|
|
774
|
+
pub fn create(path: impl AsRef<Path>, dimension: usize) -> Result<Self> {
|
|
775
|
+
ensure_dimension(dimension)?;
|
|
776
|
+
let lock = acquire_exclusive_lock(path.as_ref())?;
|
|
777
|
+
|
|
778
|
+
let mut database = Self {
|
|
779
|
+
path: path.as_ref().to_path_buf(),
|
|
780
|
+
wal_path: wal_path(path.as_ref()),
|
|
781
|
+
dimension,
|
|
782
|
+
records: BTreeMap::new(),
|
|
783
|
+
ann: AnnCatalog::default(),
|
|
784
|
+
sparse_index: SparseIndex::default(),
|
|
785
|
+
wal_entries_replayed: 0,
|
|
786
|
+
ann_loaded_from_disk: false,
|
|
787
|
+
read_only: false,
|
|
788
|
+
_lock_file: Some(lock),
|
|
789
|
+
};
|
|
790
|
+
|
|
791
|
+
database.flush()?;
|
|
792
|
+
Ok(database)
|
|
793
|
+
}
|
|
794
|
+
|
|
795
|
+
pub fn open(path: impl AsRef<Path>) -> Result<Self> {
|
|
796
|
+
let path = path.as_ref().to_path_buf();
|
|
797
|
+
let lock = acquire_exclusive_lock(&path)?;
|
|
798
|
+
let mut file = File::open(&path)?;
|
|
799
|
+
let mut database = Self::read_from(&path, &mut file)?;
|
|
800
|
+
database._lock_file = Some(lock);
|
|
801
|
+
database.read_only = false;
|
|
802
|
+
database.replay_wal()?;
|
|
803
|
+
database.rebuild_sparse_index();
|
|
804
|
+
database.ann_loaded_from_disk = database.try_load_ann_from_disk();
|
|
805
|
+
if !database.ann_loaded_from_disk {
|
|
806
|
+
database.rebuild_ann();
|
|
807
|
+
}
|
|
808
|
+
Ok(database)
|
|
809
|
+
}
|
|
810
|
+
|
|
811
|
+
/// Open an existing database in read-only mode. Acquires a shared lock
|
|
812
|
+
/// so multiple readers can coexist. All write operations will return
|
|
813
|
+
/// `VectLiteError::ReadOnly`.
|
|
814
|
+
pub fn open_read_only(path: impl AsRef<Path>) -> Result<Self> {
|
|
815
|
+
let path = path.as_ref().to_path_buf();
|
|
816
|
+
let lock = acquire_shared_lock(&path)?;
|
|
817
|
+
let mut file = File::open(&path)?;
|
|
818
|
+
let mut database = Self::read_from(&path, &mut file)?;
|
|
819
|
+
database._lock_file = Some(lock);
|
|
820
|
+
database.read_only = true;
|
|
821
|
+
database.replay_wal()?;
|
|
822
|
+
database.rebuild_sparse_index();
|
|
823
|
+
database.ann_loaded_from_disk = database.try_load_ann_from_disk();
|
|
824
|
+
if !database.ann_loaded_from_disk {
|
|
825
|
+
database.rebuild_ann();
|
|
826
|
+
}
|
|
827
|
+
Ok(database)
|
|
828
|
+
}
|
|
829
|
+
|
|
830
|
+
pub fn is_read_only(&self) -> bool {
|
|
831
|
+
self.read_only
|
|
832
|
+
}
|
|
833
|
+
|
|
834
|
+
fn check_writable(&self) -> Result<()> {
|
|
835
|
+
if self.read_only {
|
|
836
|
+
return Err(VectLiteError::ReadOnly);
|
|
837
|
+
}
|
|
838
|
+
Ok(())
|
|
839
|
+
}
|
|
840
|
+
|
|
841
|
+
pub fn open_or_create(path: impl AsRef<Path>, dimension: usize) -> Result<Self> {
|
|
842
|
+
if path.as_ref().exists() {
|
|
843
|
+
let database = Self::open(path)?;
|
|
844
|
+
if database.dimension != dimension {
|
|
845
|
+
return Err(VectLiteError::DimensionMismatch {
|
|
846
|
+
expected: database.dimension,
|
|
847
|
+
found: dimension,
|
|
848
|
+
});
|
|
849
|
+
}
|
|
850
|
+
Ok(database)
|
|
851
|
+
} else {
|
|
852
|
+
Self::create(path, dimension)
|
|
853
|
+
}
|
|
854
|
+
}
|
|
855
|
+
|
|
856
|
+
pub fn path(&self) -> &Path {
|
|
857
|
+
&self.path
|
|
858
|
+
}
|
|
859
|
+
|
|
860
|
+
pub fn dimension(&self) -> usize {
|
|
861
|
+
self.dimension
|
|
862
|
+
}
|
|
863
|
+
|
|
864
|
+
pub fn len(&self) -> usize {
|
|
865
|
+
self.records.len()
|
|
866
|
+
}
|
|
867
|
+
|
|
868
|
+
pub fn is_empty(&self) -> bool {
|
|
869
|
+
self.records.is_empty()
|
|
870
|
+
}
|
|
871
|
+
|
|
872
|
+
pub fn insert(
|
|
873
|
+
&mut self,
|
|
874
|
+
id: impl Into<String>,
|
|
875
|
+
vector: impl Into<Vec<f32>>,
|
|
876
|
+
metadata: Metadata,
|
|
877
|
+
) -> Result<()> {
|
|
878
|
+
self.insert_with_vectors_in_namespace(
|
|
879
|
+
DEFAULT_NAMESPACE,
|
|
880
|
+
id,
|
|
881
|
+
vector,
|
|
882
|
+
NamedVectors::new(),
|
|
883
|
+
SparseVector::new(),
|
|
884
|
+
metadata,
|
|
885
|
+
)
|
|
886
|
+
}
|
|
887
|
+
|
|
888
|
+
pub fn upsert(
|
|
889
|
+
&mut self,
|
|
890
|
+
id: impl Into<String>,
|
|
891
|
+
vector: impl Into<Vec<f32>>,
|
|
892
|
+
metadata: Metadata,
|
|
893
|
+
) -> Result<()> {
|
|
894
|
+
self.insert(id, vector, metadata)
|
|
895
|
+
}
|
|
896
|
+
|
|
897
|
+
pub fn insert_in_namespace(
|
|
898
|
+
&mut self,
|
|
899
|
+
namespace: impl Into<String>,
|
|
900
|
+
id: impl Into<String>,
|
|
901
|
+
vector: impl Into<Vec<f32>>,
|
|
902
|
+
metadata: Metadata,
|
|
903
|
+
) -> Result<()> {
|
|
904
|
+
self.insert_with_vectors_in_namespace(
|
|
905
|
+
namespace,
|
|
906
|
+
id,
|
|
907
|
+
vector,
|
|
908
|
+
NamedVectors::new(),
|
|
909
|
+
SparseVector::new(),
|
|
910
|
+
metadata,
|
|
911
|
+
)
|
|
912
|
+
}
|
|
913
|
+
|
|
914
|
+
pub fn insert_with_vectors_in_namespace(
|
|
915
|
+
&mut self,
|
|
916
|
+
namespace: impl Into<String>,
|
|
917
|
+
id: impl Into<String>,
|
|
918
|
+
vector: impl Into<Vec<f32>>,
|
|
919
|
+
vectors: NamedVectors,
|
|
920
|
+
sparse: SparseVector,
|
|
921
|
+
metadata: Metadata,
|
|
922
|
+
) -> Result<()> {
|
|
923
|
+
self.check_writable()?;
|
|
924
|
+
let record = self.record_from_parts(namespace, id, vector, vectors, sparse, metadata)?;
|
|
925
|
+
let key = (record.namespace.clone(), record.id.clone());
|
|
926
|
+
if self.records.contains_key(&key) {
|
|
927
|
+
return Err(VectLiteError::DuplicateId {
|
|
928
|
+
namespace: key.0,
|
|
929
|
+
id: key.1,
|
|
930
|
+
});
|
|
931
|
+
}
|
|
932
|
+
self.apply_wal_batch(vec![WalOp::Upsert(record)])?;
|
|
933
|
+
Ok(())
|
|
934
|
+
}
|
|
935
|
+
|
|
936
|
+
pub fn insert_with_sparse_in_namespace(
|
|
937
|
+
&mut self,
|
|
938
|
+
namespace: impl Into<String>,
|
|
939
|
+
id: impl Into<String>,
|
|
940
|
+
vector: impl Into<Vec<f32>>,
|
|
941
|
+
sparse: SparseVector,
|
|
942
|
+
metadata: Metadata,
|
|
943
|
+
) -> Result<()> {
|
|
944
|
+
self.insert_with_vectors_in_namespace(
|
|
945
|
+
namespace,
|
|
946
|
+
id,
|
|
947
|
+
vector,
|
|
948
|
+
NamedVectors::new(),
|
|
949
|
+
sparse,
|
|
950
|
+
metadata,
|
|
951
|
+
)
|
|
952
|
+
}
|
|
953
|
+
|
|
954
|
+
pub fn upsert_in_namespace(
|
|
955
|
+
&mut self,
|
|
956
|
+
namespace: impl Into<String>,
|
|
957
|
+
id: impl Into<String>,
|
|
958
|
+
vector: impl Into<Vec<f32>>,
|
|
959
|
+
metadata: Metadata,
|
|
960
|
+
) -> Result<()> {
|
|
961
|
+
self.upsert_with_vectors_in_namespace(
|
|
962
|
+
namespace,
|
|
963
|
+
id,
|
|
964
|
+
vector,
|
|
965
|
+
NamedVectors::new(),
|
|
966
|
+
SparseVector::new(),
|
|
967
|
+
metadata,
|
|
968
|
+
)
|
|
969
|
+
}
|
|
970
|
+
|
|
971
|
+
pub fn upsert_with_sparse_in_namespace(
|
|
972
|
+
&mut self,
|
|
973
|
+
namespace: impl Into<String>,
|
|
974
|
+
id: impl Into<String>,
|
|
975
|
+
vector: impl Into<Vec<f32>>,
|
|
976
|
+
sparse: SparseVector,
|
|
977
|
+
metadata: Metadata,
|
|
978
|
+
) -> Result<()> {
|
|
979
|
+
self.upsert_with_vectors_in_namespace(
|
|
980
|
+
namespace,
|
|
981
|
+
id,
|
|
982
|
+
vector,
|
|
983
|
+
NamedVectors::new(),
|
|
984
|
+
sparse,
|
|
985
|
+
metadata,
|
|
986
|
+
)
|
|
987
|
+
}
|
|
988
|
+
|
|
989
|
+
pub fn upsert_with_vectors_in_namespace(
|
|
990
|
+
&mut self,
|
|
991
|
+
namespace: impl Into<String>,
|
|
992
|
+
id: impl Into<String>,
|
|
993
|
+
vector: impl Into<Vec<f32>>,
|
|
994
|
+
vectors: NamedVectors,
|
|
995
|
+
sparse: SparseVector,
|
|
996
|
+
metadata: Metadata,
|
|
997
|
+
) -> Result<()> {
|
|
998
|
+
self.check_writable()?;
|
|
999
|
+
let record = self.record_from_parts(namespace, id, vector, vectors, sparse, metadata)?;
|
|
1000
|
+
self.apply_wal_batch(vec![WalOp::Upsert(record)])?;
|
|
1001
|
+
Ok(())
|
|
1002
|
+
}
|
|
1003
|
+
|
|
1004
|
+
pub fn insert_many<I>(&mut self, records: I) -> Result<usize>
|
|
1005
|
+
where
|
|
1006
|
+
I: IntoIterator<Item = Record>,
|
|
1007
|
+
{
|
|
1008
|
+
self.check_writable()?;
|
|
1009
|
+
let records = records
|
|
1010
|
+
.into_iter()
|
|
1011
|
+
.map(|record| {
|
|
1012
|
+
self.validate_record(&record)?;
|
|
1013
|
+
let key = (record.namespace.clone(), record.id.clone());
|
|
1014
|
+
if self.records.contains_key(&key) {
|
|
1015
|
+
return Err(VectLiteError::DuplicateId {
|
|
1016
|
+
namespace: key.0,
|
|
1017
|
+
id: key.1,
|
|
1018
|
+
});
|
|
1019
|
+
}
|
|
1020
|
+
Ok(record)
|
|
1021
|
+
})
|
|
1022
|
+
.collect::<Result<Vec<_>>>()?;
|
|
1023
|
+
|
|
1024
|
+
let count = records.len();
|
|
1025
|
+
if count == 0 {
|
|
1026
|
+
return Ok(0);
|
|
1027
|
+
}
|
|
1028
|
+
|
|
1029
|
+
self.apply_wal_batch(records.into_iter().map(WalOp::Upsert).collect())?;
|
|
1030
|
+
Ok(count)
|
|
1031
|
+
}
|
|
1032
|
+
|
|
1033
|
+
pub fn upsert_many<I>(&mut self, records: I) -> Result<usize>
|
|
1034
|
+
where
|
|
1035
|
+
I: IntoIterator<Item = Record>,
|
|
1036
|
+
{
|
|
1037
|
+
self.check_writable()?;
|
|
1038
|
+
let records = records
|
|
1039
|
+
.into_iter()
|
|
1040
|
+
.map(|record| {
|
|
1041
|
+
self.validate_record(&record)?;
|
|
1042
|
+
Ok(record)
|
|
1043
|
+
})
|
|
1044
|
+
.collect::<Result<Vec<_>>>()?;
|
|
1045
|
+
|
|
1046
|
+
let count = records.len();
|
|
1047
|
+
if count == 0 {
|
|
1048
|
+
return Ok(0);
|
|
1049
|
+
}
|
|
1050
|
+
|
|
1051
|
+
self.apply_wal_batch(records.into_iter().map(WalOp::Upsert).collect())?;
|
|
1052
|
+
Ok(count)
|
|
1053
|
+
}
|
|
1054
|
+
|
|
1055
|
+
pub fn get(&self, id: &str) -> Option<&Record> {
|
|
1056
|
+
self.get_in_namespace(DEFAULT_NAMESPACE, id)
|
|
1057
|
+
}
|
|
1058
|
+
|
|
1059
|
+
pub fn get_in_namespace(&self, namespace: &str, id: &str) -> Option<&Record> {
|
|
1060
|
+
self.records.get(&(namespace.to_owned(), id.to_owned()))
|
|
1061
|
+
}
|
|
1062
|
+
|
|
1063
|
+
pub fn delete(&mut self, id: &str) -> Result<bool> {
|
|
1064
|
+
self.delete_in_namespace(DEFAULT_NAMESPACE, id)
|
|
1065
|
+
}
|
|
1066
|
+
|
|
1067
|
+
pub fn delete_in_namespace(&mut self, namespace: &str, id: &str) -> Result<bool> {
|
|
1068
|
+
self.check_writable()?;
|
|
1069
|
+
let removed = self
|
|
1070
|
+
.records
|
|
1071
|
+
.contains_key(&(namespace.to_owned(), id.to_owned()));
|
|
1072
|
+
if removed {
|
|
1073
|
+
self.apply_wal_batch(vec![WalOp::Delete {
|
|
1074
|
+
namespace: namespace.to_owned(),
|
|
1075
|
+
id: id.to_owned(),
|
|
1076
|
+
}])?;
|
|
1077
|
+
}
|
|
1078
|
+
Ok(removed)
|
|
1079
|
+
}
|
|
1080
|
+
|
|
1081
|
+
pub fn delete_many<I, S>(&mut self, ids: I) -> Result<usize>
|
|
1082
|
+
where
|
|
1083
|
+
I: IntoIterator<Item = S>,
|
|
1084
|
+
S: AsRef<str>,
|
|
1085
|
+
{
|
|
1086
|
+
self.delete_many_in_namespace(DEFAULT_NAMESPACE, ids)
|
|
1087
|
+
}
|
|
1088
|
+
|
|
1089
|
+
pub fn delete_many_in_namespace<I, S>(&mut self, namespace: &str, ids: I) -> Result<usize>
|
|
1090
|
+
where
|
|
1091
|
+
I: IntoIterator<Item = S>,
|
|
1092
|
+
S: AsRef<str>,
|
|
1093
|
+
{
|
|
1094
|
+
self.check_writable()?;
|
|
1095
|
+
let mut removed = 0;
|
|
1096
|
+
let mut ops = Vec::new();
|
|
1097
|
+
|
|
1098
|
+
for id in ids {
|
|
1099
|
+
if self
|
|
1100
|
+
.records
|
|
1101
|
+
.contains_key(&(namespace.to_owned(), id.as_ref().to_owned()))
|
|
1102
|
+
{
|
|
1103
|
+
removed += 1;
|
|
1104
|
+
ops.push(WalOp::Delete {
|
|
1105
|
+
namespace: namespace.to_owned(),
|
|
1106
|
+
id: id.as_ref().to_owned(),
|
|
1107
|
+
});
|
|
1108
|
+
}
|
|
1109
|
+
}
|
|
1110
|
+
|
|
1111
|
+
if removed > 0 {
|
|
1112
|
+
self.apply_wal_batch(ops)?;
|
|
1113
|
+
}
|
|
1114
|
+
|
|
1115
|
+
Ok(removed)
|
|
1116
|
+
}
|
|
1117
|
+
|
|
1118
|
+
pub fn search(&self, query: &[f32], options: SearchOptions) -> Result<Vec<SearchResult>> {
|
|
1119
|
+
self.search_in_namespace(DEFAULT_NAMESPACE, query, options)
|
|
1120
|
+
}
|
|
1121
|
+
|
|
1122
|
+
pub fn search_in_namespace(
|
|
1123
|
+
&self,
|
|
1124
|
+
namespace: &str,
|
|
1125
|
+
query: &[f32],
|
|
1126
|
+
options: SearchOptions,
|
|
1127
|
+
) -> Result<Vec<SearchResult>> {
|
|
1128
|
+
Ok(self
|
|
1129
|
+
.hybrid_search_in_namespace_with_stats(
|
|
1130
|
+
namespace,
|
|
1131
|
+
Some(query),
|
|
1132
|
+
None,
|
|
1133
|
+
HybridSearchOptions {
|
|
1134
|
+
top_k: options.top_k,
|
|
1135
|
+
filter: options.filter,
|
|
1136
|
+
dense_weight: 1.0,
|
|
1137
|
+
sparse_weight: 0.0,
|
|
1138
|
+
..HybridSearchOptions::default()
|
|
1139
|
+
},
|
|
1140
|
+
)?
|
|
1141
|
+
.results)
|
|
1142
|
+
}
|
|
1143
|
+
|
|
1144
|
+
pub fn search_all_namespaces(
|
|
1145
|
+
&self,
|
|
1146
|
+
query: &[f32],
|
|
1147
|
+
options: SearchOptions,
|
|
1148
|
+
) -> Result<Vec<SearchResult>> {
|
|
1149
|
+
Ok(self
|
|
1150
|
+
.hybrid_search_all_namespaces_with_stats(
|
|
1151
|
+
Some(query),
|
|
1152
|
+
None,
|
|
1153
|
+
HybridSearchOptions {
|
|
1154
|
+
top_k: options.top_k,
|
|
1155
|
+
filter: options.filter,
|
|
1156
|
+
dense_weight: 1.0,
|
|
1157
|
+
sparse_weight: 0.0,
|
|
1158
|
+
..HybridSearchOptions::default()
|
|
1159
|
+
},
|
|
1160
|
+
)?
|
|
1161
|
+
.results)
|
|
1162
|
+
}
|
|
1163
|
+
|
|
1164
|
+
pub fn hybrid_search_in_namespace(
|
|
1165
|
+
&self,
|
|
1166
|
+
namespace: &str,
|
|
1167
|
+
dense_query: Option<&[f32]>,
|
|
1168
|
+
sparse_query: Option<&SparseVector>,
|
|
1169
|
+
options: HybridSearchOptions,
|
|
1170
|
+
) -> Result<Vec<SearchResult>> {
|
|
1171
|
+
Ok(self
|
|
1172
|
+
.hybrid_search_in_namespace_with_stats(namespace, dense_query, sparse_query, options)?
|
|
1173
|
+
.results)
|
|
1174
|
+
}
|
|
1175
|
+
|
|
1176
|
+
pub fn hybrid_search_in_namespace_with_stats(
|
|
1177
|
+
&self,
|
|
1178
|
+
namespace: &str,
|
|
1179
|
+
dense_query: Option<&[f32]>,
|
|
1180
|
+
sparse_query: Option<&SparseVector>,
|
|
1181
|
+
options: HybridSearchOptions,
|
|
1182
|
+
) -> Result<SearchOutcome> {
|
|
1183
|
+
self.hybrid_search_internal(dense_query, sparse_query, options, Some(namespace))
|
|
1184
|
+
}
|
|
1185
|
+
|
|
1186
|
+
pub fn hybrid_search_all_namespaces(
|
|
1187
|
+
&self,
|
|
1188
|
+
dense_query: Option<&[f32]>,
|
|
1189
|
+
sparse_query: Option<&SparseVector>,
|
|
1190
|
+
options: HybridSearchOptions,
|
|
1191
|
+
) -> Result<Vec<SearchResult>> {
|
|
1192
|
+
Ok(self
|
|
1193
|
+
.hybrid_search_all_namespaces_with_stats(dense_query, sparse_query, options)?
|
|
1194
|
+
.results)
|
|
1195
|
+
}
|
|
1196
|
+
|
|
1197
|
+
pub fn hybrid_search_all_namespaces_with_stats(
|
|
1198
|
+
&self,
|
|
1199
|
+
dense_query: Option<&[f32]>,
|
|
1200
|
+
sparse_query: Option<&SparseVector>,
|
|
1201
|
+
options: HybridSearchOptions,
|
|
1202
|
+
) -> Result<SearchOutcome> {
|
|
1203
|
+
self.hybrid_search_internal(dense_query, sparse_query, options, None)
|
|
1204
|
+
}
|
|
1205
|
+
|
|
1206
|
+
pub fn namespaces(&self) -> Vec<String> {
|
|
1207
|
+
self.records
|
|
1208
|
+
.keys()
|
|
1209
|
+
.map(|(namespace, _)| namespace.clone())
|
|
1210
|
+
.collect::<std::collections::BTreeSet<_>>()
|
|
1211
|
+
.into_iter()
|
|
1212
|
+
.collect()
|
|
1213
|
+
}
|
|
1214
|
+
|
|
1215
|
+
pub fn wal_path(&self) -> &Path {
|
|
1216
|
+
&self.wal_path
|
|
1217
|
+
}
|
|
1218
|
+
|
|
1219
|
+
pub fn apply_operations(&mut self, operations: Vec<WriteOperation>) -> Result<()> {
|
|
1220
|
+
self.check_writable()?;
|
|
1221
|
+
let ops = operations
|
|
1222
|
+
.into_iter()
|
|
1223
|
+
.map(|operation| match operation {
|
|
1224
|
+
WriteOperation::Insert(record) => {
|
|
1225
|
+
self.validate_record(&record)?;
|
|
1226
|
+
let key = (record.namespace.clone(), record.id.clone());
|
|
1227
|
+
if self.records.contains_key(&key) {
|
|
1228
|
+
return Err(VectLiteError::DuplicateId {
|
|
1229
|
+
namespace: key.0,
|
|
1230
|
+
id: key.1,
|
|
1231
|
+
});
|
|
1232
|
+
}
|
|
1233
|
+
Ok(WalOp::Upsert(record))
|
|
1234
|
+
}
|
|
1235
|
+
WriteOperation::Upsert(record) => {
|
|
1236
|
+
self.validate_record(&record)?;
|
|
1237
|
+
Ok(WalOp::Upsert(record))
|
|
1238
|
+
}
|
|
1239
|
+
WriteOperation::Delete { namespace, id } => Ok(WalOp::Delete { namespace, id }),
|
|
1240
|
+
})
|
|
1241
|
+
.collect::<Result<Vec<_>>>()?;
|
|
1242
|
+
self.apply_wal_batch(ops)
|
|
1243
|
+
}
|
|
1244
|
+
|
|
1245
|
+
fn hybrid_search_internal(
|
|
1246
|
+
&self,
|
|
1247
|
+
dense_query: Option<&[f32]>,
|
|
1248
|
+
sparse_query: Option<&SparseVector>,
|
|
1249
|
+
options: HybridSearchOptions,
|
|
1250
|
+
namespace: Option<&str>,
|
|
1251
|
+
) -> Result<SearchOutcome> {
|
|
1252
|
+
if let Some(query) = dense_query {
|
|
1253
|
+
self.validate_vector(query)?;
|
|
1254
|
+
}
|
|
1255
|
+
if dense_query.is_none() && sparse_query.is_none() {
|
|
1256
|
+
return Err(VectLiteError::InvalidFormat(
|
|
1257
|
+
"search requires a dense query, a sparse query, or both".to_owned(),
|
|
1258
|
+
));
|
|
1259
|
+
}
|
|
1260
|
+
if let Some(mmr_lambda) = options.mmr_lambda {
|
|
1261
|
+
if !(0.0..=1.0).contains(&mmr_lambda) {
|
|
1262
|
+
return Err(VectLiteError::InvalidFormat(
|
|
1263
|
+
"mmr_lambda must be between 0.0 and 1.0".to_owned(),
|
|
1264
|
+
));
|
|
1265
|
+
}
|
|
1266
|
+
}
|
|
1267
|
+
if let Some(vector_name) = options.vector_name.as_deref() {
|
|
1268
|
+
if vector_name.is_empty() {
|
|
1269
|
+
return Err(VectLiteError::InvalidFormat(
|
|
1270
|
+
"vector_name must not be empty".to_owned(),
|
|
1271
|
+
));
|
|
1272
|
+
}
|
|
1273
|
+
}
|
|
1274
|
+
|
|
1275
|
+
let search_start = Instant::now();
|
|
1276
|
+
|
|
1277
|
+
let top_k = if options.top_k == 0 {
|
|
1278
|
+
self.records.len()
|
|
1279
|
+
} else {
|
|
1280
|
+
options.top_k
|
|
1281
|
+
};
|
|
1282
|
+
let fetch_k = resolve_fetch_k(
|
|
1283
|
+
top_k,
|
|
1284
|
+
options.fetch_k,
|
|
1285
|
+
self.records.len(),
|
|
1286
|
+
options.mmr_lambda,
|
|
1287
|
+
);
|
|
1288
|
+
let vector_name = options.vector_name.as_deref();
|
|
1289
|
+
|
|
1290
|
+
let dense_start = Instant::now();
|
|
1291
|
+
let ann_candidates = dense_query
|
|
1292
|
+
.and_then(|query| self.ann_candidate_keys(namespace, vector_name, query, fetch_k));
|
|
1293
|
+
let dense_us = dense_start.elapsed().as_micros() as u64;
|
|
1294
|
+
|
|
1295
|
+
let sparse_start = Instant::now();
|
|
1296
|
+
let sparse_candidates = sparse_query
|
|
1297
|
+
.map(|query| self.sparse_candidate_keys(namespace, query, fetch_k))
|
|
1298
|
+
.unwrap_or_default();
|
|
1299
|
+
let sparse_us = sparse_start.elapsed().as_micros() as u64;
|
|
1300
|
+
|
|
1301
|
+
let candidate_keys = if dense_query.is_some() && ann_candidates.is_none() {
|
|
1302
|
+
None
|
|
1303
|
+
} else {
|
|
1304
|
+
merge_candidate_keys(
|
|
1305
|
+
ann_candidates.as_deref(),
|
|
1306
|
+
Some(sparse_candidates.as_slice()),
|
|
1307
|
+
)
|
|
1308
|
+
};
|
|
1309
|
+
let mut stats = SearchStats {
|
|
1310
|
+
used_ann: ann_candidates.is_some(),
|
|
1311
|
+
ann_candidate_count: ann_candidates.as_ref().map_or(0, Vec::len),
|
|
1312
|
+
fetch_k,
|
|
1313
|
+
sparse_candidate_count: sparse_candidates.len(),
|
|
1314
|
+
ann_loaded_from_disk: self.ann_loaded_from_disk,
|
|
1315
|
+
wal_entries_replayed: self.wal_entries_replayed,
|
|
1316
|
+
fusion: options.fusion.label().to_owned(),
|
|
1317
|
+
..SearchStats::default()
|
|
1318
|
+
};
|
|
1319
|
+
|
|
1320
|
+
let mut results = self.collect_results(
|
|
1321
|
+
dense_query,
|
|
1322
|
+
sparse_query,
|
|
1323
|
+
&options,
|
|
1324
|
+
namespace,
|
|
1325
|
+
candidate_keys.as_deref(),
|
|
1326
|
+
);
|
|
1327
|
+
stats.considered_count = results.len();
|
|
1328
|
+
|
|
1329
|
+
if ann_candidates.is_some() && results.len() < fetch_k {
|
|
1330
|
+
stats.exact_fallback = true;
|
|
1331
|
+
results = self.collect_results(dense_query, sparse_query, &options, namespace, None);
|
|
1332
|
+
stats.considered_count = results.len();
|
|
1333
|
+
}
|
|
1334
|
+
|
|
1335
|
+
let fusion_start = Instant::now();
|
|
1336
|
+
apply_rank_metadata(&mut results);
|
|
1337
|
+
apply_fusion_strategy(
|
|
1338
|
+
&mut results,
|
|
1339
|
+
&options.fusion,
|
|
1340
|
+
options.dense_weight,
|
|
1341
|
+
options.sparse_weight,
|
|
1342
|
+
);
|
|
1343
|
+
sort_scored_records(&mut results);
|
|
1344
|
+
let fusion_us = fusion_start.elapsed().as_micros() as u64;
|
|
1345
|
+
|
|
1346
|
+
let mmr_applied = options.mmr_lambda.is_some() && top_k > 1 && results.len() > 1;
|
|
1347
|
+
let results = if let Some(mmr_lambda) = options.mmr_lambda {
|
|
1348
|
+
apply_mmr(
|
|
1349
|
+
results,
|
|
1350
|
+
top_k,
|
|
1351
|
+
mmr_lambda,
|
|
1352
|
+
options.dense_weight,
|
|
1353
|
+
options.sparse_weight,
|
|
1354
|
+
vector_name,
|
|
1355
|
+
)
|
|
1356
|
+
} else {
|
|
1357
|
+
let mut results = results;
|
|
1358
|
+
results.truncate(top_k);
|
|
1359
|
+
results
|
|
1360
|
+
};
|
|
1361
|
+
stats.mmr_applied = mmr_applied;
|
|
1362
|
+
|
|
1363
|
+
let total_us = search_start.elapsed().as_micros() as u64;
|
|
1364
|
+
stats.timings = SearchTimings {
|
|
1365
|
+
dense_us,
|
|
1366
|
+
sparse_us,
|
|
1367
|
+
fusion_us,
|
|
1368
|
+
total_us,
|
|
1369
|
+
};
|
|
1370
|
+
|
|
1371
|
+
Ok(SearchOutcome {
|
|
1372
|
+
results: results
|
|
1373
|
+
.into_iter()
|
|
1374
|
+
.map(ScoredRecord::into_search_result)
|
|
1375
|
+
.collect(),
|
|
1376
|
+
stats,
|
|
1377
|
+
})
|
|
1378
|
+
}
|
|
1379
|
+
|
|
1380
|
+
pub fn flush(&mut self) -> Result<()> {
|
|
1381
|
+
self.check_writable()?;
|
|
1382
|
+
self.compact_inner()
|
|
1383
|
+
}
|
|
1384
|
+
|
|
1385
|
+
/// Bulk-ingest many records efficiently. WAL writes happen in batches of
|
|
1386
|
+
/// `batch_size`, but the ANN index and sparse index are only rebuilt once
|
|
1387
|
+
/// at the very end, making this much faster than `upsert_many` for large
|
|
1388
|
+
/// imports.
|
|
1389
|
+
pub fn bulk_ingest<I>(&mut self, records: I, batch_size: usize) -> Result<usize>
|
|
1390
|
+
where
|
|
1391
|
+
I: IntoIterator<Item = Record>,
|
|
1392
|
+
{
|
|
1393
|
+
self.check_writable()?;
|
|
1394
|
+
let batch_size = batch_size.max(1);
|
|
1395
|
+
let mut total = 0_usize;
|
|
1396
|
+
let mut batch = Vec::with_capacity(batch_size);
|
|
1397
|
+
|
|
1398
|
+
for record in records {
|
|
1399
|
+
self.validate_record(&record)?;
|
|
1400
|
+
batch.push(WalOp::Upsert(record));
|
|
1401
|
+
|
|
1402
|
+
if batch.len() >= batch_size {
|
|
1403
|
+
total += batch.len();
|
|
1404
|
+
self.append_wal_batch(&batch)?;
|
|
1405
|
+
self.apply_ops_in_memory(batch);
|
|
1406
|
+
batch = Vec::with_capacity(batch_size);
|
|
1407
|
+
}
|
|
1408
|
+
}
|
|
1409
|
+
|
|
1410
|
+
if !batch.is_empty() {
|
|
1411
|
+
total += batch.len();
|
|
1412
|
+
self.append_wal_batch(&batch)?;
|
|
1413
|
+
self.apply_ops_in_memory(batch);
|
|
1414
|
+
}
|
|
1415
|
+
|
|
1416
|
+
if total > 0 {
|
|
1417
|
+
self.rebuild_sparse_index();
|
|
1418
|
+
self.rebuild_ann();
|
|
1419
|
+
self.ann_loaded_from_disk = false;
|
|
1420
|
+
self.persist_ann_to_disk()?;
|
|
1421
|
+
}
|
|
1422
|
+
|
|
1423
|
+
Ok(total)
|
|
1424
|
+
}
|
|
1425
|
+
|
|
1426
|
+
pub fn compact(&mut self) -> Result<()> {
|
|
1427
|
+
self.check_writable()?;
|
|
1428
|
+
self.compact_inner()
|
|
1429
|
+
}
|
|
1430
|
+
|
|
1431
|
+
fn compact_inner(&mut self) -> Result<()> {
|
|
1432
|
+
if let Some(parent) = self.path.parent() {
|
|
1433
|
+
if !parent.as_os_str().is_empty() {
|
|
1434
|
+
fs::create_dir_all(parent)?;
|
|
1435
|
+
}
|
|
1436
|
+
}
|
|
1437
|
+
|
|
1438
|
+
let temp_path = temp_path(&self.path);
|
|
1439
|
+
let mut file = File::create(&temp_path)?;
|
|
1440
|
+
{
|
|
1441
|
+
let mut writer = BufWriter::new(&mut file);
|
|
1442
|
+
self.write_to(&mut writer)?;
|
|
1443
|
+
writer.flush()?;
|
|
1444
|
+
}
|
|
1445
|
+
file.sync_all()?;
|
|
1446
|
+
|
|
1447
|
+
if self.path.exists() {
|
|
1448
|
+
fs::remove_file(&self.path)?;
|
|
1449
|
+
}
|
|
1450
|
+
fs::rename(temp_path, &self.path)?;
|
|
1451
|
+
self.clear_wal()?;
|
|
1452
|
+
self.wal_entries_replayed = 0;
|
|
1453
|
+
self.persist_ann_to_disk()?;
|
|
1454
|
+
|
|
1455
|
+
Ok(())
|
|
1456
|
+
}
|
|
1457
|
+
|
|
1458
|
+
/// Create an atomic snapshot of the database at `dest`. The snapshot is a
|
|
1459
|
+
/// self-contained `.vdb` file (WAL is folded in). The current database is
|
|
1460
|
+
/// not modified. Works in both read-only and read-write mode.
|
|
1461
|
+
pub fn snapshot(&self, dest: impl AsRef<Path>) -> Result<()> {
|
|
1462
|
+
let dest = dest.as_ref();
|
|
1463
|
+
if let Some(parent) = dest.parent() {
|
|
1464
|
+
if !parent.as_os_str().is_empty() {
|
|
1465
|
+
fs::create_dir_all(parent)?;
|
|
1466
|
+
}
|
|
1467
|
+
}
|
|
1468
|
+
let mut file = File::create(dest)?;
|
|
1469
|
+
{
|
|
1470
|
+
let mut writer = BufWriter::new(&mut file);
|
|
1471
|
+
self.write_to(&mut writer)?;
|
|
1472
|
+
writer.flush()?;
|
|
1473
|
+
}
|
|
1474
|
+
file.sync_all()?;
|
|
1475
|
+
Ok(())
|
|
1476
|
+
}
|
|
1477
|
+
|
|
1478
|
+
/// Back up the database to `dest` directory. Creates a complete copy
|
|
1479
|
+
/// including the `.vdb` file and ANN sidecar files. The backup is
|
|
1480
|
+
/// compacted (WAL folded in). Works in both read-only and read-write mode.
|
|
1481
|
+
pub fn backup(&self, dest: impl AsRef<Path>) -> Result<()> {
|
|
1482
|
+
let dest = dest.as_ref();
|
|
1483
|
+
fs::create_dir_all(dest)?;
|
|
1484
|
+
|
|
1485
|
+
let file_name = self.path.file_name().ok_or_else(|| {
|
|
1486
|
+
VectLiteError::InvalidFormat("database path has no file name".to_owned())
|
|
1487
|
+
})?;
|
|
1488
|
+
let dest_vdb = dest.join(file_name);
|
|
1489
|
+
self.snapshot(&dest_vdb)?;
|
|
1490
|
+
|
|
1491
|
+
// Copy ANN sidecar files
|
|
1492
|
+
if let Some(parent) = self.path.parent() {
|
|
1493
|
+
if let Some(stem) = self.path.file_name().and_then(|n| n.to_str()) {
|
|
1494
|
+
if let Ok(entries) = fs::read_dir(parent) {
|
|
1495
|
+
for entry in entries.flatten() {
|
|
1496
|
+
if let Some(fname) = entry.file_name().to_str() {
|
|
1497
|
+
if fname.starts_with(&format!("{stem}.ann.")) {
|
|
1498
|
+
let _ = fs::copy(entry.path(), dest.join(fname));
|
|
1499
|
+
}
|
|
1500
|
+
}
|
|
1501
|
+
}
|
|
1502
|
+
}
|
|
1503
|
+
// Copy ann manifest
|
|
1504
|
+
let manifest = ann_manifest_path(&self.path);
|
|
1505
|
+
if manifest.exists() {
|
|
1506
|
+
if let Some(manifest_name) = manifest.file_name() {
|
|
1507
|
+
let _ = fs::copy(&manifest, dest.join(manifest_name));
|
|
1508
|
+
}
|
|
1509
|
+
}
|
|
1510
|
+
}
|
|
1511
|
+
}
|
|
1512
|
+
|
|
1513
|
+
Ok(())
|
|
1514
|
+
}
|
|
1515
|
+
|
|
1516
|
+
/// Restore a database from a backup directory. Opens the `.vdb` file
|
|
1517
|
+
/// found in `source` and returns a new writable Database.
|
|
1518
|
+
pub fn restore(source: impl AsRef<Path>, dest: impl AsRef<Path>) -> Result<Self> {
|
|
1519
|
+
let source = source.as_ref();
|
|
1520
|
+
let dest = dest.as_ref();
|
|
1521
|
+
|
|
1522
|
+
// Find the .vdb file in the source directory
|
|
1523
|
+
let mut vdb_file = None;
|
|
1524
|
+
for entry in fs::read_dir(source)? {
|
|
1525
|
+
let entry = entry?;
|
|
1526
|
+
let path = entry.path();
|
|
1527
|
+
if path.extension().and_then(|ext| ext.to_str()) == Some("vdb") {
|
|
1528
|
+
vdb_file = Some(path);
|
|
1529
|
+
break;
|
|
1530
|
+
}
|
|
1531
|
+
}
|
|
1532
|
+
let source_vdb = vdb_file.ok_or_else(|| {
|
|
1533
|
+
VectLiteError::InvalidFormat("no .vdb file found in backup directory".to_owned())
|
|
1534
|
+
})?;
|
|
1535
|
+
|
|
1536
|
+
// Copy the vdb file
|
|
1537
|
+
if let Some(parent) = dest.parent() {
|
|
1538
|
+
if !parent.as_os_str().is_empty() {
|
|
1539
|
+
fs::create_dir_all(parent)?;
|
|
1540
|
+
}
|
|
1541
|
+
}
|
|
1542
|
+
fs::copy(&source_vdb, dest)?;
|
|
1543
|
+
|
|
1544
|
+
// Copy ANN sidecar files
|
|
1545
|
+
if let Some(stem) = source_vdb.file_name().and_then(|n| n.to_str()) {
|
|
1546
|
+
for entry in fs::read_dir(source)?.flatten() {
|
|
1547
|
+
if let Some(fname) = entry.file_name().to_str() {
|
|
1548
|
+
if fname.starts_with(&format!("{stem}.ann.")) || fname.ends_with(".ann") {
|
|
1549
|
+
if let Some(dest_parent) = dest.parent() {
|
|
1550
|
+
let _ = fs::copy(entry.path(), dest_parent.join(fname));
|
|
1551
|
+
}
|
|
1552
|
+
}
|
|
1553
|
+
}
|
|
1554
|
+
}
|
|
1555
|
+
}
|
|
1556
|
+
|
|
1557
|
+
Self::open(dest)
|
|
1558
|
+
}
|
|
1559
|
+
|
|
1560
|
+
fn apply_wal_batch(&mut self, ops: Vec<WalOp>) -> Result<()> {
|
|
1561
|
+
if ops.is_empty() {
|
|
1562
|
+
return Ok(());
|
|
1563
|
+
}
|
|
1564
|
+
|
|
1565
|
+
self.append_wal_batch(&ops)?;
|
|
1566
|
+
self.apply_ops_in_memory(ops);
|
|
1567
|
+
self.rebuild_sparse_index();
|
|
1568
|
+
self.rebuild_ann();
|
|
1569
|
+
self.ann_loaded_from_disk = false;
|
|
1570
|
+
self.persist_ann_to_disk()?;
|
|
1571
|
+
Ok(())
|
|
1572
|
+
}
|
|
1573
|
+
|
|
1574
|
+
fn apply_ops_in_memory(&mut self, ops: Vec<WalOp>) {
|
|
1575
|
+
for op in ops {
|
|
1576
|
+
match op {
|
|
1577
|
+
WalOp::Upsert(record) => {
|
|
1578
|
+
self.records
|
|
1579
|
+
.insert((record.namespace.clone(), record.id.clone()), record);
|
|
1580
|
+
}
|
|
1581
|
+
WalOp::Delete { namespace, id } => {
|
|
1582
|
+
self.records.remove(&(namespace, id));
|
|
1583
|
+
}
|
|
1584
|
+
}
|
|
1585
|
+
}
|
|
1586
|
+
}
|
|
1587
|
+
|
|
1588
|
+
fn append_wal_batch(&self, ops: &[WalOp]) -> Result<()> {
|
|
1589
|
+
if let Some(parent) = self.wal_path.parent() {
|
|
1590
|
+
if !parent.as_os_str().is_empty() {
|
|
1591
|
+
fs::create_dir_all(parent)?;
|
|
1592
|
+
}
|
|
1593
|
+
}
|
|
1594
|
+
|
|
1595
|
+
let new_file = !self.wal_path.exists();
|
|
1596
|
+
let mut file = OpenOptions::new()
|
|
1597
|
+
.create(true)
|
|
1598
|
+
.append(true)
|
|
1599
|
+
.open(&self.wal_path)?;
|
|
1600
|
+
|
|
1601
|
+
if new_file {
|
|
1602
|
+
file.write_all(WAL_MAGIC)?;
|
|
1603
|
+
}
|
|
1604
|
+
|
|
1605
|
+
let mut buffer = Vec::new();
|
|
1606
|
+
write_u32(&mut buffer, u32_from_usize(ops.len())?)?;
|
|
1607
|
+
for op in ops {
|
|
1608
|
+
write_wal_op(&mut buffer, op)?;
|
|
1609
|
+
}
|
|
1610
|
+
|
|
1611
|
+
write_u32(&mut file, u32_from_usize(buffer.len())?)?;
|
|
1612
|
+
file.write_all(&buffer)?;
|
|
1613
|
+
file.sync_all()?;
|
|
1614
|
+
Ok(())
|
|
1615
|
+
}
|
|
1616
|
+
|
|
1617
|
+
fn replay_wal(&mut self) -> Result<()> {
|
|
1618
|
+
if !self.wal_path.exists() {
|
|
1619
|
+
self.wal_entries_replayed = 0;
|
|
1620
|
+
return Ok(());
|
|
1621
|
+
}
|
|
1622
|
+
|
|
1623
|
+
let mut reader = BufReader::new(File::open(&self.wal_path)?);
|
|
1624
|
+
let mut magic = [0_u8; 4];
|
|
1625
|
+
if let Err(err) = reader.read_exact(&mut magic) {
|
|
1626
|
+
if err.kind() == ErrorKind::UnexpectedEof {
|
|
1627
|
+
self.wal_entries_replayed = 0;
|
|
1628
|
+
return Ok(());
|
|
1629
|
+
}
|
|
1630
|
+
return Err(err.into());
|
|
1631
|
+
}
|
|
1632
|
+
if &magic != WAL_MAGIC {
|
|
1633
|
+
return Err(VectLiteError::InvalidFormat(
|
|
1634
|
+
"invalid WAL header".to_owned(),
|
|
1635
|
+
));
|
|
1636
|
+
}
|
|
1637
|
+
|
|
1638
|
+
let mut replayed = 0;
|
|
1639
|
+
loop {
|
|
1640
|
+
let batch_len = match read_u32(&mut reader) {
|
|
1641
|
+
Ok(batch_len) => usize_from_u32(batch_len)?,
|
|
1642
|
+
Err(err) if err.kind() == ErrorKind::UnexpectedEof => break,
|
|
1643
|
+
Err(err) => return Err(err.into()),
|
|
1644
|
+
};
|
|
1645
|
+
|
|
1646
|
+
let mut batch = vec![0_u8; batch_len];
|
|
1647
|
+
match reader.read_exact(&mut batch) {
|
|
1648
|
+
Ok(()) => {
|
|
1649
|
+
let mut batch_reader = &batch[..];
|
|
1650
|
+
let op_count = usize_from_u32(read_u32(&mut batch_reader)?)?;
|
|
1651
|
+
let mut ops = Vec::with_capacity(op_count);
|
|
1652
|
+
for _ in 0..op_count {
|
|
1653
|
+
ops.push(read_wal_op(&mut batch_reader, self.dimension)?);
|
|
1654
|
+
}
|
|
1655
|
+
self.apply_ops_in_memory(ops);
|
|
1656
|
+
replayed += op_count;
|
|
1657
|
+
}
|
|
1658
|
+
Err(err) if err.kind() == ErrorKind::UnexpectedEof => break,
|
|
1659
|
+
Err(err) => return Err(err.into()),
|
|
1660
|
+
}
|
|
1661
|
+
}
|
|
1662
|
+
|
|
1663
|
+
self.wal_entries_replayed = replayed;
|
|
1664
|
+
Ok(())
|
|
1665
|
+
}
|
|
1666
|
+
|
|
1667
|
+
fn clear_wal(&self) -> Result<()> {
|
|
1668
|
+
if self.wal_path.exists() {
|
|
1669
|
+
fs::remove_file(&self.wal_path)?;
|
|
1670
|
+
}
|
|
1671
|
+
Ok(())
|
|
1672
|
+
}
|
|
1673
|
+
|
|
1674
|
+
fn read_from(path: &Path, reader: &mut impl Read) -> Result<Self> {
|
|
1675
|
+
let mut magic = [0_u8; 4];
|
|
1676
|
+
reader.read_exact(&mut magic)?;
|
|
1677
|
+
if &magic != MAGIC {
|
|
1678
|
+
return Err(VectLiteError::InvalidFormat(
|
|
1679
|
+
"missing VDB1 magic header".to_owned(),
|
|
1680
|
+
));
|
|
1681
|
+
}
|
|
1682
|
+
|
|
1683
|
+
let version = read_u16(reader)?;
|
|
1684
|
+
if !(1..=VERSION).contains(&version) {
|
|
1685
|
+
return Err(VectLiteError::InvalidFormat(format!(
|
|
1686
|
+
"unsupported version {version}"
|
|
1687
|
+
)));
|
|
1688
|
+
}
|
|
1689
|
+
|
|
1690
|
+
let dimension = usize_from_u32(read_u32(reader)?)?;
|
|
1691
|
+
ensure_dimension(dimension)?;
|
|
1692
|
+
|
|
1693
|
+
let record_count = usize_from_u64(read_u64(reader)?)?;
|
|
1694
|
+
let mut records = BTreeMap::new();
|
|
1695
|
+
|
|
1696
|
+
for _ in 0..record_count {
|
|
1697
|
+
let namespace = if version >= 2 {
|
|
1698
|
+
read_string(reader)?
|
|
1699
|
+
} else {
|
|
1700
|
+
DEFAULT_NAMESPACE.to_owned()
|
|
1701
|
+
};
|
|
1702
|
+
let id = read_string(reader)?;
|
|
1703
|
+
let metadata_count = usize_from_u32(read_u32(reader)?)?;
|
|
1704
|
+
let mut metadata = Metadata::new();
|
|
1705
|
+
for _ in 0..metadata_count {
|
|
1706
|
+
let key = read_string(reader)?;
|
|
1707
|
+
let value = read_metadata_value(reader)?;
|
|
1708
|
+
metadata.insert(key, value);
|
|
1709
|
+
}
|
|
1710
|
+
|
|
1711
|
+
let vector_len = usize_from_u32(read_u32(reader)?)?;
|
|
1712
|
+
if vector_len != dimension {
|
|
1713
|
+
return Err(VectLiteError::InvalidFormat(format!(
|
|
1714
|
+
"record {id} has vector length {vector_len}, expected {dimension}"
|
|
1715
|
+
)));
|
|
1716
|
+
}
|
|
1717
|
+
|
|
1718
|
+
let mut vector = Vec::with_capacity(vector_len);
|
|
1719
|
+
for _ in 0..vector_len {
|
|
1720
|
+
vector.push(read_f32(reader)?);
|
|
1721
|
+
}
|
|
1722
|
+
|
|
1723
|
+
let vectors = if version >= 4 {
|
|
1724
|
+
read_named_vectors(reader, dimension)?
|
|
1725
|
+
} else {
|
|
1726
|
+
NamedVectors::new()
|
|
1727
|
+
};
|
|
1728
|
+
|
|
1729
|
+
let sparse = if version >= 3 {
|
|
1730
|
+
read_sparse_vector(reader)?
|
|
1731
|
+
} else {
|
|
1732
|
+
SparseVector::new()
|
|
1733
|
+
};
|
|
1734
|
+
|
|
1735
|
+
let record = Record {
|
|
1736
|
+
namespace: namespace.clone(),
|
|
1737
|
+
id: id.clone(),
|
|
1738
|
+
vector,
|
|
1739
|
+
vectors,
|
|
1740
|
+
sparse,
|
|
1741
|
+
metadata,
|
|
1742
|
+
};
|
|
1743
|
+
records.insert((namespace, id), record);
|
|
1744
|
+
}
|
|
1745
|
+
|
|
1746
|
+
Ok(Self {
|
|
1747
|
+
path: path.to_path_buf(),
|
|
1748
|
+
wal_path: wal_path(path),
|
|
1749
|
+
dimension,
|
|
1750
|
+
records,
|
|
1751
|
+
ann: AnnCatalog::default(),
|
|
1752
|
+
sparse_index: SparseIndex::default(),
|
|
1753
|
+
wal_entries_replayed: 0,
|
|
1754
|
+
ann_loaded_from_disk: false,
|
|
1755
|
+
read_only: false,
|
|
1756
|
+
_lock_file: None,
|
|
1757
|
+
})
|
|
1758
|
+
}
|
|
1759
|
+
|
|
1760
|
+
fn write_to(&self, writer: &mut impl Write) -> Result<()> {
|
|
1761
|
+
writer.write_all(MAGIC)?;
|
|
1762
|
+
write_u16(writer, VERSION)?;
|
|
1763
|
+
write_u32(writer, u32_from_usize(self.dimension)?)?;
|
|
1764
|
+
write_u64(writer, u64_from_usize(self.records.len())?)?;
|
|
1765
|
+
|
|
1766
|
+
for record in self.records.values() {
|
|
1767
|
+
write_string(writer, &record.namespace)?;
|
|
1768
|
+
write_string(writer, &record.id)?;
|
|
1769
|
+
write_u32(writer, u32_from_usize(record.metadata.len())?)?;
|
|
1770
|
+
for (key, value) in &record.metadata {
|
|
1771
|
+
write_string(writer, key)?;
|
|
1772
|
+
write_metadata_value(writer, value)?;
|
|
1773
|
+
}
|
|
1774
|
+
|
|
1775
|
+
write_u32(writer, u32_from_usize(record.vector.len())?)?;
|
|
1776
|
+
for value in &record.vector {
|
|
1777
|
+
write_f32(writer, *value)?;
|
|
1778
|
+
}
|
|
1779
|
+
write_named_vectors(writer, &record.vectors)?;
|
|
1780
|
+
write_sparse_vector(writer, &record.sparse)?;
|
|
1781
|
+
}
|
|
1782
|
+
|
|
1783
|
+
Ok(())
|
|
1784
|
+
}
|
|
1785
|
+
|
|
1786
|
+
fn validate_vector(&self, vector: &[f32]) -> Result<()> {
|
|
1787
|
+
if vector.len() != self.dimension {
|
|
1788
|
+
return Err(VectLiteError::DimensionMismatch {
|
|
1789
|
+
expected: self.dimension,
|
|
1790
|
+
found: vector.len(),
|
|
1791
|
+
});
|
|
1792
|
+
}
|
|
1793
|
+
|
|
1794
|
+
Ok(())
|
|
1795
|
+
}
|
|
1796
|
+
|
|
1797
|
+
fn validate_record(&self, record: &Record) -> Result<()> {
|
|
1798
|
+
self.validate_vector(&record.vector)?;
|
|
1799
|
+
|
|
1800
|
+
for (vector_name, vector) in &record.vectors {
|
|
1801
|
+
if vector_name.is_empty() {
|
|
1802
|
+
return Err(VectLiteError::InvalidFormat(
|
|
1803
|
+
"named vectors must not use an empty name".to_owned(),
|
|
1804
|
+
));
|
|
1805
|
+
}
|
|
1806
|
+
self.validate_vector(vector)?;
|
|
1807
|
+
}
|
|
1808
|
+
|
|
1809
|
+
Ok(())
|
|
1810
|
+
}
|
|
1811
|
+
|
|
1812
|
+
fn rebuild_ann(&mut self) {
|
|
1813
|
+
self.ann = AnnCatalog::default();
|
|
1814
|
+
let mut global_by_vector: BTreeMap<String, Vec<(RecordKey, &Vec<f32>)>> = BTreeMap::new();
|
|
1815
|
+
let mut by_namespace: BTreeMap<String, BTreeMap<String, Vec<(RecordKey, &Vec<f32>)>>> =
|
|
1816
|
+
BTreeMap::new();
|
|
1817
|
+
|
|
1818
|
+
for (key, record) in &self.records {
|
|
1819
|
+
for (vector_name, vector) in record.dense_vectors() {
|
|
1820
|
+
global_by_vector
|
|
1821
|
+
.entry(vector_name.to_owned())
|
|
1822
|
+
.or_default()
|
|
1823
|
+
.push((key.clone(), vector));
|
|
1824
|
+
by_namespace
|
|
1825
|
+
.entry(record.namespace.clone())
|
|
1826
|
+
.or_default()
|
|
1827
|
+
.entry(vector_name.to_owned())
|
|
1828
|
+
.or_default()
|
|
1829
|
+
.push((key.clone(), vector));
|
|
1830
|
+
}
|
|
1831
|
+
}
|
|
1832
|
+
|
|
1833
|
+
self.ann.global = global_by_vector
|
|
1834
|
+
.into_iter()
|
|
1835
|
+
.filter_map(|(vector_name, records)| {
|
|
1836
|
+
if records.len() < ANN_MIN_POINTS {
|
|
1837
|
+
None
|
|
1838
|
+
} else {
|
|
1839
|
+
Some((vector_name, build_ann_index(records)))
|
|
1840
|
+
}
|
|
1841
|
+
})
|
|
1842
|
+
.collect();
|
|
1843
|
+
|
|
1844
|
+
self.ann.namespaces = by_namespace
|
|
1845
|
+
.into_iter()
|
|
1846
|
+
.filter_map(|(namespace, indexes)| {
|
|
1847
|
+
let indexes = indexes
|
|
1848
|
+
.into_iter()
|
|
1849
|
+
.filter_map(|(vector_name, records)| {
|
|
1850
|
+
if records.len() < ANN_MIN_POINTS {
|
|
1851
|
+
None
|
|
1852
|
+
} else {
|
|
1853
|
+
Some((vector_name, build_ann_index(records)))
|
|
1854
|
+
}
|
|
1855
|
+
})
|
|
1856
|
+
.collect::<BTreeMap<_, _>>();
|
|
1857
|
+
|
|
1858
|
+
if indexes.is_empty() {
|
|
1859
|
+
None
|
|
1860
|
+
} else {
|
|
1861
|
+
Some((namespace, indexes))
|
|
1862
|
+
}
|
|
1863
|
+
})
|
|
1864
|
+
.collect();
|
|
1865
|
+
}
|
|
1866
|
+
|
|
1867
|
+
fn try_load_ann_from_disk(&mut self) -> bool {
|
|
1868
|
+
let Some(parent) = self.path.parent() else {
|
|
1869
|
+
return false;
|
|
1870
|
+
};
|
|
1871
|
+
|
|
1872
|
+
let Ok(entries) = read_ann_manifest(&ann_manifest_path(&self.path)) else {
|
|
1873
|
+
return false;
|
|
1874
|
+
};
|
|
1875
|
+
|
|
1876
|
+
let expected = self.expected_ann_entries();
|
|
1877
|
+
if expected.len() != entries.len() {
|
|
1878
|
+
return false;
|
|
1879
|
+
}
|
|
1880
|
+
|
|
1881
|
+
let mut loaded_global = BTreeMap::new();
|
|
1882
|
+
let mut loaded_namespaces: BTreeMap<String, BTreeMap<String, AnnIndex>> = BTreeMap::new();
|
|
1883
|
+
|
|
1884
|
+
for expected_entry in expected {
|
|
1885
|
+
let Some(manifest_entry) = entries.iter().find(|entry| {
|
|
1886
|
+
entry.namespace == expected_entry.namespace
|
|
1887
|
+
&& entry.vector_name == expected_entry.vector_name
|
|
1888
|
+
}) else {
|
|
1889
|
+
return false;
|
|
1890
|
+
};
|
|
1891
|
+
|
|
1892
|
+
if manifest_entry.record_count != expected_entry.record_count
|
|
1893
|
+
|| manifest_entry.key_signature != expected_entry.key_signature
|
|
1894
|
+
{
|
|
1895
|
+
return false;
|
|
1896
|
+
}
|
|
1897
|
+
|
|
1898
|
+
let Some(index) = load_ann_index(
|
|
1899
|
+
parent,
|
|
1900
|
+
&ann_basename(
|
|
1901
|
+
&self.path,
|
|
1902
|
+
expected_entry.namespace.as_deref(),
|
|
1903
|
+
&expected_entry.vector_name,
|
|
1904
|
+
),
|
|
1905
|
+
expected_entry.keys.clone(),
|
|
1906
|
+
) else {
|
|
1907
|
+
return false;
|
|
1908
|
+
};
|
|
1909
|
+
|
|
1910
|
+
if let Some(namespace) = expected_entry.namespace {
|
|
1911
|
+
loaded_namespaces
|
|
1912
|
+
.entry(namespace)
|
|
1913
|
+
.or_default()
|
|
1914
|
+
.insert(expected_entry.vector_name, index);
|
|
1915
|
+
} else {
|
|
1916
|
+
loaded_global.insert(expected_entry.vector_name, index);
|
|
1917
|
+
}
|
|
1918
|
+
}
|
|
1919
|
+
|
|
1920
|
+
self.ann = AnnCatalog {
|
|
1921
|
+
global: loaded_global,
|
|
1922
|
+
namespaces: loaded_namespaces,
|
|
1923
|
+
};
|
|
1924
|
+
true
|
|
1925
|
+
}
|
|
1926
|
+
|
|
1927
|
+
fn persist_ann_to_disk(&self) -> Result<()> {
|
|
1928
|
+
let Some(parent) = self.path.parent() else {
|
|
1929
|
+
return Ok(());
|
|
1930
|
+
};
|
|
1931
|
+
if !parent.exists() {
|
|
1932
|
+
return Ok(());
|
|
1933
|
+
}
|
|
1934
|
+
|
|
1935
|
+
let entries = self.expected_ann_entries();
|
|
1936
|
+
for entry in &entries {
|
|
1937
|
+
let basename = ann_basename(&self.path, entry.namespace.as_deref(), &entry.vector_name);
|
|
1938
|
+
let graph_path = parent.join(format!("{basename}.hnsw.graph"));
|
|
1939
|
+
let data_path = parent.join(format!("{basename}.hnsw.data"));
|
|
1940
|
+
if graph_path.exists() {
|
|
1941
|
+
fs::remove_file(&graph_path)?;
|
|
1942
|
+
}
|
|
1943
|
+
if data_path.exists() {
|
|
1944
|
+
fs::remove_file(&data_path)?;
|
|
1945
|
+
}
|
|
1946
|
+
|
|
1947
|
+
let index = match &entry.namespace {
|
|
1948
|
+
Some(namespace) => self
|
|
1949
|
+
.ann
|
|
1950
|
+
.namespaces
|
|
1951
|
+
.get(namespace)
|
|
1952
|
+
.and_then(|indexes| indexes.get(&entry.vector_name)),
|
|
1953
|
+
None => self.ann.global.get(&entry.vector_name),
|
|
1954
|
+
};
|
|
1955
|
+
if let Some(index) = index {
|
|
1956
|
+
index.hnsw.file_dump(parent, &basename).map_err(|err| {
|
|
1957
|
+
VectLiteError::InvalidFormat(format!("failed to persist ANN index: {err}"))
|
|
1958
|
+
})?;
|
|
1959
|
+
}
|
|
1960
|
+
}
|
|
1961
|
+
|
|
1962
|
+
write_ann_manifest(&ann_manifest_path(&self.path), &entries)
|
|
1963
|
+
}
|
|
1964
|
+
|
|
1965
|
+
fn expected_ann_entries(&self) -> Vec<AnnManifestEntry> {
|
|
1966
|
+
let mut global: BTreeMap<String, Vec<RecordKey>> = BTreeMap::new();
|
|
1967
|
+
let mut by_namespace: BTreeMap<String, BTreeMap<String, Vec<RecordKey>>> = BTreeMap::new();
|
|
1968
|
+
|
|
1969
|
+
for (key, record) in &self.records {
|
|
1970
|
+
for (vector_name, _) in record.dense_vectors() {
|
|
1971
|
+
global
|
|
1972
|
+
.entry(vector_name.to_owned())
|
|
1973
|
+
.or_default()
|
|
1974
|
+
.push(key.clone());
|
|
1975
|
+
by_namespace
|
|
1976
|
+
.entry(record.namespace.clone())
|
|
1977
|
+
.or_default()
|
|
1978
|
+
.entry(vector_name.to_owned())
|
|
1979
|
+
.or_default()
|
|
1980
|
+
.push(key.clone());
|
|
1981
|
+
}
|
|
1982
|
+
}
|
|
1983
|
+
|
|
1984
|
+
let mut entries = Vec::new();
|
|
1985
|
+
|
|
1986
|
+
for (vector_name, keys) in global {
|
|
1987
|
+
if keys.len() < ANN_MIN_POINTS {
|
|
1988
|
+
continue;
|
|
1989
|
+
}
|
|
1990
|
+
entries.push(AnnManifestEntry {
|
|
1991
|
+
namespace: None,
|
|
1992
|
+
vector_name,
|
|
1993
|
+
record_count: keys.len(),
|
|
1994
|
+
key_signature: record_key_signature(&keys),
|
|
1995
|
+
keys,
|
|
1996
|
+
});
|
|
1997
|
+
}
|
|
1998
|
+
|
|
1999
|
+
for (namespace, indexes) in by_namespace {
|
|
2000
|
+
for (vector_name, keys) in indexes {
|
|
2001
|
+
if keys.len() < ANN_MIN_POINTS {
|
|
2002
|
+
continue;
|
|
2003
|
+
}
|
|
2004
|
+
entries.push(AnnManifestEntry {
|
|
2005
|
+
namespace: Some(namespace.clone()),
|
|
2006
|
+
vector_name,
|
|
2007
|
+
record_count: keys.len(),
|
|
2008
|
+
key_signature: record_key_signature(&keys),
|
|
2009
|
+
keys,
|
|
2010
|
+
});
|
|
2011
|
+
}
|
|
2012
|
+
}
|
|
2013
|
+
|
|
2014
|
+
entries
|
|
2015
|
+
}
|
|
2016
|
+
|
|
2017
|
+
fn rebuild_sparse_index(&mut self) {
|
|
2018
|
+
self.sparse_index = SparseIndex::default();
|
|
2019
|
+
self.sparse_index.doc_count = self.records.len();
|
|
2020
|
+
|
|
2021
|
+
let mut total_doc_len = 0.0_f32;
|
|
2022
|
+
for (key, record) in &self.records {
|
|
2023
|
+
let doc_len = record
|
|
2024
|
+
.sparse
|
|
2025
|
+
.values()
|
|
2026
|
+
.copied()
|
|
2027
|
+
.filter(|weight| *weight > 0.0)
|
|
2028
|
+
.sum::<f32>();
|
|
2029
|
+
self.sparse_index.doc_lengths.insert(key.clone(), doc_len);
|
|
2030
|
+
total_doc_len += doc_len;
|
|
2031
|
+
|
|
2032
|
+
for (term, weight) in &record.sparse {
|
|
2033
|
+
if *weight <= 0.0 {
|
|
2034
|
+
continue;
|
|
2035
|
+
}
|
|
2036
|
+
self.sparse_index
|
|
2037
|
+
.postings
|
|
2038
|
+
.entry(term.clone())
|
|
2039
|
+
.or_default()
|
|
2040
|
+
.push(SparsePosting {
|
|
2041
|
+
key: key.clone(),
|
|
2042
|
+
term_weight: *weight,
|
|
2043
|
+
});
|
|
2044
|
+
}
|
|
2045
|
+
}
|
|
2046
|
+
|
|
2047
|
+
self.sparse_index.avg_doc_len = if self.sparse_index.doc_count == 0 {
|
|
2048
|
+
0.0
|
|
2049
|
+
} else {
|
|
2050
|
+
total_doc_len / self.sparse_index.doc_count as f32
|
|
2051
|
+
};
|
|
2052
|
+
}
|
|
2053
|
+
|
|
2054
|
+
fn collect_results(
|
|
2055
|
+
&self,
|
|
2056
|
+
dense_query: Option<&[f32]>,
|
|
2057
|
+
sparse_query: Option<&SparseVector>,
|
|
2058
|
+
options: &HybridSearchOptions,
|
|
2059
|
+
namespace: Option<&str>,
|
|
2060
|
+
candidate_keys: Option<&[RecordKey]>,
|
|
2061
|
+
) -> Vec<ScoredRecord<'_>> {
|
|
2062
|
+
let record_iter: Box<dyn Iterator<Item = &Record> + '_> = match candidate_keys {
|
|
2063
|
+
Some(keys) => Box::new(keys.iter().filter_map(|key| self.records.get(key))),
|
|
2064
|
+
None => Box::new(self.records.values()),
|
|
2065
|
+
};
|
|
2066
|
+
|
|
2067
|
+
record_iter
|
|
2068
|
+
.filter(|record| {
|
|
2069
|
+
namespace
|
|
2070
|
+
.map(|namespace| record.namespace == namespace)
|
|
2071
|
+
.unwrap_or(true)
|
|
2072
|
+
&& (dense_query.is_none()
|
|
2073
|
+
|| record.vector_for(options.vector_name.as_deref()).is_some())
|
|
2074
|
+
&& options
|
|
2075
|
+
.filter
|
|
2076
|
+
.as_ref()
|
|
2077
|
+
.map(|filter| filter.matches(&record.metadata))
|
|
2078
|
+
.unwrap_or(true)
|
|
2079
|
+
})
|
|
2080
|
+
.map(|record| {
|
|
2081
|
+
let (dense_score, resolved_vector_name) =
|
|
2082
|
+
if !options.multi_vector_queries.is_empty() {
|
|
2083
|
+
// Multi-vector weighted search
|
|
2084
|
+
let mut weighted_sum = 0.0_f32;
|
|
2085
|
+
for (name, (query, weight)) in &options.multi_vector_queries {
|
|
2086
|
+
if let Some(vector) = record.vector_for(Some(name.as_str())) {
|
|
2087
|
+
weighted_sum += weight * cosine_similarity(query, vector);
|
|
2088
|
+
}
|
|
2089
|
+
}
|
|
2090
|
+
(weighted_sum, None)
|
|
2091
|
+
} else {
|
|
2092
|
+
let score = dense_query
|
|
2093
|
+
.and_then(|query| {
|
|
2094
|
+
record
|
|
2095
|
+
.vector_for(options.vector_name.as_deref())
|
|
2096
|
+
.map(|vector| cosine_similarity(query, vector))
|
|
2097
|
+
})
|
|
2098
|
+
.unwrap_or(0.0);
|
|
2099
|
+
(score, options.vector_name.clone())
|
|
2100
|
+
};
|
|
2101
|
+
let sparse_score = sparse_query
|
|
2102
|
+
.map(|query| {
|
|
2103
|
+
self.bm25_score((record.namespace.clone(), record.id.clone()), query)
|
|
2104
|
+
})
|
|
2105
|
+
.unwrap_or(0.0);
|
|
2106
|
+
let record_key = (record.namespace.clone(), record.id.clone());
|
|
2107
|
+
let mut bm25_term_scores = BTreeMap::<String, f32>::new();
|
|
2108
|
+
let matched_terms = sparse_query
|
|
2109
|
+
.map(|query| {
|
|
2110
|
+
query
|
|
2111
|
+
.keys()
|
|
2112
|
+
.filter(|term| record.sparse.contains_key(*term))
|
|
2113
|
+
.map(|term| {
|
|
2114
|
+
let score = self.bm25_term_score(
|
|
2115
|
+
&record_key,
|
|
2116
|
+
term,
|
|
2117
|
+
*record.sparse.get(term).unwrap_or(&0.0),
|
|
2118
|
+
);
|
|
2119
|
+
bm25_term_scores.insert(term.clone(), score);
|
|
2120
|
+
term.clone()
|
|
2121
|
+
})
|
|
2122
|
+
.collect::<Vec<_>>()
|
|
2123
|
+
})
|
|
2124
|
+
.unwrap_or_default();
|
|
2125
|
+
|
|
2126
|
+
ScoredRecord {
|
|
2127
|
+
record,
|
|
2128
|
+
score: (options.dense_weight * dense_score)
|
|
2129
|
+
+ (options.sparse_weight * sparse_score),
|
|
2130
|
+
dense_score,
|
|
2131
|
+
sparse_score,
|
|
2132
|
+
vector_name: resolved_vector_name,
|
|
2133
|
+
matched_terms,
|
|
2134
|
+
dense_rank: None,
|
|
2135
|
+
sparse_rank: None,
|
|
2136
|
+
bm25_term_scores,
|
|
2137
|
+
}
|
|
2138
|
+
})
|
|
2139
|
+
.collect()
|
|
2140
|
+
}
|
|
2141
|
+
|
|
2142
|
+
fn ann_candidate_keys(
|
|
2143
|
+
&self,
|
|
2144
|
+
namespace: Option<&str>,
|
|
2145
|
+
vector_name: Option<&str>,
|
|
2146
|
+
query: &[f32],
|
|
2147
|
+
top_k: usize,
|
|
2148
|
+
) -> Option<Vec<RecordKey>> {
|
|
2149
|
+
let index = match namespace {
|
|
2150
|
+
Some(namespace) => self
|
|
2151
|
+
.ann
|
|
2152
|
+
.namespaces
|
|
2153
|
+
.get(namespace)
|
|
2154
|
+
.and_then(|indexes| indexes.get(vector_name.unwrap_or(DEFAULT_VECTOR_NAME))),
|
|
2155
|
+
None => self
|
|
2156
|
+
.ann
|
|
2157
|
+
.global
|
|
2158
|
+
.get(vector_name.unwrap_or(DEFAULT_VECTOR_NAME)),
|
|
2159
|
+
}?;
|
|
2160
|
+
if index.keys.len() < ANN_SEARCH_MIN_POINTS {
|
|
2161
|
+
return None;
|
|
2162
|
+
}
|
|
2163
|
+
|
|
2164
|
+
let candidate_count = candidate_count(top_k, index.keys.len());
|
|
2165
|
+
if candidate_count == 0 {
|
|
2166
|
+
return None;
|
|
2167
|
+
}
|
|
2168
|
+
|
|
2169
|
+
let ef_search = candidate_count.max(ANN_EF_CONSTRUCTION);
|
|
2170
|
+
let neighbours = index.hnsw.search(query, candidate_count, ef_search);
|
|
2171
|
+
Some(
|
|
2172
|
+
neighbours
|
|
2173
|
+
.into_iter()
|
|
2174
|
+
.filter_map(|neighbour| index.keys.get(neighbour.d_id).cloned())
|
|
2175
|
+
.collect(),
|
|
2176
|
+
)
|
|
2177
|
+
}
|
|
2178
|
+
|
|
2179
|
+
fn sparse_candidate_keys(
|
|
2180
|
+
&self,
|
|
2181
|
+
namespace: Option<&str>,
|
|
2182
|
+
sparse_query: &SparseVector,
|
|
2183
|
+
top_k: usize,
|
|
2184
|
+
) -> Vec<RecordKey> {
|
|
2185
|
+
let mut scores = BTreeMap::<RecordKey, f32>::new();
|
|
2186
|
+
for (term, query_weight) in sparse_query {
|
|
2187
|
+
let Some(postings) = self.sparse_index.postings.get(term) else {
|
|
2188
|
+
continue;
|
|
2189
|
+
};
|
|
2190
|
+
for posting in postings {
|
|
2191
|
+
if namespace
|
|
2192
|
+
.map(|namespace| posting.key.0 == namespace)
|
|
2193
|
+
.unwrap_or(true)
|
|
2194
|
+
{
|
|
2195
|
+
let key = posting.key.clone();
|
|
2196
|
+
*scores.entry(key.clone()).or_insert(0.0) +=
|
|
2197
|
+
*query_weight * self.bm25_term_score(&key, term, posting.term_weight);
|
|
2198
|
+
}
|
|
2199
|
+
}
|
|
2200
|
+
}
|
|
2201
|
+
|
|
2202
|
+
let mut scored = scores.into_iter().collect::<Vec<_>>();
|
|
2203
|
+
scored.sort_by(|left, right| {
|
|
2204
|
+
right
|
|
2205
|
+
.1
|
|
2206
|
+
.total_cmp(&left.1)
|
|
2207
|
+
.then_with(|| left.0.0.cmp(&right.0.0))
|
|
2208
|
+
.then_with(|| left.0.1.cmp(&right.0.1))
|
|
2209
|
+
});
|
|
2210
|
+
let limit = candidate_count(top_k, scored.len());
|
|
2211
|
+
scored.into_iter().take(limit).map(|(key, _)| key).collect()
|
|
2212
|
+
}
|
|
2213
|
+
|
|
2214
|
+
fn bm25_score(&self, key: RecordKey, sparse_query: &SparseVector) -> f32 {
|
|
2215
|
+
sparse_query
|
|
2216
|
+
.iter()
|
|
2217
|
+
.map(|(term, query_weight)| {
|
|
2218
|
+
let doc_weight = self
|
|
2219
|
+
.records
|
|
2220
|
+
.get(&key)
|
|
2221
|
+
.and_then(|record| record.sparse.get(term))
|
|
2222
|
+
.copied()
|
|
2223
|
+
.unwrap_or(0.0);
|
|
2224
|
+
if doc_weight <= 0.0 {
|
|
2225
|
+
0.0
|
|
2226
|
+
} else {
|
|
2227
|
+
*query_weight * self.bm25_term_score(&key, term, doc_weight)
|
|
2228
|
+
}
|
|
2229
|
+
})
|
|
2230
|
+
.sum()
|
|
2231
|
+
}
|
|
2232
|
+
|
|
2233
|
+
fn bm25_term_score(&self, key: &RecordKey, term: &str, doc_weight: f32) -> f32 {
|
|
2234
|
+
if doc_weight <= 0.0 || self.sparse_index.doc_count == 0 {
|
|
2235
|
+
return 0.0;
|
|
2236
|
+
}
|
|
2237
|
+
|
|
2238
|
+
let df = self
|
|
2239
|
+
.sparse_index
|
|
2240
|
+
.postings
|
|
2241
|
+
.get(term)
|
|
2242
|
+
.map_or(0, |postings| postings.len()) as f32;
|
|
2243
|
+
if df == 0.0 {
|
|
2244
|
+
return 0.0;
|
|
2245
|
+
}
|
|
2246
|
+
|
|
2247
|
+
let idf = (((self.sparse_index.doc_count as f32 - df + 0.5) / (df + 0.5)) + 1.0).ln();
|
|
2248
|
+
let doc_len = self
|
|
2249
|
+
.sparse_index
|
|
2250
|
+
.doc_lengths
|
|
2251
|
+
.get(key)
|
|
2252
|
+
.copied()
|
|
2253
|
+
.unwrap_or(0.0);
|
|
2254
|
+
let norm = if self.sparse_index.avg_doc_len > 0.0 {
|
|
2255
|
+
1.0 - BM25_B + BM25_B * (doc_len / self.sparse_index.avg_doc_len)
|
|
2256
|
+
} else {
|
|
2257
|
+
1.0
|
|
2258
|
+
};
|
|
2259
|
+
|
|
2260
|
+
idf * ((doc_weight * (BM25_K1 + 1.0)) / (doc_weight + BM25_K1 * norm))
|
|
2261
|
+
}
|
|
2262
|
+
|
|
2263
|
+
#[cfg(test)]
|
|
2264
|
+
fn has_ann_index(&self, namespace: Option<&str>, vector_name: Option<&str>) -> bool {
|
|
2265
|
+
match namespace {
|
|
2266
|
+
Some(namespace) => self.ann.namespaces.get(namespace).is_some_and(|indexes| {
|
|
2267
|
+
indexes.contains_key(vector_name.unwrap_or(DEFAULT_VECTOR_NAME))
|
|
2268
|
+
}),
|
|
2269
|
+
None => self
|
|
2270
|
+
.ann
|
|
2271
|
+
.global
|
|
2272
|
+
.contains_key(vector_name.unwrap_or(DEFAULT_VECTOR_NAME)),
|
|
2273
|
+
}
|
|
2274
|
+
}
|
|
2275
|
+
|
|
2276
|
+
fn record_from_parts(
|
|
2277
|
+
&self,
|
|
2278
|
+
namespace: impl Into<String>,
|
|
2279
|
+
id: impl Into<String>,
|
|
2280
|
+
vector: impl Into<Vec<f32>>,
|
|
2281
|
+
vectors: NamedVectors,
|
|
2282
|
+
sparse: SparseVector,
|
|
2283
|
+
metadata: Metadata,
|
|
2284
|
+
) -> Result<Record> {
|
|
2285
|
+
let vector = vector.into();
|
|
2286
|
+
self.validate_vector(&vector)?;
|
|
2287
|
+
|
|
2288
|
+
for (vector_name, named_vector) in &vectors {
|
|
2289
|
+
if vector_name.is_empty() {
|
|
2290
|
+
return Err(VectLiteError::InvalidFormat(
|
|
2291
|
+
"named vectors must not use an empty name".to_owned(),
|
|
2292
|
+
));
|
|
2293
|
+
}
|
|
2294
|
+
self.validate_vector(named_vector)?;
|
|
2295
|
+
}
|
|
2296
|
+
|
|
2297
|
+
Ok(Record {
|
|
2298
|
+
namespace: namespace.into(),
|
|
2299
|
+
id: id.into(),
|
|
2300
|
+
vector,
|
|
2301
|
+
vectors,
|
|
2302
|
+
sparse,
|
|
2303
|
+
metadata,
|
|
2304
|
+
})
|
|
2305
|
+
}
|
|
2306
|
+
}
|
|
2307
|
+
|
|
2308
|
+
impl ScoredRecord<'_> {
|
|
2309
|
+
fn into_search_result(self) -> SearchResult {
|
|
2310
|
+
SearchResult {
|
|
2311
|
+
namespace: self.record.namespace.clone(),
|
|
2312
|
+
id: self.record.id.clone(),
|
|
2313
|
+
score: self.score,
|
|
2314
|
+
dense_score: self.dense_score,
|
|
2315
|
+
sparse_score: self.sparse_score,
|
|
2316
|
+
vector_name: self.vector_name,
|
|
2317
|
+
matched_terms: self.matched_terms,
|
|
2318
|
+
dense_rank: self.dense_rank,
|
|
2319
|
+
sparse_rank: self.sparse_rank,
|
|
2320
|
+
metadata: self.record.metadata.clone(),
|
|
2321
|
+
bm25_term_scores: self.bm25_term_scores,
|
|
2322
|
+
}
|
|
2323
|
+
}
|
|
2324
|
+
}
|
|
2325
|
+
|
|
2326
|
+
impl FusionStrategy {
|
|
2327
|
+
fn label(&self) -> &'static str {
|
|
2328
|
+
match self {
|
|
2329
|
+
Self::Linear => "linear",
|
|
2330
|
+
Self::Rrf { .. } => "rrf",
|
|
2331
|
+
}
|
|
2332
|
+
}
|
|
2333
|
+
}
|
|
2334
|
+
|
|
2335
|
+
fn ensure_dimension(dimension: usize) -> Result<()> {
|
|
2336
|
+
if dimension == 0 {
|
|
2337
|
+
return Err(VectLiteError::InvalidFormat(
|
|
2338
|
+
"dimension must be greater than zero".to_owned(),
|
|
2339
|
+
));
|
|
2340
|
+
}
|
|
2341
|
+
|
|
2342
|
+
Ok(())
|
|
2343
|
+
}
|
|
2344
|
+
|
|
2345
|
+
fn cosine_similarity(left: &[f32], right: &[f32]) -> f32 {
|
|
2346
|
+
let mut dot = 0.0_f32;
|
|
2347
|
+
let mut left_norm = 0.0_f32;
|
|
2348
|
+
let mut right_norm = 0.0_f32;
|
|
2349
|
+
|
|
2350
|
+
for (left_value, right_value) in left.iter().zip(right.iter()) {
|
|
2351
|
+
dot += left_value * right_value;
|
|
2352
|
+
left_norm += left_value * left_value;
|
|
2353
|
+
right_norm += right_value * right_value;
|
|
2354
|
+
}
|
|
2355
|
+
|
|
2356
|
+
if left_norm == 0.0 || right_norm == 0.0 {
|
|
2357
|
+
0.0
|
|
2358
|
+
} else {
|
|
2359
|
+
dot / (left_norm.sqrt() * right_norm.sqrt())
|
|
2360
|
+
}
|
|
2361
|
+
}
|
|
2362
|
+
|
|
2363
|
+
fn sparse_dot_product(left: &SparseVector, right: &SparseVector) -> f32 {
|
|
2364
|
+
let (small, large) = if left.len() <= right.len() {
|
|
2365
|
+
(left, right)
|
|
2366
|
+
} else {
|
|
2367
|
+
(right, left)
|
|
2368
|
+
};
|
|
2369
|
+
|
|
2370
|
+
small.iter().fold(0.0_f32, |acc, (term, weight)| {
|
|
2371
|
+
acc + (*weight * large.get(term).copied().unwrap_or(0.0))
|
|
2372
|
+
})
|
|
2373
|
+
}
|
|
2374
|
+
|
|
2375
|
+
fn build_ann_index(records: Vec<(RecordKey, &Vec<f32>)>) -> AnnIndex {
|
|
2376
|
+
let max_layer = compute_hnsw_layers(records.len());
|
|
2377
|
+
let mut hnsw = Hnsw::<f32, DistCosine>::new(
|
|
2378
|
+
ANN_M,
|
|
2379
|
+
records.len(),
|
|
2380
|
+
max_layer,
|
|
2381
|
+
ANN_EF_CONSTRUCTION,
|
|
2382
|
+
DistCosine {},
|
|
2383
|
+
);
|
|
2384
|
+
|
|
2385
|
+
let mut keys = Vec::with_capacity(records.len());
|
|
2386
|
+
for (origin_id, (key, vector)) in records.into_iter().enumerate() {
|
|
2387
|
+
hnsw.insert((vector.as_slice(), origin_id));
|
|
2388
|
+
keys.push(key);
|
|
2389
|
+
}
|
|
2390
|
+
hnsw.set_searching_mode(true);
|
|
2391
|
+
|
|
2392
|
+
AnnIndex { hnsw, keys }
|
|
2393
|
+
}
|
|
2394
|
+
|
|
2395
|
+
fn compute_hnsw_layers(record_count: usize) -> usize {
|
|
2396
|
+
let _ = record_count;
|
|
2397
|
+
16
|
|
2398
|
+
}
|
|
2399
|
+
|
|
2400
|
+
fn candidate_count(top_k: usize, total: usize) -> usize {
|
|
2401
|
+
if total <= 256 {
|
|
2402
|
+
return total;
|
|
2403
|
+
}
|
|
2404
|
+
|
|
2405
|
+
let requested = top_k.max(1);
|
|
2406
|
+
requested
|
|
2407
|
+
.saturating_mul(ANN_OVERSAMPLE)
|
|
2408
|
+
.max(ANN_MIN_CANDIDATES)
|
|
2409
|
+
.min(total)
|
|
2410
|
+
}
|
|
2411
|
+
|
|
2412
|
+
fn wal_path(path: &Path) -> PathBuf {
|
|
2413
|
+
let mut wal = path.as_os_str().to_os_string();
|
|
2414
|
+
wal.push(".wal");
|
|
2415
|
+
PathBuf::from(wal)
|
|
2416
|
+
}
|
|
2417
|
+
|
|
2418
|
+
fn lock_path(path: &Path) -> PathBuf {
|
|
2419
|
+
let mut lock = path.as_os_str().to_os_string();
|
|
2420
|
+
lock.push(".lock");
|
|
2421
|
+
PathBuf::from(lock)
|
|
2422
|
+
}
|
|
2423
|
+
|
|
2424
|
+
fn acquire_exclusive_lock(path: &Path) -> Result<File> {
|
|
2425
|
+
if let Some(parent) = path.parent() {
|
|
2426
|
+
if !parent.as_os_str().is_empty() && !parent.exists() {
|
|
2427
|
+
fs::create_dir_all(parent)?;
|
|
2428
|
+
}
|
|
2429
|
+
}
|
|
2430
|
+
let file = OpenOptions::new()
|
|
2431
|
+
.create(true)
|
|
2432
|
+
.truncate(false)
|
|
2433
|
+
.read(true)
|
|
2434
|
+
.write(true)
|
|
2435
|
+
.open(lock_path(path))?;
|
|
2436
|
+
file.try_lock_exclusive().map_err(|err| {
|
|
2437
|
+
VectLiteError::LockContention(format!(
|
|
2438
|
+
"could not acquire exclusive lock on '{}': {err}",
|
|
2439
|
+
path.display()
|
|
2440
|
+
))
|
|
2441
|
+
})?;
|
|
2442
|
+
Ok(file)
|
|
2443
|
+
}
|
|
2444
|
+
|
|
2445
|
+
fn acquire_shared_lock(path: &Path) -> Result<File> {
|
|
2446
|
+
let lock_file = lock_path(path);
|
|
2447
|
+
if !lock_file.exists() {
|
|
2448
|
+
// Lock file may not exist yet for read-only opens on existing dbs
|
|
2449
|
+
if let Some(parent) = lock_file.parent() {
|
|
2450
|
+
if !parent.as_os_str().is_empty() && !parent.exists() {
|
|
2451
|
+
fs::create_dir_all(parent)?;
|
|
2452
|
+
}
|
|
2453
|
+
}
|
|
2454
|
+
}
|
|
2455
|
+
let file = OpenOptions::new()
|
|
2456
|
+
.create(true)
|
|
2457
|
+
.truncate(false)
|
|
2458
|
+
.read(true)
|
|
2459
|
+
.write(true)
|
|
2460
|
+
.open(&lock_file)?;
|
|
2461
|
+
file.try_lock_shared().map_err(|err| {
|
|
2462
|
+
VectLiteError::LockContention(format!(
|
|
2463
|
+
"could not acquire shared lock on '{}': {err}",
|
|
2464
|
+
path.display()
|
|
2465
|
+
))
|
|
2466
|
+
})?;
|
|
2467
|
+
Ok(file)
|
|
2468
|
+
}
|
|
2469
|
+
|
|
2470
|
+
fn ann_manifest_path(path: &Path) -> PathBuf {
|
|
2471
|
+
let mut manifest = path.as_os_str().to_os_string();
|
|
2472
|
+
manifest.push(".ann");
|
|
2473
|
+
PathBuf::from(manifest)
|
|
2474
|
+
}
|
|
2475
|
+
|
|
2476
|
+
fn ann_basename(path: &Path, namespace: Option<&str>, vector_name: &str) -> String {
|
|
2477
|
+
let stem = path
|
|
2478
|
+
.file_name()
|
|
2479
|
+
.and_then(|name| name.to_str())
|
|
2480
|
+
.unwrap_or("vectlite");
|
|
2481
|
+
format!(
|
|
2482
|
+
"{stem}.ann.{}.{}",
|
|
2483
|
+
hex_encode(namespace.unwrap_or(DEFAULT_NAMESPACE).as_bytes()),
|
|
2484
|
+
hex_encode(vector_name.as_bytes())
|
|
2485
|
+
)
|
|
2486
|
+
}
|
|
2487
|
+
|
|
2488
|
+
fn hex_encode(bytes: &[u8]) -> String {
|
|
2489
|
+
let mut out = String::with_capacity(bytes.len() * 2);
|
|
2490
|
+
for byte in bytes {
|
|
2491
|
+
out.push_str(&format!("{byte:02x}"));
|
|
2492
|
+
}
|
|
2493
|
+
out
|
|
2494
|
+
}
|
|
2495
|
+
|
|
2496
|
+
fn record_key_signature(keys: &[RecordKey]) -> u64 {
|
|
2497
|
+
let mut state = 0xcbf29ce484222325_u64;
|
|
2498
|
+
for (namespace, id) in keys {
|
|
2499
|
+
for byte in namespace
|
|
2500
|
+
.as_bytes()
|
|
2501
|
+
.iter()
|
|
2502
|
+
.chain(std::iter::once(&0xff))
|
|
2503
|
+
.chain(id.as_bytes().iter())
|
|
2504
|
+
.chain(std::iter::once(&0xfe))
|
|
2505
|
+
{
|
|
2506
|
+
state ^= *byte as u64;
|
|
2507
|
+
state = state.wrapping_mul(0x100000001b3);
|
|
2508
|
+
}
|
|
2509
|
+
}
|
|
2510
|
+
state
|
|
2511
|
+
}
|
|
2512
|
+
|
|
2513
|
+
fn load_ann_index(directory: &Path, basename: &str, keys: Vec<RecordKey>) -> Option<AnnIndex> {
|
|
2514
|
+
let reloader = Box::leak(Box::new(HnswIo::new(directory, basename)));
|
|
2515
|
+
let mut hnsw = reloader.load_hnsw_with_dist(DistCosine {}).ok()?;
|
|
2516
|
+
hnsw.set_searching_mode(true);
|
|
2517
|
+
Some(AnnIndex { hnsw, keys })
|
|
2518
|
+
}
|
|
2519
|
+
|
|
2520
|
+
fn write_ann_manifest(path: &Path, entries: &[AnnManifestEntry]) -> Result<()> {
|
|
2521
|
+
let mut file = File::create(path)?;
|
|
2522
|
+
file.write_all(b"ANN1")?;
|
|
2523
|
+
write_u32(&mut file, u32_from_usize(entries.len())?)?;
|
|
2524
|
+
for entry in entries {
|
|
2525
|
+
write_u8(&mut file, u8::from(entry.namespace.is_some()))?;
|
|
2526
|
+
if let Some(namespace) = &entry.namespace {
|
|
2527
|
+
write_string(&mut file, namespace)?;
|
|
2528
|
+
}
|
|
2529
|
+
write_string(&mut file, &entry.vector_name)?;
|
|
2530
|
+
write_u64(&mut file, u64_from_usize(entry.record_count)?)?;
|
|
2531
|
+
write_u64(&mut file, entry.key_signature)?;
|
|
2532
|
+
}
|
|
2533
|
+
file.sync_all()?;
|
|
2534
|
+
Ok(())
|
|
2535
|
+
}
|
|
2536
|
+
|
|
2537
|
+
fn read_ann_manifest(path: &Path) -> Result<Vec<AnnManifestEntry>> {
|
|
2538
|
+
let mut file = BufReader::new(File::open(path)?);
|
|
2539
|
+
let mut magic = [0_u8; 4];
|
|
2540
|
+
file.read_exact(&mut magic)?;
|
|
2541
|
+
if &magic != b"ANN1" {
|
|
2542
|
+
return Err(VectLiteError::InvalidFormat(
|
|
2543
|
+
"invalid ANN manifest".to_owned(),
|
|
2544
|
+
));
|
|
2545
|
+
}
|
|
2546
|
+
|
|
2547
|
+
let count = usize_from_u32(read_u32(&mut file)?)?;
|
|
2548
|
+
let mut entries = Vec::with_capacity(count);
|
|
2549
|
+
for _ in 0..count {
|
|
2550
|
+
let has_namespace = read_u8(&mut file)? != 0;
|
|
2551
|
+
let namespace = if has_namespace {
|
|
2552
|
+
Some(read_string(&mut file)?)
|
|
2553
|
+
} else {
|
|
2554
|
+
None
|
|
2555
|
+
};
|
|
2556
|
+
let vector_name = read_string(&mut file)?;
|
|
2557
|
+
let record_count = usize_from_u64(read_u64(&mut file)?)?;
|
|
2558
|
+
let key_signature = read_u64(&mut file)?;
|
|
2559
|
+
entries.push(AnnManifestEntry {
|
|
2560
|
+
namespace,
|
|
2561
|
+
vector_name,
|
|
2562
|
+
record_count,
|
|
2563
|
+
key_signature,
|
|
2564
|
+
keys: Vec::new(),
|
|
2565
|
+
});
|
|
2566
|
+
}
|
|
2567
|
+
Ok(entries)
|
|
2568
|
+
}
|
|
2569
|
+
|
|
2570
|
+
fn resolve_fetch_k(
|
|
2571
|
+
top_k: usize,
|
|
2572
|
+
requested_fetch_k: usize,
|
|
2573
|
+
total_records: usize,
|
|
2574
|
+
mmr_lambda: Option<f32>,
|
|
2575
|
+
) -> usize {
|
|
2576
|
+
if total_records == 0 {
|
|
2577
|
+
return 0;
|
|
2578
|
+
}
|
|
2579
|
+
|
|
2580
|
+
let default_fetch_k = if mmr_lambda.is_some() {
|
|
2581
|
+
top_k.max(1).saturating_mul(4)
|
|
2582
|
+
} else {
|
|
2583
|
+
top_k.max(1)
|
|
2584
|
+
};
|
|
2585
|
+
|
|
2586
|
+
requested_fetch_k
|
|
2587
|
+
.max(top_k.max(1))
|
|
2588
|
+
.max(default_fetch_k)
|
|
2589
|
+
.min(total_records)
|
|
2590
|
+
}
|
|
2591
|
+
|
|
2592
|
+
fn sort_scored_records(results: &mut [ScoredRecord<'_>]) {
|
|
2593
|
+
results.sort_by(|left, right| {
|
|
2594
|
+
right
|
|
2595
|
+
.score
|
|
2596
|
+
.total_cmp(&left.score)
|
|
2597
|
+
.then_with(|| left.record.namespace.cmp(&right.record.namespace))
|
|
2598
|
+
.then_with(|| left.record.id.cmp(&right.record.id))
|
|
2599
|
+
});
|
|
2600
|
+
}
|
|
2601
|
+
|
|
2602
|
+
fn apply_rank_metadata(results: &mut [ScoredRecord<'_>]) {
|
|
2603
|
+
let mut dense_order = (0..results.len()).collect::<Vec<_>>();
|
|
2604
|
+
dense_order.sort_by(|left, right| {
|
|
2605
|
+
results[*right]
|
|
2606
|
+
.dense_score
|
|
2607
|
+
.total_cmp(&results[*left].dense_score)
|
|
2608
|
+
.then_with(|| {
|
|
2609
|
+
results[*left]
|
|
2610
|
+
.record
|
|
2611
|
+
.namespace
|
|
2612
|
+
.cmp(&results[*right].record.namespace)
|
|
2613
|
+
})
|
|
2614
|
+
.then_with(|| results[*left].record.id.cmp(&results[*right].record.id))
|
|
2615
|
+
});
|
|
2616
|
+
for (rank, index) in dense_order.into_iter().enumerate() {
|
|
2617
|
+
if results[index].dense_score > 0.0 {
|
|
2618
|
+
results[index].dense_rank = Some(rank + 1);
|
|
2619
|
+
}
|
|
2620
|
+
}
|
|
2621
|
+
|
|
2622
|
+
let mut sparse_order = (0..results.len()).collect::<Vec<_>>();
|
|
2623
|
+
sparse_order.sort_by(|left, right| {
|
|
2624
|
+
results[*right]
|
|
2625
|
+
.sparse_score
|
|
2626
|
+
.total_cmp(&results[*left].sparse_score)
|
|
2627
|
+
.then_with(|| {
|
|
2628
|
+
results[*left]
|
|
2629
|
+
.record
|
|
2630
|
+
.namespace
|
|
2631
|
+
.cmp(&results[*right].record.namespace)
|
|
2632
|
+
})
|
|
2633
|
+
.then_with(|| results[*left].record.id.cmp(&results[*right].record.id))
|
|
2634
|
+
});
|
|
2635
|
+
for (rank, index) in sparse_order.into_iter().enumerate() {
|
|
2636
|
+
if results[index].sparse_score > 0.0 {
|
|
2637
|
+
results[index].sparse_rank = Some(rank + 1);
|
|
2638
|
+
}
|
|
2639
|
+
}
|
|
2640
|
+
}
|
|
2641
|
+
|
|
2642
|
+
fn apply_fusion_strategy(
|
|
2643
|
+
results: &mut [ScoredRecord<'_>],
|
|
2644
|
+
fusion: &FusionStrategy,
|
|
2645
|
+
dense_weight: f32,
|
|
2646
|
+
sparse_weight: f32,
|
|
2647
|
+
) {
|
|
2648
|
+
match fusion {
|
|
2649
|
+
FusionStrategy::Linear => {
|
|
2650
|
+
for result in results {
|
|
2651
|
+
result.score =
|
|
2652
|
+
(dense_weight * result.dense_score) + (sparse_weight * result.sparse_score);
|
|
2653
|
+
}
|
|
2654
|
+
}
|
|
2655
|
+
FusionStrategy::Rrf { rank_constant } => {
|
|
2656
|
+
let rank_constant = (*rank_constant).max(1) as f32;
|
|
2657
|
+
for result in results {
|
|
2658
|
+
let dense_component = result
|
|
2659
|
+
.dense_rank
|
|
2660
|
+
.map(|rank| dense_weight / (rank_constant + rank as f32))
|
|
2661
|
+
.unwrap_or(0.0);
|
|
2662
|
+
let sparse_component = result
|
|
2663
|
+
.sparse_rank
|
|
2664
|
+
.map(|rank| sparse_weight / (rank_constant + rank as f32))
|
|
2665
|
+
.unwrap_or(0.0);
|
|
2666
|
+
result.score = dense_component + sparse_component;
|
|
2667
|
+
}
|
|
2668
|
+
}
|
|
2669
|
+
}
|
|
2670
|
+
}
|
|
2671
|
+
|
|
2672
|
+
fn merge_candidate_keys(
|
|
2673
|
+
dense_candidates: Option<&[RecordKey]>,
|
|
2674
|
+
sparse_candidates: Option<&[RecordKey]>,
|
|
2675
|
+
) -> Option<Vec<RecordKey>> {
|
|
2676
|
+
let mut merged = BTreeSet::new();
|
|
2677
|
+
if let Some(dense_candidates) = dense_candidates {
|
|
2678
|
+
merged.extend(dense_candidates.iter().cloned());
|
|
2679
|
+
}
|
|
2680
|
+
if let Some(sparse_candidates) = sparse_candidates {
|
|
2681
|
+
merged.extend(sparse_candidates.iter().cloned());
|
|
2682
|
+
}
|
|
2683
|
+
|
|
2684
|
+
if merged.is_empty() {
|
|
2685
|
+
None
|
|
2686
|
+
} else {
|
|
2687
|
+
Some(merged.into_iter().collect())
|
|
2688
|
+
}
|
|
2689
|
+
}
|
|
2690
|
+
|
|
2691
|
+
fn apply_mmr<'a>(
|
|
2692
|
+
candidates: Vec<ScoredRecord<'a>>,
|
|
2693
|
+
top_k: usize,
|
|
2694
|
+
mmr_lambda: f32,
|
|
2695
|
+
dense_weight: f32,
|
|
2696
|
+
sparse_weight: f32,
|
|
2697
|
+
vector_name: Option<&str>,
|
|
2698
|
+
) -> Vec<ScoredRecord<'a>> {
|
|
2699
|
+
let limit = top_k.min(candidates.len());
|
|
2700
|
+
if limit <= 1 {
|
|
2701
|
+
return candidates.into_iter().take(limit).collect();
|
|
2702
|
+
}
|
|
2703
|
+
|
|
2704
|
+
let mut selected: Vec<ScoredRecord<'a>> = Vec::with_capacity(limit);
|
|
2705
|
+
let mut used = vec![false; candidates.len()];
|
|
2706
|
+
|
|
2707
|
+
while selected.len() < limit {
|
|
2708
|
+
let mut best: Option<(usize, f32)> = None;
|
|
2709
|
+
|
|
2710
|
+
for (index, candidate) in candidates.iter().enumerate() {
|
|
2711
|
+
if used[index] {
|
|
2712
|
+
continue;
|
|
2713
|
+
}
|
|
2714
|
+
|
|
2715
|
+
let diversity_penalty = selected
|
|
2716
|
+
.iter()
|
|
2717
|
+
.map(|selected_candidate| {
|
|
2718
|
+
record_similarity(
|
|
2719
|
+
candidate.record,
|
|
2720
|
+
selected_candidate.record,
|
|
2721
|
+
dense_weight,
|
|
2722
|
+
sparse_weight,
|
|
2723
|
+
vector_name,
|
|
2724
|
+
)
|
|
2725
|
+
})
|
|
2726
|
+
.fold(0.0_f32, f32::max);
|
|
2727
|
+
|
|
2728
|
+
let mmr_score = if selected.is_empty() {
|
|
2729
|
+
candidate.score
|
|
2730
|
+
} else {
|
|
2731
|
+
(mmr_lambda * candidate.score) - ((1.0 - mmr_lambda) * diversity_penalty)
|
|
2732
|
+
};
|
|
2733
|
+
|
|
2734
|
+
let replace_best = match best {
|
|
2735
|
+
Some((best_index, best_score)) => {
|
|
2736
|
+
mmr_score > best_score
|
|
2737
|
+
|| (mmr_score == best_score
|
|
2738
|
+
&& candidate.score > candidates[best_index].score)
|
|
2739
|
+
|| (mmr_score == best_score
|
|
2740
|
+
&& candidate.score == candidates[best_index].score
|
|
2741
|
+
&& candidate.record.namespace < candidates[best_index].record.namespace)
|
|
2742
|
+
|| (mmr_score == best_score
|
|
2743
|
+
&& candidate.score == candidates[best_index].score
|
|
2744
|
+
&& candidate.record.namespace
|
|
2745
|
+
== candidates[best_index].record.namespace
|
|
2746
|
+
&& candidate.record.id < candidates[best_index].record.id)
|
|
2747
|
+
}
|
|
2748
|
+
None => true,
|
|
2749
|
+
};
|
|
2750
|
+
|
|
2751
|
+
if replace_best {
|
|
2752
|
+
best = Some((index, mmr_score));
|
|
2753
|
+
}
|
|
2754
|
+
}
|
|
2755
|
+
|
|
2756
|
+
let Some((best_index, _)) = best else {
|
|
2757
|
+
break;
|
|
2758
|
+
};
|
|
2759
|
+
used[best_index] = true;
|
|
2760
|
+
selected.push(candidates[best_index].clone());
|
|
2761
|
+
}
|
|
2762
|
+
|
|
2763
|
+
selected
|
|
2764
|
+
}
|
|
2765
|
+
|
|
2766
|
+
fn record_similarity(
|
|
2767
|
+
left: &Record,
|
|
2768
|
+
right: &Record,
|
|
2769
|
+
dense_weight: f32,
|
|
2770
|
+
sparse_weight: f32,
|
|
2771
|
+
vector_name: Option<&str>,
|
|
2772
|
+
) -> f32 {
|
|
2773
|
+
let dense_score = match (left.vector_for(vector_name), right.vector_for(vector_name)) {
|
|
2774
|
+
(Some(left), Some(right)) => cosine_similarity(left, right),
|
|
2775
|
+
_ => 0.0,
|
|
2776
|
+
};
|
|
2777
|
+
|
|
2778
|
+
(dense_weight * dense_score) + (sparse_weight * sparse_dot_product(&left.sparse, &right.sparse))
|
|
2779
|
+
}
|
|
2780
|
+
|
|
2781
|
+
fn temp_path(path: &Path) -> PathBuf {
|
|
2782
|
+
let mut temp = path.as_os_str().to_os_string();
|
|
2783
|
+
temp.push(".tmp");
|
|
2784
|
+
PathBuf::from(temp)
|
|
2785
|
+
}
|
|
2786
|
+
|
|
2787
|
+
fn write_metadata_value(writer: &mut impl Write, value: &MetadataValue) -> Result<()> {
|
|
2788
|
+
write_u8(writer, value.type_tag())?;
|
|
2789
|
+
match value {
|
|
2790
|
+
MetadataValue::String(value) => write_string(writer, value)?,
|
|
2791
|
+
MetadataValue::Integer(value) => write_i64(writer, *value)?,
|
|
2792
|
+
MetadataValue::Float(value) => write_f64(writer, *value)?,
|
|
2793
|
+
MetadataValue::Boolean(value) => write_u8(writer, u8::from(*value))?,
|
|
2794
|
+
MetadataValue::Null => {}
|
|
2795
|
+
MetadataValue::List(values) => {
|
|
2796
|
+
write_u32(writer, u32_from_usize(values.len())?)?;
|
|
2797
|
+
for item in values {
|
|
2798
|
+
write_metadata_value(writer, item)?;
|
|
2799
|
+
}
|
|
2800
|
+
}
|
|
2801
|
+
MetadataValue::Map(entries) => {
|
|
2802
|
+
write_u32(writer, u32_from_usize(entries.len())?)?;
|
|
2803
|
+
for (key, val) in entries {
|
|
2804
|
+
write_string(writer, key)?;
|
|
2805
|
+
write_metadata_value(writer, val)?;
|
|
2806
|
+
}
|
|
2807
|
+
}
|
|
2808
|
+
}
|
|
2809
|
+
Ok(())
|
|
2810
|
+
}
|
|
2811
|
+
|
|
2812
|
+
fn read_metadata_value(reader: &mut impl Read) -> Result<MetadataValue> {
|
|
2813
|
+
let tag = read_u8(reader)?;
|
|
2814
|
+
let value = match tag {
|
|
2815
|
+
TYPE_STRING => MetadataValue::String(read_string(reader)?),
|
|
2816
|
+
TYPE_INTEGER => MetadataValue::Integer(read_i64(reader)?),
|
|
2817
|
+
TYPE_FLOAT => MetadataValue::Float(read_f64(reader)?),
|
|
2818
|
+
TYPE_BOOLEAN => MetadataValue::Boolean(read_u8(reader)? != 0),
|
|
2819
|
+
TYPE_NULL => MetadataValue::Null,
|
|
2820
|
+
TYPE_LIST => {
|
|
2821
|
+
let count = usize_from_u32(read_u32(reader)?)?;
|
|
2822
|
+
let mut items = Vec::with_capacity(count);
|
|
2823
|
+
for _ in 0..count {
|
|
2824
|
+
items.push(read_metadata_value(reader)?);
|
|
2825
|
+
}
|
|
2826
|
+
MetadataValue::List(items)
|
|
2827
|
+
}
|
|
2828
|
+
TYPE_MAP => {
|
|
2829
|
+
let count = usize_from_u32(read_u32(reader)?)?;
|
|
2830
|
+
let mut entries = BTreeMap::new();
|
|
2831
|
+
for _ in 0..count {
|
|
2832
|
+
let key = read_string(reader)?;
|
|
2833
|
+
let val = read_metadata_value(reader)?;
|
|
2834
|
+
entries.insert(key, val);
|
|
2835
|
+
}
|
|
2836
|
+
MetadataValue::Map(entries)
|
|
2837
|
+
}
|
|
2838
|
+
other => {
|
|
2839
|
+
return Err(VectLiteError::InvalidFormat(format!(
|
|
2840
|
+
"unknown metadata value tag {other}"
|
|
2841
|
+
)));
|
|
2842
|
+
}
|
|
2843
|
+
};
|
|
2844
|
+
Ok(value)
|
|
2845
|
+
}
|
|
2846
|
+
|
|
2847
|
+
fn write_sparse_vector(writer: &mut impl Write, sparse: &SparseVector) -> Result<()> {
|
|
2848
|
+
write_u32(writer, u32_from_usize(sparse.len())?)?;
|
|
2849
|
+
for (term, weight) in sparse {
|
|
2850
|
+
write_string(writer, term)?;
|
|
2851
|
+
write_f32(writer, *weight)?;
|
|
2852
|
+
}
|
|
2853
|
+
Ok(())
|
|
2854
|
+
}
|
|
2855
|
+
|
|
2856
|
+
fn read_sparse_vector(reader: &mut impl Read) -> Result<SparseVector> {
|
|
2857
|
+
let entry_count = usize_from_u32(read_u32(reader)?)?;
|
|
2858
|
+
let mut sparse = SparseVector::new();
|
|
2859
|
+
|
|
2860
|
+
for _ in 0..entry_count {
|
|
2861
|
+
let term = read_string(reader)?;
|
|
2862
|
+
let weight = read_f32(reader)?;
|
|
2863
|
+
sparse.insert(term, weight);
|
|
2864
|
+
}
|
|
2865
|
+
|
|
2866
|
+
Ok(sparse)
|
|
2867
|
+
}
|
|
2868
|
+
|
|
2869
|
+
fn write_named_vectors(writer: &mut impl Write, vectors: &NamedVectors) -> Result<()> {
|
|
2870
|
+
write_u32(writer, u32_from_usize(vectors.len())?)?;
|
|
2871
|
+
for (name, vector) in vectors {
|
|
2872
|
+
write_string(writer, name)?;
|
|
2873
|
+
write_u32(writer, u32_from_usize(vector.len())?)?;
|
|
2874
|
+
for value in vector {
|
|
2875
|
+
write_f32(writer, *value)?;
|
|
2876
|
+
}
|
|
2877
|
+
}
|
|
2878
|
+
Ok(())
|
|
2879
|
+
}
|
|
2880
|
+
|
|
2881
|
+
fn read_named_vectors(reader: &mut impl Read, dimension: usize) -> Result<NamedVectors> {
|
|
2882
|
+
let vector_count = usize_from_u32(read_u32(reader)?)?;
|
|
2883
|
+
let mut vectors = NamedVectors::new();
|
|
2884
|
+
|
|
2885
|
+
for _ in 0..vector_count {
|
|
2886
|
+
let name = read_string(reader)?;
|
|
2887
|
+
if name.is_empty() {
|
|
2888
|
+
return Err(VectLiteError::InvalidFormat(
|
|
2889
|
+
"named vectors must not use an empty name".to_owned(),
|
|
2890
|
+
));
|
|
2891
|
+
}
|
|
2892
|
+
|
|
2893
|
+
let vector_len = usize_from_u32(read_u32(reader)?)?;
|
|
2894
|
+
if vector_len != dimension {
|
|
2895
|
+
return Err(VectLiteError::InvalidFormat(format!(
|
|
2896
|
+
"named vector {name} has length {vector_len}, expected {dimension}"
|
|
2897
|
+
)));
|
|
2898
|
+
}
|
|
2899
|
+
|
|
2900
|
+
let mut vector = Vec::with_capacity(vector_len);
|
|
2901
|
+
for _ in 0..vector_len {
|
|
2902
|
+
vector.push(read_f32(reader)?);
|
|
2903
|
+
}
|
|
2904
|
+
|
|
2905
|
+
vectors.insert(name, vector);
|
|
2906
|
+
}
|
|
2907
|
+
|
|
2908
|
+
Ok(vectors)
|
|
2909
|
+
}
|
|
2910
|
+
|
|
2911
|
+
fn write_wal_op(writer: &mut impl Write, op: &WalOp) -> Result<()> {
|
|
2912
|
+
match op {
|
|
2913
|
+
WalOp::Upsert(record) => {
|
|
2914
|
+
write_u8(writer, 1)?;
|
|
2915
|
+
write_string(writer, &record.namespace)?;
|
|
2916
|
+
write_string(writer, &record.id)?;
|
|
2917
|
+
write_u32(writer, u32_from_usize(record.metadata.len())?)?;
|
|
2918
|
+
for (key, value) in &record.metadata {
|
|
2919
|
+
write_string(writer, key)?;
|
|
2920
|
+
write_metadata_value(writer, value)?;
|
|
2921
|
+
}
|
|
2922
|
+
write_u32(writer, u32_from_usize(record.vector.len())?)?;
|
|
2923
|
+
for value in &record.vector {
|
|
2924
|
+
write_f32(writer, *value)?;
|
|
2925
|
+
}
|
|
2926
|
+
write_named_vectors(writer, &record.vectors)?;
|
|
2927
|
+
write_sparse_vector(writer, &record.sparse)?;
|
|
2928
|
+
}
|
|
2929
|
+
WalOp::Delete { namespace, id } => {
|
|
2930
|
+
write_u8(writer, 2)?;
|
|
2931
|
+
write_string(writer, namespace)?;
|
|
2932
|
+
write_string(writer, id)?;
|
|
2933
|
+
}
|
|
2934
|
+
}
|
|
2935
|
+
Ok(())
|
|
2936
|
+
}
|
|
2937
|
+
|
|
2938
|
+
fn read_wal_op(reader: &mut impl Read, dimension: usize) -> Result<WalOp> {
|
|
2939
|
+
match read_u8(reader)? {
|
|
2940
|
+
1 => {
|
|
2941
|
+
let namespace = read_string(reader)?;
|
|
2942
|
+
let id = read_string(reader)?;
|
|
2943
|
+
let metadata_count = usize_from_u32(read_u32(reader)?)?;
|
|
2944
|
+
let mut metadata = Metadata::new();
|
|
2945
|
+
for _ in 0..metadata_count {
|
|
2946
|
+
let key = read_string(reader)?;
|
|
2947
|
+
let value = read_metadata_value(reader)?;
|
|
2948
|
+
metadata.insert(key, value);
|
|
2949
|
+
}
|
|
2950
|
+
let vector_len = usize_from_u32(read_u32(reader)?)?;
|
|
2951
|
+
if vector_len != dimension {
|
|
2952
|
+
return Err(VectLiteError::InvalidFormat(format!(
|
|
2953
|
+
"wal record {id} has vector length {vector_len}, expected {dimension}"
|
|
2954
|
+
)));
|
|
2955
|
+
}
|
|
2956
|
+
let mut vector = Vec::with_capacity(vector_len);
|
|
2957
|
+
for _ in 0..vector_len {
|
|
2958
|
+
vector.push(read_f32(reader)?);
|
|
2959
|
+
}
|
|
2960
|
+
let vectors = read_named_vectors(reader, dimension)?;
|
|
2961
|
+
let sparse = read_sparse_vector(reader)?;
|
|
2962
|
+
Ok(WalOp::Upsert(Record {
|
|
2963
|
+
namespace,
|
|
2964
|
+
id,
|
|
2965
|
+
vector,
|
|
2966
|
+
vectors,
|
|
2967
|
+
sparse,
|
|
2968
|
+
metadata,
|
|
2969
|
+
}))
|
|
2970
|
+
}
|
|
2971
|
+
2 => Ok(WalOp::Delete {
|
|
2972
|
+
namespace: read_string(reader)?,
|
|
2973
|
+
id: read_string(reader)?,
|
|
2974
|
+
}),
|
|
2975
|
+
other => Err(VectLiteError::InvalidFormat(format!(
|
|
2976
|
+
"unknown WAL op tag {other}"
|
|
2977
|
+
))),
|
|
2978
|
+
}
|
|
2979
|
+
}
|
|
2980
|
+
|
|
2981
|
+
fn write_string(writer: &mut impl Write, value: &str) -> Result<()> {
|
|
2982
|
+
write_u32(writer, u32_from_usize(value.len())?)?;
|
|
2983
|
+
writer.write_all(value.as_bytes())?;
|
|
2984
|
+
Ok(())
|
|
2985
|
+
}
|
|
2986
|
+
|
|
2987
|
+
fn read_string(reader: &mut impl Read) -> Result<String> {
|
|
2988
|
+
let length = usize_from_u32(read_u32(reader)?)?;
|
|
2989
|
+
let mut buffer = vec![0_u8; length];
|
|
2990
|
+
reader.read_exact(&mut buffer)?;
|
|
2991
|
+
String::from_utf8(buffer)
|
|
2992
|
+
.map_err(|err| VectLiteError::InvalidFormat(format!("invalid utf-8 string: {err}")))
|
|
2993
|
+
}
|
|
2994
|
+
|
|
2995
|
+
fn write_u8(writer: &mut impl Write, value: u8) -> io::Result<()> {
|
|
2996
|
+
writer.write_all(&[value])
|
|
2997
|
+
}
|
|
2998
|
+
|
|
2999
|
+
fn read_u8(reader: &mut impl Read) -> io::Result<u8> {
|
|
3000
|
+
let mut buffer = [0_u8; 1];
|
|
3001
|
+
reader.read_exact(&mut buffer)?;
|
|
3002
|
+
Ok(buffer[0])
|
|
3003
|
+
}
|
|
3004
|
+
|
|
3005
|
+
fn write_u16(writer: &mut impl Write, value: u16) -> io::Result<()> {
|
|
3006
|
+
writer.write_all(&value.to_le_bytes())
|
|
3007
|
+
}
|
|
3008
|
+
|
|
3009
|
+
fn read_u16(reader: &mut impl Read) -> io::Result<u16> {
|
|
3010
|
+
let mut buffer = [0_u8; 2];
|
|
3011
|
+
reader.read_exact(&mut buffer)?;
|
|
3012
|
+
Ok(u16::from_le_bytes(buffer))
|
|
3013
|
+
}
|
|
3014
|
+
|
|
3015
|
+
fn write_u32(writer: &mut impl Write, value: u32) -> io::Result<()> {
|
|
3016
|
+
writer.write_all(&value.to_le_bytes())
|
|
3017
|
+
}
|
|
3018
|
+
|
|
3019
|
+
fn read_u32(reader: &mut impl Read) -> io::Result<u32> {
|
|
3020
|
+
let mut buffer = [0_u8; 4];
|
|
3021
|
+
reader.read_exact(&mut buffer)?;
|
|
3022
|
+
Ok(u32::from_le_bytes(buffer))
|
|
3023
|
+
}
|
|
3024
|
+
|
|
3025
|
+
fn write_u64(writer: &mut impl Write, value: u64) -> io::Result<()> {
|
|
3026
|
+
writer.write_all(&value.to_le_bytes())
|
|
3027
|
+
}
|
|
3028
|
+
|
|
3029
|
+
fn read_u64(reader: &mut impl Read) -> io::Result<u64> {
|
|
3030
|
+
let mut buffer = [0_u8; 8];
|
|
3031
|
+
reader.read_exact(&mut buffer)?;
|
|
3032
|
+
Ok(u64::from_le_bytes(buffer))
|
|
3033
|
+
}
|
|
3034
|
+
|
|
3035
|
+
fn write_i64(writer: &mut impl Write, value: i64) -> io::Result<()> {
|
|
3036
|
+
writer.write_all(&value.to_le_bytes())
|
|
3037
|
+
}
|
|
3038
|
+
|
|
3039
|
+
fn read_i64(reader: &mut impl Read) -> io::Result<i64> {
|
|
3040
|
+
let mut buffer = [0_u8; 8];
|
|
3041
|
+
reader.read_exact(&mut buffer)?;
|
|
3042
|
+
Ok(i64::from_le_bytes(buffer))
|
|
3043
|
+
}
|
|
3044
|
+
|
|
3045
|
+
fn write_f64(writer: &mut impl Write, value: f64) -> io::Result<()> {
|
|
3046
|
+
writer.write_all(&value.to_le_bytes())
|
|
3047
|
+
}
|
|
3048
|
+
|
|
3049
|
+
fn read_f64(reader: &mut impl Read) -> io::Result<f64> {
|
|
3050
|
+
let mut buffer = [0_u8; 8];
|
|
3051
|
+
reader.read_exact(&mut buffer)?;
|
|
3052
|
+
Ok(f64::from_le_bytes(buffer))
|
|
3053
|
+
}
|
|
3054
|
+
|
|
3055
|
+
fn write_f32(writer: &mut impl Write, value: f32) -> io::Result<()> {
|
|
3056
|
+
writer.write_all(&value.to_le_bytes())
|
|
3057
|
+
}
|
|
3058
|
+
|
|
3059
|
+
fn read_f32(reader: &mut impl Read) -> io::Result<f32> {
|
|
3060
|
+
let mut buffer = [0_u8; 4];
|
|
3061
|
+
reader.read_exact(&mut buffer)?;
|
|
3062
|
+
Ok(f32::from_le_bytes(buffer))
|
|
3063
|
+
}
|
|
3064
|
+
|
|
3065
|
+
fn u32_from_usize(value: usize) -> Result<u32> {
|
|
3066
|
+
u32::try_from(value)
|
|
3067
|
+
.map_err(|_| VectLiteError::InvalidFormat("value exceeds the u32 storage limit".to_owned()))
|
|
3068
|
+
}
|
|
3069
|
+
|
|
3070
|
+
fn u64_from_usize(value: usize) -> Result<u64> {
|
|
3071
|
+
u64::try_from(value)
|
|
3072
|
+
.map_err(|_| VectLiteError::InvalidFormat("value exceeds the u64 storage limit".to_owned()))
|
|
3073
|
+
}
|
|
3074
|
+
|
|
3075
|
+
fn usize_from_u32(value: u32) -> Result<usize> {
|
|
3076
|
+
usize::try_from(value).map_err(|_| {
|
|
3077
|
+
VectLiteError::InvalidFormat("u32 value cannot fit into usize on this platform".to_owned())
|
|
3078
|
+
})
|
|
3079
|
+
}
|
|
3080
|
+
|
|
3081
|
+
fn usize_from_u64(value: u64) -> Result<usize> {
|
|
3082
|
+
usize::try_from(value).map_err(|_| {
|
|
3083
|
+
VectLiteError::InvalidFormat("u64 value cannot fit into usize on this platform".to_owned())
|
|
3084
|
+
})
|
|
3085
|
+
}
|
|
3086
|
+
|
|
3087
|
+
#[cfg(test)]
|
|
3088
|
+
mod tests {
|
|
3089
|
+
use super::{
|
|
3090
|
+
Database, HybridSearchOptions, Metadata, MetadataFilter, MetadataValue, NamedVectors,
|
|
3091
|
+
Record, SearchOptions, SparseVector,
|
|
3092
|
+
};
|
|
3093
|
+
use std::path::{Path, PathBuf};
|
|
3094
|
+
use std::time::{SystemTime, UNIX_EPOCH};
|
|
3095
|
+
|
|
3096
|
+
#[test]
|
|
3097
|
+
fn roundtrip_persists_records() {
|
|
3098
|
+
let path = temp_file("roundtrip");
|
|
3099
|
+
let mut metadata = Metadata::new();
|
|
3100
|
+
metadata.insert("source".to_owned(), MetadataValue::from("blog"));
|
|
3101
|
+
|
|
3102
|
+
{
|
|
3103
|
+
let mut database = Database::create(&path, 3).expect("create database");
|
|
3104
|
+
database
|
|
3105
|
+
.insert("doc1", vec![1.0, 0.0, 0.0], metadata.clone())
|
|
3106
|
+
.expect("insert record");
|
|
3107
|
+
}
|
|
3108
|
+
|
|
3109
|
+
let reopened = Database::open(&path).expect("reopen database");
|
|
3110
|
+
assert_eq!(reopened.dimension(), 3);
|
|
3111
|
+
assert_eq!(reopened.len(), 1);
|
|
3112
|
+
|
|
3113
|
+
let record = reopened.get("doc1").expect("record exists");
|
|
3114
|
+
assert_eq!(record.namespace, "");
|
|
3115
|
+
assert_eq!(record.id, "doc1");
|
|
3116
|
+
assert_eq!(record.vector, vec![1.0, 0.0, 0.0]);
|
|
3117
|
+
assert!(record.sparse.is_empty());
|
|
3118
|
+
assert_eq!(record.metadata, metadata);
|
|
3119
|
+
|
|
3120
|
+
cleanup(&path);
|
|
3121
|
+
}
|
|
3122
|
+
|
|
3123
|
+
#[test]
|
|
3124
|
+
fn search_orders_by_similarity_and_filters_metadata() {
|
|
3125
|
+
let path = temp_file("search");
|
|
3126
|
+
let mut database = Database::create(&path, 2).expect("create database");
|
|
3127
|
+
|
|
3128
|
+
let mut docs_metadata = Metadata::new();
|
|
3129
|
+
docs_metadata.insert("source".to_owned(), MetadataValue::from("notes"));
|
|
3130
|
+
docs_metadata.insert("title".to_owned(), MetadataValue::from("auth flow"));
|
|
3131
|
+
|
|
3132
|
+
let mut blog_metadata = Metadata::new();
|
|
3133
|
+
blog_metadata.insert("source".to_owned(), MetadataValue::from("blog"));
|
|
3134
|
+
blog_metadata.insert("title".to_owned(), MetadataValue::from("shipping"));
|
|
3135
|
+
|
|
3136
|
+
database
|
|
3137
|
+
.insert("doc1", vec![1.0, 0.0], docs_metadata)
|
|
3138
|
+
.expect("insert doc1");
|
|
3139
|
+
database
|
|
3140
|
+
.insert("doc2", vec![0.8, 0.2], blog_metadata)
|
|
3141
|
+
.expect("insert doc2");
|
|
3142
|
+
database
|
|
3143
|
+
.insert("doc3", vec![0.0, 1.0], Metadata::new())
|
|
3144
|
+
.expect("insert doc3");
|
|
3145
|
+
|
|
3146
|
+
let results = database
|
|
3147
|
+
.search(
|
|
3148
|
+
&[1.0, 0.0],
|
|
3149
|
+
SearchOptions {
|
|
3150
|
+
top_k: 5,
|
|
3151
|
+
filter: Some(MetadataFilter::and(vec![
|
|
3152
|
+
MetadataFilter::eq("source", "notes"),
|
|
3153
|
+
MetadataFilter::contains("title", "auth"),
|
|
3154
|
+
])),
|
|
3155
|
+
},
|
|
3156
|
+
)
|
|
3157
|
+
.expect("search database");
|
|
3158
|
+
|
|
3159
|
+
assert_eq!(results.len(), 1);
|
|
3160
|
+
assert_eq!(results[0].id, "doc1");
|
|
3161
|
+
assert!(results[0].score > 0.99);
|
|
3162
|
+
|
|
3163
|
+
cleanup(&path);
|
|
3164
|
+
}
|
|
3165
|
+
|
|
3166
|
+
#[test]
|
|
3167
|
+
fn delete_removes_records_from_disk() {
|
|
3168
|
+
let path = temp_file("delete");
|
|
3169
|
+
{
|
|
3170
|
+
let mut database = Database::create(&path, 2).expect("create database");
|
|
3171
|
+
database
|
|
3172
|
+
.insert("doc1", vec![1.0, 0.0], Metadata::new())
|
|
3173
|
+
.expect("insert record");
|
|
3174
|
+
database.delete("doc1").expect("delete record");
|
|
3175
|
+
}
|
|
3176
|
+
|
|
3177
|
+
let reopened = Database::open(&path).expect("reopen database");
|
|
3178
|
+
assert!(reopened.get("doc1").is_none());
|
|
3179
|
+
assert_eq!(reopened.len(), 0);
|
|
3180
|
+
|
|
3181
|
+
cleanup(&path);
|
|
3182
|
+
}
|
|
3183
|
+
|
|
3184
|
+
#[test]
|
|
3185
|
+
fn batch_upsert_persists_records_in_one_call() {
|
|
3186
|
+
let path = temp_file("batch-upsert");
|
|
3187
|
+
{
|
|
3188
|
+
let mut database = Database::create(&path, 2).expect("create database");
|
|
3189
|
+
|
|
3190
|
+
let inserted = database
|
|
3191
|
+
.upsert_many(vec![
|
|
3192
|
+
Record {
|
|
3193
|
+
namespace: "".to_owned(),
|
|
3194
|
+
id: "doc1".to_owned(),
|
|
3195
|
+
vector: vec![1.0, 0.0],
|
|
3196
|
+
vectors: NamedVectors::new(),
|
|
3197
|
+
sparse: SparseVector::new(),
|
|
3198
|
+
metadata: Metadata::new(),
|
|
3199
|
+
},
|
|
3200
|
+
Record {
|
|
3201
|
+
namespace: "".to_owned(),
|
|
3202
|
+
id: "doc2".to_owned(),
|
|
3203
|
+
vector: vec![0.0, 1.0],
|
|
3204
|
+
vectors: NamedVectors::new(),
|
|
3205
|
+
sparse: SparseVector::new(),
|
|
3206
|
+
metadata: Metadata::new(),
|
|
3207
|
+
},
|
|
3208
|
+
])
|
|
3209
|
+
.expect("batch upsert");
|
|
3210
|
+
|
|
3211
|
+
assert_eq!(inserted, 2);
|
|
3212
|
+
}
|
|
3213
|
+
|
|
3214
|
+
let reopened = Database::open(&path).expect("reopen database");
|
|
3215
|
+
assert_eq!(reopened.len(), 2);
|
|
3216
|
+
|
|
3217
|
+
cleanup(&path);
|
|
3218
|
+
}
|
|
3219
|
+
|
|
3220
|
+
#[test]
|
|
3221
|
+
fn extended_filters_match_expected_records() {
|
|
3222
|
+
let path = temp_file("extended-filters");
|
|
3223
|
+
let mut database = Database::create(&path, 2).expect("create database");
|
|
3224
|
+
|
|
3225
|
+
let mut metadata = Metadata::new();
|
|
3226
|
+
metadata.insert("source".to_owned(), MetadataValue::from("blog"));
|
|
3227
|
+
metadata.insert("priority".to_owned(), MetadataValue::from(10));
|
|
3228
|
+
|
|
3229
|
+
database
|
|
3230
|
+
.insert("doc1", vec![1.0, 0.0], metadata)
|
|
3231
|
+
.expect("insert record");
|
|
3232
|
+
|
|
3233
|
+
let results = database
|
|
3234
|
+
.search(
|
|
3235
|
+
&[1.0, 0.0],
|
|
3236
|
+
SearchOptions {
|
|
3237
|
+
top_k: 10,
|
|
3238
|
+
filter: Some(MetadataFilter::and(vec![
|
|
3239
|
+
MetadataFilter::ne("source", "notes"),
|
|
3240
|
+
MetadataFilter::gte("priority", 10.0),
|
|
3241
|
+
MetadataFilter::lte("priority", 10.0),
|
|
3242
|
+
])),
|
|
3243
|
+
},
|
|
3244
|
+
)
|
|
3245
|
+
.expect("search database");
|
|
3246
|
+
|
|
3247
|
+
assert_eq!(results.len(), 1);
|
|
3248
|
+
assert_eq!(results[0].namespace, "");
|
|
3249
|
+
assert_eq!(results[0].id, "doc1");
|
|
3250
|
+
assert!(results[0].dense_score > 0.99);
|
|
3251
|
+
assert_eq!(results[0].sparse_score, 0.0);
|
|
3252
|
+
|
|
3253
|
+
cleanup(&path);
|
|
3254
|
+
}
|
|
3255
|
+
|
|
3256
|
+
#[test]
|
|
3257
|
+
fn namespaces_isolate_same_ids() {
|
|
3258
|
+
let path = temp_file("namespaces");
|
|
3259
|
+
let mut database = Database::create(&path, 2).expect("create database");
|
|
3260
|
+
|
|
3261
|
+
database
|
|
3262
|
+
.insert_in_namespace("docs", "doc1", vec![1.0, 0.0], Metadata::new())
|
|
3263
|
+
.expect("insert docs");
|
|
3264
|
+
database
|
|
3265
|
+
.insert_in_namespace("notes", "doc1", vec![0.0, 1.0], Metadata::new())
|
|
3266
|
+
.expect("insert notes");
|
|
3267
|
+
|
|
3268
|
+
assert!(database.get_in_namespace("docs", "doc1").is_some());
|
|
3269
|
+
assert!(database.get_in_namespace("notes", "doc1").is_some());
|
|
3270
|
+
assert!(database.get("doc1").is_none());
|
|
3271
|
+
|
|
3272
|
+
let docs = database
|
|
3273
|
+
.search_in_namespace("docs", &[1.0, 0.0], SearchOptions::default())
|
|
3274
|
+
.expect("search docs");
|
|
3275
|
+
let all = database
|
|
3276
|
+
.search_all_namespaces(&[1.0, 0.0], SearchOptions::default())
|
|
3277
|
+
.expect("search all");
|
|
3278
|
+
|
|
3279
|
+
assert_eq!(docs.len(), 1);
|
|
3280
|
+
assert_eq!(docs[0].namespace, "docs");
|
|
3281
|
+
assert_eq!(all.len(), 2);
|
|
3282
|
+
assert_eq!(
|
|
3283
|
+
database.namespaces(),
|
|
3284
|
+
vec!["docs".to_owned(), "notes".to_owned()]
|
|
3285
|
+
);
|
|
3286
|
+
|
|
3287
|
+
cleanup(&path);
|
|
3288
|
+
}
|
|
3289
|
+
|
|
3290
|
+
#[test]
|
|
3291
|
+
fn hybrid_search_combines_dense_and_sparse_scores() {
|
|
3292
|
+
let path = temp_file("hybrid");
|
|
3293
|
+
let mut database = Database::create(&path, 2).expect("create database");
|
|
3294
|
+
|
|
3295
|
+
let mut sparse_auth = SparseVector::new();
|
|
3296
|
+
sparse_auth.insert("auth".to_owned(), 1.0);
|
|
3297
|
+
sparse_auth.insert("sso".to_owned(), 0.5);
|
|
3298
|
+
|
|
3299
|
+
let mut sparse_billing = SparseVector::new();
|
|
3300
|
+
sparse_billing.insert("billing".to_owned(), 1.0);
|
|
3301
|
+
|
|
3302
|
+
database
|
|
3303
|
+
.upsert_with_sparse_in_namespace(
|
|
3304
|
+
"docs",
|
|
3305
|
+
"doc1",
|
|
3306
|
+
vec![1.0, 0.0],
|
|
3307
|
+
sparse_auth,
|
|
3308
|
+
Metadata::new(),
|
|
3309
|
+
)
|
|
3310
|
+
.expect("insert doc1");
|
|
3311
|
+
database
|
|
3312
|
+
.upsert_with_sparse_in_namespace(
|
|
3313
|
+
"docs",
|
|
3314
|
+
"doc2",
|
|
3315
|
+
vec![1.0, 0.0],
|
|
3316
|
+
sparse_billing,
|
|
3317
|
+
Metadata::new(),
|
|
3318
|
+
)
|
|
3319
|
+
.expect("insert doc2");
|
|
3320
|
+
|
|
3321
|
+
let mut query_sparse = SparseVector::new();
|
|
3322
|
+
query_sparse.insert("auth".to_owned(), 1.0);
|
|
3323
|
+
|
|
3324
|
+
let results = database
|
|
3325
|
+
.hybrid_search_in_namespace(
|
|
3326
|
+
"docs",
|
|
3327
|
+
Some(&[1.0, 0.0]),
|
|
3328
|
+
Some(&query_sparse),
|
|
3329
|
+
HybridSearchOptions {
|
|
3330
|
+
top_k: 10,
|
|
3331
|
+
filter: None,
|
|
3332
|
+
dense_weight: 1.0,
|
|
3333
|
+
sparse_weight: 1.0,
|
|
3334
|
+
..HybridSearchOptions::default()
|
|
3335
|
+
},
|
|
3336
|
+
)
|
|
3337
|
+
.expect("hybrid search");
|
|
3338
|
+
|
|
3339
|
+
assert_eq!(results.len(), 2);
|
|
3340
|
+
assert_eq!(results[0].id, "doc1");
|
|
3341
|
+
assert!(results[0].sparse_score > results[1].sparse_score);
|
|
3342
|
+
|
|
3343
|
+
cleanup(&path);
|
|
3344
|
+
}
|
|
3345
|
+
|
|
3346
|
+
#[test]
|
|
3347
|
+
fn named_vectors_roundtrip_and_search() {
|
|
3348
|
+
let path = temp_file("named-vectors");
|
|
3349
|
+
{
|
|
3350
|
+
let mut database = Database::create(&path, 2).expect("create database");
|
|
3351
|
+
|
|
3352
|
+
let mut doc1_vectors = NamedVectors::new();
|
|
3353
|
+
doc1_vectors.insert("title".to_owned(), vec![1.0, 0.0]);
|
|
3354
|
+
doc1_vectors.insert("body".to_owned(), vec![0.0, 1.0]);
|
|
3355
|
+
|
|
3356
|
+
let mut doc2_vectors = NamedVectors::new();
|
|
3357
|
+
doc2_vectors.insert("title".to_owned(), vec![0.0, 1.0]);
|
|
3358
|
+
doc2_vectors.insert("body".to_owned(), vec![1.0, 0.0]);
|
|
3359
|
+
|
|
3360
|
+
database
|
|
3361
|
+
.upsert_with_vectors_in_namespace(
|
|
3362
|
+
"docs",
|
|
3363
|
+
"doc1",
|
|
3364
|
+
vec![0.2, 0.8],
|
|
3365
|
+
doc1_vectors,
|
|
3366
|
+
SparseVector::new(),
|
|
3367
|
+
Metadata::new(),
|
|
3368
|
+
)
|
|
3369
|
+
.expect("insert doc1");
|
|
3370
|
+
database
|
|
3371
|
+
.upsert_with_vectors_in_namespace(
|
|
3372
|
+
"docs",
|
|
3373
|
+
"doc2",
|
|
3374
|
+
vec![0.8, 0.2],
|
|
3375
|
+
doc2_vectors,
|
|
3376
|
+
SparseVector::new(),
|
|
3377
|
+
Metadata::new(),
|
|
3378
|
+
)
|
|
3379
|
+
.expect("insert doc2");
|
|
3380
|
+
}
|
|
3381
|
+
|
|
3382
|
+
let reopened = Database::open(&path).expect("reopen database");
|
|
3383
|
+
let record = reopened
|
|
3384
|
+
.get_in_namespace("docs", "doc1")
|
|
3385
|
+
.expect("expected stored record");
|
|
3386
|
+
assert_eq!(record.vectors.len(), 2);
|
|
3387
|
+
assert_eq!(record.vectors.get("title"), Some(&vec![1.0, 0.0]));
|
|
3388
|
+
|
|
3389
|
+
let title_results = reopened
|
|
3390
|
+
.hybrid_search_in_namespace(
|
|
3391
|
+
"docs",
|
|
3392
|
+
Some(&[1.0, 0.0]),
|
|
3393
|
+
None,
|
|
3394
|
+
HybridSearchOptions {
|
|
3395
|
+
top_k: 2,
|
|
3396
|
+
filter: None,
|
|
3397
|
+
dense_weight: 1.0,
|
|
3398
|
+
sparse_weight: 0.0,
|
|
3399
|
+
vector_name: Some("title".to_owned()),
|
|
3400
|
+
..HybridSearchOptions::default()
|
|
3401
|
+
},
|
|
3402
|
+
)
|
|
3403
|
+
.expect("search title vector");
|
|
3404
|
+
|
|
3405
|
+
let default_results = reopened
|
|
3406
|
+
.search_in_namespace("docs", &[1.0, 0.0], SearchOptions::default())
|
|
3407
|
+
.expect("search default vector");
|
|
3408
|
+
|
|
3409
|
+
assert_eq!(
|
|
3410
|
+
title_results
|
|
3411
|
+
.iter()
|
|
3412
|
+
.map(|result| result.id.as_str())
|
|
3413
|
+
.collect::<Vec<_>>(),
|
|
3414
|
+
vec!["doc1", "doc2"]
|
|
3415
|
+
);
|
|
3416
|
+
assert_eq!(
|
|
3417
|
+
default_results
|
|
3418
|
+
.iter()
|
|
3419
|
+
.map(|result| result.id.as_str())
|
|
3420
|
+
.collect::<Vec<_>>(),
|
|
3421
|
+
vec!["doc2", "doc1"]
|
|
3422
|
+
);
|
|
3423
|
+
|
|
3424
|
+
cleanup(&path);
|
|
3425
|
+
}
|
|
3426
|
+
|
|
3427
|
+
#[test]
|
|
3428
|
+
fn ann_index_is_built_for_larger_collections() {
|
|
3429
|
+
let path = temp_file("ann");
|
|
3430
|
+
let mut database = Database::create(&path, 128).expect("create database");
|
|
3431
|
+
|
|
3432
|
+
for i in 0..128 {
|
|
3433
|
+
let mut vector = vec![0.0_f32; 128];
|
|
3434
|
+
vector[i] = 1.0;
|
|
3435
|
+
database
|
|
3436
|
+
.insert_in_namespace("docs", format!("doc{i}"), vector, Metadata::new())
|
|
3437
|
+
.expect("insert record");
|
|
3438
|
+
}
|
|
3439
|
+
|
|
3440
|
+
assert!(database.has_ann_index(None, None));
|
|
3441
|
+
assert!(database.has_ann_index(Some("docs"), None));
|
|
3442
|
+
|
|
3443
|
+
let mut query = vec![0.0_f32; 128];
|
|
3444
|
+
query[42] = 1.0;
|
|
3445
|
+
let outcome = database
|
|
3446
|
+
.hybrid_search_in_namespace_with_stats(
|
|
3447
|
+
"docs",
|
|
3448
|
+
Some(&query),
|
|
3449
|
+
None,
|
|
3450
|
+
HybridSearchOptions {
|
|
3451
|
+
top_k: 10,
|
|
3452
|
+
filter: None,
|
|
3453
|
+
dense_weight: 1.0,
|
|
3454
|
+
sparse_weight: 0.0,
|
|
3455
|
+
..HybridSearchOptions::default()
|
|
3456
|
+
},
|
|
3457
|
+
)
|
|
3458
|
+
.expect("search database");
|
|
3459
|
+
|
|
3460
|
+
assert!(outcome.stats.used_ann);
|
|
3461
|
+
assert!(outcome.stats.ann_candidate_count >= 10);
|
|
3462
|
+
assert!(!outcome.stats.exact_fallback);
|
|
3463
|
+
assert_eq!(outcome.results.len(), 10);
|
|
3464
|
+
|
|
3465
|
+
cleanup(&path);
|
|
3466
|
+
}
|
|
3467
|
+
|
|
3468
|
+
#[test]
|
|
3469
|
+
fn mmr_diversifies_near_duplicate_results() {
|
|
3470
|
+
let path = temp_file("mmr");
|
|
3471
|
+
let mut database = Database::create(&path, 2).expect("create database");
|
|
3472
|
+
|
|
3473
|
+
database
|
|
3474
|
+
.insert("doc1", vec![1.0, 0.0], Metadata::new())
|
|
3475
|
+
.expect("insert doc1");
|
|
3476
|
+
database
|
|
3477
|
+
.insert("doc2", vec![0.99, 0.01], Metadata::new())
|
|
3478
|
+
.expect("insert doc2");
|
|
3479
|
+
database
|
|
3480
|
+
.insert("doc3", vec![0.7, 0.7], Metadata::new())
|
|
3481
|
+
.expect("insert doc3");
|
|
3482
|
+
|
|
3483
|
+
let plain_results = database
|
|
3484
|
+
.search(
|
|
3485
|
+
&[1.0, 0.0],
|
|
3486
|
+
SearchOptions {
|
|
3487
|
+
top_k: 2,
|
|
3488
|
+
filter: None,
|
|
3489
|
+
},
|
|
3490
|
+
)
|
|
3491
|
+
.expect("search database");
|
|
3492
|
+
let mmr_outcome = database
|
|
3493
|
+
.hybrid_search_in_namespace_with_stats(
|
|
3494
|
+
"",
|
|
3495
|
+
Some(&[1.0, 0.0]),
|
|
3496
|
+
None,
|
|
3497
|
+
HybridSearchOptions {
|
|
3498
|
+
top_k: 2,
|
|
3499
|
+
filter: None,
|
|
3500
|
+
dense_weight: 1.0,
|
|
3501
|
+
sparse_weight: 0.0,
|
|
3502
|
+
fetch_k: 3,
|
|
3503
|
+
mmr_lambda: Some(0.3),
|
|
3504
|
+
..HybridSearchOptions::default()
|
|
3505
|
+
},
|
|
3506
|
+
)
|
|
3507
|
+
.expect("mmr search");
|
|
3508
|
+
|
|
3509
|
+
assert_eq!(
|
|
3510
|
+
plain_results
|
|
3511
|
+
.iter()
|
|
3512
|
+
.map(|result| result.id.as_str())
|
|
3513
|
+
.collect::<Vec<_>>(),
|
|
3514
|
+
vec!["doc1", "doc2"]
|
|
3515
|
+
);
|
|
3516
|
+
assert_eq!(
|
|
3517
|
+
mmr_outcome
|
|
3518
|
+
.results
|
|
3519
|
+
.iter()
|
|
3520
|
+
.map(|result| result.id.as_str())
|
|
3521
|
+
.collect::<Vec<_>>(),
|
|
3522
|
+
vec!["doc1", "doc3"]
|
|
3523
|
+
);
|
|
3524
|
+
assert!(mmr_outcome.stats.mmr_applied);
|
|
3525
|
+
assert_eq!(mmr_outcome.stats.fetch_k, 3);
|
|
3526
|
+
assert_eq!(mmr_outcome.stats.considered_count, 3);
|
|
3527
|
+
|
|
3528
|
+
cleanup(&path);
|
|
3529
|
+
}
|
|
3530
|
+
|
|
3531
|
+
fn temp_file(name: &str) -> PathBuf {
|
|
3532
|
+
let nanos = SystemTime::now()
|
|
3533
|
+
.duration_since(UNIX_EPOCH)
|
|
3534
|
+
.expect("system time")
|
|
3535
|
+
.as_nanos();
|
|
3536
|
+
std::env::temp_dir().join(format!(
|
|
3537
|
+
"vectlite-{name}-{}-{nanos}.vdb",
|
|
3538
|
+
std::process::id()
|
|
3539
|
+
))
|
|
3540
|
+
}
|
|
3541
|
+
|
|
3542
|
+
fn cleanup(path: &Path) {
|
|
3543
|
+
let _ = std::fs::remove_file(path);
|
|
3544
|
+
}
|
|
3545
|
+
}
|