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.
@@ -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
+ }