From 4efb698dbff723083764abc64296b41ff4602c22 Mon Sep 17 00:00:00 2001 From: "Yujie Zhang (from Dev Box)" Date: Thu, 20 Aug 2026 17:51:15 +0800 Subject: [PATCH 1/4] Return indexed vectors from disk search Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/search/provider/disk_provider.rs | 379 +++++++++++++++--- 1 file changed, 317 insertions(+), 62 deletions(-) diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 100936fbab..75966ca167 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -238,6 +238,7 @@ where { // Borrowed from `search_internal` so the strategy can be passed by value io_tracker: &'a IOTracker, + cache_indexed_vectors: bool, /// Consumed only by `default_post_processor()` → `RerankAndFilter`. /// `FlatScan` and `InlineFilter` filter earlier in their pipelines and /// pass `AcceptAll` here to avoid a redundant second pass. @@ -319,14 +320,135 @@ impl<'a> DeterminantDiversityAndFilter<'a> { } } +type IndexedVector = Option>; +type DistanceCacheEntry = + (f32, AssociatedData, IndexedVector); +type IndexedVectorOutput<'a, VectorData> = Option<&'a mut [IndexedVector]>; +type SearchPayload = (u32, AssociatedData, IndexedVector); + +struct SearchOutput<'a, A, V> { + output: search_output_buffer::IdDistanceAssociatedData<'a, u32, A>, + indexed_vectors: IndexedVectorOutput<'a, V>, +} + +impl<'a, A, V> SearchOutput<'a, A, V> { + fn new( + ids: &'a mut [u32], + distances: &'a mut [f32], + associated_data: &'a mut [A], + indexed_vectors: IndexedVectorOutput<'a, V>, + ) -> Self { + if let Some(vectors) = &indexed_vectors { + assert_eq!(ids.len(), vectors.len()); + } + Self { + output: search_output_buffer::IdDistanceAssociatedData::new( + ids, + distances, + associated_data, + ), + indexed_vectors, + } + } +} + +impl search_output_buffer::SearchOutputBuffer> + for SearchOutput<'_, A, V> +{ + fn size_hint(&self) -> Option { + search_output_buffer::SearchOutputBuffer::size_hint(&self.output) + } + + fn push( + &mut self, + neighbor: Neighbor>, + ) -> search_output_buffer::BufferState { + let position = search_output_buffer::SearchOutputBuffer::current_len(&self.output); + let ((id, data, vector), distance) = neighbor.as_tuple(); + let state = search_output_buffer::SearchOutputBuffer::push( + &mut self.output, + Neighbor::new((id, data), distance), + ); + if search_output_buffer::SearchOutputBuffer::current_len(&self.output) != position { + if let Some(vectors) = &mut self.indexed_vectors { + vectors[position] = vector; + } + } + state + } + + fn current_len(&self) -> usize { + search_output_buffer::SearchOutputBuffer::current_len(&self.output) + } + + fn extend(&mut self, iter: Iter) -> usize + where + Iter: IntoIterator>>, + { + let start = self.current_len(); + for neighbor in iter { + let previous = self.current_len(); + let state = self.push(neighbor); + if self.current_len() == previous || state.is_full() { + break; + } + } + self.current_len() - start + } +} + +fn extend_output( + accessor: &mut DiskAccessor<'_, Data, VP>, + reranked: I, + output: &mut B, +) -> ANNResult +where + Data: GraphDataType, + VP: VertexProvider, + I: IntoIterator>, + B: search_output_buffer::SearchOutputBuffer< + SearchPayload, + > + Send + + ?Sized, +{ + if !accessor.cache_indexed_vectors { + return Ok(output.extend(reranked.into_iter().map(|candidate| { + let ((id, data), distance) = candidate.as_tuple(); + Neighbor::new((id, data, None), distance) + }))); + } + + let mut count = 0; + for candidate in reranked { + if output.size_hint() == Some(0) { + break; + } + let ((id, data), distance) = candidate.as_tuple(); + let vector = match accessor + .scratch + .distance_cache + .remove(&id) + .and_then(|(_, _, vector)| vector) + { + Some(vector) => vector, + None => Box::from(accessor.scratch.vertex_provider.get_vector(&id)?), + }; + count += 1; + if output + .push(Neighbor::new((id, data, Some(vector)), distance)) + .is_full() + { + break; + } + } + Ok(count) +} + impl SearchPostProcess< DiskAccessor<'_, Data, VP>, &[Data::VectorDataType], - ( - as DataProvider>::InternalId, - Data::AssociatedDataType, - ), + SearchPayload, > for RerankAndFilter<'_> where Data: GraphDataType, @@ -342,8 +464,9 @@ where ) -> Result where I: Iterator> + Send, - B: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> - + Send + B: search_output_buffer::SearchOutputBuffer< + SearchPayload, + > + Send + ?Sized, { let provider = accessor.provider; @@ -384,7 +507,7 @@ where reranked.sort_unstable_by(neighbor::ord::fast_distance); // Store the reranked results. - Ok(output.extend(reranked)) + extend_output(accessor, reranked, output) } } @@ -392,10 +515,7 @@ impl SearchPostProcess< DiskAccessor<'_, Data, VP>, &[Data::VectorDataType], - ( - as DataProvider>::InternalId, - Data::AssociatedDataType, - ), + SearchPayload, > for DeterminantDiversityAndFilter<'_> where Data: GraphDataType, @@ -411,8 +531,9 @@ where ) -> Result where I: Iterator> + Send, - B: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> - + Send + B: search_output_buffer::SearchOutputBuffer< + SearchPayload, + > + Send + ?Sized, { let provider = accessor.provider; @@ -459,11 +580,15 @@ where &self.params, )?; - Ok(output.extend(reranked.into_iter().map(|idx| { - let id = candidate_ids[idx]; - let distance = candidate_distances[idx]; - Neighbor::new((id, associated_data[idx]), distance) - }))) + extend_output( + accessor, + reranked.into_iter().map(|idx| { + let id = candidate_ids[idx]; + let distance = candidate_distances[idx]; + Neighbor::new((id, associated_data[idx]), distance) + }), + output, + ) } } @@ -471,10 +596,7 @@ impl SearchPostProcess< DiskAccessor<'_, Data, VP>, &[Data::VectorDataType], - ( - as DataProvider>::InternalId, - Data::AssociatedDataType, - ), + SearchPayload, > for DiskSearchPostProcessor<'_> where Data: GraphDataType, @@ -490,8 +612,9 @@ where ) -> Result where I: Iterator> + Send, - B: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> - + Send + B: search_output_buffer::SearchOutputBuffer< + SearchPayload, + > + Send + ?Sized, { match self { @@ -527,6 +650,7 @@ where query, self.vertex_provider_factory, self.scratch_pool, + self.cache_indexed_vectors, ) } } @@ -536,10 +660,7 @@ impl<'this, Data, ProviderFactory> 'this, DiskProvider, &'this [Data::VectorDataType], - ( - as DataProvider>::InternalId, - Data::AssociatedDataType, - ), + SearchPayload, > for DiskSearchStrategy<'this, Data, ProviderFactory> where Data: GraphDataType, @@ -559,7 +680,8 @@ where Data: GraphDataType, VP: VertexProvider, { - distance_cache: HashMap, + distance_cache: + HashMap>, pq_scratch: PQScratch, vertex_provider: VP, } @@ -621,6 +743,7 @@ where io_tracker: &'a IOTracker, scratch: PoolOption>, query: &'a [Data::VectorDataType], + cache_indexed_vectors: bool, } impl DiskAccessor<'_, Data, VP> @@ -732,6 +855,7 @@ where query: &'a [Data::VectorDataType], vertex_provider_factory: &'a VPF, scratch_pool: &'a Arc>>, + cache_indexed_vectors: bool, ) -> ANNResult where VPF: VertexProviderFactory, @@ -770,6 +894,7 @@ where io_tracker, scratch, query, + cache_indexed_vectors, }) } @@ -786,14 +911,16 @@ where ); self.io_tracker.add_io_count(ids.len()); for id in ids { + let vector = scratch.vertex_provider.get_vector(id)?; let distance = self .provider .distance_comparer - .evaluate_similarity(self.query, scratch.vertex_provider.get_vector(id)?); + .evaluate_similarity(self.query, vector); let associated_data = *scratch.vertex_provider.get_associated_data(id)?; + let indexed_vector = self.cache_indexed_vectors.then(|| Box::from(vector)); scratch .distance_cache - .insert(*id, (distance, associated_data)); + .insert(*id, (distance, associated_data, indexed_vector)); } Ok(()) } @@ -850,6 +977,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, @@ -920,14 +1059,15 @@ where }) } - /// Helper method to create a `DiskSearchStrategy` with common parameters. fn search_strategy<'a>( &'a self, io_tracker: &'a IOTracker, postprocess_filter: PostprocessStrategy<'a>, + cache_indexed_vectors: bool, ) -> DiskSearchStrategy<'a, Data, ProviderFactory> { DiskSearchStrategy { io_tracker, + cache_indexed_vectors, postprocess_filter, vertex_provider_factory: &self.vertex_provider_factory, scratch_pool: &self.scratch_pool, @@ -951,7 +1091,9 @@ where output: &mut OB, ) -> ANNResult where - OB: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> + Send, + OB: search_output_buffer::SearchOutputBuffer< + SearchPayload, + > + Send, { let provider = self.index.provider(); let mut accessor = strategy @@ -1033,7 +1175,9 @@ where output: &mut OB, ) -> ANNResult where - OB: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> + Send, + OB: search_output_buffer::SearchOutputBuffer< + SearchPayload, + > + Send, { let filtered_strategy = labeled::Filtered::new(strategy, label_provider); let search = InlineFilterSearch::new(knn, adaptive_l); @@ -1095,6 +1239,61 @@ where Ok(search_result) } + 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 result_count = return_list_size as usize; + let mut query_stats = QueryStatistics::default(); + let mut indices = vec![0u32; result_count]; + let mut distances = vec![0f32; result_count]; + let mut associated_data = vec![Data::AssociatedDataType::default(); result_count]; + let mut indexed_vectors = std::iter::repeat_with(|| None) + .take(result_count) + .collect::>(); + + let stats = self.search_internal_impl( + query, + result_count, + search_list_size, + beam_width, + &mut query_stats, + &mut indices, + &mut distances, + &mut associated_data, + Some(&mut indexed_vectors), + &mode, + )?; + + let results = indices + .into_iter() + .zip(distances) + .zip(associated_data) + .zip(indexed_vectors) + .take(stats.result_count as usize) + .map(|(((vertex_id, distance), data), indexed_vector)| { + Ok(SearchResultItemWithIndexedVector { + vertex_id, + data, + distance, + indexed_vector: indexed_vector.ok_or_else(|| { + diskann_error!( + ErrorKind::IndexError, + "missing indexed vector for vertex {vertex_id}" + ) + })?, + }) + }) + .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)] @@ -1109,6 +1308,34 @@ where distances: &mut [f32], associated_data: &mut [Data::AssociatedDataType], mode: &SearchMode<'_>, + ) -> ANNResult { + self.search_internal_impl( + query, + k_value, + search_list_size, + beam_width, + query_stats, + indices, + distances, + associated_data, + None, + mode, + ) + } + + #[allow(clippy::too_many_arguments)] + fn search_internal_impl( + &self, + query: &[Data::VectorDataType], + k_value: usize, + search_list_size: u32, + beam_width: Option, + query_stats: &mut QueryStatistics, + indices: &mut [u32], + distances: &mut [f32], + associated_data: &mut [Data::AssociatedDataType], + indexed_vectors: IndexedVectorOutput<'_, Data::VectorDataType>, + mode: &SearchMode<'_>, ) -> ANNResult { let l = search_list_size as usize; @@ -1119,10 +1346,12 @@ where )); } - let mut result_output_buffer = search_output_buffer::IdDistanceAssociatedData::new( + let cache_indexed_vectors = indexed_vectors.is_some(); + let mut result_output_buffer = SearchOutput::new( &mut indices[..k_value], &mut distances[..k_value], &mut associated_data[..k_value], + indexed_vectors.map(|vectors| &mut vectors[..k_value]), ); let timer = Instant::now(); @@ -1140,7 +1369,11 @@ where // as the post-processor over the L candidate pool. let stats = match mode { SearchMode::FlatScan { filter } => { - let strategy = self.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll); + let strategy = self.search_strategy( + &io_tracker, + PostprocessStrategy::AcceptAll, + cache_indexed_vectors, + ); self.runtime.block_on(self.flat_search( &strategy, query, @@ -1155,6 +1388,7 @@ where filter .as_deref() .map_or(PostprocessStrategy::AcceptAll, PostprocessStrategy::Apply), + cache_indexed_vectors, ); let knn_search = Knn::new(l, beam_width) .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?; @@ -1170,7 +1404,11 @@ where // Strategy is passed by value into `filter_search` so that the // `labeled::Filtered` wrapper can own it; `io_tracker` keeps // its counters reachable from this scope. - let strategy = self.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll); + let strategy = self.search_strategy( + &io_tracker, + PostprocessStrategy::AcceptAll, + cache_indexed_vectors, + ); let knn_search = Knn::new(l, beam_width)?; self.runtime.block_on(self.filter_search( strategy, @@ -1188,7 +1426,8 @@ where let postprocess_config = filter .as_deref() .map_or(PostprocessStrategy::AcceptAll, PostprocessStrategy::Apply); - let strategy = self.search_strategy(&io_tracker, postprocess_config); + let strategy = + self.search_strategy(&io_tracker, postprocess_config, cache_indexed_vectors); let knn_search = Knn::new(l, beam_width)?; let processor = DiskSearchPostProcessor::DeterminantDiversity( DeterminantDiversityAndFilter::new(postprocess_config, *params), @@ -1372,6 +1611,30 @@ mod disk_provider_tests { k: 10, l: 20, }); + + let query = vec![0.1f32; 128]; + let plain = search_engine + .search(&query, 10, 20, None, SearchMode::graph()) + .unwrap(); + let indexed = search_engine + .search_with_indexed_vectors(&query, 10, 20, None, SearchMode::graph()) + .unwrap(); + let source = + read_bin::(&mut storage_provider.open_reader(TEST_DATA_FILE).unwrap()).unwrap(); + + assert!(plain + .results + .iter() + .zip(&indexed.results) + .all(|(left, right)| (left.vertex_id, left.distance) + == (right.vertex_id, right.distance))); + assert_eq!(indexed.results.len(), indexed.stats.result_count as usize); + for item in &indexed.results { + assert_eq!( + item.indexed_vector.as_ref(), + source.row(item.vertex_id as usize) + ); + } } fn get_truth_associated_data( @@ -1818,13 +2081,11 @@ mod disk_provider_tests { let mut distances = vec![0f32; k]; let mut associated_data = vec![(); k]; - let mut result_output_buffer = search_output_buffer::IdDistanceAssociatedData::new( - &mut indices, - &mut distances, - &mut associated_data, - ); + let mut result_output_buffer = + SearchOutput::new(&mut indices, &mut distances, &mut associated_data, None); let io_tracker = IOTracker::default(); - let strategy = search_engine.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll); + let strategy = + search_engine.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll, false); let mut search_record = VisitedSearchRecord::new(0); let search_params = Knn::new(10, Some(4)).unwrap(); let recorded_search = @@ -2054,13 +2315,11 @@ mod disk_provider_tests { let mut distances = vec![0f32; original_k]; let mut associated_data = vec![(); original_k]; - let mut result_output_buffer = search_output_buffer::IdDistanceAssociatedData::new( - &mut indices, - &mut distances, - &mut associated_data, - ); + let mut result_output_buffer = + SearchOutput::new(&mut indices, &mut distances, &mut associated_data, None); let io_tracker = IOTracker::default(); - let strategy = search_engine.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll); + let strategy = + search_engine.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll, false); // Create diverse search parameters with attribute provider let diverse_params = DiverseSearchParams::new( @@ -2106,13 +2365,11 @@ mod disk_provider_tests { let mut indices2 = vec![0u32; original_k]; let mut distances2 = vec![0f32; original_k]; let mut associated_data2 = vec![(); original_k]; - let mut result_output_buffer2 = search_output_buffer::IdDistanceAssociatedData::new( - &mut indices2, - &mut distances2, - &mut associated_data2, - ); + let mut result_output_buffer2 = + SearchOutput::new(&mut indices2, &mut distances2, &mut associated_data2, None); let io_tracker2 = IOTracker::default(); - let strategy2 = search_engine.search_strategy(&io_tracker2, PostprocessStrategy::AcceptAll); + let strategy2 = + search_engine.search_strategy(&io_tracker2, PostprocessStrategy::AcceptAll, false); let search_params2 = Knn::new(search_list_size as usize, None).unwrap(); let diverse_search2 = @@ -2520,14 +2777,12 @@ mod disk_provider_tests { let mut distances = vec![0f32; k]; let mut associated_data = vec![(); k]; - let mut result_output_buffer = search_output_buffer::IdDistanceAssociatedData::new( - &mut indices, - &mut distances, - &mut associated_data, - ); + let mut result_output_buffer = + SearchOutput::new(&mut indices, &mut distances, &mut associated_data, None); let io_tracker = IOTracker::default(); - let strategy = search_engine.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll); + let strategy = + search_engine.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll, false); let mut search_record = VisitedSearchRecord::new(0); let search_params = Knn::new(10, Some(4)).unwrap(); From 4d6d5444ed3c74073086023fb22169ddc5d880ba Mon Sep 17 00:00:00 2001 From: "Yujie Zhang (from Dev Box)" Date: Fri, 21 Aug 2026 00:32:37 +0800 Subject: [PATCH 2/4] Release indexed-vector cache after each search Co-Authored-By: Claude Opus 4.8 (1M context) --- diskann-disk/src/search/provider/disk_provider.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 75966ca167..15d50883b0 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -926,6 +926,16 @@ where } } +impl Drop for DiskAccessor<'_, Data, VP> +where + Data: GraphDataType, + VP: VertexProvider, +{ + fn drop(&mut self) { + self.scratch.distance_cache.clear(); + } +} + /// [`DiskIndexSearcher`] is a helper class to make it easy to construct index /// and do repeated search operations. It is a wrapper around the index. /// This is useful for drivers such as search_disk_index.exe in tools. From aa952c2845e14b6e1fde4a7b522a572d76fef396 Mon Sep 17 00:00:00 2001 From: "Yujie Zhang (from Dev Box)" Date: Fri, 21 Aug 2026 17:12:17 +0800 Subject: [PATCH 3/4] Address indexed vector review feedback Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/search/provider/disk_provider.rs | 227 +++++++++++++----- 1 file changed, 167 insertions(+), 60 deletions(-) diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 15d50883b0..8749e73394 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -21,7 +21,8 @@ use diskann::{ ext::labeled::{self, QueryLabelProvider}, glue::{self, DefaultPostProcessor, SearchPostProcess, SearchStrategy}, search::{AdaptiveL, InlineFilterSearch, Knn}, - search_output_buffer, DiskANNIndex, + search_output_buffer::{self, BufferState, IdDistanceAssociatedData}, + DiskANNIndex, }, neighbor::{self, Neighbor, NeighborPriorityQueue}, provider::{DataProvider, DefaultContext, HasId, NoopGuard}, @@ -327,7 +328,7 @@ type IndexedVectorOutput<'a, VectorData> = Option<&'a mut [IndexedVector = (u32, AssociatedData, IndexedVector); struct SearchOutput<'a, A, V> { - output: search_output_buffer::IdDistanceAssociatedData<'a, u32, A>, + output: IdDistanceAssociatedData<'a, u32, A>, indexed_vectors: IndexedVectorOutput<'a, V>, } @@ -337,18 +338,20 @@ impl<'a, A, V> SearchOutput<'a, A, V> { distances: &'a mut [f32], associated_data: &'a mut [A], indexed_vectors: IndexedVectorOutput<'a, V>, - ) -> Self { - if let Some(vectors) = &indexed_vectors { - assert_eq!(ids.len(), vectors.len()); + ) -> ANNResult { + if indexed_vectors + .as_ref() + .is_some_and(|vectors| ids.len() != vectors.len()) + { + return Err(diskann_error!( + ErrorKind::IndexError, + "ids and indexed_vectors must have the same length", + )); } - Self { - output: search_output_buffer::IdDistanceAssociatedData::new( - ids, - distances, - associated_data, - ), + Ok(Self { + output: IdDistanceAssociatedData::new(ids, distances, associated_data), indexed_vectors, - } + }) } } @@ -359,20 +362,19 @@ impl search_output_buffer::SearchOutputBuffer>, - ) -> search_output_buffer::BufferState { - let position = search_output_buffer::SearchOutputBuffer::current_len(&self.output); + fn push(&mut self, neighbor: Neighbor>) -> BufferState { + let position = self.current_len(); + if position == self.output.capacity() { + return BufferState::Full; + } + let ((id, data, vector), distance) = neighbor.as_tuple(); let state = search_output_buffer::SearchOutputBuffer::push( &mut self.output, Neighbor::new((id, data), distance), ); - if search_output_buffer::SearchOutputBuffer::current_len(&self.output) != position { - if let Some(vectors) = &mut self.indexed_vectors { - vectors[position] = vector; - } + if let Some(vectors) = &mut self.indexed_vectors { + vectors[position] = vector; } state } @@ -387,9 +389,7 @@ impl search_output_buffer::SearchOutputBuffer { pub distance: f32, } -pub struct SearchResultWithIndexedVectors { +pub struct SearchResultWithVectors { pub results: Vec>, pub stats: SearchResultStats, } @@ -1256,8 +1256,7 @@ where search_list_size: u32, beam_width: Option, mode: SearchMode<'_>, - ) -> ANNResult> - { + ) -> ANNResult> { let result_count = return_list_size as usize; let mut query_stats = QueryStatistics::default(); let mut indices = vec![0u32; result_count]; @@ -1301,7 +1300,7 @@ where }) .collect::>>()?; - Ok(SearchResultWithIndexedVectors { results, stats }) + Ok(SearchResultWithVectors { results, stats }) } /// Perform a raw search on the disk index. @@ -1362,7 +1361,7 @@ where &mut distances[..k_value], &mut associated_data[..k_value], indexed_vectors.map(|vectors| &mut vectors[..k_value]), - ); + )?; let timer = Instant::now(); @@ -1499,7 +1498,7 @@ mod disk_provider_tests { DynWriteProvider, StorageReadProvider, VirtualStorageProvider, }; use diskann_providers::utils::{create_thread_pool, PQPathNames, ParallelIteratorInPool}; - use diskann_utils::{io::read_bin, test_data_root}; + use diskann_utils::{io::read_bin, test_data_root, views::Matrix}; use diskann_vector::distance::Metric; use rayon::prelude::IndexedParallelIterator; use rstest::rstest; @@ -1566,6 +1565,7 @@ mod disk_provider_tests { thread_num: 1, query_file_path: TEST_QUERY_10PTS_100DIM, truth_result_file_path: TEST_TRUTH_RESULT_10PTS_100DIM, + source_data_file_path: None, k: 10, l: 20, }); @@ -1576,6 +1576,7 @@ mod disk_provider_tests { thread_num: 5, query_file_path: TEST_QUERY_10PTS_100DIM, truth_result_file_path: TEST_TRUTH_RESULT_10PTS_100DIM, + source_data_file_path: None, k: 10, l: 20, }); @@ -1608,6 +1609,7 @@ mod disk_provider_tests { thread_num: 1, query_file_path: TEST_QUERY_10PTS_128DIM, truth_result_file_path: TEST_TRUTH_RESULT_10PTS_128DIM, + source_data_file_path: Some(TEST_DATA_FILE), k: 10, l: 20, }); @@ -1618,33 +1620,10 @@ mod disk_provider_tests { thread_num: 5, query_file_path: TEST_QUERY_10PTS_128DIM, truth_result_file_path: TEST_TRUTH_RESULT_10PTS_128DIM, + source_data_file_path: Some(TEST_DATA_FILE), k: 10, l: 20, }); - - let query = vec![0.1f32; 128]; - let plain = search_engine - .search(&query, 10, 20, None, SearchMode::graph()) - .unwrap(); - let indexed = search_engine - .search_with_indexed_vectors(&query, 10, 20, None, SearchMode::graph()) - .unwrap(); - let source = - read_bin::(&mut storage_provider.open_reader(TEST_DATA_FILE).unwrap()).unwrap(); - - assert!(plain - .results - .iter() - .zip(&indexed.results) - .all(|(left, right)| (left.vertex_id, left.distance) - == (right.vertex_id, right.distance))); - assert_eq!(indexed.results.len(), indexed.stats.result_count as usize); - for item in &indexed.results { - assert_eq!( - item.indexed_vector.as_ref(), - source.row(item.vertex_id as usize) - ); - } } fn get_truth_associated_data( @@ -1808,6 +1787,58 @@ mod disk_provider_tests { result.into_inner().into_vec() } + fn load_source_data( + storage_provider: &StorageReader, + path: &str, + ) -> Matrix { + read_bin(&mut storage_provider.open_reader(path).unwrap()).unwrap() + } + + fn assert_indexed_results_match( + indexed: &SearchResultWithVectors<(), f32>, + expected_result_count: u32, + expected_io_operations: u32, + expected_results: impl IntoIterator, + source: Option<&Matrix>, + vector_dimension: usize, + ) { + assert_eq!(indexed.stats.result_count, expected_result_count); + assert_eq!(indexed.results.len(), expected_result_count as usize); + assert_eq!( + indexed.stats.query_statistics.total_io_operations, + expected_io_operations + ); + for (actual, expected) in indexed.results.iter().zip(expected_results) { + assert_eq!((actual.vertex_id, actual.distance), expected); + if let Some(source) = source { + assert_eq!( + actual.indexed_vector.as_ref(), + source.row(actual.vertex_id as usize) + ); + } else { + assert_eq!(actual.indexed_vector.len(), vector_dimension); + } + } + } + + fn assert_indexed_search_results_match( + result: &SearchResult<()>, + indexed: &SearchResultWithVectors<(), f32>, + source: &Matrix, + ) { + assert_indexed_results_match( + indexed, + result.stats.result_count, + result.stats.query_statistics.total_io_operations, + result + .results + .iter() + .map(|item| (item.vertex_id, item.distance)), + Some(source), + source.ncols(), + ); + } + struct TestDiskSearchParams<'a, StorageType> { storage_provider: &'a StorageType, index_search_engine: &'a DiskIndexSearcher< @@ -1820,6 +1851,7 @@ mod disk_provider_tests { thread_num: u64, query_file_path: &'a str, truth_result_file_path: &'a str, + source_data_file_path: Option<&'a str>, k: usize, l: usize, } @@ -1852,6 +1884,9 @@ mod disk_provider_tests { .unwrap(); let truth_result = load_query_result(params.storage_provider, params.truth_result_file_path); + let source = params + .source_data_file_path + .map(|path| load_source_data(params.storage_provider, path)); let pool = create_thread_pool(params.thread_num.into_usize()).unwrap(); queries @@ -1890,6 +1925,25 @@ mod disk_provider_tests { "Expected vertices loaded to be greater than 0" ); + let indexed = params + .index_search_engine + .search_with_indexed_vectors( + query, + params.k as u32, + params.l as u32, + None, + SearchMode::graph(), + ) + .unwrap(); + assert_indexed_results_match( + &indexed, + result_unwrapped.result_count, + result_unwrapped.query_statistics.total_io_operations, + indices.iter().copied().zip(distances.iter().copied()), + source.as_ref(), + query.len(), + ); + // Compare res with truth_slice using assert_eq! assert_eq!( indices, truth_slice, @@ -2067,6 +2121,17 @@ mod disk_provider_tests { Err(_) => {} Ok(_) => panic!("Expected error when search_list_size < return_list_size"), } + + match search_engine.search_with_indexed_vectors( + &query, + return_list_size, + search_list_size, + None, + SearchMode::graph(), + ) { + Err(error) => assert_eq!(error_kind(&error), ErrorKind::IndexError), + Ok(_) => panic!("Expected error when search_list_size < return_list_size"), + } } #[test] @@ -2092,7 +2157,7 @@ mod disk_provider_tests { let mut associated_data = vec![(); k]; let mut result_output_buffer = - SearchOutput::new(&mut indices, &mut distances, &mut associated_data, None); + SearchOutput::new(&mut indices, &mut distances, &mut associated_data, None).unwrap(); let io_tracker = IOTracker::default(); let strategy = search_engine.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll, false); @@ -2193,6 +2258,18 @@ mod disk_provider_tests { SearchMode::diverse_graph(params), ) .unwrap(); + let indexed = search_engine + .search_with_indexed_vectors( + &query_vector, + return_list_size, + search_list_size, + Some(4), + SearchMode::diverse_graph(params), + ) + .unwrap(); + let source = load_source_data(storage_provider.as_ref(), TEST_DATA_FILE); + assert_indexed_search_results_match(&result, &indexed, &source); + let det_div_ids: Vec = result.results.iter().map(|r| r.vertex_id).collect(); assert_eq!( @@ -2326,7 +2403,7 @@ mod disk_provider_tests { let mut associated_data = vec![(); original_k]; let mut result_output_buffer = - SearchOutput::new(&mut indices, &mut distances, &mut associated_data, None); + SearchOutput::new(&mut indices, &mut distances, &mut associated_data, None).unwrap(); let io_tracker = IOTracker::default(); let strategy = search_engine.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll, false); @@ -2376,7 +2453,7 @@ mod disk_provider_tests { let mut distances2 = vec![0f32; original_k]; let mut associated_data2 = vec![(); original_k]; let mut result_output_buffer2 = - SearchOutput::new(&mut indices2, &mut distances2, &mut associated_data2, None); + SearchOutput::new(&mut indices2, &mut distances2, &mut associated_data2, None).unwrap(); let io_tracker2 = IOTracker::default(); let strategy2 = search_engine.search_strategy(&io_tracker2, PostprocessStrategy::AcceptAll, false); @@ -2493,6 +2570,7 @@ mod disk_provider_tests { 10, vec![152, 118, 72, 170, 87, 141, 79, 207, 124, 86], vec![256101.7, 256675.3, 256709.69, 256712.5, 256760.08, 256958.5, 257006.1, 257025.7, 257105.67, 257107.67], + false, )] // This case validates post-filtering using 2 ids which are not present in the unfiltered result set. // It is expected that the post-filtering will return an empty result @@ -2502,24 +2580,29 @@ mod disk_provider_tests { 0, vec![0; 10], vec![0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + false, )] // This case validates pre-filtering using 2 ids which are not present in the unfiltered result set. - // It is expected that the pre-filtering will do search over matching ids + // It is expected that the pre-filtering will do search over matching ids. + // It also verifies indexed-vector output for FlatScan. #[case( |id: &u32| *id == 0 || *id == 1, true, 2, vec![1, 0, 0, 0, 0, 0, 0, 0, 0, 0], vec![257247.28, 258179.28, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + true, )] // This case validates post-filtering using 3 ids from the unfiltered result set. - // It is expected that the post-filtering will filter out non-matching ids + // It is expected that the post-filtering will filter out non-matching ids. + // It also verifies indexed-vector output and partial-result alignment for Graph. #[case( |id: &u32| *id == 72 || *id == 87 || *id == 170, false, 3, vec![72, 170, 87, 0, 0, 0, 0, 0, 0, 0], vec![256709.69, 256712.5, 256760.08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + true, )] // This case validates pre-filtering using 3 ids from the unfiltered result set. // It is expected that the pre-filtering will do search over matching ids @@ -2529,6 +2612,7 @@ mod disk_provider_tests { 3, vec![72, 170, 87, 0, 0, 0, 0, 0, 0, 0], vec![256709.69, 256712.5, 256760.08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + false, )] fn test_search_with_vector_filter( #[case] vector_filter: fn(&u32) -> bool, @@ -2536,6 +2620,7 @@ mod disk_provider_tests { #[case] expected_result_count: u32, #[case] expected_indices: Vec, #[case] expected_distances: Vec, + #[case] check_indexed_vectors: bool, ) { // Exact distances can vary slightly depending on the architecture used // to compute distances due to different unrolling strategies and SIMD widthd. @@ -2641,6 +2726,14 @@ mod disk_provider_tests { check_distances(&actual_distances, &expected_distances), "Expected distances to match" ); + + if check_indexed_vectors { + let indexed = search_engine + .search_with_indexed_vectors(&query, 10, 10, None, make_mode()) + .unwrap(); + let source = load_source_data(storage_provider.as_ref(), TEST_DATA_FILE); + assert_indexed_search_results_match(&result_with_filter_unwrapped, &indexed, &source); + } } // =========================================================================== @@ -2740,6 +2833,20 @@ mod disk_provider_tests { ), ) .expect("inline filter search with AdaptiveL must succeed"); + let indexed = search_engine + .search_with_indexed_vectors( + &query, + 10, + 10, + None, + SearchMode::inline_filter( + predicate, + Some(AdaptiveL::new(5, 16.0).expect("valid AdaptiveL")), + ), + ) + .expect("indexed-vector inline filter search with AdaptiveL must succeed"); + let source = load_source_data(storage_provider.as_ref(), TEST_DATA_FILE); + assert_indexed_search_results_match(&result, &indexed, &source); // `result.results` is pre-allocated to `return_list_size`; only the // first `result_count` entries are populated. The trailing entries @@ -2788,7 +2895,7 @@ mod disk_provider_tests { let mut associated_data = vec![(); k]; let mut result_output_buffer = - SearchOutput::new(&mut indices, &mut distances, &mut associated_data, None); + SearchOutput::new(&mut indices, &mut distances, &mut associated_data, None).unwrap(); let io_tracker = IOTracker::default(); let strategy = From e554e3f768e03d939fa7956b5bbb3f27b694bee5 Mon Sep 17 00:00:00 2001 From: "Yujie Zhang (from Dev Box)" Date: Fri, 21 Aug 2026 18:50:03 +0800 Subject: [PATCH 4/4] Document indexed vector API and bump version Co-Authored-By: Claude Opus 4.8 (1M context) --- Cargo.lock | 34 +++++++++---------- Cargo.toml | 32 ++++++++--------- .../src/search/provider/disk_provider.rs | 6 ++++ 3 files changed, 39 insertions(+), 33 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2588eee92a..9789e45509 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -433,7 +433,7 @@ dependencies = [ [[package]] name = "diskann" -version = "0.55.0" +version = "0.56.0" dependencies = [ "anyhow", "bytemuck", @@ -457,7 +457,7 @@ dependencies = [ [[package]] name = "diskann-benchmark" -version = "0.55.0" +version = "0.56.0" dependencies = [ "anyhow", "bf-tree", @@ -496,7 +496,7 @@ dependencies = [ [[package]] name = "diskann-benchmark-core" -version = "0.55.0" +version = "0.56.0" dependencies = [ "anyhow", "diskann", @@ -513,7 +513,7 @@ dependencies = [ [[package]] name = "diskann-benchmark-runner" -version = "0.55.0" +version = "0.56.0" dependencies = [ "anyhow", "clap", @@ -527,7 +527,7 @@ dependencies = [ [[package]] name = "diskann-benchmark-simd" -version = "0.55.0" +version = "0.56.0" dependencies = [ "anyhow", "diskann-benchmark-runner", @@ -544,7 +544,7 @@ dependencies = [ [[package]] name = "diskann-bftree" -version = "0.55.0" +version = "0.56.0" dependencies = [ "bf-tree", "bytemuck", @@ -566,7 +566,7 @@ dependencies = [ [[package]] name = "diskann-disk" -version = "0.55.0" +version = "0.56.0" dependencies = [ "anyhow", "approx", @@ -621,7 +621,7 @@ dependencies = [ [[package]] name = "diskann-inmem" -version = "0.55.0" +version = "0.56.0" dependencies = [ "anyhow", "bytemuck", @@ -645,7 +645,7 @@ dependencies = [ [[package]] name = "diskann-label-filter" -version = "0.55.0" +version = "0.56.0" dependencies = [ "anyhow", "bf-tree", @@ -668,7 +668,7 @@ dependencies = [ [[package]] name = "diskann-linalg" -version = "0.55.0" +version = "0.56.0" dependencies = [ "approx", "cfg-if", @@ -682,7 +682,7 @@ dependencies = [ [[package]] name = "diskann-providers" -version = "0.55.0" +version = "0.56.0" dependencies = [ "approx", "arc-swap", @@ -716,7 +716,7 @@ dependencies = [ [[package]] name = "diskann-quantization" -version = "0.55.0" +version = "0.56.0" dependencies = [ "bytemuck", "cfg-if", @@ -737,7 +737,7 @@ dependencies = [ [[package]] name = "diskann-record" -version = "0.55.0" +version = "0.56.0" dependencies = [ "anyhow", "serde", @@ -747,7 +747,7 @@ dependencies = [ [[package]] name = "diskann-tools" -version = "0.55.0" +version = "0.56.0" dependencies = [ "anyhow", "bincode", @@ -780,7 +780,7 @@ dependencies = [ [[package]] name = "diskann-utils" -version = "0.55.0" +version = "0.56.0" dependencies = [ "bytemuck", "cfg-if", @@ -796,7 +796,7 @@ dependencies = [ [[package]] name = "diskann-vector" -version = "0.55.0" +version = "0.56.0" dependencies = [ "approx", "bytemuck", @@ -811,7 +811,7 @@ dependencies = [ [[package]] name = "diskann-wide" -version = "0.55.0" +version = "0.56.0" dependencies = [ "bytemuck", "cfg-if", diff --git a/Cargo.toml b/Cargo.toml index 73d7d2d611..59e02da4ba 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,7 +40,7 @@ default-members = [ exclude = ["vectorset"] [workspace.package] -version = "0.55.0" # Obeying semver +version = "0.56.0" # Obeying semver description = "DiskANN3 is a composable library for bringing scalable, accurate and cost-effective vector indexing to multiple databases." authors = ["Microsoft"] repository = "https://github.com/microsoft/DiskANN" @@ -52,24 +52,24 @@ undocumented_unsafe_blocks = "warn" [workspace.dependencies] # Base And Numerics -diskann-wide = { path = "diskann-wide", version = "0.55.0" } -diskann-vector = { path = "diskann-vector", version = "0.55.0" } -diskann-linalg = { path = "diskann-linalg", version = "0.55.0" } -diskann-utils = { path = "diskann-utils", default-features = false, version = "0.55.0" } -diskann-quantization = { path = "diskann-quantization", default-features = false, version = "0.55.0" } +diskann-wide = { path = "diskann-wide", version = "0.56.0" } +diskann-vector = { path = "diskann-vector", version = "0.56.0" } +diskann-linalg = { path = "diskann-linalg", version = "0.56.0" } +diskann-utils = { path = "diskann-utils", default-features = false, version = "0.56.0" } +diskann-quantization = { path = "diskann-quantization", default-features = false, version = "0.56.0" } # Algorithm -diskann = { path = "diskann", version = "0.55.0" } +diskann = { path = "diskann", version = "0.56.0" } # Providers -diskann-providers = { path = "diskann-providers", default-features = false, version = "0.55.0" } -diskann-inmem = { path = "diskann-inmem", default-features = false, version = "0.55.0" } -diskann-disk = { path = "diskann-disk", version = "0.55.0" } -diskann-label-filter = { path = "diskann-label-filter", version = "0.55.0" } +diskann-providers = { path = "diskann-providers", default-features = false, version = "0.56.0" } +diskann-inmem = { path = "diskann-inmem", default-features = false, version = "0.56.0" } +diskann-disk = { path = "diskann-disk", version = "0.56.0" } +diskann-label-filter = { path = "diskann-label-filter", version = "0.56.0" } # Infra -diskann-benchmark-runner = { path = "diskann-benchmark-runner", version = "0.55.0" } -diskann-benchmark-core = { path = "diskann-benchmark-core", version = "0.55.0" } -diskann-tools = { path = "diskann-tools", version = "0.55.0" } -diskann-bftree = { path = "diskann-bftree", version = "0.55.0" } -diskann-record = { path = "diskann-record", version = "0.55.0" } +diskann-benchmark-runner = { path = "diskann-benchmark-runner", version = "0.56.0" } +diskann-benchmark-core = { path = "diskann-benchmark-core", version = "0.56.0" } +diskann-tools = { path = "diskann-tools", version = "0.56.0" } +diskann-bftree = { path = "diskann-bftree", version = "0.56.0" } +diskann-record = { path = "diskann-record", version = "0.56.0" } # External dependencies (shared versions) anyhow = "1.0.98" diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 8749e73394..37a3171dfa 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -987,11 +987,13 @@ pub struct SearchResultItem { pub distance: f32, } +/// Search results that include the native indexed vector for each neighbor. pub struct SearchResultWithVectors { pub results: Vec>, pub stats: SearchResultStats, } +/// A nearest-neighbor result with its native indexed vector. pub struct SearchResultItemWithIndexedVector { pub vertex_id: u32, pub data: AssociatedData, @@ -1249,6 +1251,10 @@ where Ok(search_result) } + /// Perform a search on the disk index and return each result with its native indexed vector. + /// + /// Unlike [`Self::search`], this returns only valid results, so the result list length equals + /// `stats.result_count`. pub fn search_with_indexed_vectors( &self, query: &[Data::VectorDataType],