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 e533165958c62c1eb775a6267564518b7202e905 Mon Sep 17 00:00:00 2001 From: "Yujie Zhang (from Dev Box)" Date: Thu, 20 Aug 2026 18:44:22 +0800 Subject: [PATCH 2/4] Benchmark indexed-vector disk search Add opt-in API latency and payload metrics plus a manual paired benchmark workflow for legacy and indexed-vector disk search. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/disk-benchmarks.yml | 255 +++++++- ...ndex-api-call-metrics-indexed-vectors.json | 30 + .../disk-index-api-call-metrics-legacy.json | 30 + .../src/disk_index/benchmarks.rs | 61 ++ diskann-benchmark/src/disk_index/search.rs | 573 ++++++++++++++++-- diskann-benchmark/src/inputs/disk.rs | 110 ++++ 6 files changed, 993 insertions(+), 66 deletions(-) create mode 100644 diskann-benchmark/example/disk-index-api-call-metrics-indexed-vectors.json create mode 100644 diskann-benchmark/example/disk-index-api-call-metrics-legacy.json diff --git a/.github/workflows/disk-benchmarks.yml b/.github/workflows/disk-benchmarks.yml index 36753c846e..ed71974dc7 100644 --- a/.github/workflows/disk-benchmarks.yml +++ b/.github/workflows/disk-benchmarks.yml @@ -11,6 +11,19 @@ name: Disk Benchmarks on: workflow_dispatch: inputs: + benchmark_mode: + description: 'Benchmark mode to run' + required: true + default: regression + type: choice + options: + - regression + - indexed-vector-api + search_l: + description: 'Search-list size for indexed-vector API mode' + required: true + default: '2000' + type: string baseline_ref: description: 'A branch, commit SHA, or tag name to compare the current branch with' required: true @@ -97,6 +110,7 @@ jobs: lfs: true - name: Checkout baseline (${{ inputs.baseline_ref || 'main' }}) + if: ${{ github.event_name != 'workflow_dispatch' || inputs.benchmark_mode != 'indexed-vector-api' }} uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ inputs.baseline_ref || 'main' }} @@ -112,11 +126,13 @@ jobs: extract-to: diskann_rust/target/tmp - name: Copy dataset to baseline + if: ${{ github.event_name != 'workflow_dispatch' || inputs.benchmark_mode != 'indexed-vector-api' }} run: | mkdir -p baseline/target/tmp cp -r diskann_rust/target/tmp/${{ matrix.data_dir }} baseline/target/tmp/ - name: Run baseline benchmark + if: ${{ github.event_name != 'workflow_dispatch' || inputs.benchmark_mode != 'indexed-vector-api' }} working-directory: baseline run: | cargo run -p diskann-benchmark --features disk-index --release -- \ @@ -124,6 +140,7 @@ jobs: --output-file target/tmp/${{ matrix.dataset }}_baseline.json - name: Run current branch benchmark + if: ${{ github.event_name != 'workflow_dispatch' || inputs.benchmark_mode != 'indexed-vector-api' }} working-directory: diskann_rust run: | cargo run -p diskann-benchmark --features disk-index --release -- \ @@ -131,6 +148,7 @@ jobs: --output-file target/tmp/${{ matrix.dataset }}_target.json - name: Validate benchmark results + if: ${{ github.event_name != 'workflow_dispatch' || inputs.benchmark_mode != 'indexed-vector-api' }} working-directory: diskann_rust run: | cargo run -p diskann-benchmark --features disk-index --release -- \ @@ -140,12 +158,245 @@ jobs: --before ../baseline/target/tmp/${{ matrix.dataset }}_baseline.json \ --after target/tmp/${{ matrix.dataset }}_target.json + - name: Build current benchmark binary + if: ${{ github.event_name == 'workflow_dispatch' && inputs.benchmark_mode == 'indexed-vector-api' }} + working-directory: diskann_rust + run: cargo build -p diskann-benchmark --features disk-index --release + + - name: Build benchmark index + if: ${{ github.event_name == 'workflow_dispatch' && inputs.benchmark_mode == 'indexed-vector-api' }} + working-directory: diskann_rust + run: | + target/release/diskann-benchmark \ + run --input-file ${{ env.PERF_INPUTS }}/${{ matrix.config }} \ + --output-file target/tmp/${{ matrix.dataset }}_api_index_build.json + + - name: Generate paired API benchmark configs + if: ${{ github.event_name == 'workflow_dispatch' && inputs.benchmark_mode == 'indexed-vector-api' }} + working-directory: diskann_rust + env: + SEARCH_L: ${{ inputs.search_l }} + SOURCE_CONFIG: ${{ env.PERF_INPUTS }}/${{ matrix.config }} + LEGACY_CONFIG: target/tmp/${{ matrix.dataset }}_legacy_config.json + INDEXED_CONFIG: target/tmp/${{ matrix.dataset }}_indexed_vectors_config.json + run: | + python3 - <<'PY' + import copy + import json + import os + from pathlib import Path + + search_l = int(os.environ["SEARCH_L"]) + if not 0 < search_l <= 2**32 - 1: + raise ValueError(f"search_l must be in [1, {2**32 - 1}], got {search_l}") + + with Path(os.environ["SOURCE_CONFIG"]).open(encoding="utf-8") as stream: + original = json.load(stream) + + if len(original.get("jobs", [])) != 1: + raise ValueError("expected exactly one benchmark job in the source config") + content = original["jobs"][0].get("content", {}) + source = content.get("source", {}) + if source.get("disk-index-source") != "Build": + raise ValueError("expected the source config to contain a Build disk index") + + recall_at = content.get("search_phase", {}).get("recall_at") + if not isinstance(recall_at, int) or search_l < recall_at: + raise ValueError(f"search_l ({search_l}) must be at least recall_at ({recall_at})") + + common = copy.deepcopy(original) + common_content = common["jobs"][0]["content"] + common_content["source"] = { + "disk-index-source": "Load", + "data_type": source["data_type"], + "load_path": source["save_path"], + } + search_phase = common_content["search_phase"] + search_phase["collect_api_metrics"] = True + search_phase["search_list"] = [search_l] + + configs = { + "legacy": Path(os.environ["LEGACY_CONFIG"]), + "indexed-vectors": Path(os.environ["INDEXED_CONFIG"]), + } + generated = {} + for search_api, path in configs.items(): + config = copy.deepcopy(common) + config["jobs"][0]["content"]["search_phase"]["search_api"] = search_api + path.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8") + generated[search_api] = config + + comparable_legacy = copy.deepcopy(generated["legacy"]) + comparable_legacy["jobs"][0]["content"]["search_phase"]["search_api"] = "indexed-vectors" + if comparable_legacy != generated["indexed-vectors"]: + raise AssertionError("generated API configs differ by more than search_api") + PY + + - name: Run legacy API benchmark + if: ${{ github.event_name == 'workflow_dispatch' && inputs.benchmark_mode == 'indexed-vector-api' }} + working-directory: diskann_rust + run: | + target/release/diskann-benchmark \ + run --input-file target/tmp/${{ matrix.dataset }}_legacy_config.json \ + --output-file target/tmp/${{ matrix.dataset }}_legacy_result.json + + - name: Run indexed-vector API benchmark + if: ${{ github.event_name == 'workflow_dispatch' && inputs.benchmark_mode == 'indexed-vector-api' }} + working-directory: diskann_rust + run: | + target/release/diskann-benchmark \ + run --input-file target/tmp/${{ matrix.dataset }}_indexed_vectors_config.json \ + --output-file target/tmp/${{ matrix.dataset }}_indexed_vectors_result.json + + - name: Summarize API benchmark results + if: ${{ github.event_name == 'workflow_dispatch' && inputs.benchmark_mode == 'indexed-vector-api' }} + working-directory: diskann_rust + env: + SEARCH_L: ${{ inputs.search_l }} + LEGACY_RESULT: target/tmp/${{ matrix.dataset }}_legacy_result.json + INDEXED_RESULT: target/tmp/${{ matrix.dataset }}_indexed_vectors_result.json + run: | + python3 - <<'PY' + import json + import math + import os + from pathlib import Path + + expected_l = int(os.environ["SEARCH_L"]) + + def finite_number(value, label, *, positive=False): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + raise ValueError(f"{label} must be a finite number, got {value!r}") + if positive and value <= 0: + raise ValueError(f"{label} must be positive, got {value!r}") + if not positive and value < 0: + raise ValueError(f"{label} must be non-negative, got {value!r}") + return value + + def search_completed_peak(span_metrics, label): + roots = span_metrics.get("spans") + if not isinstance(roots, list): + raise ValueError(f"{label}.span_metrics.spans must be a list") + stack = list(roots) + while stack: + span = stack.pop() + if not isinstance(span, dict): + raise ValueError(f"{label}.span_metrics contains a non-object span") + children = span.get("children", []) + if not isinstance(children, list): + raise ValueError(f"{label}.span_metrics span children must be a list") + stack.extend(children) + if str(span.get("span_name", "")).endswith("-search_completed"): + metrics = span.get("metrics", {}) + return finite_number( + metrics.get("peak_memory_usage"), + f"{label}.search_completed.peak_memory_usage", + ) + raise ValueError(f"{label} result has no search_completed span") + + def load_result(path, expected_api): + with Path(path).open(encoding="utf-8") as stream: + records = json.load(stream) + label = expected_api + if not isinstance(records, list) or len(records) != 1: + raise ValueError(f"{label} output must contain exactly one result record") + record = records[0] + if not isinstance(record, dict) or record.get("input", {}).get("type") != "disk-index": + raise ValueError(f"{label} output does not contain a disk-index input") + results = record.get("results") + if not isinstance(results, dict) or results.get("build") is not None: + raise ValueError(f"{label} output must contain Load-only disk-index results") + search = results.get("search") + if not isinstance(search, dict) or search.get("search_api") != expected_api: + raise ValueError(f"{label} output has the wrong search_api") + per_l = search.get("search_results_per_l") + if not isinstance(per_l, list) or len(per_l) != 1: + raise ValueError(f"{label} output must contain exactly one search result") + result = per_l[0] + if not isinstance(result, dict) or result.get("search_l") != expected_l: + actual_l = result.get("search_l") if isinstance(result, dict) else None + raise ValueError(f"{label} output search_l is {actual_l!r}, expected {expected_l}") + recall = finite_number(result.get("recall"), f"{label}.recall") + if recall > 100: + raise ValueError(f"{label}.recall must be at most 100, got {recall!r}") + for field in ( + "qps", + "mean_public_api_call_latency_us", + "p95_public_api_call_latency_us", + "p999_public_api_call_latency_us", + ): + finite_number(result.get(field), f"{label}.{field}", positive=True) + for field in ( + "mean_latency", + "p95_latency", + "p999_latency", + "mean_returned_vector_payload_bytes", + "max_returned_vector_payload_bytes", + ): + finite_number(result.get(field), f"{label}.{field}") + result["search_completed_peak_memory_gb"] = search_completed_peak( + search.get("span_metrics", {}), label + ) + return result + + rows = [ + ("legacy", load_result(os.environ["LEGACY_RESULT"], "legacy")), + ( + "indexed-vectors", + load_result(os.environ["INDEXED_RESULT"], "indexed-vectors"), + ), + ] + if rows[0][1]["max_returned_vector_payload_bytes"] != 0: + raise ValueError("legacy API returned a nonzero vector payload") + if rows[1][1]["max_returned_vector_payload_bytes"] <= 0: + raise ValueError("indexed-vectors API returned an empty vector payload") + + print(f"Indexed-vector API A/B summary (L={expected_l})") + for name, result in rows: + print(f"\n{name}:") + print(f" recall: {result['recall']:.6f}") + print(f" qps: {result['qps']:.3f}") + print( + " internal latency us (mean/p95/p999): " + f"{result['mean_latency']:.3f} / {result['p95_latency']:.3f} / " + f"{result['p999_latency']:.3f}" + ) + print( + " public API call latency us (mean/p95/p999): " + f"{result['mean_public_api_call_latency_us']:.3f} / " + f"{result['p95_public_api_call_latency_us']:.3f} / " + f"{result['p999_public_api_call_latency_us']:.3f}" + ) + print( + " returned vector payload bytes (mean/max): " + f"{result['mean_returned_vector_payload_bytes']:.3f} / " + f"{result['max_returned_vector_payload_bytes']:.0f}" + ) + print( + " search_completed peak_memory_usage GB: " + f"{result['search_completed_peak_memory_gb']:.6f}" + ) + PY + - name: Upload benchmark results uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 - if: always() # Upload even if validation fails + if: ${{ always() && (github.event_name != 'workflow_dispatch' || inputs.benchmark_mode != 'indexed-vector-api') }} # Upload even if validation fails with: name: benchmark-results-${{ matrix.dataset }} path: | diskann_rust/target/tmp/${{ matrix.dataset }}_target.json baseline/target/tmp/${{ matrix.dataset }}_baseline.json - retention-days: 30 \ No newline at end of file + retention-days: 30 + + - name: Upload API benchmark artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + if: ${{ always() && github.event_name == 'workflow_dispatch' && inputs.benchmark_mode == 'indexed-vector-api' }} + with: + name: indexed-vector-api-${{ matrix.dataset }} + path: | + diskann_rust/target/tmp/${{ matrix.dataset }}_api_index_build.json + diskann_rust/target/tmp/${{ matrix.dataset }}_legacy_config.json + diskann_rust/target/tmp/${{ matrix.dataset }}_indexed_vectors_config.json + diskann_rust/target/tmp/${{ matrix.dataset }}_legacy_result.json + diskann_rust/target/tmp/${{ matrix.dataset }}_indexed_vectors_result.json + retention-days: 30 diff --git a/diskann-benchmark/example/disk-index-api-call-metrics-indexed-vectors.json b/diskann-benchmark/example/disk-index-api-call-metrics-indexed-vectors.json new file mode 100644 index 0000000000..b040d3983c --- /dev/null +++ b/diskann-benchmark/example/disk-index-api-call-metrics-indexed-vectors.json @@ -0,0 +1,30 @@ +{ + "search_directories": [ + "test_data/disk_index_search" + ], + "jobs": [ + { + "type": "disk-index", + "content": { + "source": { + "disk-index-source": "Load", + "data_type": "float32", + "load_path": "test_data/disk_index_search/disk_index_sift_learn_R4_L50_A1.2_truth_search" + }, + "search_phase": { + "queries": "disk_index_sample_query_10pts.fbin", + "groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin", + "search_list": [10], + "beam_width": 4, + "recall_at": 10, + "num_threads": 1, + "search_api": "indexed-vectors", + "collect_api_metrics": true, + "is_flat_search": true, + "distance": "squared_l2", + "vector_filters_file": null + } + } + } + ] +} diff --git a/diskann-benchmark/example/disk-index-api-call-metrics-legacy.json b/diskann-benchmark/example/disk-index-api-call-metrics-legacy.json new file mode 100644 index 0000000000..a8594710c2 --- /dev/null +++ b/diskann-benchmark/example/disk-index-api-call-metrics-legacy.json @@ -0,0 +1,30 @@ +{ + "search_directories": [ + "test_data/disk_index_search" + ], + "jobs": [ + { + "type": "disk-index", + "content": { + "source": { + "disk-index-source": "Load", + "data_type": "float32", + "load_path": "test_data/disk_index_search/disk_index_sift_learn_R4_L50_A1.2_truth_search" + }, + "search_phase": { + "queries": "disk_index_sample_query_10pts.fbin", + "groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin", + "search_list": [10], + "beam_width": 4, + "recall_at": 10, + "num_threads": 1, + "search_api": "legacy", + "collect_api_metrics": true, + "is_flat_search": true, + "distance": "squared_l2", + "vector_filters_file": null + } + } + } + ] +} diff --git a/diskann-benchmark/src/disk_index/benchmarks.rs b/diskann-benchmark/src/disk_index/benchmarks.rs index 1acc56abc9..bbcaa3305a 100644 --- a/diskann-benchmark/src/disk_index/benchmarks.rs +++ b/diskann-benchmark/src/disk_index/benchmarks.rs @@ -304,6 +304,13 @@ where let mut passed = true; let mut comparisons = Vec::new(); + anyhow::ensure!( + before.search.search_api == after.search.search_api, + "search_api mismatch: before={} after={}", + before.search.search_api, + after.search.search_api, + ); + // Check build time if both sides have it if let (Some(b_build), Some(a_build)) = (&before.build, &after.build) { comparisons.push(check_metric( @@ -403,3 +410,57 @@ where } } } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn output(search_api: Option<&str>) -> DiskIndexStats { + let mut search = json!({ + "num_threads": 1, + "beam_width": 4, + "recall_at": 10, + "is_flat_search": false, + "distance": "squared_l2", + "uses_vector_filters": false, + "num_nodes_to_cache": null, + "search_results_per_l": [{ + "search_l": 10, + "qps": 100.0, + "mean_latency": 10.0, + "p95_latency": 12, + "p999_latency": 14, + "mean_ios": 1.0, + "mean_io_time": 2.0, + "mean_cpu_time": 3.0, + "mean_pq_preprocess_time": 4.0, + "mean_comparisons": 5.0, + "mean_hops": 6.0, + "cache_hit_percentage": 90.0, + "recall": 0.9 + }], + "span_metrics": {"span_data": []} + }); + if let Some(search_api) = search_api { + search["search_api"] = json!(search_api); + } + serde_json::from_value(json!({"build": null, "search": search})).unwrap() + } + + #[test] + fn regression_check_rejects_search_api_mismatch() { + let before = output(None); + let after = output(Some("indexed-vectors")); + let input = ::example(); + let tolerances = DiskIndexTolerance::example(); + + let error = DiskIndex::::new() + .check(&tolerances, &input, &before, &after) + .unwrap_err(); + assert_eq!( + error.to_string(), + "search_api mismatch: before=legacy after=indexed-vectors" + ); + } +} diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index 8cd0f5e606..901b81271a 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -4,7 +4,10 @@ */ use rayon::prelude::*; -use std::{collections::HashSet, fmt, sync::atomic::AtomicBool, time::Instant}; +use std::{ + collections::HashSet, fmt, hint::black_box, mem::size_of, sync::atomic::AtomicBool, + time::Instant, +}; use opentelemetry::{global, trace::Span, trace::Tracer}; use opentelemetry_sdk::trace::SdkTracerProvider; @@ -36,12 +39,14 @@ use serde::{Deserialize, Serialize}; use crate::{ disk_index::json_spancollector::JsonSpanCollector, - inputs::disk::{DiskIndexLoad, DiskSearchPhase}, + inputs::disk::{DiskIndexLoad, DiskSearchApi, DiskSearchPhase}, utils::{datafiles, SimilarityMeasure}, }; #[derive(Serialize, Deserialize, Debug)] pub(super) struct DiskSearchStats { + #[serde(default)] + pub(super) search_api: DiskSearchApi, pub(super) num_threads: usize, pub(super) beam_width: usize, pub(super) recall_at: u32, @@ -60,6 +65,16 @@ pub(super) struct DiskSearchResult { pub(super) mean_latency: f64, pub(super) p95_latency: MicroSeconds, pub(super) p999_latency: MicroSeconds, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(super) mean_public_api_call_latency_us: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(super) p95_public_api_call_latency_us: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(super) p999_public_api_call_latency_us: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(super) mean_returned_vector_payload_bytes: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(super) max_returned_vector_payload_bytes: Option, pub(super) mean_ios: f64, pub(super) mean_io_time: f64, pub(super) mean_cpu_time: f64, @@ -70,11 +85,66 @@ pub(super) struct DiskSearchResult { pub(super) recall: f32, } +#[derive(Debug, PartialEq)] +struct PublicApiMetrics { + mean_call_latency_us: f64, + p95_call_latency_us: MicroSeconds, + p999_call_latency_us: MicroSeconds, + mean_payload_bytes: f64, + max_payload_bytes: u64, +} + +fn percentile_from_sorted(values: &[u64], numerator: usize, denominator: usize) -> u64 { + let index = values + .len() + .saturating_mul(numerator) + .checked_div(denominator) + .unwrap_or(0) + .min(values.len() - 1); + values[index] +} + +fn aggregate_public_api_metrics( + mut call_latencies_us: Vec, + payload_bytes: &[u64], +) -> anyhow::Result { + anyhow::ensure!( + !call_latencies_us.is_empty(), + "cannot aggregate zero queries" + ); + anyhow::ensure!( + call_latencies_us.len() == payload_bytes.len(), + "latency and payload sample counts differ" + ); + + let mean_call_latency_us = call_latencies_us + .iter() + .map(|&value| value as f64) + .sum::() + / call_latencies_us.len() as f64; + call_latencies_us.sort_unstable(); + + Ok(PublicApiMetrics { + mean_call_latency_us, + p95_call_latency_us: MicroSeconds::new(percentile_from_sorted(&call_latencies_us, 95, 100)), + p999_call_latency_us: MicroSeconds::new(percentile_from_sorted( + &call_latencies_us, + 999, + 1000, + )), + mean_payload_bytes: payload_bytes.iter().map(|&value| value as f64).sum::() + / payload_bytes.len() as f64, + max_payload_bytes: payload_bytes.iter().copied().max().unwrap_or(0), + }) +} + impl DiskSearchResult { - pub(super) fn new( + #[allow(clippy::too_many_arguments)] + fn new( statistics: &[QueryStatistics], result_ids: &[u32], result_counts: &[u32], + public_api_metrics: Option, search_l: u32, total_time_as_secs: f32, num_queries: usize, @@ -142,6 +212,20 @@ impl DiskSearchResult { 0.999, |s| s.total_execution_time_us, ) as u64), + mean_public_api_call_latency_us: public_api_metrics + .as_ref() + .map(|metrics| metrics.mean_call_latency_us), + p95_public_api_call_latency_us: public_api_metrics + .as_ref() + .map(|metrics| metrics.p95_call_latency_us), + p999_public_api_call_latency_us: public_api_metrics + .as_ref() + .map(|metrics| metrics.p999_call_latency_us), + mean_returned_vector_payload_bytes: public_api_metrics + .as_ref() + .map(|metrics| metrics.mean_payload_bytes), + max_returned_vector_payload_bytes: public_api_metrics + .map(|metrics| metrics.max_payload_bytes), mean_ios: statistics::get_mean_stats(statistics, |s| s.total_io_operations), mean_io_time: statistics::get_mean_stats(statistics, |s| s.io_time_us as f64), mean_cpu_time: statistics::get_mean_stats(statistics, |stats| stats.cpu_time_us as f64), @@ -248,7 +332,6 @@ where let mut result_ids: Vec = vec![0; (search_params.recall_at as usize) * num_queries]; let mut result_dists: Vec = vec![0.0; (search_params.recall_at as usize) * num_queries]; - let start = Instant::now(); let mut l_span = { @@ -257,71 +340,287 @@ where tracer.start(span_name) }; - let zipped = queries - .par_row_iter() - .zip(vector_filters.par_iter()) - .zip(result_ids.par_chunks_mut(search_params.recall_at as usize)) - .zip(result_dists.par_chunks_mut(search_params.recall_at as usize)) - .zip(statistics_vec.par_iter_mut()) - .zip(result_counts.par_iter_mut()); - - zipped.for_each_in_pool( - pool.as_ref(), - |(((((q, vf), id_chunk), dist_chunk), stats), rc)| { - // Construct the SearchMode from the JSON-driven - // `adaptive_l` is now encapsulated in `DiskSearchMode`, so the - // benchmark only supplies the per-query filter and post-processor. - let has_filter = search_params.vector_filters_file.is_some(); - let mode: SearchMode<'_> = search_params.search_mode.search_mode( - has_filter, - vf, - search_params.post_processor.as_ref(), + let public_api_samples = match (search_params.collect_api_metrics, search_params.search_api) + { + (false, DiskSearchApi::Legacy) => { + // Keep the pre-existing hot loop unchanged for old benchmark JSON. + let zipped = queries + .par_row_iter() + .zip(vector_filters.par_iter()) + .zip(result_ids.par_chunks_mut(search_params.recall_at as usize)) + .zip(result_dists.par_chunks_mut(search_params.recall_at as usize)) + .zip(statistics_vec.par_iter_mut()) + .zip(result_counts.par_iter_mut()); + + zipped.for_each_in_pool( + pool.as_ref(), + |(((((q, vf), id_chunk), dist_chunk), stats), rc)| { + // Construct the SearchMode from the JSON-driven + // `adaptive_l` is now encapsulated in `DiskSearchMode`, so the + // benchmark only supplies the per-query filter and post-processor. + let has_filter = search_params.vector_filters_file.is_some(); + let mode: SearchMode<'_> = search_params.search_mode.search_mode( + has_filter, + vf, + search_params.post_processor.as_ref(), + ); + + match searcher.search( + q, + search_params.recall_at, + l, + Some(search_params.beam_width), + mode, + ) { + Ok(search_result) => { + *stats = search_result.stats.query_statistics; + let base_count = (search_result.stats.result_count as usize) + .min(search_params.recall_at as usize) + .min(search_result.results.len()); + + *rc = base_count as u32; + id_chunk.fill(0); + dist_chunk.fill(0.0); + + for (i, result_item) in + search_result.results.iter().take(base_count).enumerate() + { + id_chunk[i] = result_item.vertex_id; + dist_chunk[i] = result_item.distance; + } + } + Err(e) => { + eprintln!("Search failed for query: {:?}", e); + *rc = 0; + id_chunk.fill(0); + dist_chunk.fill(0.0); + has_any_search_failed + .store(true, std::sync::atomic::Ordering::Release); + } + } + }, ); - - match searcher.search( - q, - search_params.recall_at, - l, - Some(search_params.beam_width), - mode, - ) { - Ok(search_result) => { - *stats = search_result.stats.query_statistics; - let base_count = (search_result.stats.result_count as usize) - .min(search_params.recall_at as usize) - .min(search_result.results.len()); - - *rc = base_count as u32; - id_chunk.fill(0); - dist_chunk.fill(0.0); - - for (i, result_item) in - search_result.results.iter().take(base_count).enumerate() - { - id_chunk[i] = result_item.vertex_id; - dist_chunk[i] = result_item.distance; + None + } + (false, DiskSearchApi::IndexedVectors) => { + let zipped = queries + .par_row_iter() + .zip(vector_filters.par_iter()) + .zip(result_ids.par_chunks_mut(search_params.recall_at as usize)) + .zip(result_dists.par_chunks_mut(search_params.recall_at as usize)) + .zip(statistics_vec.par_iter_mut()) + .zip(result_counts.par_iter_mut()); + + zipped.for_each_in_pool( + pool.as_ref(), + |(((((q, vf), id_chunk), dist_chunk), stats), rc)| { + let has_filter = search_params.vector_filters_file.is_some(); + let mode: SearchMode<'_> = search_params.search_mode.search_mode( + has_filter, + vf, + search_params.post_processor.as_ref(), + ); + + match searcher.search_with_indexed_vectors( + q, + search_params.recall_at, + l, + Some(search_params.beam_width), + mode, + ) { + Ok(search_result) => { + *stats = search_result.stats.query_statistics; + let base_count = (search_result.stats.result_count as usize) + .min(search_params.recall_at as usize) + .min(search_result.results.len()); + + *rc = base_count as u32; + id_chunk.fill(0); + dist_chunk.fill(0.0); + + for (i, result_item) in + search_result.results.iter().take(base_count).enumerate() + { + id_chunk[i] = result_item.vertex_id; + dist_chunk[i] = result_item.distance; + } + for result_item in &search_result.results { + black_box(result_item.indexed_vector.as_ref()); + } + } + Err(e) => { + eprintln!("Search failed for query: {:?}", e); + *rc = 0; + id_chunk.fill(0); + dist_chunk.fill(0.0); + has_any_search_failed + .store(true, std::sync::atomic::Ordering::Release); + } } - } - Err(e) => { - eprintln!("Search failed for query: {:?}", e); - *rc = 0; - id_chunk.fill(0); - dist_chunk.fill(0.0); - has_any_search_failed.store(true, std::sync::atomic::Ordering::Release); - } - } - }, - ); + }, + ); + None + } + (true, DiskSearchApi::Legacy) => { + let mut public_api_call_latencies_us = vec![0u64; num_queries]; + let returned_vector_payload_bytes = vec![0u64; num_queries]; + let zipped = queries + .par_row_iter() + .zip(vector_filters.par_iter()) + .zip(result_ids.par_chunks_mut(search_params.recall_at as usize)) + .zip(result_dists.par_chunks_mut(search_params.recall_at as usize)) + .zip(statistics_vec.par_iter_mut()) + .zip(result_counts.par_iter_mut()) + .zip(public_api_call_latencies_us.par_iter_mut()); + + zipped.for_each_in_pool( + pool.as_ref(), + |((((((q, vf), id_chunk), dist_chunk), stats), rc), call_latency_us)| { + let has_filter = search_params.vector_filters_file.is_some(); + let mode: SearchMode<'_> = search_params.search_mode.search_mode( + has_filter, + vf, + search_params.post_processor.as_ref(), + ); + + let api_start = Instant::now(); + let search_result = searcher.search( + q, + search_params.recall_at, + l, + Some(search_params.beam_width), + mode, + ); + let api_elapsed = api_start.elapsed(); + *call_latency_us = api_elapsed.as_micros().min(u64::MAX as u128) as u64; + + match search_result { + Ok(search_result) => { + *stats = search_result.stats.query_statistics; + let base_count = (search_result.stats.result_count as usize) + .min(search_params.recall_at as usize) + .min(search_result.results.len()); + + *rc = base_count as u32; + id_chunk.fill(0); + dist_chunk.fill(0.0); + + for (i, result_item) in + search_result.results.iter().take(base_count).enumerate() + { + id_chunk[i] = result_item.vertex_id; + dist_chunk[i] = result_item.distance; + } + } + Err(e) => { + eprintln!("Search failed for query: {:?}", e); + *rc = 0; + id_chunk.fill(0); + dist_chunk.fill(0.0); + has_any_search_failed + .store(true, std::sync::atomic::Ordering::Release); + } + } + }, + ); + Some((public_api_call_latencies_us, returned_vector_payload_bytes)) + } + (true, DiskSearchApi::IndexedVectors) => { + let mut public_api_call_latencies_us = vec![0u64; num_queries]; + let mut returned_vector_payload_bytes = vec![0u64; num_queries]; + let zipped = queries + .par_row_iter() + .zip(vector_filters.par_iter()) + .zip(result_ids.par_chunks_mut(search_params.recall_at as usize)) + .zip(result_dists.par_chunks_mut(search_params.recall_at as usize)) + .zip(statistics_vec.par_iter_mut()) + .zip(result_counts.par_iter_mut()) + .zip(public_api_call_latencies_us.par_iter_mut()) + .zip(returned_vector_payload_bytes.par_iter_mut()); + + zipped.for_each_in_pool( + pool.as_ref(), + |( + ((((((q, vf), id_chunk), dist_chunk), stats), rc), call_latency_us), + payload_bytes, + )| { + let has_filter = search_params.vector_filters_file.is_some(); + let mode: SearchMode<'_> = search_params.search_mode.search_mode( + has_filter, + vf, + search_params.post_processor.as_ref(), + ); + + let api_start = Instant::now(); + let search_result = searcher.search_with_indexed_vectors( + q, + search_params.recall_at, + l, + Some(search_params.beam_width), + mode, + ); + let api_elapsed = api_start.elapsed(); + *call_latency_us = api_elapsed.as_micros().min(u64::MAX as u128) as u64; + + match search_result { + Ok(search_result) => { + *stats = search_result.stats.query_statistics; + let base_count = (search_result.stats.result_count as usize) + .min(search_params.recall_at as usize) + .min(search_result.results.len()); + + *rc = base_count as u32; + id_chunk.fill(0); + dist_chunk.fill(0.0); + + for (i, result_item) in + search_result.results.iter().take(base_count).enumerate() + { + id_chunk[i] = result_item.vertex_id; + dist_chunk[i] = result_item.distance; + } + *payload_bytes = search_result + .results + .iter() + .map(|result_item| { + black_box(result_item.indexed_vector.as_ref()); + (result_item.indexed_vector.len() as u64) + * (size_of::() as u64) + }) + .sum(); + // The owned indexed vectors are dropped with `search_result` before + // this per-query closure returns; none are cloned or serialized. + } + Err(e) => { + eprintln!("Search failed for query: {:?}", e); + *rc = 0; + id_chunk.fill(0); + dist_chunk.fill(0.0); + has_any_search_failed + .store(true, std::sync::atomic::Ordering::Release); + } + } + }, + ); + Some((public_api_call_latencies_us, returned_vector_payload_bytes)) + } + }; let total_time = start.elapsed(); if has_any_search_failed.load(std::sync::atomic::Ordering::Acquire) { anyhow::bail!("One or more searches failed. See logs for details."); } + let public_api_metrics = public_api_samples + .map(|(call_latencies_us, payload_bytes)| { + aggregate_public_api_metrics(call_latencies_us, &payload_bytes) + }) + .transpose()?; let search_result = DiskSearchResult::new( &statistics_vec, &result_ids, &result_counts, + public_api_metrics, l, total_time.as_secs_f32(), num_queries, @@ -346,6 +645,7 @@ where global::set_tracer_provider(previous_tracer_provider); Ok(DiskSearchStats { + search_api: search_params.search_api, num_threads: search_params.num_threads, beam_width: search_params.beam_width, recall_at: search_params.recall_at, @@ -401,13 +701,31 @@ impl fmt::Display for DiskSearchStats { let fmt_us = |v: f64| -> String { format!("{:.1}us", v) }; let fmt_pct = |v: f64| -> String { format!("{:.1}%", v) }; - let cols: [(&str, usize); 14] = [ + let show_api_metrics = self.search_results_per_l.iter().all(|result| { + result.mean_public_api_call_latency_us.is_some() + && result.p95_public_api_call_latency_us.is_some() + && result.p999_public_api_call_latency_us.is_some() + && result.mean_returned_vector_payload_bytes.is_some() + && result.max_returned_vector_payload_bytes.is_some() + }); + let mut cols = vec![ ("L", 2), ("KNN", 3), ("QPS", 8), - ("Mean Latency", 13), - ("95% Latency", 13), - ("99.9 Latency", 13), + ("Internal Mean", 13), + ("Internal P95", 13), + ("Internal P999", 13), + ]; + if show_api_metrics { + cols.extend([ + ("API Call Mean", 13), + ("API Call P95", 13), + ("API Call P999", 13), + ("Vector B/q Mean", 15), + ("Vector B/q Max", 14), + ]); + } + cols.extend([ ("IOs", 6), ("IO (us)", 10), ("CPU (us)", 10), @@ -416,7 +734,7 @@ impl fmt::Display for DiskSearchStats { ("Mean Hops", 10), ("Cache Hit %", 12), ("Recall", 7), - ]; + ]); // Build header with exact widths let mut header = String::new(); @@ -430,6 +748,7 @@ impl fmt::Display for DiskSearchStats { // Summary writeln!(f, "Search Stats")?; + writeln!(f, "Search API, : {}", self.search_api)?; writeln!(f, "Threads, : {}", self.num_threads)?; writeln!(f, "Beam width, : {}", self.beam_width)?; writeln!(f, "Recall at, : {}", self.recall_at)?; @@ -451,13 +770,24 @@ impl fmt::Display for DiskSearchStats { for r in &self.search_results_per_l { // Prepare values as strings with numeric formatting - let vals: [String; 14] = [ + let mut vals = vec![ format!("{}", r.search_l), format!("{}", self.recall_at), format!("{:.1}", r.qps), fmt_us(r.mean_latency), format!("{}", r.p95_latency), format!("{}", r.p999_latency), + ]; + if show_api_metrics { + vals.extend([ + fmt_us(r.mean_public_api_call_latency_us.unwrap()), + format!("{}", r.p95_public_api_call_latency_us.unwrap()), + format!("{}", r.p999_public_api_call_latency_us.unwrap()), + format!("{:.1}", r.mean_returned_vector_payload_bytes.unwrap()), + format!("{}", r.max_returned_vector_payload_bytes.unwrap()), + ]); + } + vals.extend([ format!("{:.1}", r.mean_ios), fmt_us(r.mean_io_time), fmt_us(r.mean_cpu_time), @@ -466,7 +796,7 @@ impl fmt::Display for DiskSearchStats { format!("{:.1}", r.mean_hops), fmt_pct(r.cache_hit_percentage), format!("{:.3}", r.recall), - ]; + ]); // Right align each value to the column width, one space between columns let mut line = String::new(); @@ -482,3 +812,118 @@ impl fmt::Display for DiskSearchStats { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn output_without_api_metrics() -> serde_json::Value { + json!({ + "num_threads": 1, + "beam_width": 4, + "recall_at": 10, + "is_flat_search": false, + "distance": "squared_l2", + "uses_vector_filters": false, + "num_nodes_to_cache": null, + "search_results_per_l": [{ + "search_l": 10, + "qps": 100.0, + "mean_latency": 10.0, + "p95_latency": 12, + "p999_latency": 14, + "mean_ios": 1.0, + "mean_io_time": 2.0, + "mean_cpu_time": 3.0, + "mean_pq_preprocess_time": 4.0, + "mean_comparisons": 5.0, + "mean_hops": 6.0, + "cache_hit_percentage": 90.0, + "recall": 0.9 + }], + "span_metrics": {"span_data": []} + }) + } + + #[test] + fn aggregates_public_api_call_latency_and_payload_metrics() { + let metrics = + aggregate_public_api_metrics(vec![10, 20, 30, 40, 50], &[128, 256, 384, 512, 640]) + .unwrap(); + assert_eq!( + metrics, + PublicApiMetrics { + mean_call_latency_us: 30.0, + p95_call_latency_us: MicroSeconds::new(50), + p999_call_latency_us: MicroSeconds::new(50), + mean_payload_bytes: 384.0, + max_payload_bytes: 640, + } + ); + } + + #[test] + fn percentiles_keep_existing_rank_semantics() { + let metrics = aggregate_public_api_metrics( + (0..100).rev().collect(), + &(0..100).map(|_| 0).collect::>(), + ) + .unwrap(); + assert_eq!(metrics.p95_call_latency_us, MicroSeconds::new(95)); + assert_eq!(metrics.p999_call_latency_us, MicroSeconds::new(99)); + } + + #[test] + fn legacy_payload_aggregation_is_zero() { + let metrics = aggregate_public_api_metrics(vec![4, 8], &[0, 0]).unwrap(); + assert_eq!(metrics.mean_payload_bytes, 0.0); + assert_eq!(metrics.max_payload_bytes, 0); + } + + #[test] + fn metric_aggregation_validates_samples() { + assert!(aggregate_public_api_metrics(vec![], &[]).is_err()); + assert!(aggregate_public_api_metrics(vec![1], &[]).is_err()); + } + + #[test] + fn old_output_schema_deserializes_with_legacy_and_no_api_metrics() { + let stats: DiskSearchStats = serde_json::from_value(output_without_api_metrics()).unwrap(); + assert_eq!(stats.search_api, DiskSearchApi::Legacy); + let result = &stats.search_results_per_l[0]; + assert_eq!(result.mean_public_api_call_latency_us, None); + assert_eq!(result.p95_public_api_call_latency_us, None); + assert_eq!(result.p999_public_api_call_latency_us, None); + assert_eq!(result.mean_returned_vector_payload_bytes, None); + assert_eq!(result.max_returned_vector_payload_bytes, None); + + let serialized = serde_json::to_value(stats).unwrap(); + assert!(serialized["search_results_per_l"][0] + .get("mean_public_api_call_latency_us") + .is_none()); + } + + #[test] + fn collected_api_metrics_deserialize_as_some() { + let mut output = output_without_api_metrics(); + output["search_api"] = json!("indexed-vectors"); + let result = &mut output["search_results_per_l"][0]; + result["mean_public_api_call_latency_us"] = json!(11.0); + result["p95_public_api_call_latency_us"] = json!(12); + result["p999_public_api_call_latency_us"] = json!(13); + result["mean_returned_vector_payload_bytes"] = json!(512.0); + result["max_returned_vector_payload_bytes"] = json!(640); + + let stats: DiskSearchStats = serde_json::from_value(output).unwrap(); + assert_eq!(stats.search_api, DiskSearchApi::IndexedVectors); + let result = &stats.search_results_per_l[0]; + assert_eq!(result.mean_public_api_call_latency_us, Some(11.0)); + assert_eq!( + result.p95_public_api_call_latency_us, + Some(MicroSeconds::new(12)) + ); + assert_eq!(result.mean_returned_vector_payload_bytes, Some(512.0)); + assert_eq!(result.max_returned_vector_payload_bytes, Some(640)); + } +} diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index 032665cd7a..257f673c21 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -146,6 +146,23 @@ impl fmt::Display for DiskSearchMode { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "kebab-case")] +pub(crate) enum DiskSearchApi { + #[default] + Legacy, + IndexedVectors, +} + +impl fmt::Display for DiskSearchApi { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Legacy => write!(f, "legacy"), + Self::IndexedVectors => write!(f, "indexed-vectors"), + } + } +} + /// Search phase configuration #[derive(Debug, Deserialize, Serialize)] pub(crate) struct DiskSearchPhase { @@ -155,6 +172,10 @@ pub(crate) struct DiskSearchPhase { pub(crate) beam_width: usize, pub(crate) search_list: Vec, pub(crate) recall_at: u32, + #[serde(default)] + pub(crate) search_api: DiskSearchApi, + #[serde(default)] + pub(crate) collect_api_metrics: bool, #[cfg(feature = "disk-index")] #[serde(default)] pub(crate) search_mode: DiskSearchMode, @@ -353,6 +374,8 @@ impl Example for DiskIndexOperation { beam_width: 16, recall_at: 10, num_threads: 8, + search_api: DiskSearchApi::default(), + collect_api_metrics: false, #[cfg(feature = "disk-index")] search_mode: DiskSearchMode { is_flat_search: false, @@ -483,6 +506,8 @@ impl DiskSearchPhase { write_field!(f, "Beam Width", self.beam_width)?; write_field!(f, "Recall@", self.recall_at)?; write_field!(f, "Threads", self.num_threads)?; + write_field!(f, "Search API", self.search_api)?; + write_field!(f, "Collect API Metrics", self.collect_api_metrics)?; #[cfg(feature = "disk-index")] write_field!(f, "Search Mode", self.search_mode)?; #[cfg(not(feature = "disk-index"))] @@ -514,3 +539,88 @@ impl fmt::Display for DiskSearchPhase { self.summarize_fields(f) } } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn search_phase_json() -> serde_json::Value { + json!({ + "queries": "queries.fbin", + "groundtruth": "groundtruth.bin", + "num_threads": 1, + "beam_width": 4, + "search_list": [10], + "recall_at": 10, + "is_flat_search": false, + "distance": "squared_l2" + }) + } + + #[test] + fn disk_search_api_defaults_to_legacy_for_old_inputs() { + let phase: DiskSearchPhase = serde_json::from_value(search_phase_json()).unwrap(); + assert_eq!(phase.search_api, DiskSearchApi::Legacy); + assert!(!phase.collect_api_metrics); + assert_eq!(DiskSearchApi::default(), DiskSearchApi::Legacy); + assert_eq!( + DiskIndexOperation::example().search_phase.search_api, + DiskSearchApi::Legacy + ); + assert!(phase.to_string().contains("Search API: legacy")); + } + + #[test] + fn disk_search_api_uses_kebab_case_serde_and_display() { + let mut value = search_phase_json(); + value["search_api"] = json!("indexed-vectors"); + let phase: DiskSearchPhase = serde_json::from_value(value).unwrap(); + assert_eq!(phase.search_api, DiskSearchApi::IndexedVectors); + assert_eq!(phase.search_api.to_string(), "indexed-vectors"); + assert_eq!( + serde_json::to_value(phase.search_api).unwrap(), + json!("indexed-vectors") + ); + } + + #[test] + fn disk_search_api_rejects_unknown_values() { + let mut value = search_phase_json(); + value["search_api"] = json!("return-vectors"); + let error = serde_json::from_value::(value).unwrap_err(); + assert!(error.to_string().contains("unknown variant")); + assert!(error.to_string().contains("indexed-vectors")); + } + + fn example_operation(contents: &str) -> DiskIndexOperation { + let value: serde_json::Value = serde_json::from_str(contents).unwrap(); + serde_json::from_value(value["jobs"][0]["content"].clone()).unwrap() + } + + #[test] + fn api_call_metric_examples_deserialize_and_only_differ_by_selector() { + let legacy_contents = include_str!("../../example/disk-index-api-call-metrics-legacy.json"); + let indexed_contents = + include_str!("../../example/disk-index-api-call-metrics-indexed-vectors.json"); + let legacy = example_operation(legacy_contents); + let indexed = example_operation(indexed_contents); + + assert!(matches!(legacy.source, DiskIndexSource::Load(_))); + assert!(matches!(indexed.source, DiskIndexSource::Load(_))); + assert_eq!(legacy.search_phase.search_list, vec![10]); + assert_eq!(indexed.search_phase.search_list, vec![10]); + assert!(legacy.search_phase.collect_api_metrics); + assert!(indexed.search_phase.collect_api_metrics); + assert_eq!(legacy.search_phase.search_api, DiskSearchApi::Legacy); + assert_eq!( + indexed.search_phase.search_api, + DiskSearchApi::IndexedVectors + ); + + let mut legacy_json: serde_json::Value = serde_json::from_str(legacy_contents).unwrap(); + let indexed_json: serde_json::Value = serde_json::from_str(indexed_contents).unwrap(); + legacy_json["jobs"][0]["content"]["search_phase"]["search_api"] = json!("indexed-vectors"); + assert_eq!(legacy_json, indexed_json); + } +} From c517a8f0b69b655c6d1f7d0b950d75364b238c78 Mon Sep 17 00:00:00 2001 From: "Yujie Zhang (from Dev Box)" Date: Thu, 20 Aug 2026 23:22:07 +0800 Subject: [PATCH 3/4] Benchmark configurable result count Decouple returned K from recall@K so indexed-vector benchmarks can exercise K=1000 with existing top-100 ground truth. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/disk-benchmarks.yml | 30 ++++++++--- .../src/disk_index/benchmarks.rs | 12 +++++ diskann-benchmark/src/disk_index/search.rs | 52 +++++++++++-------- diskann-benchmark/src/inputs/disk.rs | 23 ++++++-- 4 files changed, 85 insertions(+), 32 deletions(-) diff --git a/.github/workflows/disk-benchmarks.yml b/.github/workflows/disk-benchmarks.yml index ed71974dc7..dd10012304 100644 --- a/.github/workflows/disk-benchmarks.yml +++ b/.github/workflows/disk-benchmarks.yml @@ -20,10 +20,15 @@ on: - regression - indexed-vector-api search_l: - description: 'Search-list size for indexed-vector API mode' + description: 'Search-list size L for indexed-vector API mode' required: true default: '2000' type: string + return_k: + description: 'Result count K for indexed-vector API mode' + required: true + default: '100' + type: string baseline_ref: description: 'A branch, commit SHA, or tag name to compare the current branch with' required: true @@ -176,6 +181,7 @@ jobs: working-directory: diskann_rust env: SEARCH_L: ${{ inputs.search_l }} + RETURN_K: ${{ inputs.return_k }} SOURCE_CONFIG: ${{ env.PERF_INPUTS }}/${{ matrix.config }} LEGACY_CONFIG: target/tmp/${{ matrix.dataset }}_legacy_config.json INDEXED_CONFIG: target/tmp/${{ matrix.dataset }}_indexed_vectors_config.json @@ -187,8 +193,12 @@ jobs: from pathlib import Path search_l = int(os.environ["SEARCH_L"]) - if not 0 < search_l <= 2**32 - 1: - raise ValueError(f"search_l must be in [1, {2**32 - 1}], got {search_l}") + return_k = int(os.environ["RETURN_K"]) + if not 0 < return_k <= search_l <= 2**32 - 1: + raise ValueError( + f"expected 0 < return_k <= search_l <= {2**32 - 1}, " + f"got return_k={return_k}, search_l={search_l}" + ) with Path(os.environ["SOURCE_CONFIG"]).open(encoding="utf-8") as stream: original = json.load(stream) @@ -201,8 +211,8 @@ jobs: raise ValueError("expected the source config to contain a Build disk index") recall_at = content.get("search_phase", {}).get("recall_at") - if not isinstance(recall_at, int) or search_l < recall_at: - raise ValueError(f"search_l ({search_l}) must be at least recall_at ({recall_at})") + if not isinstance(recall_at, int) or return_k < recall_at: + raise ValueError(f"return_k ({return_k}) must be at least recall_at ({recall_at})") common = copy.deepcopy(original) common_content = common["jobs"][0]["content"] @@ -213,6 +223,7 @@ jobs: } search_phase = common_content["search_phase"] search_phase["collect_api_metrics"] = True + search_phase["return_list_size"] = return_k search_phase["search_list"] = [search_l] configs = { @@ -253,6 +264,7 @@ jobs: working-directory: diskann_rust env: SEARCH_L: ${{ inputs.search_l }} + RETURN_K: ${{ inputs.return_k }} LEGACY_RESULT: target/tmp/${{ matrix.dataset }}_legacy_result.json INDEXED_RESULT: target/tmp/${{ matrix.dataset }}_indexed_vectors_result.json run: | @@ -263,6 +275,7 @@ jobs: from pathlib import Path expected_l = int(os.environ["SEARCH_L"]) + expected_k = int(os.environ["RETURN_K"]) def finite_number(value, label, *, positive=False): if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): @@ -309,6 +322,11 @@ jobs: search = results.get("search") if not isinstance(search, dict) or search.get("search_api") != expected_api: raise ValueError(f"{label} output has the wrong search_api") + if search.get("return_list_size") != expected_k: + raise ValueError( + f"{label} output return_list_size is {search.get('return_list_size')!r}, " + f"expected {expected_k}" + ) per_l = search.get("search_results_per_l") if not isinstance(per_l, list) or len(per_l) != 1: raise ValueError(f"{label} output must contain exactly one search result") @@ -351,7 +369,7 @@ jobs: if rows[1][1]["max_returned_vector_payload_bytes"] <= 0: raise ValueError("indexed-vectors API returned an empty vector payload") - print(f"Indexed-vector API A/B summary (L={expected_l})") + print(f"Indexed-vector API A/B summary (K={expected_k}, L={expected_l})") for name, result in rows: print(f"\n{name}:") print(f" recall: {result['recall']:.6f}") diff --git a/diskann-benchmark/src/disk_index/benchmarks.rs b/diskann-benchmark/src/disk_index/benchmarks.rs index bbcaa3305a..0eb37e8357 100644 --- a/diskann-benchmark/src/disk_index/benchmarks.rs +++ b/diskann-benchmark/src/disk_index/benchmarks.rs @@ -310,6 +310,18 @@ where before.search.search_api, after.search.search_api, ); + let before_k = before + .search + .return_list_size + .unwrap_or(before.search.recall_at); + let after_k = after + .search + .return_list_size + .unwrap_or(after.search.recall_at); + anyhow::ensure!( + before_k == after_k, + "return_list_size mismatch: before={before_k} after={after_k}", + ); // Check build time if both sides have it if let (Some(b_build), Some(a_build)) = (&before.build, &after.build) { diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index 901b81271a..4bc0b851d5 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -50,6 +50,8 @@ pub(super) struct DiskSearchStats { pub(super) num_threads: usize, pub(super) beam_width: usize, pub(super) recall_at: u32, + #[serde(default)] + pub(super) return_list_size: Option, pub(crate) is_flat_search: bool, pub(crate) distance: SimilarityMeasure, pub(crate) uses_vector_filters: bool, @@ -148,6 +150,7 @@ impl DiskSearchResult { search_l: u32, total_time_as_secs: f32, num_queries: usize, + result_dim: u32, gt_context: &GroundTruthContext, ) -> anyhow::Result { let total_ios = statistics::get_sum_stats(statistics, |stats| stats.total_io_operations); @@ -161,7 +164,7 @@ impl DiskSearchResult { let recall = if let Some(var_gt) = >_context.gt_ids_variable_length { let ours: Vec> = result_ids - .chunks_exact(gt_context.recall_at as usize) + .chunks_exact(result_dim as usize) .enumerate() .map(|(qi, chunk)| { let written = result_counts[qi] as usize; @@ -187,7 +190,7 @@ impl DiskSearchResult { gt_context.gt_dists.as_ref(), gt_context.gt_dim, result_ids, - gt_context.recall_at, + result_dim, KRecallAtN::new(gt_context.recall_at, gt_context.recall_at)?, )?; recall_value as f32 @@ -288,6 +291,7 @@ where search_params.recall_at, storage_provider, )?; + let return_list_size = search_params.return_list_size(); // Setup disk index components let pivot_path = get_pq_pivot_file(&index_load.load_path); @@ -329,9 +333,8 @@ where let mut statistics_vec: Vec = vec![QueryStatistics::default(); num_queries]; let mut result_counts: Vec = vec![0; num_queries]; - let mut result_ids: Vec = vec![0; (search_params.recall_at as usize) * num_queries]; - let mut result_dists: Vec = - vec![0.0; (search_params.recall_at as usize) * num_queries]; + let mut result_ids: Vec = vec![0; (return_list_size as usize) * num_queries]; + let mut result_dists: Vec = vec![0.0; (return_list_size as usize) * num_queries]; let start = Instant::now(); let mut l_span = { @@ -347,8 +350,8 @@ where let zipped = queries .par_row_iter() .zip(vector_filters.par_iter()) - .zip(result_ids.par_chunks_mut(search_params.recall_at as usize)) - .zip(result_dists.par_chunks_mut(search_params.recall_at as usize)) + .zip(result_ids.par_chunks_mut(return_list_size as usize)) + .zip(result_dists.par_chunks_mut(return_list_size as usize)) .zip(statistics_vec.par_iter_mut()) .zip(result_counts.par_iter_mut()); @@ -367,7 +370,7 @@ where match searcher.search( q, - search_params.recall_at, + return_list_size, l, Some(search_params.beam_width), mode, @@ -375,7 +378,7 @@ where Ok(search_result) => { *stats = search_result.stats.query_statistics; let base_count = (search_result.stats.result_count as usize) - .min(search_params.recall_at as usize) + .min(return_list_size as usize) .min(search_result.results.len()); *rc = base_count as u32; @@ -406,8 +409,8 @@ where let zipped = queries .par_row_iter() .zip(vector_filters.par_iter()) - .zip(result_ids.par_chunks_mut(search_params.recall_at as usize)) - .zip(result_dists.par_chunks_mut(search_params.recall_at as usize)) + .zip(result_ids.par_chunks_mut(return_list_size as usize)) + .zip(result_dists.par_chunks_mut(return_list_size as usize)) .zip(statistics_vec.par_iter_mut()) .zip(result_counts.par_iter_mut()); @@ -423,7 +426,7 @@ where match searcher.search_with_indexed_vectors( q, - search_params.recall_at, + return_list_size, l, Some(search_params.beam_width), mode, @@ -431,7 +434,7 @@ where Ok(search_result) => { *stats = search_result.stats.query_statistics; let base_count = (search_result.stats.result_count as usize) - .min(search_params.recall_at as usize) + .min(return_list_size as usize) .min(search_result.results.len()); *rc = base_count as u32; @@ -467,8 +470,8 @@ where let zipped = queries .par_row_iter() .zip(vector_filters.par_iter()) - .zip(result_ids.par_chunks_mut(search_params.recall_at as usize)) - .zip(result_dists.par_chunks_mut(search_params.recall_at as usize)) + .zip(result_ids.par_chunks_mut(return_list_size as usize)) + .zip(result_dists.par_chunks_mut(return_list_size as usize)) .zip(statistics_vec.par_iter_mut()) .zip(result_counts.par_iter_mut()) .zip(public_api_call_latencies_us.par_iter_mut()); @@ -486,7 +489,7 @@ where let api_start = Instant::now(); let search_result = searcher.search( q, - search_params.recall_at, + return_list_size, l, Some(search_params.beam_width), mode, @@ -498,7 +501,7 @@ where Ok(search_result) => { *stats = search_result.stats.query_statistics; let base_count = (search_result.stats.result_count as usize) - .min(search_params.recall_at as usize) + .min(return_list_size as usize) .min(search_result.results.len()); *rc = base_count as u32; @@ -531,8 +534,8 @@ where let zipped = queries .par_row_iter() .zip(vector_filters.par_iter()) - .zip(result_ids.par_chunks_mut(search_params.recall_at as usize)) - .zip(result_dists.par_chunks_mut(search_params.recall_at as usize)) + .zip(result_ids.par_chunks_mut(return_list_size as usize)) + .zip(result_dists.par_chunks_mut(return_list_size as usize)) .zip(statistics_vec.par_iter_mut()) .zip(result_counts.par_iter_mut()) .zip(public_api_call_latencies_us.par_iter_mut()) @@ -554,7 +557,7 @@ where let api_start = Instant::now(); let search_result = searcher.search_with_indexed_vectors( q, - search_params.recall_at, + return_list_size, l, Some(search_params.beam_width), mode, @@ -566,7 +569,7 @@ where Ok(search_result) => { *stats = search_result.stats.query_statistics; let base_count = (search_result.stats.result_count as usize) - .min(search_params.recall_at as usize) + .min(return_list_size as usize) .min(search_result.results.len()); *rc = base_count as u32; @@ -624,6 +627,7 @@ where l, total_time.as_secs_f32(), num_queries, + return_list_size, >_context, )?; @@ -649,6 +653,7 @@ where num_threads: search_params.num_threads, beam_width: search_params.beam_width, recall_at: search_params.recall_at, + return_list_size: Some(return_list_size), is_flat_search: search_params.search_mode.is_flat_search, distance: search_params.distance, uses_vector_filters: search_params.vector_filters_file.is_some(), @@ -752,6 +757,11 @@ impl fmt::Display for DiskSearchStats { writeln!(f, "Threads, : {}", self.num_threads)?; writeln!(f, "Beam width, : {}", self.beam_width)?; writeln!(f, "Recall at, : {}", self.recall_at)?; + writeln!( + f, + "Return K, : {}", + self.return_list_size.unwrap_or(self.recall_at) + )?; writeln!(f, "Flat search, : {}", self.is_flat_search)?; writeln!(f, "Distance, : {}", self.distance)?; writeln!(f, "Vector filters, : {}", self.uses_vector_filters)?; diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index 257f673c21..ab9293c035 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -172,6 +172,8 @@ pub(crate) struct DiskSearchPhase { pub(crate) beam_width: usize, pub(crate) search_list: Vec, pub(crate) recall_at: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) return_list_size: Option, #[serde(default)] pub(crate) search_api: DiskSearchApi, #[serde(default)] @@ -284,6 +286,10 @@ impl DiskIndexBuild { } impl DiskSearchPhase { + pub(crate) fn return_list_size(&self) -> u32 { + self.return_list_size.unwrap_or(self.recall_at) + } + pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.queries .resolve(checker) @@ -306,22 +312,26 @@ impl DiskSearchPhase { .context("invalid disk search mode")?; // basic numeric sanity checks + if self.recall_at == 0 { + anyhow::bail!("recall_at must be positive"); + } + let return_list_size = self.return_list_size(); + if return_list_size < self.recall_at { + anyhow::bail!("return_list_size must be at least recall_at"); + } if self.search_list.is_empty() { anyhow::bail!("search_list must have at least one value"); } if self .search_list .iter() - .any(|&l| l == 0 || l < self.recall_at) + .any(|&l| l == 0 || l < return_list_size) { - anyhow::bail!("search_list must contain positive values only"); + anyhow::bail!("search_list values must be at least return_list_size"); } if self.beam_width == 0 { anyhow::bail!("beam_width must be positive"); } - if self.recall_at == 0 { - anyhow::bail!("recall_at must be positive"); - } if self.num_threads == 0 { anyhow::bail!("num_threads must be positive"); } @@ -373,6 +383,7 @@ impl Example for DiskIndexOperation { search_list: vec![64, 128, 256, 512], beam_width: 16, recall_at: 10, + return_list_size: None, num_threads: 8, search_api: DiskSearchApi::default(), collect_api_metrics: false, @@ -505,6 +516,7 @@ impl DiskSearchPhase { } write_field!(f, "Beam Width", self.beam_width)?; write_field!(f, "Recall@", self.recall_at)?; + write_field!(f, "Return K", self.return_list_size())?; write_field!(f, "Threads", self.num_threads)?; write_field!(f, "Search API", self.search_api)?; write_field!(f, "Collect API Metrics", self.collect_api_metrics)?; @@ -563,6 +575,7 @@ mod tests { let phase: DiskSearchPhase = serde_json::from_value(search_phase_json()).unwrap(); assert_eq!(phase.search_api, DiskSearchApi::Legacy); assert!(!phase.collect_api_metrics); + assert_eq!(phase.return_list_size(), phase.recall_at); assert_eq!(DiskSearchApi::default(), DiskSearchApi::Legacy); assert_eq!( DiskIndexOperation::example().search_phase.search_api, From bb7b7ba1ff8924bd9d54f5d392ba17c2e789e874 Mon Sep 17 00:00:00 2001 From: "Yujie Zhang (from Dev Box)" Date: Thu, 20 Aug 2026 23:48:51 +0800 Subject: [PATCH 4/4] Clear indexed-vector cache after each query Release per-query vector boxes before pooled scratch is returned so the memory benchmark can measure the cleanup effect. 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.