1use std::collections::HashMap;
14use std::path::{Path, PathBuf};
15use std::sync::atomic::{AtomicUsize, Ordering};
16use std::sync::{Arc, Mutex};
17
18use crate::blob_store::BlobStore;
19use crate::index::{SparseIndex, SparseVector};
20use crate::mmap_index::{self, MmapPostingData};
21use crate::segments::{self, IndexMeta, Segment, SegmentMeta};
22use crate::wand::Postings;
23
24const MMAP_FILE: &str = "sparse.mmap";
25const VECTORS_FILE: &str = "sparse_vectors.bin";
26const DIMS_FILE: &str = "sparse_dims.bin";
27const LEGACY_FILE: &str = "sparse.bin";
29
30const STALE_FILES: &[&str] = &[MMAP_FILE, VECTORS_FILE, DIMS_FILE, LEGACY_FILE];
35
36const BLOB_PREFIX: &str = "Sparse_";
38
39static CACHE_SEQ: AtomicUsize = AtomicUsize::new(0);
41
42enum StorageBackend {
44 Filesystem,
46 Store {
48 store: Arc<dyn BlobStore>,
49 index_name: String,
50 },
51}
52
53struct Inner {
54 index: SparseIndex,
57 mmap: Option<MmapPostingData>,
60 postings_loaded: bool,
62 vectors_loaded: bool,
64 num_vectors: usize,
67 dirty: bool,
68 segments: Vec<Segment>,
71 meta: IndexMeta,
72 written: u64,
74}
75
76impl Inner {
79 fn segmented(&self) -> bool {
80 self.mmap.is_none() || !self.meta.segments.is_empty()
81 }
82}
83
84fn empty_inner() -> Inner {
86 Inner {
87 index: SparseIndex::new(),
88 mmap: None,
89 postings_loaded: true,
90 vectors_loaded: true,
91 num_vectors: 0,
92 dirty: false,
93 segments: Vec::new(),
94 meta: IndexMeta::default(),
95 written: 0,
96 }
97}
98
99pub struct SparseHandle {
100 inner: Mutex<Inner>,
101 path: PathBuf,
102 backend: StorageBackend,
103}
104
105impl SparseHandle {
106 pub fn create(path: &str) -> Result<Self, String> {
112 std::fs::create_dir_all(Path::new(path))
113 .map_err(|e| format!("cannot create directory {path}: {e}"))?;
114 let handle = Self {
115 inner: Mutex::new(empty_inner()),
116 path: PathBuf::from(path),
117 backend: StorageBackend::Filesystem,
118 };
119 handle.commit_inner()?;
120 Ok(handle)
121 }
122
123 pub fn open(path: &str) -> Result<Self, String> {
126 Self::open_backed(Path::new(path), StorageBackend::Filesystem)
127 }
128
129 fn open_backed(base: &Path, backend: StorageBackend) -> Result<Self, String> {
132 if base.join(segments::META_FILE).exists() {
133 Self::open_segmented(base, backend)
134 } else if base.join(MMAP_FILE).exists() {
135 Self::open_mmap(base, backend)
136 } else {
137 Self::open_legacy(base)
138 }
139 }
140
141 fn open_segmented(base: &Path, backend: StorageBackend) -> Result<Self, String> {
144 let meta = IndexMeta::read(base)?;
145 let mut segments = Vec::with_capacity(meta.segments.len());
146 for sm in &meta.segments {
147 segments.push(Segment::open(base, sm.clone())?);
148 }
149 let num_vectors = meta.live_vectors();
150 Ok(Self {
151 inner: Mutex::new(Inner {
152 num_vectors,
153 segments,
154 meta,
155 ..empty_inner()
156 }),
157 path: base.to_path_buf(),
158 backend,
159 })
160 }
161
162 pub fn create_with_store(
172 store: Arc<dyn BlobStore>,
173 index_name: &str,
174 cache_base: &Path,
175 ) -> Result<Self, String> {
176 let blob_name = format!("{BLOB_PREFIX}{index_name}");
177 let cache_dir = Self::make_cache_dir(cache_base, &blob_name)?;
178
179 let handle = Self {
180 inner: Mutex::new(empty_inner()),
181 path: cache_dir,
182 backend: StorageBackend::Store {
183 store,
184 index_name: blob_name,
185 },
186 };
187 handle.commit_inner()?;
188 Ok(handle)
189 }
190
191 pub fn open_with_store(
196 store: Arc<dyn BlobStore>,
197 index_name: &str,
198 cache_base: &Path,
199 ) -> Result<Self, String> {
200 let blob_name = format!("{BLOB_PREFIX}{index_name}");
201 let cache_dir = Self::make_cache_dir(cache_base, &blob_name)?;
202
203 let files = store
205 .list(&blob_name)
206 .map_err(|e| format!("cannot list blobs for {blob_name}: {e}"))?;
207
208 for file_name in &files {
209 let data = store
210 .load(&blob_name, file_name)
211 .map_err(|e| format!("cannot load {blob_name}/{file_name}: {e}"))?;
212 mmap_index::write_file_atomic(&cache_dir.join(file_name), &data)?;
215 }
216
217 let backend = StorageBackend::Store {
218 store,
219 index_name: blob_name,
220 };
221
222 if cache_dir.join(segments::META_FILE).exists() || cache_dir.join(MMAP_FILE).exists() {
224 Self::open_backed(&cache_dir, backend)
225 } else {
226 let handle = Self {
228 inner: Mutex::new(empty_inner()),
229 path: cache_dir,
230 backend,
231 };
232 handle.commit_inner()?;
233 Ok(handle)
234 }
235 }
236
237 fn make_cache_dir(base: &Path, index_name: &str) -> Result<PathBuf, String> {
243 let seq = CACHE_SEQ.fetch_add(1, Ordering::Relaxed);
244 let pid = std::process::id();
245 let dir = base
246 .join(format!("{pid}"))
247 .join(format!("{index_name}_{seq}"));
248 std::fs::create_dir_all(&dir)
249 .map_err(|e| format!("cannot create cache dir {}: {e}", dir.display()))?;
250 Ok(dir)
251 }
252
253 fn open_mmap(base: &Path, backend: StorageBackend) -> Result<Self, String> {
260 let mmap = MmapPostingData::open(&base.join(MMAP_FILE))?;
261
262 let (dim_map, dim_reverse): (HashMap<u32, usize>, Vec<u32>) = if mmap.has_global_dims() {
265 let reverse: Vec<u32> = mmap.tokens().map(|(t, _)| t).collect();
266 let map = reverse.iter().enumerate().map(|(i, &t)| (t, i)).collect();
267 (map, reverse)
268 } else {
269 let dims_data = std::fs::read(base.join(DIMS_FILE))
270 .map_err(|e| format!("cannot read {DIMS_FILE}: {e}"))?;
271 bincode::deserialize(&dims_data)
272 .map_err(|e| format!("cannot deserialize dims: {e}"))?
273 };
274
275 let num_dims = mmap.num_dims();
276 let num_vectors = mmap.num_vectors();
277 let empty_postings: Vec<Postings> = (0..num_dims).map(|_| Postings::new()).collect();
278 let index =
280 SparseIndex::from_parts(dim_map, dim_reverse, empty_postings, HashMap::new());
281
282 Ok(Self {
283 inner: Mutex::new(Inner {
284 index,
285 mmap: Some(mmap),
286 postings_loaded: false,
287 vectors_loaded: false,
288 num_vectors,
289 ..empty_inner()
290 }),
291 path: base.to_path_buf(),
292 backend,
293 })
294 }
295
296 fn open_legacy(base: &Path) -> Result<Self, String> {
298 let data_path = base.join(LEGACY_FILE);
299 let data = std::fs::read(&data_path)
300 .map_err(|e| format!("cannot read {}: {e}", data_path.display()))?;
301 let index: SparseIndex = bincode::deserialize(&data)
302 .map_err(|e| format!("cannot deserialize sparse index: {e}"))?;
303 let num_vectors = index.len();
304 Ok(Self {
305 inner: Mutex::new(Inner { index, num_vectors, ..empty_inner() }),
306 path: base.to_path_buf(),
307 backend: StorageBackend::Filesystem,
308 })
309 }
310
311 fn ensure_postings_loaded(inner: &mut Inner) {
317 if inner.postings_loaded {
318 return;
319 }
320 if let Some(ref mmap) = inner.mmap {
321 if mmap.has_global_dims() {
326 let tokens: Vec<u32> = inner.index.dim_reverse().to_vec();
327 let postings = inner.index.postings_mut();
328 for (i, pl) in postings.iter_mut().enumerate() {
329 *pl = match tokens.get(i) {
333 Some(&token) => mmap.load_postings_of_token(token),
334 None => Postings::new(),
335 };
336 }
337 } else {
338 let postings = inner.index.postings_mut();
339 for (i, pl) in postings.iter_mut().enumerate() {
340 *pl = mmap.load_postings(i);
341 }
342 }
343 }
344 inner.postings_loaded = true;
345 }
346
347 fn ensure_vectors_loaded(inner: &mut Inner, path: &Path) -> Result<(), String> {
349 if inner.vectors_loaded {
350 return Ok(());
351 }
352 if inner.segmented() {
356 inner.vectors_loaded = true;
357 return Ok(());
358 }
359 let vectors_path = path.join(VECTORS_FILE);
360 if !vectors_path.exists()
365 && inner.mmap.as_ref().is_some_and(|m| m.has_global_dims())
366 {
367 inner.vectors_loaded = true;
368 return Ok(());
369 }
370 let data = std::fs::read(&vectors_path)
371 .map_err(|e| format!("cannot read {}: {e}", vectors_path.display()))?;
372 let vectors: HashMap<u64, SparseVector> = bincode::deserialize(&data)
373 .map_err(|e| format!("cannot deserialize vectors: {e}"))?;
374 inner.index.set_vectors(vectors);
375 inner.vectors_loaded = true;
376 Ok(())
377 }
378
379 pub fn insert(&self, node_id: u64, vector: &SparseVector) -> Result<(), String> {
384 let mut inner = self.inner.lock().map_err(|_| "lock poisoned".to_string())?;
385 Self::ensure_vectors_loaded(&mut inner, &self.path)?;
386 Self::ensure_postings_loaded(&mut inner);
387 Self::tombstone_committed(&mut inner, node_id)?;
391 inner.index.insert(node_id, vector);
392 inner.num_vectors = Self::count(&inner);
393 inner.dirty = true;
394 Ok(())
395 }
396
397 pub fn remove(&self, node_id: u64) -> Result<bool, String> {
398 let mut inner = self.inner.lock().map_err(|_| "lock poisoned".to_string())?;
399 Self::ensure_vectors_loaded(&mut inner, &self.path)?;
400 Self::ensure_postings_loaded(&mut inner);
401 let from_segments = Self::tombstone_committed(&mut inner, node_id)?;
402 let from_ram = inner.index.remove(node_id);
403 let removed = from_segments || from_ram;
404 if removed {
405 inner.num_vectors = Self::count(&inner);
406 inner.dirty = true;
407 }
408 Ok(removed)
409 }
410
411 fn tombstone_committed(inner: &mut Inner, node_id: u64) -> Result<bool, String> {
414 let mut hit = false;
415 for seg in &mut inner.segments {
416 if seg.holds(node_id)? && seg.tombstone(node_id) {
417 hit = true;
418 }
419 }
420 if hit {
421 for (sm, seg) in inner.meta.segments.iter_mut().zip(inner.segments.iter()) {
424 sm.deleted = seg.meta.deleted.clone();
425 }
426 }
427 Ok(hit)
428 }
429
430 fn count(inner: &Inner) -> usize {
432 inner.meta.live_vectors() + inner.index.len()
433 }
434
435 pub fn search(&self, query: &SparseVector, limit: usize) -> Vec<(u64, f32)> {
436 let inner = self.inner.lock().unwrap();
437 if !inner.segments.is_empty() {
438 return Self::search_segments(&inner, query, limit, None);
439 }
440 if !inner.dirty {
441 if let Some(ref mmap) = inner.mmap {
442 return mmap_index::search_mmap(
443 mmap,
444 inner.index.dim_map(),
445 query,
446 limit,
447 &|_| true,
448 );
449 }
450 }
451 inner.index.search(query, limit)
452 }
453
454 fn search_segments(
471 inner: &Inner,
472 query: &SparseVector,
473 limit: usize,
474 allowed: Option<&[u64]>,
475 ) -> Vec<(u64, f32)> {
476 let no_dims = HashMap::new();
477 let mut all: Vec<(u64, f32)> = Vec::new();
478 for seg in &inner.segments {
479 if seg.data.num_vectors() == 0 {
480 continue;
481 }
482 let hits = match allowed {
483 Some(ids) if seg.meta.deleted.is_empty() => {
484 mmap_index::search_mmap_allowed(&seg.data, &no_dims, query, limit, ids)
485 }
486 Some(ids) => {
487 let live: Vec<u64> = ids.iter().copied().filter(|&id| seg.is_live(id)).collect();
488 mmap_index::search_mmap_allowed(&seg.data, &no_dims, query, limit, &live)
489 }
490 None => mmap_index::search_mmap(
491 &seg.data, &no_dims, query, limit, &|id| seg.is_live(id)),
492 };
493 all.extend(hits);
494 }
495 if !inner.index.is_empty() {
496 all.extend(match allowed {
497 Some(ids) => inner.index.search_filtered(query, limit, ids),
498 None => inner.index.search(query, limit),
499 });
500 }
501 all.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal).then(a.0.cmp(&b.0)));
502 all.truncate(limit);
503 all
504 }
505
506 pub fn search_filtered(
521 &self,
522 query: &SparseVector,
523 limit: usize,
524 allowed_ids: &[u64],
525 ) -> Vec<(u64, f32)> {
526 let inner = self.inner.lock().unwrap();
527 if !inner.segments.is_empty() {
528 return Self::search_segments(&inner, query, limit, Some(allowed_ids));
529 }
530 if !inner.dirty {
531 if let Some(ref mmap) = inner.mmap {
532 return mmap_index::search_mmap_allowed(
533 mmap,
534 inner.index.dim_map(),
535 query,
536 limit,
537 allowed_ids,
538 );
539 }
540 }
541 inner.index.search_filtered(query, limit, allowed_ids)
542 }
543
544 pub fn compact(&self) -> Result<(), String> {
553 let mut inner = self.inner.lock().map_err(|_| "lock poisoned".to_string())?;
554 if inner.dirty {
555 drop(inner);
556 self.commit_inner()?;
557 inner = self.inner.lock().map_err(|_| "lock poisoned".to_string())?;
558 }
559 if inner.segments.len() < 2 {
560 return Ok(());
561 }
562
563 inner.written += 1;
564 let new_id = segments::new_segment_id(inner.written);
565 let sources: Vec<&Segment> = inner.segments.iter().collect();
566 let merged = segments::merge_segments(&self.path, &sources, &new_id)?;
567 let dropped: Vec<String> = inner.segments.iter()
568 .flat_map(|s| [segments::segment_file(&s.meta.id), segments::ids_file(&s.meta.id)])
569 .collect();
570
571 inner.meta.segments = vec![merged.clone()];
574 inner.meta.write(&self.path)?;
575 inner.segments = vec![Segment::open(&self.path, merged)?];
576 inner.num_vectors = Self::count(&inner);
577
578 if let StorageBackend::Store { ref store, ref index_name } = self.backend {
579 for file in [segments::segment_file(&new_id), segments::ids_file(&new_id), segments::META_FILE.to_string()] {
580 let data = std::fs::read(self.path.join(&file))
581 .map_err(|e| format!("cannot read cache {file}: {e}"))?;
582 store.save(index_name, &file, &data)
583 .map_err(|e| format!("cannot save {index_name}/{file} to store: {e}"))?;
584 }
585 }
586 for file in dropped {
588 let _ = std::fs::remove_file(self.path.join(&file));
589 if let StorageBackend::Store { ref store, ref index_name } = self.backend {
590 let _ = store.delete(index_name, &file);
591 }
592 }
593 Ok(())
594 }
595
596 fn max_segments() -> usize {
625 static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
626 *CAP.get_or_init(|| {
627 std::env::var("LUCIVY_SPARSE_MAX_SEGMENTS").ok()
628 .and_then(|v| v.parse().ok())
629 .unwrap_or(8)
630 })
631 }
632
633 pub fn num_segments(&self) -> usize {
636 self.inner.lock().map(|i| i.segments.len()).unwrap_or(0)
637 }
638
639 pub fn len(&self) -> usize {
640 self.inner.lock().unwrap().num_vectors
641 }
642
643 pub fn is_empty(&self) -> bool {
644 self.len() == 0
645 }
646
647 pub fn commit_inner(&self) -> Result<(), String> {
658 let mut inner = self.inner.lock().map_err(|_| "lock poisoned".to_string())?;
659
660 let converting = inner.mmap.is_some();
661 if converting {
662 Self::ensure_postings_loaded(&mut inner);
665 Self::ensure_vectors_loaded(&mut inner, &self.path)?;
666 } else if !inner.dirty && !inner.meta.segments.is_empty() {
667 return Ok(());
670 }
671
672 let mut written: Vec<String> = Vec::new();
674 let has_delta = !inner.index.is_empty()
675 || inner.index.postings().iter().any(|p| !p.is_empty());
676 if has_delta {
677 inner.written += 1;
678 let id = segments::new_segment_id(inner.written);
679 let file = segments::segment_file(&id);
680 mmap_index::write_mmap_file(
681 &self.path.join(&file),
682 inner.index.postings(),
683 inner.index.dim_reverse(),
684 inner.index.len() as u32,
685 )?;
686 let mut ids: Vec<u64> = if converting {
692 let mut set = std::collections::HashSet::new();
693 for p in inner.index.postings() {
694 for x in p.as_slice() { set.insert(x.id); }
695 }
696 set.into_iter().collect()
697 } else {
698 inner.index.vectors().keys().copied().collect()
699 };
700 ids.sort_unstable();
701 let ids_name = segments::ids_file(&id);
702 mmap_index::write_file_atomic(&self.path.join(&ids_name), &segments::encode_ids(&ids))?;
703 inner.meta.segments.push(SegmentMeta {
704 id,
705 num_vectors: ids.len() as u32,
706 deleted: Vec::new(),
707 });
708 written.push(file);
709 written.push(ids_name);
710 }
711
712 inner.meta.version = segments::META_VERSION;
716 inner.meta.write(&self.path)?;
717 written.push(segments::META_FILE.to_string());
718
719 let metas: Vec<SegmentMeta> = inner.meta.segments.clone();
721 let mut opened = Vec::with_capacity(metas.len());
722 for sm in metas {
723 opened.push(Segment::open(&self.path, sm)?);
724 }
725 inner.segments = opened;
726 inner.index = SparseIndex::new();
727 inner.mmap = None;
728 inner.postings_loaded = true;
729 inner.vectors_loaded = true;
730 inner.dirty = false;
731 inner.num_vectors = Self::count(&inner);
732
733 if let StorageBackend::Store { ref store, ref index_name } = self.backend {
735 for file in &written {
736 let data = std::fs::read(self.path.join(file))
737 .map_err(|e| format!("cannot read cache {file}: {e}"))?;
738 store
739 .save(index_name, file, &data)
740 .map_err(|e| format!("cannot save {index_name}/{file} to store: {e}"))?;
741 }
742 }
743 for &stale in STALE_FILES {
747 let path = self.path.join(stale);
748 if path.exists() {
749 let _ = std::fs::remove_file(&path);
750 if let StorageBackend::Store { ref store, ref index_name } = self.backend {
751 let _ = store.delete(index_name, stale);
752 }
753 }
754 }
755
756 let cap = Self::max_segments();
759 let pile = inner.segments.len();
760 drop(inner);
761 if cap > 0 && pile > cap {
762 self.compact()?;
763 }
764 Ok(())
765 }
766}
767
768impl Drop for SparseHandle {
769 fn drop(&mut self) {
770 if let StorageBackend::Store { .. } = &self.backend {
773 let _ = std::fs::remove_dir_all(&self.path);
774 }
775 }
776}
777
778#[cfg(test)]
779mod tests {
780 use super::*;
781 use crate::blob_store::MemBlobStore;
782
783 fn tmp_path(name: &str) -> PathBuf {
784 std::env::temp_dir().join(name)
785 }
786
787 fn cleanup(path: &Path) {
788 let _ = std::fs::remove_dir_all(path);
789 }
790
791 #[test]
796 fn create_writes_a_manifest_and_a_segment_per_commit() {
797 let p = tmp_path("sparse_mmap_create_test");
798 cleanup(&p);
799 let path = p.to_str().unwrap();
800
801 let handle = SparseHandle::create(path).unwrap();
804 assert!(p.join(crate::segments::META_FILE).exists());
805 assert_eq!(segment_files(&p).len(), 0);
806 assert!(!p.join(MMAP_FILE).exists(), "the single-file format is not written any more");
807
808 handle.insert(1, &SparseVector::new(vec![7], vec![1.0])).unwrap();
809 handle.commit_inner().unwrap();
810 assert_eq!(segment_files(&p).len(), 1);
811
812 handle.insert(2, &SparseVector::new(vec![7], vec![1.0])).unwrap();
814 handle.commit_inner().unwrap();
815 assert_eq!(segment_files(&p).len(), 2);
816
817 let handle2 = SparseHandle::open(path).unwrap();
818 assert_eq!(handle2.len(), 2);
819 assert_eq!(handle2.search(&SparseVector::new(vec![7], vec![1.0]), 10).len(), 2);
820
821 cleanup(&p);
822 }
823
824 fn segment_files(base: &Path) -> Vec<String> {
826 let mut names: Vec<String> = std::fs::read_dir(base).unwrap()
827 .filter_map(|e| e.ok().map(|e| e.file_name().to_string_lossy().into_owned()))
828 .filter(|n| n.starts_with("seg_") && n.ends_with(".mmap"))
829 .collect();
830 names.sort();
831 names
832 }
833
834 #[test]
835 fn persistence_roundtrip_mmap() {
836 let p = tmp_path("sparse_mmap_roundtrip_test");
837 cleanup(&p);
838 let path = p.to_str().unwrap();
839
840 let handle = SparseHandle::create(path).unwrap();
841 handle
842 .insert(42, &SparseVector::new(vec![1, 2], vec![0.5, 0.3]))
843 .unwrap();
844 handle
845 .insert(99, &SparseVector::new(vec![2, 3], vec![0.8, 0.2]))
846 .unwrap();
847 handle.commit_inner().unwrap();
848 drop(handle);
849
850 let handle2 = SparseHandle::open(path).unwrap();
852 assert_eq!(handle2.len(), 2);
853
854 let results = handle2.search(&SparseVector::new(vec![2], vec![1.0]), 10);
856 assert_eq!(results.len(), 2);
857 assert_eq!(results[0].0, 99);
858 assert!((results[0].1 - 0.8).abs() < 1e-6);
859 assert_eq!(results[1].0, 42);
860 assert!((results[1].1 - 0.3).abs() < 1e-6);
861
862 cleanup(&p);
863 }
864
865 #[test]
866 fn mmap_search_filtered() {
867 let p = tmp_path("sparse_mmap_filtered_test");
868 cleanup(&p);
869 let path = p.to_str().unwrap();
870
871 let handle = SparseHandle::create(path).unwrap();
872 handle
873 .insert(1, &SparseVector::new(vec![1, 2], vec![0.5, 0.3]))
874 .unwrap();
875 handle
876 .insert(2, &SparseVector::new(vec![1, 3], vec![0.9, 0.1]))
877 .unwrap();
878 handle
879 .insert(3, &SparseVector::new(vec![1], vec![0.7]))
880 .unwrap();
881 handle.commit_inner().unwrap();
882 drop(handle);
883
884 let handle2 = SparseHandle::open(path).unwrap();
885 let results = handle2.search_filtered(&SparseVector::new(vec![1], vec![1.0]), 10, &[1, 3]);
886 assert_eq!(results.len(), 2);
887 assert_eq!(results[0].0, 3); assert_eq!(results[1].0, 1); cleanup(&p);
891 }
892
893 #[test]
894 fn mutation_after_mmap_open() {
895 let p = tmp_path("sparse_mmap_mutation_test");
896 cleanup(&p);
897 let path = p.to_str().unwrap();
898
899 let handle = SparseHandle::create(path).unwrap();
900 handle
901 .insert(1, &SparseVector::new(vec![10], vec![1.0]))
902 .unwrap();
903 handle.commit_inner().unwrap();
904 drop(handle);
905
906 let handle2 = SparseHandle::open(path).unwrap();
908 handle2
909 .insert(2, &SparseVector::new(vec![10], vec![2.0]))
910 .unwrap();
911
912 let results = handle2.search(&SparseVector::new(vec![10], vec![1.0]), 10);
913 assert_eq!(results.len(), 2);
914 assert_eq!(results[0].0, 2); assert_eq!(results[1].0, 1); handle2.commit_inner().unwrap();
919 drop(handle2);
920
921 let handle3 = SparseHandle::open(path).unwrap();
922 let results = handle3.search(&SparseVector::new(vec![10], vec![1.0]), 10);
923 assert_eq!(results.len(), 2);
924 assert_eq!(results[0].0, 2);
925
926 cleanup(&p);
927 }
928
929 #[test]
930 fn legacy_fallback() {
931 let p = tmp_path("sparse_mmap_legacy_test");
932 cleanup(&p);
933 let path = p.to_str().unwrap();
934
935 std::fs::create_dir_all(&p).unwrap();
937 let mut index = SparseIndex::new();
938 index.insert(7, &SparseVector::new(vec![1], vec![0.42]));
939 let data = bincode::serialize(&index).unwrap();
940 std::fs::write(p.join(LEGACY_FILE), data).unwrap();
941
942 let handle = SparseHandle::open(path).unwrap();
944 assert_eq!(handle.len(), 1);
945 let results = handle.search(&SparseVector::new(vec![1], vec![1.0]), 10);
946 assert_eq!(results[0].0, 7);
947
948 handle.commit_inner().unwrap();
950 assert!(p.join(crate::segments::META_FILE).exists());
951 assert_eq!(segment_files(&p).len(), 1);
952 assert!(!p.join(LEGACY_FILE).exists());
953 assert!(!p.join(MMAP_FILE).exists());
954 assert_eq!(handle.search(&SparseVector::new(vec![1], vec![1.0]), 10)[0].0, 7);
955
956 cleanup(&p);
957 }
958
959 #[test]
960 fn many_docs_mmap_roundtrip() {
961 let p = tmp_path("sparse_mmap_many_docs_test");
962 cleanup(&p);
963 let path = p.to_str().unwrap();
964
965 let handle = SparseHandle::create(path).unwrap();
966 for i in 0..500u64 {
967 let token = (i % 50) as u32;
968 let weight = (i as f32) / 500.0;
969 handle
970 .insert(
971 i,
972 &SparseVector::new(vec![token, token + 50], vec![weight, weight * 0.5]),
973 )
974 .unwrap();
975 }
976 handle.commit_inner().unwrap();
977 drop(handle);
978
979 let handle2 = SparseHandle::open(path).unwrap();
980 assert_eq!(handle2.len(), 500);
981
982 let results = handle2.search(&SparseVector::new(vec![0, 50], vec![1.0, 1.0]), 5);
983 assert_eq!(results.len(), 5);
984 assert_eq!(results[0].0, 450);
986
987 cleanup(&p);
988 }
989
990 fn test_cache_base() -> PathBuf {
995 std::env::temp_dir().join("sparse_test_cache")
996 }
997
998 #[test]
999 fn blob_store_create_and_search() {
1000 let store = Arc::new(MemBlobStore::new());
1001 let cb = test_cache_base();
1002 let handle = SparseHandle::create_with_store(store.clone(), "test_idx", &cb).unwrap();
1003
1004 handle
1005 .insert(42, &SparseVector::new(vec![1, 2], vec![0.5, 0.3]))
1006 .unwrap();
1007 handle
1008 .insert(99, &SparseVector::new(vec![2, 3], vec![0.8, 0.2]))
1009 .unwrap();
1010 handle.commit_inner().unwrap();
1011
1012 let names = store.list("Sparse_test_idx").unwrap();
1014 assert!(names.iter().any(|n| n == crate::segments::META_FILE), "{names:?}");
1015 assert_eq!(names.iter().filter(|n| n.ends_with(".mmap")).count(), 1, "{names:?}");
1016 assert_eq!(names.iter().filter(|n| n.ends_with(".ids")).count(), 1, "{names:?}");
1017
1018 let results = handle.search(&SparseVector::new(vec![2], vec![1.0]), 10);
1020 assert_eq!(results.len(), 2);
1021 assert_eq!(results[0].0, 99);
1022 }
1023
1024 #[test]
1025 fn blob_store_close_reopen() {
1026 let store = Arc::new(MemBlobStore::new());
1027 let cb = test_cache_base();
1028
1029 {
1031 let handle = SparseHandle::create_with_store(store.clone(), "reopen_idx", &cb).unwrap();
1032 handle
1033 .insert(1, &SparseVector::new(vec![10], vec![1.0]))
1034 .unwrap();
1035 handle
1036 .insert(2, &SparseVector::new(vec![10, 20], vec![0.5, 0.8]))
1037 .unwrap();
1038 handle.commit_inner().unwrap();
1039 }
1040 let handle2 = SparseHandle::open_with_store(store.clone(), "reopen_idx", &cb).unwrap();
1044 assert_eq!(handle2.len(), 2);
1045
1046 let results = handle2.search(&SparseVector::new(vec![10], vec![1.0]), 10);
1047 assert_eq!(results.len(), 2);
1048 assert_eq!(results[0].0, 1); assert_eq!(results[1].0, 2); }
1051
1052 #[test]
1053 fn blob_store_mutation_after_reopen() {
1054 let store = Arc::new(MemBlobStore::new());
1055 let cb = test_cache_base();
1056
1057 {
1058 let handle = SparseHandle::create_with_store(store.clone(), "mut_idx", &cb).unwrap();
1059 handle
1060 .insert(1, &SparseVector::new(vec![5], vec![1.0]))
1061 .unwrap();
1062 handle.commit_inner().unwrap();
1063 }
1064
1065 let handle2 = SparseHandle::open_with_store(store.clone(), "mut_idx", &cb).unwrap();
1066 handle2
1067 .insert(2, &SparseVector::new(vec![5], vec![2.0]))
1068 .unwrap();
1069 handle2.commit_inner().unwrap();
1070
1071 drop(handle2);
1073 let handle3 = SparseHandle::open_with_store(store.clone(), "mut_idx", &cb).unwrap();
1074 assert_eq!(handle3.len(), 2);
1075
1076 let results = handle3.search(&SparseVector::new(vec![5], vec![1.0]), 10);
1077 assert_eq!(results.len(), 2);
1078 assert_eq!(results[0].0, 2); assert_eq!(results[1].0, 1); }
1081
1082 #[test]
1083 fn blob_store_delete_and_reopen() {
1084 let store = Arc::new(MemBlobStore::new());
1085 let cb = test_cache_base();
1086
1087 {
1088 let handle = SparseHandle::create_with_store(store.clone(), "del_idx", &cb).unwrap();
1089 handle
1090 .insert(1, &SparseVector::new(vec![1], vec![1.0]))
1091 .unwrap();
1092 handle
1093 .insert(2, &SparseVector::new(vec![1], vec![2.0]))
1094 .unwrap();
1095 handle.commit_inner().unwrap();
1096 }
1097
1098 let handle2 = SparseHandle::open_with_store(store.clone(), "del_idx", &cb).unwrap();
1100 assert_eq!(handle2.len(), 2);
1101 handle2.remove(1).unwrap();
1102 assert_eq!(handle2.len(), 1);
1103 handle2.commit_inner().unwrap();
1104 drop(handle2);
1105
1106 let handle3 = SparseHandle::open_with_store(store.clone(), "del_idx", &cb).unwrap();
1108 assert_eq!(handle3.len(), 1);
1109
1110 let results = handle3.search(&SparseVector::new(vec![1], vec![1.0]), 10);
1111 assert_eq!(results.len(), 1);
1112 assert_eq!(results[0].0, 2);
1113 }
1114
1115 #[test]
1116 fn blob_store_multiple_indexes_isolated() {
1117 let store = Arc::new(MemBlobStore::new());
1118 let cb = test_cache_base();
1119
1120 let h1 = SparseHandle::create_with_store(store.clone(), "idx_a", &cb).unwrap();
1121 let h2 = SparseHandle::create_with_store(store.clone(), "idx_b", &cb).unwrap();
1122
1123 h1.insert(1, &SparseVector::new(vec![1], vec![1.0]))
1124 .unwrap();
1125 h1.insert(2, &SparseVector::new(vec![1], vec![0.5]))
1126 .unwrap();
1127 h2.insert(10, &SparseVector::new(vec![1], vec![3.0]))
1128 .unwrap();
1129
1130 h1.commit_inner().unwrap();
1131 h2.commit_inner().unwrap();
1132
1133 assert_eq!(h1.len(), 2);
1134 assert_eq!(h2.len(), 1);
1135
1136 assert_eq!(store.list("Sparse_idx_a").unwrap().len(), 3);
1138 assert_eq!(store.list("Sparse_idx_b").unwrap().len(), 3);
1139 }
1140
1141 #[test]
1142 fn blob_store_survives_cache_cleanup() {
1143 let store = Arc::new(MemBlobStore::new());
1144 let cb = test_cache_base();
1145
1146 {
1147 let handle = SparseHandle::create_with_store(store.clone(), "surv_idx", &cb).unwrap();
1148 for i in 0..50u64 {
1149 handle
1150 .insert(i, &SparseVector::new(vec![(i % 10) as u32], vec![i as f32]))
1151 .unwrap();
1152 }
1153 handle.commit_inner().unwrap();
1154 }
1155 let handle2 = SparseHandle::open_with_store(store.clone(), "surv_idx", &cb).unwrap();
1159 assert_eq!(handle2.len(), 50);
1160
1161 let results = handle2.search(&SparseVector::new(vec![0], vec![1.0]), 5);
1162 assert_eq!(results.len(), 4);
1165 assert_eq!(results[0].0, 40);
1166 }
1167
1168 #[test]
1169 fn blob_store_search_filtered_after_reopen() {
1170 let store = Arc::new(MemBlobStore::new());
1171 let cb = test_cache_base();
1172
1173 {
1174 let handle = SparseHandle::create_with_store(store.clone(), "filt_idx", &cb).unwrap();
1175 handle
1176 .insert(1, &SparseVector::new(vec![1, 2], vec![0.5, 0.3]))
1177 .unwrap();
1178 handle
1179 .insert(2, &SparseVector::new(vec![1, 3], vec![0.9, 0.1]))
1180 .unwrap();
1181 handle
1182 .insert(3, &SparseVector::new(vec![1], vec![0.7]))
1183 .unwrap();
1184 handle.commit_inner().unwrap();
1185 }
1186
1187 let handle2 = SparseHandle::open_with_store(store.clone(), "filt_idx", &cb).unwrap();
1188 let results =
1189 handle2.search_filtered(&SparseVector::new(vec![1], vec![1.0]), 10, &[1, 3]);
1190 assert_eq!(results.len(), 2);
1191 assert_eq!(results[0].0, 3); assert_eq!(results[1].0, 1); }
1194}