diff --git a/lib/segment/src/index/sparse_index/sparse_vector_index.rs b/lib/segment/src/index/sparse_index/sparse_vector_index.rs index f53e218f91..15b71e947b 100644 --- a/lib/segment/src/index/sparse_index/sparse_vector_index.rs +++ b/lib/segment/src/index/sparse_index/sparse_vector_index.rs @@ -1,4 +1,5 @@ use std::collections::HashSet; +use std::fs::create_dir_all; use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicBool; use std::sync::Arc; @@ -33,13 +34,16 @@ pub struct SparseVectorIndex { } impl SparseVectorIndex { - /// Create new sparse vector index + /// Open a sparse vector index at a given path pub fn open( id_tracker: Arc>, vector_storage: Arc>, payload_index: Arc>, path: &Path, ) -> OperationResult { + // create directory if it does not exist + create_dir_all(path)?; + let searches_telemetry = SparseSearchesTelemetry::new(); let max_point_id = 0; let inverted_index = TInvertedIndex::open(path)?; @@ -158,13 +162,14 @@ impl VectorIndex for SparseVectorIndex( - rnd: &mut R, - max_dim: usize, - stopped: &AtomicBool, -) -> SparseVectorIndex { - // test params - let num_vectors = 1000; +const MAX_SPARSE_DIM: usize = 512; +/// Helper to open a test sparse vector index +fn fixture_open_sparse_index( + index_dir: &Path, + num_vectors: usize, // used to size the id tracker +) -> OperationResult> { // temp dirs let payload_dir = Builder::new().prefix("payload_dir").tempdir().unwrap(); - let index_dir = Builder::new().prefix("index_dir").tempdir().unwrap(); let storage_dir = Builder::new().prefix("storage_dir").tempdir().unwrap(); // setup @@ -51,36 +49,54 @@ fn fixture_sparse_index( id_tracker.clone(), payload_dir.path(), true, - ) - .unwrap(); + )?; let wrapped_payload_index = Arc::new(AtomicRefCell::new(payload_index)); let db = open_db(storage_dir.path(), &[DB_VECTOR_CF]).unwrap(); - let vector_storage = - open_simple_sparse_vector_storage(db, DB_VECTOR_CF, Distance::Dot).unwrap(); + let vector_storage = open_simple_sparse_vector_storage(db, DB_VECTOR_CF, Distance::Dot)?; - let mut sparse_vector_index: SparseVectorIndex = SparseVectorIndex::open( + let sparse_vector_index: SparseVectorIndex = SparseVectorIndex::open( id_tracker, vector_storage.clone(), wrapped_payload_index, - index_dir.path(), - ) - .unwrap(); + index_dir, + )?; + + Ok(sparse_vector_index) +} + +/// Prepares a sparse vector index with random sparse vectors +fn fixture_sparse_index_ram( + rnd: &mut R, + max_dim: usize, + stopped: &AtomicBool, +) -> SparseVectorIndex { + // test params + let num_vectors = 1000; + + let index_dir = Builder::new().prefix("index_dir").tempdir().unwrap(); + let mut sparse_vector_index = fixture_open_sparse_index(index_dir.path(), num_vectors).unwrap(); + let mut borrowed_storage = sparse_vector_index.vector_storage.borrow_mut(); // add points to storage for idx in 0..num_vectors { let vec = &random_sparse_vector(rnd, max_dim); - vector_storage - .borrow_mut() + borrowed_storage .insert_vector(idx as PointOffsetType, vec.into()) .unwrap(); } + drop(borrowed_storage); + + // assert all points are in storage assert_eq!( - vector_storage.borrow().available_vector_count(), + sparse_vector_index + .vector_storage + .borrow() + .available_vector_count(), num_vectors ); - // build index + // build index to refresh RAM index sparse_vector_index.build_index(stopped).unwrap(); assert_eq!(sparse_vector_index.indexed_vector_count(), num_vectors); sparse_vector_index @@ -91,8 +107,7 @@ fn sparse_vector_index_ram_no_filter_search() { let stopped = AtomicBool::new(false); let mut rnd = StdRng::seed_from_u64(42); - let sparse_vector_index: SparseVectorIndex = - fixture_sparse_index(&mut rnd, MAX_SPARSE_DIM, &stopped); + let sparse_vector_index = fixture_sparse_index_ram(&mut rnd, MAX_SPARSE_DIM, &stopped); // random query vectors let attempts = 100; @@ -140,13 +155,8 @@ fn sparse_vector_index_ram_no_filter_search() { } } -#[test] -fn sparse_vector_index_ram_consistent_with_storage() { - let stopped = AtomicBool::new(false); - let mut rnd = StdRng::seed_from_u64(42); - - let sparse_vector_index: SparseVectorIndex = - fixture_sparse_index(&mut rnd, MAX_SPARSE_DIM, &stopped); +/// Checks that the sparse vector index is consistent with the underlying storage +fn check_index_storage_consistency(sparse_vector_index: &SparseVectorIndex) { let borrowed_vector_storage = sparse_vector_index.vector_storage.borrow(); let point_count = borrowed_vector_storage.available_vector_count(); for id in 0..point_count as PointOffsetType { @@ -170,20 +180,59 @@ fn sparse_vector_index_ram_consistent_with_storage() { let top = sparse_vector_index.max_result_count(vector); let query_vector: QueryVector = vector.to_owned().into(); let results = sparse_vector_index - .search(&[&query_vector], None, top, None, &stopped) + .search(&[&query_vector], None, top, None, &false.into()) .unwrap(); assert!(results[0].iter().any(|s| s.idx == id)); } } +#[test] +fn sparse_vector_index_consistent_with_storage() { + let stopped = AtomicBool::new(false); + let mut rnd = StdRng::seed_from_u64(42); + + let sparse_vector_ram_index = fixture_sparse_index_ram(&mut rnd, MAX_SPARSE_DIM, &stopped); + + // check consistency with underlying RAM inverted index + check_index_storage_consistency(&sparse_vector_ram_index); + + let mmap_index_dir = Builder::new().prefix("mmap_index_dir").tempdir().unwrap(); + // copy index mmap and save to disk + let _mmap_inverted_index = InvertedIndexMmap::convert_and_save( + &sparse_vector_ram_index.inverted_index, + &mmap_index_dir, + ) + .unwrap(); + + // load index from memmap file + let sparse_vector_mmap_index: SparseVectorIndex = SparseVectorIndex::open( + sparse_vector_ram_index.id_tracker.clone(), + sparse_vector_ram_index.vector_storage.clone(), + sparse_vector_ram_index.payload_index.clone(), + mmap_index_dir.path(), + ) + .unwrap(); + + // check consistency with underlying mmap inverted index + check_index_storage_consistency(&sparse_vector_mmap_index); +} + +#[test] +fn sparse_vector_index_load_missing_mmap() { + let index_dir = Builder::new().prefix("index_dir").tempdir().unwrap(); + let sparse_vector_index: OperationResult> = + fixture_open_sparse_index(index_dir.path(), 0); + // fails to open index if mmap file is missing + assert!(sparse_vector_index.is_err()) +} + #[test] fn sparse_vector_index_ram_deleted_points_search() { let stopped = AtomicBool::new(false); let top = 10; let mut rnd = StdRng::seed_from_u64(42); - let mut sparse_vector_index: SparseVectorIndex = - fixture_sparse_index(&mut rnd, MAX_SPARSE_DIM, &stopped); + let mut sparse_vector_index = fixture_sparse_index_ram(&mut rnd, MAX_SPARSE_DIM, &stopped); // sanity check (all indexed, no deleted points) assert_eq!( @@ -270,8 +319,7 @@ fn sparse_vector_index_ram_filtered_search() { let field_value = "important value"; // setup index - let sparse_vector_index: SparseVectorIndex = - fixture_sparse_index(&mut rnd, MAX_SPARSE_DIM, &stopped); + let sparse_vector_index = fixture_sparse_index_ram(&mut rnd, MAX_SPARSE_DIM, &stopped); // query index by payload let filter = Filter::new_must(Condition::Field(FieldCondition::new_match(