From bcf01c1b1fc4a09dadb079975d0b33e9fc0a128a Mon Sep 17 00:00:00 2001 From: "Yujie Zhang (from Dev Box)" Date: Thu, 20 Aug 2026 23:45:05 +0800 Subject: [PATCH 1/2] Return indexed vectors after disk search Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/search/provider/disk_provider.rs | 118 ++++++++++++++++++ 1 file changed, 118 insertions(+) diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 100936fbab..9a4f5d8874 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,51 @@ 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 ids = results + .iter() + .map(|result| result.vertex_id) + .collect::>(); + let graph_header = &self.index.provider().graph_header; + let mut vertex_provider = self + .vertex_provider_factory + .create_vertex_provider(ids.len(), graph_header)?; + ensure_vertex_loaded(&mut vertex_provider, &ids)?; + + let results = results + .into_iter() + .map(|result| { + Ok(SearchResultItemWithIndexedVector { + indexed_vector: Box::from(vertex_provider.get_vector(&result.vertex_id)?), + 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 +1592,67 @@ 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 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< From 8642d25627df46cf0efec1b37451738b93a76de0 Mon Sep 17 00:00:00 2001 From: "Yujie Zhang (from Dev Box)" Date: Fri, 21 Aug 2026 00:42:16 +0800 Subject: [PATCH 2/2] Reuse pooled storage for final vectors Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/search/provider/disk_provider.rs | 57 +++++++++++++------ 1 file changed, 41 insertions(+), 16 deletions(-) diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 9a4f5d8874..6628b526d5 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -1127,28 +1127,44 @@ where }); } - let ids = results - .iter() - .map(|result| result.vertex_id) - .collect::>(); - let graph_header = &self.index.provider().graph_header; - let mut vertex_provider = self - .vertex_provider_factory - .create_vertex_provider(ids.len(), graph_header)?; - ensure_vertex_loaded(&mut vertex_provider, &ids)?; + 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() - .map(|result| { - Ok(SearchResultItemWithIndexedVector { - indexed_vector: Box::from(vertex_provider.get_vector(&result.vertex_id)?), + .zip(indexed_vectors) + .map( + |(result, indexed_vector)| SearchResultItemWithIndexedVector { + indexed_vector, vertex_id: result.vertex_id, data: result.data, distance: result.distance, - }) - }) - .collect::>()?; - + }, + ) + .collect(); Ok(SearchResultWithIndexedVectors { results, stats }) } @@ -1632,6 +1648,15 @@ mod disk_provider_tests { .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,