diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 100936fbab..6628b526d5 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -850,6 +850,18 @@ pub struct SearchResultItem { pub distance: f32, } +pub struct SearchResultWithIndexedVectors { + pub results: Vec>, + pub stats: SearchResultStats, +} + +pub struct SearchResultItemWithIndexedVector { + pub vertex_id: u32, + pub data: AssociatedData, + pub distance: f32, + pub indexed_vector: Box<[VectorData]>, +} + impl DiskIndexSearcher where Data: GraphDataType, @@ -1095,6 +1107,67 @@ where Ok(search_result) } + /// Perform a disk search and return valid hits with their indexed vectors. + pub fn search_with_indexed_vectors( + &self, + query: &[Data::VectorDataType], + return_list_size: u32, + search_list_size: u32, + beam_width: Option, + mode: SearchMode<'_>, + ) -> ANNResult> + { + let SearchResult { mut results, stats } = + self.search(query, return_list_size, search_list_size, beam_width, mode)?; + results.truncate(stats.result_count as usize); + if results.is_empty() { + return Ok(SearchResultWithIndexedVectors { + results: Vec::new(), + stats, + }); + } + + let provider = self.index.provider(); + let mut scratch = PoolOption::try_pooled( + &self.scratch_pool, + &DiskSearchScratchArgs { + graph_degree: provider.graph_header.max_degree::()?, + pq_dim: provider.pq_data.get_dim(), + num_pq_chunks: provider.pq_data.get_num_chunks(), + num_pq_centers: provider.pq_data.get_num_centers(), + vertex_factory: &self.vertex_provider_factory, + graph_header: &provider.graph_header, + }, + )?; + + const BATCH_SIZE: usize = 128; + let mut indexed_vectors = Vec::with_capacity(results.len()); + for batch in results.chunks(BATCH_SIZE) { + let ids = batch + .iter() + .map(|result| result.vertex_id) + .collect::>(); + ensure_vertex_loaded(&mut scratch.vertex_provider, &ids)?; + for id in ids { + indexed_vectors.push(Box::from(scratch.vertex_provider.get_vector(&id)?)); + } + } + + let results = results + .into_iter() + .zip(indexed_vectors) + .map( + |(result, indexed_vector)| SearchResultItemWithIndexedVector { + indexed_vector, + vertex_id: result.vertex_id, + data: result.data, + distance: result.distance, + }, + ) + .collect(); + Ok(SearchResultWithIndexedVectors { results, stats }) + } + /// Perform a raw search on the disk index. /// This is a lower-level API that allows more control over the search parameters and output buffers. #[allow(clippy::too_many_arguments)] @@ -1535,6 +1608,76 @@ mod disk_provider_tests { result.into_inner().into_vec() } + #[rstest] + #[case(CachingStrategy::None)] + #[case(CachingStrategy::StaticCacheWithBfsNodes(32))] + fn test_search_with_indexed_vectors(#[case] caching_strategy: CachingStrategy) { + let storage_provider = Arc::new(VirtualStorageProvider::new_overlay(test_data_root())); + let searcher = create_disk_index_searcher::( + CreateDiskIndexSearcherParams { + max_thread_num: 1, + pq_pivot_file_path: TEST_PQ_PIVOT_128DIM, + pq_compressed_file_path: TEST_PQ_COMPRESSED_128DIM, + index_path: TEST_INDEX_128DIM, + index_path_prefix: TEST_INDEX_PREFIX_128DIM, + caching_strategy, + ..Default::default() + }, + &storage_provider, + ); + let query = vec![0.1f32; 128]; + let legacy = searcher + .search(&query, 10, 20, None, SearchMode::graph()) + .unwrap(); + let indexed = searcher + .search_with_indexed_vectors(&query, 10, 20, None, SearchMode::graph()) + .unwrap(); + + assert_eq!(indexed.results.len(), indexed.stats.result_count as usize); + assert!(indexed.results.iter().zip(&legacy.results).all(|(a, b)| ( + a.vertex_id, + a.distance + ) == ( + b.vertex_id, + b.distance + ))); + let source = + read_bin::(&mut storage_provider.open_reader(TEST_DATA_FILE).unwrap()).unwrap(); + assert!(indexed + .results + .iter() + .all(|result| result.indexed_vector.as_ref() == source.row(result.vertex_id as usize))); + + let batched = searcher + .search_with_indexed_vectors(&query, 129, 129, None, SearchMode::graph()) + .unwrap(); + assert_eq!(batched.results.len(), 129); + assert!(batched + .results + .iter() + .all(|result| result.indexed_vector.as_ref() == source.row(result.vertex_id as usize))); + + let partial = searcher + .search_with_indexed_vectors( + &query, + 10, + 10, + None, + SearchMode::flat_filtered(|id| matches!(*id, 72 | 87 | 170)), + ) + .unwrap(); + assert_eq!(partial.results.len(), 3); + assert_eq!(partial.stats.result_count, 3); + + let empty = searcher + .search_with_indexed_vectors(&query, 10, 10, None, SearchMode::flat_filtered(|_| false)) + .unwrap(); + assert!(empty.results.is_empty()); + assert!(searcher + .search_with_indexed_vectors(&query, 2, 1, None, SearchMode::graph()) + .is_err()); + } + struct TestDiskSearchParams<'a, StorageType> { storage_provider: &'a StorageType, index_search_engine: &'a DiskIndexSearcher<