1mod filter;
11mod options;
12
13pub use crate::omen::Metric;
14pub use filter::MetadataFilter;
15pub use options::VectorStoreOptions;
16
17use super::hnsw::HNSWParams;
18use super::hnsw_index::HNSWIndex;
19use super::types::Vector;
20use super::QuantizationMode;
21use crate::compression::{QuantizationBits, RaBitQParams};
22use crate::distance::l2_distance;
23use crate::omen::{MetadataIndex, OmenFile};
24use crate::text::{
25 weighted_reciprocal_rank_fusion, weighted_reciprocal_rank_fusion_with_subscores, HybridResult,
26 TextIndex, TextSearchConfig, DEFAULT_RRF_K,
27};
28use anyhow::Result;
29use rayon::prelude::*;
30use rustc_hash::FxHashMap;
31use serde_json::Value as JsonValue;
32use std::collections::HashMap;
33use std::path::{Path, PathBuf};
34
35const DEFAULT_HNSW_M: usize = 16;
41const DEFAULT_HNSW_EF_CONSTRUCTION: usize = 100;
43const DEFAULT_HNSW_EF_SEARCH: usize = 100;
45const DEFAULT_OVERSAMPLE_FACTOR: f32 = 3.0;
47
48#[inline]
56fn compute_effective_ef(ef: Option<usize>, stored_ef: usize, k: usize) -> usize {
57 ef.unwrap_or(stored_ef).max(k)
58}
59
60#[cfg(debug_assertions)]
64fn debug_assert_mapping_consistency(
65 id_to_index: &FxHashMap<String, usize>,
66 index_to_id: &FxHashMap<usize, String>,
67) {
68 debug_assert_eq!(
70 id_to_index.len(),
71 index_to_id.len(),
72 "ID mapping size mismatch: id_to_index={}, index_to_id={}",
73 id_to_index.len(),
74 index_to_id.len()
75 );
76
77 for (id, &idx) in id_to_index {
79 debug_assert_eq!(
80 index_to_id.get(&idx),
81 Some(id),
82 "Mapping inconsistency: id_to_index[{id}]={idx} but index_to_id[{idx}]={:?}",
83 index_to_id.get(&idx)
84 );
85 }
86}
87
88#[cfg(not(debug_assertions))]
89#[inline]
90fn debug_assert_mapping_consistency(
91 _id_to_index: &FxHashMap<String, usize>,
92 _index_to_id: &FxHashMap<usize, String>,
93) {
94 }
96
97#[cfg(test)]
98mod tests;
99
100fn default_oversample_for_quantization(mode: Option<&QuantizationMode>) -> f32 {
110 match mode {
111 None => 1.0,
112 Some(QuantizationMode::Binary) => 5.0, Some(QuantizationMode::SQ8) => 2.0,
114 Some(QuantizationMode::RaBitQ(params)) => match params.bits_per_dim.to_u8() {
115 2 => 4.0, 8 => 2.0, _ => 3.0, },
119 }
120}
121
122fn quantization_mode_from_id(mode_id: u64) -> Option<QuantizationMode> {
126 match mode_id {
127 1 => Some(QuantizationMode::SQ8),
128 2 => Some(QuantizationMode::RaBitQ(RaBitQParams {
129 bits_per_dim: QuantizationBits::Bits4,
130 ..RaBitQParams::default()
131 })),
132 3 => Some(QuantizationMode::RaBitQ(RaBitQParams {
133 bits_per_dim: QuantizationBits::Bits2,
134 ..RaBitQParams::default()
135 })),
136 4 => Some(QuantizationMode::RaBitQ(RaBitQParams {
137 bits_per_dim: QuantizationBits::Bits8,
138 ..RaBitQParams::default()
139 })),
140 5 => Some(QuantizationMode::Binary),
141 _ => None, }
143}
144
145fn quantization_mode_to_id(mode: &QuantizationMode) -> u64 {
149 match mode {
150 QuantizationMode::Binary => 5,
151 QuantizationMode::SQ8 => 1,
152 QuantizationMode::RaBitQ(p) => match p.bits_per_dim.to_u8() {
153 2 => 3,
154 8 => 4,
155 _ => 2, },
157 }
158}
159
160fn create_hnsw_index(
164 dimensions: usize,
165 hnsw_m: usize,
166 hnsw_ef_construction: usize,
167 hnsw_ef_search: usize,
168 distance_metric: Metric,
169 quantization_mode: Option<&QuantizationMode>,
170 training_vectors: &[Vec<f32>],
171) -> Result<HNSWIndex> {
172 use super::hnsw_index::HNSWQuantization;
173
174 let m = hnsw_m.max(DEFAULT_HNSW_M);
176 let ef_construction = hnsw_ef_construction.max(DEFAULT_HNSW_EF_CONSTRUCTION);
177 let ef_search = hnsw_ef_search.max(DEFAULT_HNSW_EF_SEARCH);
178
179 let quantization = match quantization_mode {
181 Some(QuantizationMode::Binary) => HNSWQuantization::Binary,
182 Some(QuantizationMode::SQ8) => HNSWQuantization::SQ8,
183 Some(QuantizationMode::RaBitQ(params)) => HNSWQuantization::RaBitQ(params.clone()),
184 None => HNSWQuantization::None,
185 };
186
187 HNSWIndex::builder()
188 .dimensions(dimensions)
189 .max_elements(training_vectors.len().max(10_000))
190 .m(m)
191 .ef_construction(ef_construction)
192 .ef_search(ef_search)
193 .metric(distance_metric.into())
194 .quantization(quantization)
195 .build_with_training(training_vectors)
196}
197
198fn initialize_quantized_hnsw(
200 dimensions: usize,
201 hnsw_m: usize,
202 hnsw_ef_construction: usize,
203 hnsw_ef_search: usize,
204 distance_metric: Metric,
205 quant_mode: QuantizationMode,
206 training_vectors: &[Vec<f32>],
207) -> Result<HNSWIndex> {
208 let hnsw_params = HNSWParams::default()
209 .with_m(hnsw_m)
210 .with_ef_construction(hnsw_ef_construction)
211 .with_ef_search(hnsw_ef_search);
212
213 match quant_mode {
214 QuantizationMode::Binary => {
215 let mut idx =
216 HNSWIndex::new_with_binary(dimensions, hnsw_params, distance_metric.into())?;
217 idx.train_quantizer(training_vectors)?;
218 Ok(idx)
219 }
220 QuantizationMode::SQ8 => {
221 HNSWIndex::new_with_sq8(dimensions, hnsw_params, distance_metric.into())
222 }
223 QuantizationMode::RaBitQ(params) => {
224 let mut idx = HNSWIndex::new_with_asymmetric(
225 dimensions,
226 hnsw_params,
227 distance_metric.into(),
228 params,
229 )?;
230 idx.train_quantizer(training_vectors)?;
231 Ok(idx)
232 }
233 }
234}
235
236fn initialize_standard_hnsw(
238 dimensions: usize,
239 hnsw_m: usize,
240 hnsw_ef_construction: usize,
241 hnsw_ef_search: usize,
242 distance_metric: Metric,
243 capacity: usize,
244) -> Result<HNSWIndex> {
245 HNSWIndex::new_with_params(
246 capacity,
247 dimensions,
248 hnsw_m,
249 hnsw_ef_construction,
250 hnsw_ef_search,
251 distance_metric.into(),
252 )
253}
254
255#[inline]
257fn default_metadata() -> JsonValue {
258 serde_json::json!({})
259}
260
261pub struct VectorStore {
263 pub vectors: Vec<Vector>,
265
266 pub hnsw_index: Option<HNSWIndex>,
268
269 dimensions: usize,
271
272 rescore_enabled: bool,
274
275 oversample_factor: f32,
277
278 metadata: HashMap<usize, JsonValue>,
280
281 pub id_to_index: FxHashMap<String, usize>,
283
284 index_to_id: FxHashMap<usize, String>,
286
287 deleted: HashMap<usize, bool>,
289
290 metadata_index: MetadataIndex,
292
293 storage: Option<OmenFile>,
295
296 storage_path: Option<PathBuf>,
298
299 text_index: Option<TextIndex>,
301
302 text_search_config: Option<TextSearchConfig>,
304
305 pending_quantization: Option<QuantizationMode>,
307
308 hnsw_m: usize,
310 hnsw_ef_construction: usize,
311 hnsw_ef_search: usize,
312
313 distance_metric: Metric,
315
316 next_index: usize,
318}
319
320impl VectorStore {
321 #[must_use]
327 pub fn new(dimensions: usize) -> Self {
328 Self {
329 vectors: Vec::new(),
330 hnsw_index: None,
331 dimensions,
332 rescore_enabled: false,
333 oversample_factor: DEFAULT_OVERSAMPLE_FACTOR,
334 metadata: HashMap::new(),
335 id_to_index: FxHashMap::default(),
336 index_to_id: FxHashMap::default(),
337 deleted: HashMap::new(),
338 metadata_index: MetadataIndex::new(),
339 storage: None,
340 storage_path: None,
341 text_index: None,
342 text_search_config: None,
343 pending_quantization: None,
344 hnsw_m: DEFAULT_HNSW_M,
345 hnsw_ef_construction: DEFAULT_HNSW_EF_CONSTRUCTION,
346 hnsw_ef_search: DEFAULT_HNSW_EF_SEARCH,
347 distance_metric: Metric::L2,
348 next_index: 0,
349 }
350 }
351
352 #[must_use]
356 pub fn new_with_quantization(dimensions: usize, mode: QuantizationMode) -> Self {
357 Self {
358 vectors: Vec::new(),
359 hnsw_index: None,
360 dimensions,
361 rescore_enabled: true,
362 oversample_factor: DEFAULT_OVERSAMPLE_FACTOR,
363 metadata: HashMap::new(),
364 id_to_index: FxHashMap::default(),
365 index_to_id: FxHashMap::default(),
366 deleted: HashMap::new(),
367 metadata_index: MetadataIndex::new(),
368 storage: None,
369 storage_path: None,
370 text_index: None,
371 text_search_config: None,
372 pending_quantization: Some(mode),
373 hnsw_m: DEFAULT_HNSW_M,
374 hnsw_ef_construction: DEFAULT_HNSW_EF_CONSTRUCTION,
375 hnsw_ef_search: DEFAULT_HNSW_EF_SEARCH,
376 distance_metric: Metric::L2,
377 next_index: 0,
378 }
379 }
380
381 pub fn new_with_params(
383 dimensions: usize,
384 m: usize,
385 ef_construction: usize,
386 ef_search: usize,
387 distance_metric: Metric,
388 ) -> Result<Self> {
389 let hnsw_index = Some(HNSWIndex::new_with_params(
390 1_000_000,
391 dimensions,
392 m,
393 ef_construction,
394 ef_search,
395 distance_metric.into(),
396 )?);
397
398 Ok(Self {
399 vectors: Vec::new(),
400 hnsw_index,
401 dimensions,
402 rescore_enabled: false,
403 oversample_factor: DEFAULT_OVERSAMPLE_FACTOR,
404 metadata: HashMap::new(),
405 id_to_index: FxHashMap::default(),
406 index_to_id: FxHashMap::default(),
407 deleted: HashMap::new(),
408 metadata_index: MetadataIndex::new(),
409 storage: None,
410 storage_path: None,
411 text_index: None,
412 text_search_config: None,
413 pending_quantization: None,
414 hnsw_m: m,
415 hnsw_ef_construction: ef_construction,
416 hnsw_ef_search: ef_search,
417 distance_metric,
418 next_index: 0,
419 })
420 }
421
422 pub fn open(path: impl AsRef<Path>) -> Result<Self> {
441 let path = path.as_ref();
442 let omen_path = OmenFile::compute_omen_path(path);
443 let storage = if omen_path.exists() {
444 OmenFile::open(path)?
445 } else {
446 OmenFile::create(path, 0)?
447 };
448
449 let is_quantized = storage.is_quantized()?;
451 let quantization_mode =
452 quantization_mode_from_id(storage.get_quantization_mode()?.unwrap_or(0));
453
454 let metadata = storage.load_all_metadata()?;
456 let id_to_index: FxHashMap<String, usize> =
457 storage.load_all_id_mappings()?.into_iter().collect();
458 let deleted = storage.load_all_deleted()?;
459
460 let dimensions = storage.get_config("dimensions")?.unwrap_or(0) as usize;
462
463 let header = storage.header();
465 let distance_metric = header.distance_fn;
466 let hnsw_m = header.m as usize;
467 let hnsw_ef_construction = header.ef_construction as usize;
468 let hnsw_ef_search = header.ef_search as usize;
469
470 let (vectors, real_indices) = if is_quantized {
472 (Vec::new(), std::collections::HashSet::new())
473 } else {
474 let vectors_data = storage.load_all_vectors()?;
475 let mut vectors: Vec<Vector> = Vec::new();
476 let mut real_indices: std::collections::HashSet<usize> =
477 std::collections::HashSet::new();
478
479 for (id, data) in &vectors_data {
480 while vectors.len() < *id {
481 vectors.push(Vector::new(vec![0.0; dimensions.max(1)]));
482 }
483 vectors.push(Vector::new(data.clone()));
484 real_indices.insert(*id);
485 }
486 (vectors, real_indices)
487 };
488
489 let mut deleted = deleted;
491 for idx in 0..vectors.len() {
492 if !real_indices.contains(&idx) && !deleted.contains_key(&idx) {
493 deleted.insert(idx, true);
494 }
495 }
496
497 let active_vector_count = vectors
500 .iter()
501 .enumerate()
502 .filter(|(i, _)| !deleted.contains_key(i))
503 .count();
504
505 let hnsw_index = if let Some(hnsw_bytes) = storage.get_hnsw_index() {
506 match postcard::from_bytes::<HNSWIndex>(hnsw_bytes) {
507 Ok(index) => {
508 if index.len() != active_vector_count && !vectors.is_empty() {
510 tracing::info!(
511 "HNSW index count ({}) differs from vector count ({}), rebuilding",
512 index.len(),
513 active_vector_count
514 );
515 let vector_data: Vec<Vec<f32>> =
516 vectors.iter().map(|v| v.data.clone()).collect();
517 let mut new_index = create_hnsw_index(
518 dimensions,
519 hnsw_m,
520 hnsw_ef_construction,
521 hnsw_ef_search,
522 distance_metric,
523 quantization_mode.as_ref(),
524 &vector_data,
525 )?;
526 new_index.batch_insert(&vector_data)?;
527 Some(new_index)
528 } else {
529 Some(index)
530 }
531 }
532 Err(e) => {
533 tracing::warn!("Failed to deserialize HNSW index, rebuilding: {}", e);
534 None
535 }
536 }
537 } else if !vectors.is_empty() {
538 let vector_data: Vec<Vec<f32>> = vectors.iter().map(|v| v.data.clone()).collect();
539 let mut index = create_hnsw_index(
540 dimensions,
541 hnsw_m,
542 hnsw_ef_construction,
543 hnsw_ef_search,
544 distance_metric,
545 quantization_mode.as_ref(),
546 &vector_data,
547 )?;
548 index.batch_insert(&vector_data)?;
549 Some(index)
550 } else if is_quantized && dimensions > 0 {
551 let vectors_data = storage.load_all_vectors()?;
552 if vectors_data.is_empty() {
553 None
554 } else {
555 let vector_data: Vec<Vec<f32>> =
556 vectors_data.iter().map(|(_, v)| v.clone()).collect();
557 let mut index = create_hnsw_index(
558 dimensions,
559 hnsw_m,
560 hnsw_ef_construction,
561 hnsw_ef_search,
562 distance_metric,
563 quantization_mode.as_ref(),
564 &vector_data,
565 )?;
566 index.batch_insert(&vector_data)?;
567 Some(index)
568 }
569 } else {
570 None
571 };
572
573 let text_index_path = path.join("text_index");
575 let text_index = if text_index_path.exists() {
576 Some(TextIndex::open(&text_index_path)?)
577 } else {
578 None
579 };
580
581 let index_to_id: FxHashMap<usize, String> = id_to_index
583 .iter()
584 .map(|(id, &idx)| (idx, id.clone()))
585 .collect();
586
587 let mut metadata_index = MetadataIndex::new();
589 for (&idx, meta) in &metadata {
590 if !deleted.contains_key(&idx) {
591 metadata_index.index_json(idx as u32, meta);
592 }
593 }
594
595 let rescore_enabled = hnsw_index
597 .as_ref()
598 .is_some_and(super::hnsw_index::HNSWIndex::is_asymmetric);
599
600 debug_assert_mapping_consistency(&id_to_index, &index_to_id);
602
603 let next_index = id_to_index.values().max().map_or(0, |&max| max + 1);
605
606 Ok(Self {
607 vectors,
608 hnsw_index,
609 dimensions,
610 rescore_enabled,
611 oversample_factor: DEFAULT_OVERSAMPLE_FACTOR,
612 metadata,
613 id_to_index,
614 index_to_id,
615 deleted,
616 metadata_index,
617 storage: Some(storage),
618 storage_path: Some(path.to_path_buf()),
619 text_index,
620 text_search_config: None,
621 pending_quantization: None,
622 hnsw_m: hnsw_m.max(DEFAULT_HNSW_M),
623 hnsw_ef_construction: hnsw_ef_construction.max(DEFAULT_HNSW_EF_CONSTRUCTION),
624 hnsw_ef_search: hnsw_ef_search.max(DEFAULT_HNSW_EF_SEARCH),
625 distance_metric,
626 next_index,
627 })
628 }
629
630 pub fn open_with_dimensions(path: impl AsRef<Path>, dimensions: usize) -> Result<Self> {
634 let mut store = Self::open(path)?;
635 if store.dimensions == 0 {
636 store.dimensions = dimensions;
637 if let Some(ref mut storage) = store.storage {
638 storage.put_config("dimensions", dimensions as u64)?;
639 }
640 }
641 Ok(store)
642 }
643
644 pub fn open_with_options(path: impl AsRef<Path>, options: &VectorStoreOptions) -> Result<Self> {
648 let path = path.as_ref();
649 let omen_path = OmenFile::compute_omen_path(path);
650
651 if path.exists() || omen_path.exists() {
653 let mut store = Self::open(path)?;
654
655 if store.dimensions == 0 && options.dimensions > 0 {
657 store.dimensions = options.dimensions;
658 if let Some(ref mut storage) = store.storage {
659 storage.put_config("dimensions", options.dimensions as u64)?;
660 }
661 }
662
663 if let Some(ef) = options.ef_search {
665 store.set_ef_search(ef);
666 }
667
668 return Ok(store);
669 }
670
671 let mut storage = OmenFile::create(path, options.dimensions as u32)?;
673 let dimensions = options.dimensions;
674
675 let m = options.m.unwrap_or(16);
677 let ef_construction = options.ef_construction.unwrap_or(100);
678 let ef_search = options.ef_search.unwrap_or(100);
679
680 let distance_metric = options.metric.unwrap_or(Metric::L2);
682
683 let (hnsw_index, pending_quantization) = if options.quantization.is_some() {
685 (None, options.quantization.clone())
686 } else if dimensions > 0 {
687 if options.m.is_some() || options.ef_construction.is_some() {
688 (
689 Some(HNSWIndex::new_with_params(
690 10_000,
691 dimensions,
692 m,
693 ef_construction,
694 ef_search,
695 distance_metric.into(),
696 )?),
697 None,
698 )
699 } else {
700 (None, None)
701 }
702 } else {
703 (None, None)
704 };
705
706 if dimensions > 0 {
708 storage.put_config("dimensions", dimensions as u64)?;
709 }
710
711 let text_index = if let Some(ref config) = options.text_search_config {
713 let text_path = path.join("text_index");
714 Some(TextIndex::open_with_config(&text_path, config)?)
715 } else {
716 None
717 };
718
719 let rescore_enabled = options.rescore.unwrap_or(options.quantization.is_some());
721 let oversample_factor = options
722 .oversample
723 .unwrap_or_else(|| default_oversample_for_quantization(options.quantization.as_ref()));
724
725 Ok(Self {
726 vectors: Vec::new(),
727 hnsw_index,
728 dimensions,
729 rescore_enabled,
730 oversample_factor,
731 metadata: HashMap::new(),
732 id_to_index: FxHashMap::default(),
733 index_to_id: FxHashMap::default(),
734 deleted: HashMap::new(),
735 metadata_index: MetadataIndex::new(),
736 storage: Some(storage),
737 storage_path: Some(path.to_path_buf()),
738 text_index,
739 text_search_config: options.text_search_config.clone(),
740 pending_quantization,
741 hnsw_m: m,
742 hnsw_ef_construction: ef_construction,
743 hnsw_ef_search: ef_search,
744 distance_metric,
745 next_index: 0,
746 })
747 }
748
749 pub fn build_with_options(options: &VectorStoreOptions) -> Result<Self> {
751 let dimensions = options.dimensions;
752
753 let m = options.m.unwrap_or(16);
755 let ef_construction = options.ef_construction.unwrap_or(100);
756 let ef_search = options.ef_search.unwrap_or(100);
757
758 let distance_metric = options.metric.unwrap_or(Metric::L2);
760
761 let (hnsw_index, pending_quantization) = if options.quantization.is_some() {
763 (None, options.quantization.clone())
764 } else if dimensions > 0 {
765 if options.m.is_some() || options.ef_construction.is_some() {
766 (
767 Some(HNSWIndex::new_with_params(
768 10_000,
769 dimensions,
770 m,
771 ef_construction,
772 ef_search,
773 distance_metric.into(),
774 )?),
775 None,
776 )
777 } else {
778 (None, None)
779 }
780 } else {
781 (None, None)
782 };
783
784 let text_index = if let Some(ref config) = options.text_search_config {
786 Some(TextIndex::open_in_memory_with_config(config)?)
787 } else {
788 None
789 };
790
791 let rescore_enabled = options.rescore.unwrap_or(options.quantization.is_some());
793 let oversample_factor = options
794 .oversample
795 .unwrap_or_else(|| default_oversample_for_quantization(options.quantization.as_ref()));
796
797 Ok(Self {
798 vectors: Vec::new(),
799 hnsw_index,
800 dimensions,
801 rescore_enabled,
802 oversample_factor,
803 metadata: HashMap::new(),
804 id_to_index: FxHashMap::default(),
805 index_to_id: FxHashMap::default(),
806 deleted: HashMap::new(),
807 metadata_index: MetadataIndex::new(),
808 storage: None,
809 storage_path: None,
810 text_index,
811 text_search_config: options.text_search_config.clone(),
812 pending_quantization,
813 hnsw_m: m,
814 hnsw_ef_construction: ef_construction,
815 hnsw_ef_search: ef_search,
816 distance_metric,
817 next_index: 0,
818 })
819 }
820
821 fn resolve_dimensions(&self, vector_dim: usize) -> Result<usize> {
827 if self.dimensions == 0 {
828 Ok(vector_dim)
829 } else if vector_dim != self.dimensions {
830 anyhow::bail!(
831 "Vector dimension mismatch: store expects {}, got {}",
832 self.dimensions,
833 vector_dim
834 );
835 } else {
836 Ok(self.dimensions)
837 }
838 }
839
840 fn create_initial_hnsw(
842 &mut self,
843 dimensions: usize,
844 training_vectors: &[Vec<f32>],
845 ) -> Result<HNSWIndex> {
846 self.create_initial_hnsw_with_capacity(dimensions, training_vectors, 10_000)
847 }
848
849 fn create_initial_hnsw_with_capacity(
851 &mut self,
852 dimensions: usize,
853 training_vectors: &[Vec<f32>],
854 capacity: usize,
855 ) -> Result<HNSWIndex> {
856 if let Some(quant_mode) = self.pending_quantization.take() {
857 if let Some(ref mut storage) = self.storage {
858 storage.put_quantization_mode(quantization_mode_to_id(&quant_mode))?;
859 }
860 initialize_quantized_hnsw(
861 dimensions,
862 self.hnsw_m,
863 self.hnsw_ef_construction,
864 self.hnsw_ef_search,
865 self.distance_metric,
866 quant_mode,
867 training_vectors,
868 )
869 } else {
870 initialize_standard_hnsw(
871 dimensions,
872 self.hnsw_m,
873 self.hnsw_ef_construction,
874 self.hnsw_ef_search,
875 self.distance_metric,
876 capacity,
877 )
878 }
879 }
880
881 pub fn insert(&mut self, vector: Vector) -> Result<usize> {
887 let id = self.next_index;
888
889 if self.hnsw_index.is_none() {
890 let dimensions = self.resolve_dimensions(vector.dim())?;
891 self.hnsw_index =
892 Some(self.create_initial_hnsw(dimensions, std::slice::from_ref(&vector.data))?);
893 self.dimensions = dimensions;
894 } else if vector.dim() != self.dimensions {
895 anyhow::bail!(
896 "Vector dimension mismatch: store expects {}, got {}. All vectors in same store must have same dimension.",
897 self.dimensions,
898 vector.dim()
899 );
900 }
901
902 if let Some(ref mut index) = self.hnsw_index {
903 index.insert(&vector.data)?;
904 }
905
906 if let Some(ref mut storage) = self.storage {
907 storage.put_vector(id, &vector.data)?;
908 storage.increment_count()?;
909 if id == 0 {
910 storage.put_config("dimensions", self.dimensions as u64)?;
911 }
912 }
913
914 if !self.is_quantized() || self.storage.is_none() {
915 self.vectors.push(vector);
916 }
917
918 self.next_index += 1;
920
921 Ok(id)
922 }
923
924 pub fn insert_with_metadata(
929 &mut self,
930 id: String,
931 vector: Vector,
932 metadata: JsonValue,
933 ) -> Result<usize> {
934 if self.id_to_index.contains_key(&id) {
935 anyhow::bail!("Vector with ID '{id}' already exists. Use set() to update.");
936 }
937
938 let index = self.insert(vector)?;
939
940 self.metadata.insert(index, metadata.clone());
941 self.metadata_index.index_json(index as u32, &metadata);
942 self.id_to_index.insert(id.clone(), index);
943 self.index_to_id.insert(index, id.clone());
944
945 debug_assert_mapping_consistency(&self.id_to_index, &self.index_to_id);
947
948 if let Some(ref mut storage) = self.storage {
949 storage.put_metadata(index, &metadata)?;
950 storage.put_id_mapping(&id, index)?;
951 }
952
953 Ok(index)
954 }
955
956 pub fn set(&mut self, id: String, vector: Vector, metadata: JsonValue) -> Result<usize> {
960 if let Some(&index) = self.id_to_index.get(&id) {
961 self.update_by_index(index, Some(vector), Some(metadata))?;
962 Ok(index)
963 } else {
964 self.insert_with_metadata(id, vector, metadata)
965 }
966 }
967
968 pub fn set_batch(&mut self, batch: Vec<(String, Vector, JsonValue)>) -> Result<Vec<usize>> {
972 if batch.is_empty() {
973 return Ok(Vec::new());
974 }
975
976 let mut updates: Vec<(usize, Vector, JsonValue)> = Vec::new();
978 let mut inserts: Vec<(String, Vector, JsonValue)> = Vec::new();
979
980 for (id, vector, metadata) in batch {
981 if let Some(&index) = self.id_to_index.get(&id) {
982 updates.push((index, vector, metadata));
983 } else {
984 inserts.push((id, vector, metadata));
985 }
986 }
987
988 let mut result_indices = Vec::new();
989
990 for (index, vector, metadata) in updates {
992 self.update_by_index(index, Some(vector), Some(metadata))?;
993 result_indices.push(index);
994 }
995
996 if !inserts.is_empty() {
997 let vectors_data: Vec<Vec<f32>> =
998 inserts.iter().map(|(_, v, _)| v.data.clone()).collect();
999
1000 if self.hnsw_index.is_none() {
1001 let dimensions = self.resolve_dimensions(inserts[0].1.dim())?;
1002 self.hnsw_index = Some(self.create_initial_hnsw(dimensions, &vectors_data)?);
1003 self.dimensions = dimensions;
1004 }
1005
1006 for (i, (_, vector, _)) in inserts.iter().enumerate() {
1007 if vector.dim() != self.dimensions {
1008 anyhow::bail!(
1009 "Vector {} dimension mismatch: expected {}, got {}",
1010 i,
1011 self.dimensions,
1012 vector.dim()
1013 );
1014 }
1015 }
1016
1017 let base_index = self.next_index;
1018 let insert_count = inserts.len();
1019 if let Some(ref mut index) = self.hnsw_index {
1020 index.batch_insert(&vectors_data)?;
1021 }
1022
1023 if let Some(ref mut storage) = self.storage {
1025 if base_index == 0 {
1026 storage.put_config("dimensions", self.dimensions as u64)?;
1027 }
1028
1029 let batch_items: Vec<(usize, String, Vec<f32>, serde_json::Value)> = inserts
1030 .iter()
1031 .enumerate()
1032 .map(|(i, (id, vector, metadata))| {
1033 (
1034 base_index + i,
1035 id.clone(),
1036 vector.data.clone(),
1037 metadata.clone(),
1038 )
1039 })
1040 .collect();
1041
1042 storage.put_batch(batch_items)?;
1043 }
1044
1045 let skip_ram = self.is_quantized() && self.storage.is_some();
1048 for (i, (id, vector, metadata)) in inserts.into_iter().enumerate() {
1049 let idx = base_index + i;
1050 if !skip_ram {
1051 self.vectors.push(vector);
1052 }
1053 self.metadata.insert(idx, metadata.clone());
1054 self.metadata_index.index_json(idx as u32, &metadata);
1055 self.index_to_id.insert(idx, id.clone());
1056 self.id_to_index.insert(id, idx);
1057 result_indices.push(idx);
1058 }
1059
1060 self.next_index += insert_count;
1062
1063 debug_assert_mapping_consistency(&self.id_to_index, &self.index_to_id);
1065 }
1066
1067 Ok(result_indices)
1068 }
1069
1070 pub fn enable_text_search(&mut self) -> Result<()> {
1076 self.enable_text_search_with_config(None)
1077 }
1078
1079 pub fn enable_text_search_with_config(
1081 &mut self,
1082 config: Option<TextSearchConfig>,
1083 ) -> Result<()> {
1084 if self.text_index.is_some() {
1085 return Ok(());
1086 }
1087
1088 let config = config
1089 .or_else(|| self.text_search_config.clone())
1090 .unwrap_or_default();
1091
1092 self.text_index = if let Some(ref path) = self.storage_path {
1093 let text_path = path.join("text_index");
1094 Some(TextIndex::open_with_config(&text_path, &config)?)
1095 } else {
1096 Some(TextIndex::open_in_memory_with_config(&config)?)
1097 };
1098
1099 Ok(())
1100 }
1101
1102 #[must_use]
1104 pub fn has_text_search(&self) -> bool {
1105 self.text_index.is_some()
1106 }
1107
1108 pub fn set_with_text(
1110 &mut self,
1111 id: String,
1112 vector: Vector,
1113 text: &str,
1114 metadata: JsonValue,
1115 ) -> Result<usize> {
1116 let Some(ref mut text_index) = self.text_index else {
1117 anyhow::bail!("Text search not enabled. Call enable_text_search() first.");
1118 };
1119
1120 text_index.index_document(&id, text)?;
1121 self.set(id, vector, metadata)
1122 }
1123
1124 pub fn set_batch_with_text(
1126 &mut self,
1127 batch: Vec<(String, Vector, String, JsonValue)>,
1128 ) -> Result<Vec<usize>> {
1129 let Some(ref mut text_index) = self.text_index else {
1130 anyhow::bail!("Text search not enabled. Call enable_text_search() first.");
1131 };
1132
1133 for (id, _, text, _) in &batch {
1134 text_index.index_document(id, text)?;
1135 }
1136
1137 let vector_batch: Vec<(String, Vector, JsonValue)> = batch
1138 .into_iter()
1139 .map(|(id, vector, _, metadata)| (id, vector, metadata))
1140 .collect();
1141
1142 self.set_batch(vector_batch)
1143 }
1144
1145 pub fn text_search(&self, query: &str, k: usize) -> Result<Vec<(String, f32)>> {
1147 let Some(ref text_index) = self.text_index else {
1148 anyhow::bail!("Text search not enabled. Call enable_text_search() first.");
1149 };
1150
1151 text_index.search(query, k)
1152 }
1153
1154 pub fn hybrid_search(
1156 &mut self,
1157 query_vector: &Vector,
1158 query_text: &str,
1159 k: usize,
1160 alpha: Option<f32>,
1161 ) -> Result<Vec<(String, f32, JsonValue)>> {
1162 self.hybrid_search_with_rrf_k(query_vector, query_text, k, alpha, None)
1163 }
1164
1165 pub fn hybrid_search_with_rrf_k(
1167 &mut self,
1168 query_vector: &Vector,
1169 query_text: &str,
1170 k: usize,
1171 alpha: Option<f32>,
1172 rrf_k: Option<usize>,
1173 ) -> Result<Vec<(String, f32, JsonValue)>> {
1174 if query_vector.data.len() != self.dimensions {
1175 anyhow::bail!(
1176 "Query vector dimension {} does not match store dimension {}",
1177 query_vector.data.len(),
1178 self.dimensions
1179 );
1180 }
1181 if self.text_index.is_none() {
1182 anyhow::bail!("Text search not enabled. Call enable_text_search() first.");
1183 }
1184
1185 let fetch_k = k * 2;
1186
1187 let vector_results = self.knn_search(query_vector, fetch_k)?;
1188 let vector_results: Vec<(String, f32)> = vector_results
1189 .into_iter()
1190 .filter_map(|(idx, distance)| {
1191 self.index_to_id.get(&idx).map(|id| (id.clone(), distance))
1192 })
1193 .collect();
1194
1195 let text_results = self.text_search(query_text, fetch_k)?;
1196
1197 let fused = weighted_reciprocal_rank_fusion(
1198 vector_results,
1199 text_results,
1200 k,
1201 rrf_k.unwrap_or(DEFAULT_RRF_K),
1202 alpha.unwrap_or(0.5),
1203 );
1204
1205 Ok(self.attach_metadata(fused))
1206 }
1207
1208 pub fn hybrid_search_with_filter(
1210 &mut self,
1211 query_vector: &Vector,
1212 query_text: &str,
1213 k: usize,
1214 filter: &MetadataFilter,
1215 alpha: Option<f32>,
1216 ) -> Result<Vec<(String, f32, JsonValue)>> {
1217 self.hybrid_search_with_filter_rrf_k(query_vector, query_text, k, filter, alpha, None)
1218 }
1219
1220 pub fn hybrid_search_with_filter_rrf_k(
1222 &mut self,
1223 query_vector: &Vector,
1224 query_text: &str,
1225 k: usize,
1226 filter: &MetadataFilter,
1227 alpha: Option<f32>,
1228 rrf_k: Option<usize>,
1229 ) -> Result<Vec<(String, f32, JsonValue)>> {
1230 if query_vector.data.len() != self.dimensions {
1231 anyhow::bail!(
1232 "Query vector dimension {} does not match store dimension {}",
1233 query_vector.data.len(),
1234 self.dimensions
1235 );
1236 }
1237 if self.text_index.is_none() {
1238 anyhow::bail!("Text search not enabled. Call enable_text_search() first.");
1239 }
1240
1241 let fetch_k = k * 4;
1242
1243 let vector_results = self.knn_search_with_filter(query_vector, fetch_k, filter)?;
1244 let vector_results: Vec<(String, f32)> = vector_results
1245 .into_iter()
1246 .filter_map(|(idx, distance, _)| {
1247 self.index_to_id.get(&idx).map(|id| (id.clone(), distance))
1248 })
1249 .collect();
1250
1251 let text_results = self.text_search(query_text, fetch_k)?;
1252 let text_results: Vec<(String, f32)> = text_results
1253 .into_iter()
1254 .filter(|(id, _)| {
1255 self.id_to_index
1256 .get(id)
1257 .and_then(|&idx| self.metadata.get(&idx))
1258 .is_some_and(|meta| filter.matches(meta))
1259 })
1260 .collect();
1261
1262 let fused = weighted_reciprocal_rank_fusion(
1263 vector_results,
1264 text_results,
1265 k,
1266 rrf_k.unwrap_or(DEFAULT_RRF_K),
1267 alpha.unwrap_or(0.5),
1268 );
1269
1270 Ok(self.attach_metadata(fused))
1271 }
1272
1273 fn attach_metadata(&self, results: Vec<(String, f32)>) -> Vec<(String, f32, JsonValue)> {
1275 results
1276 .into_iter()
1277 .map(|(id, score)| {
1278 let metadata = self
1279 .id_to_index
1280 .get(&id)
1281 .and_then(|&idx| self.metadata.get(&idx))
1282 .cloned()
1283 .unwrap_or_else(default_metadata);
1284 (id, score, metadata)
1285 })
1286 .collect()
1287 }
1288
1289 pub fn hybrid_search_with_subscores(
1294 &mut self,
1295 query_vector: &Vector,
1296 query_text: &str,
1297 k: usize,
1298 alpha: Option<f32>,
1299 rrf_k: Option<usize>,
1300 ) -> Result<Vec<(HybridResult, JsonValue)>> {
1301 if query_vector.data.len() != self.dimensions {
1302 anyhow::bail!(
1303 "Query vector dimension {} does not match store dimension {}",
1304 query_vector.data.len(),
1305 self.dimensions
1306 );
1307 }
1308 if self.text_index.is_none() {
1309 anyhow::bail!("Text search not enabled. Call enable_text_search() first.");
1310 }
1311
1312 let fetch_k = k * 2;
1313
1314 let vector_results = self.knn_search(query_vector, fetch_k)?;
1315 let vector_results: Vec<(String, f32)> = vector_results
1316 .into_iter()
1317 .filter_map(|(idx, distance)| {
1318 self.index_to_id.get(&idx).map(|id| (id.clone(), distance))
1319 })
1320 .collect();
1321
1322 let text_results = self.text_search(query_text, fetch_k)?;
1323
1324 let fused = weighted_reciprocal_rank_fusion_with_subscores(
1325 vector_results,
1326 text_results,
1327 k,
1328 rrf_k.unwrap_or(DEFAULT_RRF_K),
1329 alpha.unwrap_or(0.5),
1330 );
1331
1332 Ok(self.attach_metadata_to_hybrid_results(fused))
1333 }
1334
1335 pub fn hybrid_search_with_filter_subscores(
1337 &mut self,
1338 query_vector: &Vector,
1339 query_text: &str,
1340 k: usize,
1341 filter: &MetadataFilter,
1342 alpha: Option<f32>,
1343 rrf_k: Option<usize>,
1344 ) -> Result<Vec<(HybridResult, JsonValue)>> {
1345 if query_vector.data.len() != self.dimensions {
1346 anyhow::bail!(
1347 "Query vector dimension {} does not match store dimension {}",
1348 query_vector.data.len(),
1349 self.dimensions
1350 );
1351 }
1352 if self.text_index.is_none() {
1353 anyhow::bail!("Text search not enabled. Call enable_text_search() first.");
1354 }
1355
1356 let fetch_k = k * 4;
1357
1358 let vector_results = self.knn_search_with_filter(query_vector, fetch_k, filter)?;
1359 let vector_results: Vec<(String, f32)> = vector_results
1360 .into_iter()
1361 .filter_map(|(idx, distance, _)| {
1362 self.index_to_id.get(&idx).map(|id| (id.clone(), distance))
1363 })
1364 .collect();
1365
1366 let text_results = self.text_search(query_text, fetch_k)?;
1367 let text_results: Vec<(String, f32)> = text_results
1368 .into_iter()
1369 .filter(|(id, _)| {
1370 self.id_to_index
1371 .get(id)
1372 .and_then(|&idx| self.metadata.get(&idx))
1373 .is_some_and(|meta| filter.matches(meta))
1374 })
1375 .collect();
1376
1377 let fused = weighted_reciprocal_rank_fusion_with_subscores(
1378 vector_results,
1379 text_results,
1380 k,
1381 rrf_k.unwrap_or(DEFAULT_RRF_K),
1382 alpha.unwrap_or(0.5),
1383 );
1384
1385 Ok(self.attach_metadata_to_hybrid_results(fused))
1386 }
1387
1388 fn attach_metadata_to_hybrid_results(
1390 &self,
1391 results: Vec<HybridResult>,
1392 ) -> Vec<(HybridResult, JsonValue)> {
1393 results
1394 .into_iter()
1395 .map(|result| {
1396 let metadata = self
1397 .id_to_index
1398 .get(&result.id)
1399 .and_then(|&idx| self.metadata.get(&idx))
1400 .cloned()
1401 .unwrap_or_else(default_metadata);
1402 (result, metadata)
1403 })
1404 .collect()
1405 }
1406
1407 fn update_by_index(
1413 &mut self,
1414 index: usize,
1415 vector: Option<Vector>,
1416 metadata: Option<JsonValue>,
1417 ) -> Result<()> {
1418 if index >= self.next_index {
1420 anyhow::bail!("Vector index {index} does not exist");
1421 }
1422 if self.deleted.contains_key(&index) {
1423 anyhow::bail!("Vector index {index} has been deleted");
1424 }
1425
1426 if let Some(new_vector) = vector {
1427 if new_vector.dim() != self.dimensions {
1428 anyhow::bail!(
1429 "Vector dimension mismatch: expected {}, got {}",
1430 self.dimensions,
1431 new_vector.dim()
1432 );
1433 }
1434
1435 if let Some(v) = self.vectors.get_mut(index) {
1437 *v = new_vector.clone();
1438 }
1439
1440 if let Some(ref mut storage) = self.storage {
1441 storage.put_vector(index, &new_vector.data)?;
1442 }
1443 }
1444
1445 if let Some(ref new_metadata) = metadata {
1446 self.metadata_index.remove(index as u32);
1448 self.metadata_index.index_json(index as u32, new_metadata);
1449 self.metadata.insert(index, new_metadata.clone());
1450
1451 if let Some(ref mut storage) = self.storage {
1452 storage.put_metadata(index, new_metadata)?;
1453 }
1454 }
1455
1456 Ok(())
1457 }
1458
1459 pub fn update(
1461 &mut self,
1462 id: &str,
1463 vector: Option<Vector>,
1464 metadata: Option<JsonValue>,
1465 ) -> Result<()> {
1466 let index = self
1467 .id_to_index
1468 .get(id)
1469 .copied()
1470 .ok_or_else(|| anyhow::anyhow!("Vector with ID '{id}' not found"))?;
1471
1472 self.update_by_index(index, vector, metadata)
1473 }
1474
1475 pub fn delete(&mut self, id: &str) -> Result<()> {
1483 let index = self
1484 .id_to_index
1485 .get(id)
1486 .copied()
1487 .ok_or_else(|| anyhow::anyhow!("Vector with ID '{id}' not found"))?;
1488
1489 self.deleted.insert(index, true);
1490 self.metadata_index.remove(index as u32);
1491
1492 if let Some(ref mut hnsw) = self.hnsw_index {
1495 if let Err(e) = hnsw.mark_deleted(index as u32) {
1496 tracing::warn!(
1497 id = id,
1498 index = index,
1499 error = ?e,
1500 "Failed to repair HNSW graph after deletion"
1501 );
1502 }
1503 }
1504
1505 if let Some(ref mut storage) = self.storage {
1507 storage.delete(id)?;
1508 }
1509
1510 if let Some(ref mut text_index) = self.text_index {
1511 text_index.delete_document(id)?;
1512 }
1513
1514 self.id_to_index.remove(id);
1515 self.index_to_id.remove(&index);
1516
1517 debug_assert_mapping_consistency(&self.id_to_index, &self.index_to_id);
1519
1520 Ok(())
1521 }
1522
1523 pub fn delete_batch(&mut self, ids: &[String]) -> Result<usize> {
1527 let mut node_ids: Vec<u32> = Vec::with_capacity(ids.len());
1529 let mut valid_ids: Vec<String> = Vec::with_capacity(ids.len());
1530
1531 for id in ids {
1532 if let Some(&index) = self.id_to_index.get(id) {
1533 self.deleted.insert(index, true);
1534 self.metadata_index.remove(index as u32);
1535 node_ids.push(index as u32);
1536 valid_ids.push(id.clone());
1537 }
1538 }
1539
1540 if !node_ids.is_empty() {
1542 if let Some(ref mut hnsw) = self.hnsw_index {
1543 if let Err(e) = hnsw.mark_deleted_batch(&node_ids) {
1544 tracing::warn!(
1545 count = node_ids.len(),
1546 error = ?e,
1547 "Failed to batch repair HNSW graph after deletion"
1548 );
1549 }
1550 }
1551 }
1552
1553 for (id, &node_id) in valid_ids.iter().zip(node_ids.iter()) {
1555 if let Some(ref mut storage) = self.storage {
1556 if let Err(e) = storage.delete(id) {
1557 tracing::warn!(id = %id, error = ?e, "Failed to persist deletion to storage");
1558 }
1559 }
1560 if let Some(ref mut text_index) = self.text_index {
1561 if let Err(e) = text_index.delete_document(id) {
1562 tracing::warn!(id = %id, error = ?e, "Failed to delete from text index");
1563 }
1564 }
1565 self.id_to_index.remove(id);
1566 self.index_to_id.remove(&(node_id as usize));
1567 }
1568
1569 debug_assert_mapping_consistency(&self.id_to_index, &self.index_to_id);
1571
1572 Ok(valid_ids.len())
1573 }
1574
1575 pub fn delete_by_filter(&mut self, filter: &MetadataFilter) -> Result<usize> {
1586 let ids_to_delete: Vec<String> = self
1588 .id_to_index
1589 .iter()
1590 .filter_map(|(id, &idx)| {
1591 if self.deleted.contains_key(&idx) {
1592 return None;
1593 }
1594 let metadata = self.metadata.get(&idx)?;
1595 if filter.matches(metadata) {
1596 Some(id.clone())
1597 } else {
1598 None
1599 }
1600 })
1601 .collect();
1602
1603 if ids_to_delete.is_empty() {
1604 return Ok(0);
1605 }
1606
1607 self.delete_batch(&ids_to_delete)
1608 }
1609
1610 #[must_use]
1621 pub fn count_by_filter(&self, filter: &MetadataFilter) -> usize {
1622 self.id_to_index
1623 .iter()
1624 .filter(|(_, &idx)| {
1625 if self.deleted.contains_key(&idx) {
1626 return false;
1627 }
1628 self.metadata
1629 .get(&idx)
1630 .is_some_and(|metadata| filter.matches(metadata))
1631 })
1632 .count()
1633 }
1634
1635 #[must_use]
1639 pub fn get(&self, id: &str) -> Option<(Vector, JsonValue)> {
1640 let &index = self.id_to_index.get(id)?;
1641 if self.deleted.contains_key(&index) {
1642 return None;
1643 }
1644
1645 if let Some(vec) = self.vectors.get(index) {
1647 return self
1648 .metadata
1649 .get(&index)
1650 .map(|meta| (vec.clone(), meta.clone()));
1651 }
1652
1653 if let Some(ref storage) = self.storage {
1655 if let Ok(Some(vec_data)) = storage.get_vector(index) {
1656 return self
1657 .metadata
1658 .get(&index)
1659 .map(|meta| (Vector::new(vec_data), meta.clone()));
1660 }
1661 }
1662
1663 None
1664 }
1665
1666 #[must_use]
1671 pub fn get_batch(&self, ids: &[impl AsRef<str>]) -> Vec<Option<(Vector, JsonValue)>> {
1672 ids.iter().map(|id| self.get(id.as_ref())).collect()
1673 }
1674
1675 #[must_use]
1677 pub fn get_metadata_by_id(&self, id: &str) -> Option<&JsonValue> {
1678 self.id_to_index.get(id).and_then(|&index| {
1679 if self.deleted.contains_key(&index) {
1680 return None;
1681 }
1682 self.metadata.get(&index)
1683 })
1684 }
1685
1686 pub fn batch_insert(&mut self, vectors: Vec<Vector>) -> Result<Vec<usize>> {
1692 const CHUNK_SIZE: usize = 10_000;
1693
1694 if vectors.is_empty() {
1695 return Ok(Vec::new());
1696 }
1697
1698 for (i, vector) in vectors.iter().enumerate() {
1699 if vector.dim() != self.dimensions {
1700 anyhow::bail!(
1701 "Vector {} dimension mismatch: expected {}, got {}",
1702 i,
1703 self.dimensions,
1704 vector.dim()
1705 );
1706 }
1707 }
1708
1709 if self.hnsw_index.is_none() {
1710 let training: Vec<Vec<f32>> = vectors.iter().map(|v| v.data.clone()).collect();
1711 let capacity = vectors.len().max(1_000_000);
1712 self.hnsw_index = Some(self.create_initial_hnsw_with_capacity(
1713 self.dimensions,
1714 &training,
1715 capacity,
1716 )?);
1717 }
1718
1719 let mut all_ids = Vec::with_capacity(vectors.len());
1720 for chunk in vectors.chunks(CHUNK_SIZE) {
1721 let vector_data: Vec<Vec<f32>> = chunk.iter().map(|v| v.data.clone()).collect();
1722 if let Some(ref mut index) = self.hnsw_index {
1723 all_ids.extend(index.batch_insert(&vector_data)?);
1724 }
1725 }
1726
1727 self.vectors.extend(vectors);
1728 Ok(all_ids)
1729 }
1730
1731 pub fn rebuild_index(&mut self) -> Result<()> {
1733 if self.vectors.is_empty() {
1734 return Ok(());
1735 }
1736
1737 let mut index = HNSWIndex::new_with_params(
1738 self.vectors.len().max(1_000_000),
1739 self.dimensions,
1740 self.hnsw_m,
1741 self.hnsw_ef_construction,
1742 self.hnsw_ef_search,
1743 self.distance_metric.into(),
1744 )?;
1745
1746 for vector in &self.vectors {
1747 index.insert(&vector.data)?;
1748 }
1749
1750 self.hnsw_index = Some(index);
1751 Ok(())
1752 }
1753
1754 pub fn merge_from(&mut self, other: &VectorStore) -> Result<usize> {
1756 if other.dimensions != self.dimensions {
1757 anyhow::bail!(
1758 "Dimension mismatch: self={}, other={}",
1759 self.dimensions,
1760 other.dimensions
1761 );
1762 }
1763
1764 if other.vectors.is_empty() {
1765 return Ok(0);
1766 }
1767
1768 if self.hnsw_index.is_none() {
1769 let capacity = (self.vectors.len() + other.vectors.len()).max(1_000_000);
1770 self.hnsw_index = Some(HNSWIndex::new_with_params(
1771 capacity,
1772 self.dimensions,
1773 self.hnsw_m,
1774 self.hnsw_ef_construction,
1775 self.hnsw_ef_search,
1776 self.distance_metric.into(),
1777 )?);
1778 }
1779
1780 let mut merged_count = 0;
1781 let base_index = self.vectors.len();
1782
1783 for (other_idx, vector) in other.vectors.iter().enumerate() {
1784 let has_conflict = other
1785 .id_to_index
1786 .iter()
1787 .find(|(_, &idx)| idx == other_idx)
1788 .is_some_and(|(string_id, _)| self.id_to_index.contains_key(string_id));
1789
1790 if has_conflict {
1791 continue;
1792 }
1793
1794 self.vectors.push(vector.clone());
1795
1796 if let Some(meta) = other.metadata.get(&other_idx) {
1797 self.metadata
1798 .insert(base_index + merged_count, meta.clone());
1799 }
1800
1801 if let Some((string_id, _)) =
1802 other.id_to_index.iter().find(|(_, &idx)| idx == other_idx)
1803 {
1804 self.id_to_index
1805 .insert(string_id.clone(), base_index + merged_count);
1806 }
1807
1808 merged_count += 1;
1809 }
1810
1811 self.rebuild_index()?;
1814
1815 Ok(merged_count)
1816 }
1817
1818 #[inline]
1820 #[must_use]
1821 pub fn needs_index_rebuild(&self) -> bool {
1822 self.hnsw_index.is_none() && self.vectors.len() > 100
1823 }
1824
1825 pub fn ensure_index_ready(&mut self) -> Result<()> {
1827 if self.needs_index_rebuild() {
1828 self.rebuild_index()?;
1829 }
1830 Ok(())
1831 }
1832
1833 pub fn knn_search(&mut self, query: &Vector, k: usize) -> Result<Vec<(usize, f32)>> {
1839 self.knn_search_with_ef(query, k, None)
1840 }
1841
1842 pub fn knn_search_with_ef(
1844 &mut self,
1845 query: &Vector,
1846 k: usize,
1847 ef: Option<usize>,
1848 ) -> Result<Vec<(usize, f32)>> {
1849 self.ensure_index_ready()?;
1850 self.knn_search_readonly(query, k, ef)
1851 }
1852
1853 #[inline]
1855 pub fn knn_search_readonly(
1856 &self,
1857 query: &Vector,
1858 k: usize,
1859 ef: Option<usize>,
1860 ) -> Result<Vec<(usize, f32)>> {
1861 let effective_ef = compute_effective_ef(ef, self.hnsw_ef_search, k);
1864 self.knn_search_ef(query, k, effective_ef)
1865 }
1866
1867 #[inline]
1869 pub fn knn_search_ef(&self, query: &Vector, k: usize, ef: usize) -> Result<Vec<(usize, f32)>> {
1870 if query.dim() != self.dimensions {
1871 anyhow::bail!(
1872 "Query dimension mismatch: expected {}, got {}",
1873 self.dimensions,
1874 query.dim()
1875 );
1876 }
1877
1878 let has_data =
1879 !self.vectors.is_empty() || self.hnsw_index.as_ref().is_some_and(|idx| !idx.is_empty());
1880
1881 if !has_data {
1882 return Ok(Vec::new());
1883 }
1884
1885 if let Some(ref index) = self.hnsw_index {
1886 let results = if index.is_asymmetric() {
1887 let can_rescore = self.storage.is_some() || !self.vectors.is_empty();
1889 if self.rescore_enabled && can_rescore {
1890 self.knn_search_with_rescore(query, k, ef)?
1891 } else {
1892 index.search_asymmetric_ef(&query.data, k, ef)?
1893 }
1894 } else {
1895 index.search_ef(&query.data, k, ef)?
1896 };
1897
1898 if results.is_empty() && self.has_live_vectors() {
1901 return self.knn_search_brute_force(query, k);
1902 }
1903 return Ok(results);
1904 }
1905
1906 self.knn_search_brute_force(query, k)
1907 }
1908
1909 fn knn_search_with_rescore(
1911 &self,
1912 query: &Vector,
1913 k: usize,
1914 ef: usize,
1915 ) -> Result<Vec<(usize, f32)>> {
1916 let index = self
1917 .hnsw_index
1918 .as_ref()
1919 .ok_or_else(|| anyhow::anyhow!("HNSW index required for rescore"))?;
1920
1921 let oversample_k = ((k as f32) * self.oversample_factor).ceil() as usize;
1922 let candidates = index.search_asymmetric_ef(&query.data, oversample_k, ef)?;
1923
1924 if candidates.is_empty() {
1925 return Ok(Vec::new());
1926 }
1927
1928 let mut rescored: Vec<(usize, f32)> = candidates
1931 .iter()
1932 .filter_map(|&(id, _quantized_dist)| {
1933 if let Some(ref storage) = self.storage {
1936 storage
1937 .get_vector(id)
1938 .ok()
1939 .flatten()
1940 .map(|data| (id, l2_distance(&query.data, &data)))
1941 } else {
1942 self.vectors
1943 .get(id)
1944 .map(|v| (id, l2_distance(&query.data, &v.data)))
1945 }
1946 })
1947 .collect();
1948
1949 rescored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
1950 rescored.truncate(k);
1951
1952 Ok(rescored)
1953 }
1954
1955 pub fn knn_search_with_filter(
1957 &mut self,
1958 query: &Vector,
1959 k: usize,
1960 filter: &MetadataFilter,
1961 ) -> Result<Vec<(usize, f32, JsonValue)>> {
1962 self.ensure_index_ready()?;
1963 self.knn_search_with_filter_ef_readonly(query, k, filter, None)
1964 }
1965
1966 pub fn knn_search_with_filter_ef(
1968 &mut self,
1969 query: &Vector,
1970 k: usize,
1971 filter: &MetadataFilter,
1972 ef: Option<usize>,
1973 ) -> Result<Vec<(usize, f32, JsonValue)>> {
1974 self.ensure_index_ready()?;
1975 self.knn_search_with_filter_ef_readonly(query, k, filter, ef)
1976 }
1977
1978 pub fn knn_search_with_filter_ef_readonly(
1983 &self,
1984 query: &Vector,
1985 k: usize,
1986 filter: &MetadataFilter,
1987 ef: Option<usize>,
1988 ) -> Result<Vec<(usize, f32, JsonValue)>> {
1989 let effective_ef = compute_effective_ef(ef, self.hnsw_ef_search, k);
1992
1993 let filter_bitmap = filter.evaluate_bitmap(&self.metadata_index);
1995
1996 if let Some(ref hnsw) = self.hnsw_index {
1997 let metadata_map = &self.metadata;
1998 let deleted_map = &self.deleted;
1999
2000 let search_results = if let Some(ref bitmap) = filter_bitmap {
2001 let filter_fn = |node_id: u32| -> bool {
2003 let index = node_id as usize;
2004 !deleted_map.contains_key(&index) && bitmap.contains(node_id)
2005 };
2006 hnsw.search_with_filter_ef(&query.data, k, Some(effective_ef), filter_fn)?
2007 } else {
2008 let filter_fn = |node_id: u32| -> bool {
2010 let index = node_id as usize;
2011 if deleted_map.contains_key(&index) {
2012 return false;
2013 }
2014 let metadata = metadata_map
2015 .get(&index)
2016 .cloned()
2017 .unwrap_or_else(default_metadata);
2018 filter.matches(&metadata)
2019 };
2020 hnsw.search_with_filter_ef(&query.data, k, Some(effective_ef), filter_fn)?
2021 };
2022
2023 let filtered_results: Vec<(usize, f32, JsonValue)> = search_results
2024 .into_iter()
2025 .map(|(index, distance)| {
2026 let metadata = self
2027 .metadata
2028 .get(&index)
2029 .cloned()
2030 .unwrap_or_else(default_metadata);
2031 (index, distance, metadata)
2032 })
2033 .collect();
2034
2035 return Ok(filtered_results);
2036 }
2037
2038 let mut all_results: Vec<(usize, f32, JsonValue)> = self
2040 .vectors
2041 .iter()
2042 .enumerate()
2043 .filter_map(|(index, vec)| {
2044 if self.deleted.contains_key(&index) {
2045 return None;
2046 }
2047
2048 let passes_filter = if let Some(ref bitmap) = filter_bitmap {
2050 bitmap.contains(index as u32)
2051 } else {
2052 let metadata = self
2053 .metadata
2054 .get(&index)
2055 .cloned()
2056 .unwrap_or_else(default_metadata);
2057 filter.matches(&metadata)
2058 };
2059
2060 if !passes_filter {
2061 return None;
2062 }
2063
2064 let metadata = self
2065 .metadata
2066 .get(&index)
2067 .cloned()
2068 .unwrap_or_else(default_metadata);
2069 let distance = query.l2_distance(vec).unwrap_or(f32::MAX);
2070 Some((index, distance, metadata))
2071 })
2072 .collect();
2073
2074 all_results.sort_by(|a, b| a.1.total_cmp(&b.1));
2075 all_results.truncate(k);
2076
2077 Ok(all_results)
2078 }
2079
2080 pub fn search(
2082 &mut self,
2083 query: &Vector,
2084 k: usize,
2085 filter: Option<&MetadataFilter>,
2086 ) -> Result<Vec<(usize, f32, JsonValue)>> {
2087 self.search_with_options(query, k, filter, None, None)
2088 }
2089
2090 pub fn search_with_ef(
2092 &mut self,
2093 query: &Vector,
2094 k: usize,
2095 filter: Option<&MetadataFilter>,
2096 ef: Option<usize>,
2097 ) -> Result<Vec<(usize, f32, JsonValue)>> {
2098 self.search_with_options(query, k, filter, ef, None)
2099 }
2100
2101 pub fn search_with_options(
2103 &mut self,
2104 query: &Vector,
2105 k: usize,
2106 filter: Option<&MetadataFilter>,
2107 ef: Option<usize>,
2108 max_distance: Option<f32>,
2109 ) -> Result<Vec<(usize, f32, JsonValue)>> {
2110 self.ensure_index_ready()?;
2111 self.search_with_options_readonly(query, k, filter, ef, max_distance)
2112 }
2113
2114 pub fn search_with_ef_readonly(
2116 &self,
2117 query: &Vector,
2118 k: usize,
2119 filter: Option<&MetadataFilter>,
2120 ef: Option<usize>,
2121 ) -> Result<Vec<(usize, f32, JsonValue)>> {
2122 self.search_with_options_readonly(query, k, filter, ef, None)
2123 }
2124
2125 pub fn search_with_options_readonly(
2127 &self,
2128 query: &Vector,
2129 k: usize,
2130 filter: Option<&MetadataFilter>,
2131 ef: Option<usize>,
2132 max_distance: Option<f32>,
2133 ) -> Result<Vec<(usize, f32, JsonValue)>> {
2134 let mut results = if let Some(f) = filter {
2135 self.knn_search_with_filter_ef_readonly(query, k, f, ef)?
2136 } else {
2137 let results = self.knn_search_readonly(query, k, ef)?;
2138 let filtered: Vec<(usize, f32, JsonValue)> = results
2139 .into_iter()
2140 .filter_map(|(index, distance)| {
2141 if self.deleted.contains_key(&index) {
2142 return None;
2143 }
2144 let metadata = self
2145 .metadata
2146 .get(&index)
2147 .cloned()
2148 .unwrap_or_else(default_metadata);
2149 Some((index, distance, metadata))
2150 })
2151 .collect();
2152
2153 if filtered.is_empty() && self.has_live_vectors() {
2155 self.knn_search_brute_force_with_metadata(query, k)?
2156 } else {
2157 filtered
2158 }
2159 };
2160
2161 if let Some(max_dist) = max_distance {
2162 results.retain(|(_, distance, _)| *distance <= max_dist);
2163 }
2164
2165 Ok(results)
2166 }
2167
2168 fn has_live_vectors(&self) -> bool {
2170 let total = self
2171 .vectors
2172 .len()
2173 .max(self.hnsw_index.as_ref().map_or(0, HNSWIndex::len));
2174 total > self.deleted.len()
2175 }
2176
2177 fn is_quantized(&self) -> bool {
2179 self.pending_quantization.is_some()
2180 || self
2181 .hnsw_index
2182 .as_ref()
2183 .is_some_and(|idx| idx.is_asymmetric() || idx.is_sq8())
2184 }
2185
2186 fn knn_search_brute_force_with_metadata(
2188 &self,
2189 query: &Vector,
2190 k: usize,
2191 ) -> Result<Vec<(usize, f32, JsonValue)>> {
2192 let results = self.knn_search_brute_force(query, k)?;
2193 Ok(results
2194 .into_iter()
2195 .filter_map(|(index, distance)| {
2196 if self.deleted.contains_key(&index) {
2197 return None;
2198 }
2199 let metadata = self
2200 .metadata
2201 .get(&index)
2202 .cloned()
2203 .unwrap_or_else(default_metadata);
2204 Some((index, distance, metadata))
2205 })
2206 .collect())
2207 }
2208
2209 #[must_use]
2211 pub fn search_batch(
2212 &self,
2213 queries: &[Vector],
2214 k: usize,
2215 ef: Option<usize>,
2216 ) -> Vec<Result<Vec<(usize, f32)>>> {
2217 let effective_ef = compute_effective_ef(ef, self.hnsw_ef_search, k);
2220 queries
2221 .par_iter()
2222 .map(|q| self.knn_search_ef(q, k, effective_ef))
2223 .collect()
2224 }
2225
2226 #[must_use]
2228 pub fn search_batch_with_metadata(
2229 &self,
2230 queries: &[Vector],
2231 k: usize,
2232 ef: Option<usize>,
2233 ) -> Vec<Result<Vec<(usize, f32, JsonValue)>>> {
2234 queries
2235 .par_iter()
2236 .map(|q| self.search_with_ef_readonly(q, k, None, ef))
2237 .collect()
2238 }
2239
2240 pub fn knn_search_brute_force(&self, query: &Vector, k: usize) -> Result<Vec<(usize, f32)>> {
2242 if query.dim() != self.dimensions {
2243 anyhow::bail!(
2244 "Query dimension mismatch: expected {}, got {}",
2245 self.dimensions,
2246 query.dim()
2247 );
2248 }
2249
2250 let total_count = if !self.vectors.is_empty() {
2252 self.vectors.len()
2253 } else if let Some(ref idx) = self.hnsw_index {
2254 idx.len()
2255 } else {
2256 return Ok(Vec::new());
2257 };
2258
2259 if total_count == 0 {
2260 return Ok(Vec::new());
2261 }
2262
2263 let mut distances: Vec<(usize, f32)> = (0..total_count)
2264 .filter_map(|id| {
2265 let data = if let Some(vec) = self.vectors.get(id) {
2267 Some(vec.data.clone())
2268 } else if let Some(ref storage) = self.storage {
2269 storage.get_vector(id).ok().flatten()
2270 } else {
2271 None
2272 };
2273
2274 data.map(|vec_data| {
2275 let dist = l2_distance(&query.data, &vec_data);
2276 (id, dist)
2277 })
2278 })
2279 .collect();
2280
2281 distances.sort_by(|a, b| a.1.total_cmp(&b.1));
2282 Ok(distances.into_iter().take(k).collect())
2283 }
2284
2285 pub fn optimize(&mut self) -> Result<usize> {
2300 let Some(ref mut index) = self.hnsw_index else {
2301 return Ok(0);
2302 };
2303
2304 let old_to_new = index
2306 .optimize_cache_locality()
2307 .map_err(|e| anyhow::anyhow!("Optimization failed: {e}"))?;
2308
2309 if old_to_new.is_empty() {
2310 return Ok(0);
2311 }
2312
2313 let num_reordered = old_to_new.len();
2314
2315 if !self.vectors.is_empty() {
2317 let old_vectors = std::mem::take(&mut self.vectors);
2318 let mut new_vectors = Vec::with_capacity(old_vectors.len());
2319 new_vectors.resize_with(old_vectors.len(), || Vector::new(Vec::new()));
2320
2321 for (old_idx, &new_idx) in old_to_new.iter().enumerate() {
2322 let new_idx = new_idx as usize;
2323 if old_idx < old_vectors.len() && new_idx < new_vectors.len() {
2324 new_vectors[new_idx] = old_vectors[old_idx].clone();
2325 }
2326 }
2327 self.vectors = new_vectors;
2328 }
2329
2330 let mut new_id_to_index: FxHashMap<String, usize> =
2332 FxHashMap::with_capacity_and_hasher(self.id_to_index.len(), rustc_hash::FxBuildHasher);
2333 let mut new_index_to_id: FxHashMap<usize, String> =
2334 FxHashMap::with_capacity_and_hasher(self.index_to_id.len(), rustc_hash::FxBuildHasher);
2335
2336 for (string_id, &old_idx) in &self.id_to_index {
2337 if old_idx < old_to_new.len() {
2338 let new_idx = old_to_new[old_idx] as usize;
2339 new_id_to_index.insert(string_id.clone(), new_idx);
2340 new_index_to_id.insert(new_idx, string_id.clone());
2341 }
2342 }
2343
2344 self.id_to_index = new_id_to_index;
2345 self.index_to_id = new_index_to_id;
2346
2347 if !self.deleted.is_empty() {
2349 let mut new_deleted = HashMap::with_capacity(self.deleted.len());
2350 for (&old_idx, &is_deleted) in &self.deleted {
2351 if old_idx < old_to_new.len() {
2352 let new_idx = old_to_new[old_idx] as usize;
2353 new_deleted.insert(new_idx, is_deleted);
2354 }
2355 }
2356 self.deleted = new_deleted;
2357 }
2358
2359 Ok(num_reordered)
2362 }
2363
2364 #[must_use]
2370 #[allow(dead_code)] pub(crate) fn get_by_internal_index(&self, idx: usize) -> Option<&Vector> {
2372 self.vectors.get(idx)
2373 }
2374
2375 #[must_use]
2377 #[allow(dead_code)] pub(crate) fn get_by_internal_index_owned(&self, idx: usize) -> Option<Vector> {
2379 if let Some(v) = self.vectors.get(idx) {
2380 return Some(v.clone());
2381 }
2382
2383 if let Some(ref storage) = self.storage {
2384 if let Ok(Some(data)) = storage.get_vector(idx) {
2385 return Some(Vector::new(data));
2386 }
2387 }
2388
2389 None
2390 }
2391
2392 #[must_use]
2394 pub fn len(&self) -> usize {
2395 if let Some(ref index) = self.hnsw_index {
2396 let hnsw_len = index.len();
2397 if hnsw_len > 0 {
2398 return hnsw_len.saturating_sub(self.deleted.len());
2399 }
2400 }
2401 self.vectors.len().saturating_sub(self.deleted.len())
2402 }
2403
2404 #[must_use]
2408 pub fn count(&self) -> usize {
2409 self.len()
2410 }
2411
2412 #[must_use]
2414 pub fn is_empty(&self) -> bool {
2415 self.len() == 0
2416 }
2417
2418 #[must_use]
2423 pub fn ids(&self) -> Vec<String> {
2424 self.id_to_index
2425 .iter()
2426 .filter_map(|(id, &idx)| {
2427 if self.deleted.contains_key(&idx) {
2428 None
2429 } else {
2430 Some(id.clone())
2431 }
2432 })
2433 .collect()
2434 }
2435
2436 #[must_use]
2440 pub fn items(&self) -> Vec<(String, Vec<f32>, JsonValue)> {
2441 self.id_to_index
2442 .iter()
2443 .filter_map(|(id, &idx)| {
2444 if self.deleted.contains_key(&idx) {
2445 return None;
2446 }
2447
2448 let vec_data = if let Some(vec) = self.vectors.get(idx) {
2450 vec.data.clone()
2451 } else if let Some(ref storage) = self.storage {
2452 storage.get_vector(idx).ok().flatten()?
2454 } else {
2455 return None;
2456 };
2457
2458 let metadata = self.metadata.get(&idx).cloned().unwrap_or_default();
2459 Some((id.clone(), vec_data, metadata))
2460 })
2461 .collect()
2462 }
2463
2464 #[must_use]
2466 pub fn contains(&self, id: &str) -> bool {
2467 self.id_to_index
2468 .get(id)
2469 .is_some_and(|&idx| !self.deleted.contains_key(&idx))
2470 }
2471
2472 #[must_use]
2474 pub fn memory_usage(&self) -> usize {
2475 self.vectors.iter().map(|v| v.dim() * 4).sum::<usize>()
2476 }
2477
2478 #[must_use]
2480 pub fn bytes_per_vector(&self) -> f32 {
2481 if self.vectors.is_empty() {
2482 return 0.0;
2483 }
2484 self.memory_usage() as f32 / self.vectors.len() as f32
2485 }
2486
2487 pub fn set_ef_search(&mut self, ef_search: usize) {
2489 self.hnsw_ef_search = ef_search;
2490 if let Some(ref mut index) = self.hnsw_index {
2491 index.set_ef_search(ef_search);
2492 }
2493 }
2494
2495 #[must_use]
2497 pub fn get_ef_search(&self) -> Option<usize> {
2498 Some(self.hnsw_ef_search)
2500 }
2501
2502 pub fn flush(&mut self) -> Result<()> {
2510 let hnsw_bytes = self
2511 .hnsw_index
2512 .as_ref()
2513 .map(postcard::to_allocvec)
2514 .transpose()?;
2515
2516 if let Some(ref mut storage) = self.storage {
2517 storage.set_hnsw_params(
2519 self.hnsw_m as u16,
2520 self.hnsw_ef_construction as u16,
2521 self.hnsw_ef_search as u16,
2522 );
2523
2524 if let Some(bytes) = hnsw_bytes {
2525 storage.put_hnsw_index(bytes);
2526 }
2527 storage.flush()?;
2528 }
2529
2530 if let Some(ref mut text_index) = self.text_index {
2531 text_index.commit()?;
2532 }
2533
2534 Ok(())
2535 }
2536
2537 #[must_use]
2539 pub fn is_persistent(&self) -> bool {
2540 self.storage.is_some()
2541 }
2542
2543 #[must_use]
2545 pub fn storage(&self) -> Option<&OmenFile> {
2546 self.storage.as_ref()
2547 }
2548}