From 5c86c1649fdf5f091bf8432fdbb8e6efa7a6a637 Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Tue, 7 Jul 2026 17:03:49 +0800 Subject: [PATCH 01/11] consolidate disk search config into DiskSearchMode and move SearchMode construction to search.rs --- .../disk-index-determinant-diversity.json | 17 +-- .../example/disk-index-filter.json | 16 ++- diskann-benchmark/example/disk-index.json | 10 +- .../openai-100K-disk-index.json | 5 +- .../wikipedia-100K-disk-index.json | 5 +- diskann-benchmark/src/disk_index/search.rs | 63 ++++++++-- diskann-benchmark/src/inputs/disk.rs | 109 +++--------------- 7 files changed, 99 insertions(+), 126 deletions(-) diff --git a/diskann-benchmark/example/disk-index-determinant-diversity.json b/diskann-benchmark/example/disk-index-determinant-diversity.json index 2962c1d97b..a8c236a99e 100644 --- a/diskann-benchmark/example/disk-index-determinant-diversity.json +++ b/diskann-benchmark/example/disk-index-determinant-diversity.json @@ -27,14 +27,15 @@ "beam_width": 4, "recall_at": 10, "num_threads": 1, - "is_flat_search": false, - "distance": "squared_l2", - "vector_filters_file": null, - "post_processor": { - "type": "determinant-diversity", - "power": 2.0, - "eta": 1.0 - } + "search_mode": { + "is_flat_search": false, + "post_processor": { + "type": "determinant-diversity", + "power": 2.0, + "eta": 1.0 + } + }, + "distance": "squared_l2" } } } diff --git a/diskann-benchmark/example/disk-index-filter.json b/diskann-benchmark/example/disk-index-filter.json index a3f35ca91a..cdb9842d7f 100644 --- a/diskann-benchmark/example/disk-index-filter.json +++ b/diskann-benchmark/example/disk-index-filter.json @@ -27,9 +27,11 @@ "beam_width": 4, "recall_at": 10, "num_threads": 1, - "is_flat_search": false, - "distance": "squared_l2", - "vector_filters_file": "disk_index_10pts_idx_uint32_range_res_r_100000.bin" + "search_mode": { + "is_flat_search": false, + "vector_filters_file": "disk_index_10pts_idx_uint32_range_res_r_100000.bin" + }, + "distance": "squared_l2" } } }, @@ -57,9 +59,11 @@ "beam_width": 4, "recall_at": 10, "num_threads": 1, - "is_flat_search": true, - "distance": "squared_l2", - "vector_filters_file": "disk_index_10pts_idx_uint32_range_res_r_100000.bin" + "search_mode": { + "is_flat_search": true, + "vector_filters_file": "disk_index_10pts_idx_uint32_range_res_r_100000.bin" + }, + "distance": "squared_l2" } } } diff --git a/diskann-benchmark/example/disk-index.json b/diskann-benchmark/example/disk-index.json index 4d60fdb8ea..c8a0c7b7d9 100644 --- a/diskann-benchmark/example/disk-index.json +++ b/diskann-benchmark/example/disk-index.json @@ -27,9 +27,8 @@ "beam_width": 4, "recall_at": 10, "num_threads": 1, - "is_flat_search": false, - "distance": "squared_l2", - "vector_filters_file": null + "search_mode": { "is_flat_search": false }, + "distance": "squared_l2" } } }, @@ -48,9 +47,8 @@ "beam_width": 4, "recall_at": 10, "num_threads": 1, - "is_flat_search": true, - "distance": "squared_l2", - "vector_filters_file": null + "search_mode": { "is_flat_search": true }, + "distance": "squared_l2" } } } diff --git a/diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json b/diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json index 6b3e3b42d5..2417828997 100644 --- a/diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json +++ b/diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json @@ -29,9 +29,8 @@ "beam_width": 4, "recall_at": 100, "num_threads": 4, - "is_flat_search": false, - "distance": "squared_l2", - "vector_filters_file": null + "search_mode": { "is_flat_search": false }, + "distance": "squared_l2" } } } diff --git a/diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json b/diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json index 59c439017d..aa00ab4ac2 100644 --- a/diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json +++ b/diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json @@ -29,9 +29,8 @@ "beam_width": 4, "recall_at": 100, "num_threads": 4, - "is_flat_search": false, - "distance": "inner_product", - "vector_filters_file": null + "search_mode": { "is_flat_search": false }, + "distance": "inner_product" } } } diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index a1e84e79b6..342b1adde7 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -9,6 +9,7 @@ use std::{collections::HashSet, fmt, sync::atomic::AtomicBool, time::Instant}; use opentelemetry::{global, trace::Span, trace::Tracer}; use opentelemetry_sdk::trace::SdkTracerProvider; +use diskann::graph; use diskann::utils::VectorRepr; use diskann_benchmark_runner::{files::InputFile, utils::MicroSeconds}; use diskann_disk::{ @@ -36,7 +37,8 @@ use serde::{Deserialize, Serialize}; use crate::{ disk_index::json_spancollector::JsonSpanCollector, - inputs::disk::{DiskIndexLoad, DiskSearchPhase}, + inputs::disk::{DiskIndexLoad, DiskSearchMode, DiskSearchPhase}, + inputs::post_processor::TopkPostProcessor, utils::{datafiles, SimilarityMeasure}, }; @@ -158,6 +160,52 @@ impl DiskSearchResult { } } +/// Construct the disk [`SearchMode`] from the JSON-driven [`DiskSearchMode`] +/// config plus the per-query filter and post-processor supplied at search time. +fn build_search_mode<'a>( + mode: &'a DiskSearchMode, + has_vector_filters: bool, + vector_filter: &'a HashSet, + post_processor: Option<&TopkPostProcessor>, +) -> SearchMode<'a> { + let adaptive_l = mode.adaptive_l.as_ref().map(|adaptive_l| { + graph::search::AdaptiveL::new(adaptive_l.sample_count.into(), adaptive_l.scale_factor) + .expect("validated adaptive L must construct") + }); + + match ( + mode.is_flat_search, + has_vector_filters, + post_processor, + adaptive_l, + ) { + (true, false, _, _) => SearchMode::flat(), + (true, true, _, _) => { + SearchMode::flat_filtered(move |vid: &u32| vector_filter.contains(vid)) + } + (false, false, Some(TopkPostProcessor::DeterminantDiversity(params)), _) => { + SearchMode::diverse_graph(*params) + } + (false, true, Some(TopkPostProcessor::DeterminantDiversity(params)), _) => { + SearchMode::diverse_graph_filtered( + move |vid: &u32| vector_filter.contains(vid), + *params, + ) + } + (false, false, None, Some(adaptive_l)) => { + SearchMode::inline_filter(|_| true, Some(adaptive_l)) + } + (false, true, None, Some(adaptive_l)) => SearchMode::inline_filter( + move |vid: &u32| vector_filter.contains(vid), + Some(adaptive_l), + ), + (false, false, None, None) => SearchMode::graph(), + (false, true, None, None) => { + SearchMode::graph_filtered(move |vid: &u32| vector_filter.contains(vid)) + } + } +} + pub(super) fn search_disk_index( index_load: &DiskIndexLoad, search_params: &DiskSearchPhase, @@ -185,7 +233,7 @@ where let num_queries = queries.nrows(); // Load the vector filters - let vector_filters = match &search_params.vector_filters_file { + let vector_filters = match &search_params.search_mode.vector_filters_file { Some(vector_filters_file) => { let vector_filters_file = vector_filters_file.to_string_lossy().to_string(); search_index_utils::load_vector_filters(storage_provider, &vector_filters_file)? @@ -199,7 +247,7 @@ where // Prepare ground truth context let gt_context = prepare_ground_truth_context( - search_params.vector_filters_file.is_some(), + search_params.search_mode.vector_filters_file.is_some(), &search_params.groundtruth, search_params.recall_at, storage_provider, @@ -271,11 +319,12 @@ where // 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( + let has_filter = search_params.search_mode.vector_filters_file.is_some(); + let mode: SearchMode<'_> = build_search_mode( + &search_params.search_mode, has_filter, vf, - search_params.post_processor.as_ref(), + search_params.search_mode.post_processor.as_ref(), ); match searcher.search( @@ -351,7 +400,7 @@ where recall_at: search_params.recall_at, is_flat_search: search_params.search_mode.is_flat_search, distance: search_params.distance, - uses_vector_filters: search_params.vector_filters_file.is_some(), + uses_vector_filters: search_params.search_mode.vector_filters_file.is_some(), num_nodes_to_cache: search_params.num_nodes_to_cache, search_results_per_l, span_metrics, diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index 7ed521a874..44b7d61d57 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -6,27 +6,18 @@ use std::{fmt, num::NonZeroUsize, path::Path}; use anyhow::Context; -#[cfg(feature = "disk-index")] -use std::collections::HashSet; -#[cfg(feature = "disk-index")] -use diskann::graph; use diskann_benchmark_runner::{files::InputFile, utils::datatype::DataType, Checker}; #[cfg(feature = "disk-index")] -use diskann_disk::search::search_mode::SearchMode; -#[cfg(feature = "disk-index")] use diskann_disk::QuantizationType; use diskann_providers::storage::{get_compressed_pq_file, get_disk_index_file, get_pq_pivot_file}; use serde::{Deserialize, Serialize}; use crate::{ - inputs::{as_input, post_processor::TopkPostProcessor, Example}, + inputs::{as_input, graph_index::AdaptiveL, post_processor::TopkPostProcessor, Example}, utils::SimilarityMeasure, }; -#[cfg(feature = "disk-index")] -use crate::inputs::graph_index::AdaptiveL; - ////////////// // Registry // ////////////// @@ -72,69 +63,33 @@ pub(crate) struct DiskIndexBuild { pub(crate) save_path: String, } -#[cfg(feature = "disk-index")] #[derive(Debug, Serialize, Deserialize, Default)] pub(crate) struct DiskSearchMode { pub(crate) is_flat_search: bool, #[serde(default)] pub(crate) adaptive_l: Option, + #[serde(default)] + pub(crate) vector_filters_file: Option, + #[serde(default)] + pub(crate) post_processor: Option, } -#[cfg(feature = "disk-index")] impl DiskSearchMode { - pub(crate) fn search_mode<'a>( - &'a self, - has_vector_filters: bool, - vector_filter: &'a HashSet, - post_processor: Option<&TopkPostProcessor>, - ) -> SearchMode<'a> { - let adaptive_l = self.adaptive_l.as_ref().map(|adaptive_l| { - graph::search::AdaptiveL::new(adaptive_l.sample_count.into(), adaptive_l.scale_factor) - .expect("validated adaptive L must construct") - }); - - match ( - self.is_flat_search, - has_vector_filters, - post_processor, - adaptive_l, - ) { - (true, false, _, _) => SearchMode::flat(), - (true, true, _, _) => { - SearchMode::flat_filtered(move |vid: &u32| vector_filter.contains(vid)) - } - (false, false, Some(TopkPostProcessor::DeterminantDiversity(params)), _) => { - SearchMode::diverse_graph(*params) - } - (false, true, Some(TopkPostProcessor::DeterminantDiversity(params)), _) => { - SearchMode::diverse_graph_filtered( - move |vid: &u32| vector_filter.contains(vid), - *params, - ) - } - (false, false, None, Some(adaptive_l)) => { - SearchMode::inline_filter(|_| true, Some(adaptive_l)) - } - (false, true, None, Some(adaptive_l)) => SearchMode::inline_filter( - move |vid: &u32| vector_filter.contains(vid), - Some(adaptive_l), - ), - (false, false, None, None) => SearchMode::graph(), - (false, true, None, None) => { - SearchMode::graph_filtered(move |vid: &u32| vector_filter.contains(vid)) - } - } - } - pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { if let Some(adaptive_l) = self.adaptive_l.as_mut() { adaptive_l.validate(checker)?; } + if let Some(vf) = self.vector_filters_file.as_mut() { + vf.resolve(checker).context("invalid vector_filters_file")?; + } + if let Some(pp) = self.post_processor.as_mut() { + pp.validate(checker) + .context("invalid disk search post processor")?; + } Ok(()) } } -#[cfg(feature = "disk-index")] impl fmt::Display for DiskSearchMode { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let base = if self.is_flat_search { "flat" } else { "graph" }; @@ -155,21 +110,11 @@ pub(crate) struct DiskSearchPhase { pub(crate) beam_width: usize, pub(crate) search_list: Vec, pub(crate) recall_at: u32, - #[cfg(feature = "disk-index")] #[serde(default)] pub(crate) search_mode: DiskSearchMode, - // Backward compatibility for older benchmark inputs that used - // `is_flat_search` directly at the search-phase level. - #[cfg(feature = "disk-index")] - #[serde(default, skip_serializing)] - pub(crate) is_flat_search: Option, - #[cfg(not(feature = "disk-index"))] - pub(crate) is_flat_search: bool, pub(crate) distance: SimilarityMeasure, - pub(crate) vector_filters_file: Option, pub(crate) num_nodes_to_cache: Option, pub(crate) search_io_limit: Option, - pub(crate) post_processor: Option, } ///////// @@ -270,16 +215,7 @@ impl DiskSearchPhase { self.groundtruth .resolve(checker) .context("invalid groundtruth file")?; - if let Some(vf) = self.vector_filters_file.as_mut() { - vf.resolve(checker).context("invalid vector_filters_file")?; - } - #[cfg(feature = "disk-index")] - if let Some(is_flat_search) = self.is_flat_search { - self.search_mode.is_flat_search = is_flat_search; - } - - #[cfg(feature = "disk-index")] self.search_mode .validate(checker) .context("invalid disk search mode")?; @@ -315,11 +251,6 @@ impl DiskSearchPhase { } } - if let Some(pp) = self.post_processor.as_mut() { - pp.validate(checker) - .context("invalid disk search post processor")?; - } - Ok(()) } } @@ -353,20 +284,15 @@ impl Example for DiskIndexOperation { beam_width: 16, recall_at: 10, num_threads: 8, - #[cfg(feature = "disk-index")] search_mode: DiskSearchMode { is_flat_search: false, adaptive_l: None, + vector_filters_file: None, + post_processor: None, }, - #[cfg(feature = "disk-index")] - is_flat_search: None, - #[cfg(not(feature = "disk-index"))] - is_flat_search: false, distance: SimilarityMeasure::SquaredL2, - vector_filters_file: None, num_nodes_to_cache: None, search_io_limit: None, - post_processor: None, }; Self { @@ -478,12 +404,9 @@ impl DiskSearchPhase { write_field!(f, "Beam Width", self.beam_width)?; write_field!(f, "Recall@", self.recall_at)?; write_field!(f, "Threads", self.num_threads)?; - #[cfg(feature = "disk-index")] write_field!(f, "Search Mode", self.search_mode)?; - #[cfg(not(feature = "disk-index"))] - write_field!(f, "Flat Search", self.is_flat_search)?; write_field!(f, "Distance", self.distance)?; - match &self.vector_filters_file { + match &self.search_mode.vector_filters_file { Some(vf) => write_field!(f, "Vector Filters File", vf.display())?, None => write_field!(f, "Vector Filters File", "none")?, } @@ -495,7 +418,7 @@ impl DiskSearchPhase { Some(lim) => write_field!(f, "Search IO Limit", format!("{lim}"))?, None => write_field!(f, "Search IO Limit", "none (defaults to `usize::MAX`)")?, } - match &self.post_processor { + match &self.search_mode.post_processor { Some(pp) => write_field!(f, "Post Processor", pp)?, None => write_field!(f, "Post Processor", "none")?, } From 6fe7448c777ff14123100582caad8810f56c0dc2 Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Wed, 8 Jul 2026 14:15:31 +0800 Subject: [PATCH 02/11] remove unexpected output files from disk-filter test --- diskann-benchmark/src/main.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/diskann-benchmark/src/main.rs b/diskann-benchmark/src/main.rs index e1062b1da4..9933256c9c 100644 --- a/diskann-benchmark/src/main.rs +++ b/diskann-benchmark/src/main.rs @@ -697,6 +697,18 @@ mod tests { prefix_search_directories(&mut raw, &root_directory()); let tempdir = tempfile::tempdir().unwrap(); + + // Redirect each build job's `save_path` into the tempdir so the disk index + // artifacts are not written relative to the process cwd (the repo tree). + let jobs = raw["jobs"] + .as_array_mut() + .expect("\"jobs\" should be an array"); + for (i, job) in jobs.iter_mut().enumerate() { + let save_path = tempdir.path().join(format!("disk_index_filter_job_{i}")); + job["content"]["source"]["save_path"] = + serde_json::Value::String(save_path.to_str().unwrap().to_string()); + } + let input_path = tempdir.path().join("disk-index-filter.json"); save_to_file(&input_path, &raw); let output_path = tempdir.path().join("output.json"); From 29907bf78dd0defc8d5b3e7c27dd7bde62ed786f Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Tue, 14 Jul 2026 18:28:21 +0800 Subject: [PATCH 03/11] Fix disk benchmark search-mode configuration validation --- diskann-benchmark/src/disk_index/search.rs | 44 ++++++++------- diskann-benchmark/src/inputs/disk.rs | 62 ++++++++++++++++++++++ 2 files changed, 87 insertions(+), 19 deletions(-) diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index 342b1adde7..07ec82ce4d 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -164,8 +164,7 @@ impl DiskSearchResult { /// config plus the per-query filter and post-processor supplied at search time. fn build_search_mode<'a>( mode: &'a DiskSearchMode, - has_vector_filters: bool, - vector_filter: &'a HashSet, + vector_filter: Option<&'a HashSet>, post_processor: Option<&TopkPostProcessor>, ) -> SearchMode<'a> { let adaptive_l = mode.adaptive_l.as_ref().map(|adaptive_l| { @@ -175,32 +174,32 @@ fn build_search_mode<'a>( match ( mode.is_flat_search, - has_vector_filters, + vector_filter, post_processor, adaptive_l, ) { - (true, false, _, _) => SearchMode::flat(), - (true, true, _, _) => { + (true, None, _, _) => SearchMode::flat(), + (true, Some(vector_filter), _, _) => { SearchMode::flat_filtered(move |vid: &u32| vector_filter.contains(vid)) } - (false, false, Some(TopkPostProcessor::DeterminantDiversity(params)), _) => { + (false, None, Some(TopkPostProcessor::DeterminantDiversity(params)), _) => { SearchMode::diverse_graph(*params) } - (false, true, Some(TopkPostProcessor::DeterminantDiversity(params)), _) => { + (false, Some(vector_filter), Some(TopkPostProcessor::DeterminantDiversity(params)), _) => { SearchMode::diverse_graph_filtered( move |vid: &u32| vector_filter.contains(vid), *params, ) } - (false, false, None, Some(adaptive_l)) => { + (false, None, None, Some(adaptive_l)) => { SearchMode::inline_filter(|_| true, Some(adaptive_l)) } - (false, true, None, Some(adaptive_l)) => SearchMode::inline_filter( + (false, Some(vector_filter), None, Some(adaptive_l)) => SearchMode::inline_filter( move |vid: &u32| vector_filter.contains(vid), Some(adaptive_l), ), - (false, false, None, None) => SearchMode::graph(), - (false, true, None, None) => { + (false, None, None, None) => SearchMode::graph(), + (false, Some(vector_filter), None, None) => { SearchMode::graph_filtered(move |vid: &u32| vector_filter.contains(vid)) } } @@ -236,12 +235,18 @@ where let vector_filters = match &search_params.search_mode.vector_filters_file { Some(vector_filters_file) => { let vector_filters_file = vector_filters_file.to_string_lossy().to_string(); - search_index_utils::load_vector_filters(storage_provider, &vector_filters_file)? + Some(search_index_utils::load_vector_filters( + storage_provider, + &vector_filters_file, + )?) } - None => vec![HashSet::::new(); num_queries], + None => None, }; - if vector_filters.len() != num_queries { + if vector_filters + .as_ref() + .is_some_and(|filters| filters.len() != num_queries) + { anyhow::bail!("Mismatch in query and vector filter sizes"); } @@ -307,7 +312,7 @@ where let zipped = queries .par_row_iter() - .zip(vector_filters.par_iter()) + .enumerate() .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()) @@ -315,15 +320,16 @@ where zipped.for_each_in_pool( pool.as_ref(), - |(((((q, vf), id_chunk), dist_chunk), stats), rc)| { + |(((((query_index, q), 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.search_mode.vector_filters_file.is_some(); + let vector_filter = vector_filters + .as_ref() + .and_then(|filters| filters.get(query_index)); let mode: SearchMode<'_> = build_search_mode( &search_params.search_mode, - has_filter, - vf, + vector_filter, search_params.search_mode.post_processor.as_ref(), ); diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index 44b7d61d57..d89feff3cd 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -76,6 +76,8 @@ pub(crate) struct DiskSearchMode { impl DiskSearchMode { pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { + self.validate_compatibility()?; + if let Some(adaptive_l) = self.adaptive_l.as_mut() { adaptive_l.validate(checker)?; } @@ -88,6 +90,21 @@ impl DiskSearchMode { } Ok(()) } + + fn validate_compatibility(&self) -> Result<(), anyhow::Error> { + if !self.is_flat_search { + return Ok(()); + } + + match (self.adaptive_l.is_some(), self.post_processor.is_some()) { + (false, false) => Ok(()), + (true, false) => anyhow::bail!("flat disk search does not support adaptive_l"), + (false, true) => anyhow::bail!("flat disk search does not support post_processor"), + (true, true) => { + anyhow::bail!("flat disk search does not support adaptive_l or post_processor") + } + } + } } impl fmt::Display for DiskSearchMode { @@ -101,6 +118,51 @@ impl fmt::Display for DiskSearchMode { } } +#[cfg(test)] +mod tests { + use std::num::NonZeroUsize; + + use super::*; + + #[test] + fn flat_disk_search_rejects_adaptive_l() { + let mode = DiskSearchMode { + is_flat_search: true, + adaptive_l: Some(AdaptiveL { + sample_count: NonZeroUsize::MIN, + scale_factor: 1.0, + }), + vector_filters_file: None, + post_processor: None, + }; + + let err = mode + .validate_compatibility() + .expect_err("flat search with adaptive_l must be invalid"); + assert!(err.to_string().contains("does not support adaptive_l")); + } + + #[test] + fn flat_disk_search_rejects_post_processor() { + let mode: DiskSearchMode = serde_json::from_str( + r#"{ + "is_flat_search": true, + "post_processor": { + "type": "determinant-diversity", + "power": 1.0, + "eta": 0.0 + } + }"#, + ) + .expect("test post-processor configuration must deserialize"); + + let err = mode + .validate_compatibility() + .expect_err("flat search with a post_processor must be invalid"); + assert!(err.to_string().contains("does not support post_processor")); + } +} + /// Search phase configuration #[derive(Debug, Deserialize, Serialize)] pub(crate) struct DiskSearchPhase { From ca59956d36cf08244722a939ab473b99f07a1d07 Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Tue, 14 Jul 2026 18:36:36 +0800 Subject: [PATCH 04/11] force not tolerant for old is_flat_search json config --- diskann-benchmark/src/inputs/disk.rs | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index d89feff3cd..b6744a29c9 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -161,10 +161,30 @@ mod tests { .expect_err("flat search with a post_processor must be invalid"); assert!(err.to_string().contains("does not support post_processor")); } + + #[test] + fn disk_search_phase_rejects_legacy_phase_level_search_mode_fields() { + let error = serde_json::from_str::( + r#"{ + "queries": "queries.fbin", + "groundtruth": "groundtruth.bin", + "num_threads": 1, + "beam_width": 1, + "search_list": [1], + "recall_at": 1, + "distance": "squared_l2", + "is_flat_search": true + }"#, + ) + .expect_err("legacy phase-level search settings must be rejected"); + + assert!(error.to_string().contains("is_flat_search")); + } } /// Search phase configuration #[derive(Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] pub(crate) struct DiskSearchPhase { pub(crate) queries: InputFile, pub(crate) groundtruth: InputFile, From 2e088a28f723320185bda3e09bcb5e250120d5f8 Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Tue, 4 Aug 2026 17:04:51 +0800 Subject: [PATCH 05/11] update benchmark input for DiskSearchMode, change to enum structure --- .../disk-index-determinant-diversity.json | 2 +- .../example/disk-index-filter.json | 4 +- diskann-benchmark/example/disk-index.json | 4 +- .../openai-100K-disk-index.json | 2 +- .../wikipedia-100K-disk-index.json | 2 +- diskann-benchmark/src/disk_index/search.rs | 95 +++++----- diskann-benchmark/src/inputs/disk.rs | 171 ++++++++++-------- 7 files changed, 155 insertions(+), 125 deletions(-) diff --git a/diskann-benchmark/example/disk-index-determinant-diversity.json b/diskann-benchmark/example/disk-index-determinant-diversity.json index a8c236a99e..5dd12e87fc 100644 --- a/diskann-benchmark/example/disk-index-determinant-diversity.json +++ b/diskann-benchmark/example/disk-index-determinant-diversity.json @@ -28,7 +28,7 @@ "recall_at": 10, "num_threads": 1, "search_mode": { - "is_flat_search": false, + "mode": "graph", "post_processor": { "type": "determinant-diversity", "power": 2.0, diff --git a/diskann-benchmark/example/disk-index-filter.json b/diskann-benchmark/example/disk-index-filter.json index cdb9842d7f..ee55aeffa2 100644 --- a/diskann-benchmark/example/disk-index-filter.json +++ b/diskann-benchmark/example/disk-index-filter.json @@ -28,7 +28,7 @@ "recall_at": 10, "num_threads": 1, "search_mode": { - "is_flat_search": false, + "mode": "graph", "vector_filters_file": "disk_index_10pts_idx_uint32_range_res_r_100000.bin" }, "distance": "squared_l2" @@ -60,7 +60,7 @@ "recall_at": 10, "num_threads": 1, "search_mode": { - "is_flat_search": true, + "mode": "flat", "vector_filters_file": "disk_index_10pts_idx_uint32_range_res_r_100000.bin" }, "distance": "squared_l2" diff --git a/diskann-benchmark/example/disk-index.json b/diskann-benchmark/example/disk-index.json index c8a0c7b7d9..d2d3222e1c 100644 --- a/diskann-benchmark/example/disk-index.json +++ b/diskann-benchmark/example/disk-index.json @@ -27,7 +27,7 @@ "beam_width": 4, "recall_at": 10, "num_threads": 1, - "search_mode": { "is_flat_search": false }, + "search_mode": { "mode": "graph" }, "distance": "squared_l2" } } @@ -47,7 +47,7 @@ "beam_width": 4, "recall_at": 10, "num_threads": 1, - "search_mode": { "is_flat_search": true }, + "search_mode": { "mode": "flat" }, "distance": "squared_l2" } } diff --git a/diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json b/diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json index 2417828997..2e535d4f4c 100644 --- a/diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json +++ b/diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json @@ -29,7 +29,7 @@ "beam_width": 4, "recall_at": 100, "num_threads": 4, - "search_mode": { "is_flat_search": false }, + "search_mode": { "mode": "graph" }, "distance": "squared_l2" } } diff --git a/diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json b/diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json index aa00ab4ac2..3593ad8d62 100644 --- a/diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json +++ b/diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json @@ -29,7 +29,7 @@ "beam_width": 4, "recall_at": 100, "num_threads": 4, - "search_mode": { "is_flat_search": false }, + "search_mode": { "mode": "graph" }, "distance": "inner_product" } } diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index 5ae45c2daa..65aab80d20 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -165,42 +165,49 @@ impl DiskSearchResult { fn build_search_mode<'a>( mode: &'a DiskSearchMode, vector_filter: Option<&'a HashSet>, - post_processor: Option<&TopkPostProcessor>, ) -> SearchMode<'a> { - let adaptive_l = mode.adaptive_l.as_ref().map(|adaptive_l| { - graph::search::AdaptiveL::new(adaptive_l.sample_count.into(), adaptive_l.scale_factor) - .expect("validated adaptive L must construct") - }); - - match ( - mode.is_flat_search, - vector_filter, - post_processor, - adaptive_l, - ) { - (true, None, _, _) => SearchMode::flat(), - (true, Some(vector_filter), _, _) => { - SearchMode::flat_filtered(move |vid: &u32| vector_filter.contains(vid)) - } - (false, None, Some(TopkPostProcessor::DeterminantDiversity(params)), _) => { - SearchMode::diverse_graph(*params) - } - (false, Some(vector_filter), Some(TopkPostProcessor::DeterminantDiversity(params)), _) => { - SearchMode::diverse_graph_filtered( - move |vid: &u32| vector_filter.contains(vid), - *params, - ) - } - (false, None, None, Some(adaptive_l)) => { - SearchMode::inline_filter(|_| true, Some(adaptive_l)) - } - (false, Some(vector_filter), None, Some(adaptive_l)) => SearchMode::inline_filter( - move |vid: &u32| vector_filter.contains(vid), - Some(adaptive_l), - ), - (false, None, None, None) => SearchMode::graph(), - (false, Some(vector_filter), None, None) => { - SearchMode::graph_filtered(move |vid: &u32| vector_filter.contains(vid)) + match mode { + DiskSearchMode::Flat { .. } => match vector_filter { + None => SearchMode::flat(), + Some(vector_filter) => { + SearchMode::flat_filtered(move |vid: &u32| vector_filter.contains(vid)) + } + }, + DiskSearchMode::Graph { + adaptive_l, + post_processor, + .. + } => { + let adaptive_l = adaptive_l.as_ref().map(|adaptive_l| { + graph::search::AdaptiveL::new( + adaptive_l.sample_count.into(), + adaptive_l.scale_factor, + ) + .expect("validated adaptive L must construct") + }); + + match (post_processor, adaptive_l, vector_filter) { + (Some(TopkPostProcessor::DeterminantDiversity(params)), _, None) => { + SearchMode::diverse_graph(*params) + } + (Some(TopkPostProcessor::DeterminantDiversity(params)), _, Some(vector_filter)) => { + SearchMode::diverse_graph_filtered( + move |vid: &u32| vector_filter.contains(vid), + *params, + ) + } + (None, Some(adaptive_l), None) => { + SearchMode::inline_filter(|_| true, Some(adaptive_l)) + } + (None, Some(adaptive_l), Some(vector_filter)) => SearchMode::inline_filter( + move |vid: &u32| vector_filter.contains(vid), + Some(adaptive_l), + ), + (None, None, None) => SearchMode::graph(), + (None, None, Some(vector_filter)) => { + SearchMode::graph_filtered(move |vid: &u32| vector_filter.contains(vid)) + } + } } } } @@ -232,7 +239,7 @@ where let num_queries = queries.nrows(); // Load the vector filters - let vector_filters = match &search_params.search_mode.vector_filters_file { + let vector_filters = match search_params.search_mode.vector_filters_file() { Some(vector_filters_file) => { let vector_filters_file = vector_filters_file.to_string_lossy().to_string(); Some(search_index_utils::load_vector_filters( @@ -252,7 +259,7 @@ where // Prepare ground truth context let gt_context = prepare_ground_truth_context( - search_params.search_mode.vector_filters_file.is_some(), + search_params.search_mode.vector_filters_file().is_some(), &search_params.groundtruth, search_params.recall_at, storage_provider, @@ -321,17 +328,11 @@ where zipped.for_each_in_pool( pool.as_ref(), |(((((query_index, q), 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 vector_filter = vector_filters .as_ref() .and_then(|filters| filters.get(query_index)); - let mode: SearchMode<'_> = build_search_mode( - &search_params.search_mode, - vector_filter, - search_params.search_mode.post_processor.as_ref(), - ); + let mode: SearchMode<'_> = + build_search_mode(&search_params.search_mode, vector_filter); match searcher.search( q, @@ -404,9 +405,9 @@ where num_threads: search_params.num_threads, beam_width: search_params.beam_width, recall_at: search_params.recall_at, - is_flat_search: search_params.search_mode.is_flat_search, + is_flat_search: search_params.search_mode.is_flat_search(), distance: search_params.distance, - uses_vector_filters: search_params.search_mode.vector_filters_file.is_some(), + uses_vector_filters: search_params.search_mode.vector_filters_file().is_some(), num_nodes_to_cache: search_params.num_nodes_to_cache, search_results_per_l, span_metrics, diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index b6744a29c9..2223928ce0 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -63,103 +63,133 @@ pub(crate) struct DiskIndexBuild { pub(crate) save_path: String, } -#[derive(Debug, Serialize, Deserialize, Default)] -pub(crate) struct DiskSearchMode { - pub(crate) is_flat_search: bool, - #[serde(default)] - pub(crate) adaptive_l: Option, - #[serde(default)] - pub(crate) vector_filters_file: Option, - #[serde(default)] - pub(crate) post_processor: Option, +/// Disk search mode. The `flat` / `graph` split is encoded as an enum so that +/// combinations invalid for flat scan (adaptive L, post-processor) are +/// unrepresentable rather than rejected at validation time. Mirrors the spirit +/// of `diskann_disk::search::search_mode::SearchMode` at the config level. +#[derive(Debug, Serialize, Deserialize)] +#[serde(tag = "mode", rename_all = "kebab-case")] +pub(crate) enum DiskSearchMode { + /// Brute-force flat scan, optionally restricted by a per-query vector filter. + Flat { + #[serde(default)] + vector_filters_file: Option, + }, + /// Greedy graph search, optionally with inline adaptive-L, a per-query + /// vector filter, and/or a top-k post-processor. + Graph { + #[serde(default)] + adaptive_l: Option, + #[serde(default)] + vector_filters_file: Option, + #[serde(default)] + post_processor: Option, + }, +} + +impl Default for DiskSearchMode { + fn default() -> Self { + Self::Graph { + adaptive_l: None, + vector_filters_file: None, + post_processor: None, + } + } } impl DiskSearchMode { - pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { - self.validate_compatibility()?; + pub(crate) fn is_flat_search(&self) -> bool { + matches!(self, Self::Flat { .. }) + } - if let Some(adaptive_l) = self.adaptive_l.as_mut() { - adaptive_l.validate(checker)?; - } - if let Some(vf) = self.vector_filters_file.as_mut() { - vf.resolve(checker).context("invalid vector_filters_file")?; - } - if let Some(pp) = self.post_processor.as_mut() { - pp.validate(checker) - .context("invalid disk search post processor")?; + pub(crate) fn vector_filters_file(&self) -> Option<&InputFile> { + match self { + Self::Flat { + vector_filters_file, + } + | Self::Graph { + vector_filters_file, + .. + } => vector_filters_file.as_ref(), } - Ok(()) } - fn validate_compatibility(&self) -> Result<(), anyhow::Error> { - if !self.is_flat_search { - return Ok(()); + pub(crate) fn post_processor(&self) -> Option<&TopkPostProcessor> { + match self { + Self::Flat { .. } => None, + Self::Graph { post_processor, .. } => post_processor.as_ref(), } + } - match (self.adaptive_l.is_some(), self.post_processor.is_some()) { - (false, false) => Ok(()), - (true, false) => anyhow::bail!("flat disk search does not support adaptive_l"), - (false, true) => anyhow::bail!("flat disk search does not support post_processor"), - (true, true) => { - anyhow::bail!("flat disk search does not support adaptive_l or post_processor") + pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { + match self { + Self::Flat { + vector_filters_file, + } => { + if let Some(vf) = vector_filters_file.as_mut() { + vf.resolve(checker).context("invalid vector_filters_file")?; + } + } + Self::Graph { + adaptive_l, + vector_filters_file, + post_processor, + } => { + if let Some(adaptive_l) = adaptive_l.as_mut() { + adaptive_l.validate(checker)?; + } + if let Some(vf) = vector_filters_file.as_mut() { + vf.resolve(checker).context("invalid vector_filters_file")?; + } + if let Some(pp) = post_processor.as_mut() { + pp.validate(checker) + .context("invalid disk search post processor")?; + } } } + Ok(()) } } impl fmt::Display for DiskSearchMode { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let base = if self.is_flat_search { "flat" } else { "graph" }; - if self.adaptive_l.is_some() { - write!(f, "{} + adaptive-l", base) - } else { - write!(f, "{}", base) + match self { + Self::Flat { .. } => write!(f, "flat"), + Self::Graph { adaptive_l, .. } if adaptive_l.is_some() => write!(f, "graph + adaptive-l"), + Self::Graph { .. } => write!(f, "graph"), } } } #[cfg(test)] mod tests { - use std::num::NonZeroUsize; - use super::*; + // The flat/graph split is an enum, so combinations invalid for flat scan + // (adaptive L, post-processor) are unrepresentable by construction — there + // is no runtime compatibility check left to test. These round-trip tests + // just pin the JSON shape. #[test] - fn flat_disk_search_rejects_adaptive_l() { - let mode = DiskSearchMode { - is_flat_search: true, - adaptive_l: Some(AdaptiveL { - sample_count: NonZeroUsize::MIN, - scale_factor: 1.0, - }), - vector_filters_file: None, - post_processor: None, - }; - - let err = mode - .validate_compatibility() - .expect_err("flat search with adaptive_l must be invalid"); - assert!(err.to_string().contains("does not support adaptive_l")); + fn flat_disk_search_deserializes_without_graph_only_fields() { + let mode: DiskSearchMode = + serde_json::from_str(r#"{ "mode": "flat" }"#).expect("flat mode must deserialize"); + assert!(mode.is_flat_search()); + assert!(mode.post_processor().is_none()); } #[test] - fn flat_disk_search_rejects_post_processor() { + fn graph_disk_search_deserializes_with_adaptive_l() { let mode: DiskSearchMode = serde_json::from_str( - r#"{ - "is_flat_search": true, - "post_processor": { - "type": "determinant-diversity", - "power": 1.0, - "eta": 0.0 - } - }"#, + r#"{ "mode": "graph", "adaptive_l": { "sample_count": 1, "scale_factor": 2.0 } }"#, ) - .expect("test post-processor configuration must deserialize"); - - let err = mode - .validate_compatibility() - .expect_err("flat search with a post_processor must be invalid"); - assert!(err.to_string().contains("does not support post_processor")); + .expect("graph mode with adaptive_l must deserialize"); + assert!(matches!( + mode, + DiskSearchMode::Graph { + adaptive_l: Some(_), + .. + } + )); } #[test] @@ -366,8 +396,7 @@ impl Example for DiskIndexOperation { beam_width: 16, recall_at: 10, num_threads: 8, - search_mode: DiskSearchMode { - is_flat_search: false, + search_mode: DiskSearchMode::Graph { adaptive_l: None, vector_filters_file: None, post_processor: None, @@ -488,7 +517,7 @@ impl DiskSearchPhase { write_field!(f, "Threads", self.num_threads)?; write_field!(f, "Search Mode", self.search_mode)?; write_field!(f, "Distance", self.distance)?; - match &self.search_mode.vector_filters_file { + match self.search_mode.vector_filters_file() { Some(vf) => write_field!(f, "Vector Filters File", vf.display())?, None => write_field!(f, "Vector Filters File", "none")?, } @@ -500,7 +529,7 @@ impl DiskSearchPhase { Some(lim) => write_field!(f, "Search IO Limit", format!("{lim}"))?, None => write_field!(f, "Search IO Limit", "none (defaults to `usize::MAX`)")?, } - match &self.search_mode.post_processor { + match self.search_mode.post_processor() { Some(pp) => write_field!(f, "Post Processor", pp)?, None => write_field!(f, "Post Processor", "none")?, } From acd030b648aa72df44e9d84307df98aa97a9e7b4 Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Tue, 4 Aug 2026 17:09:59 +0800 Subject: [PATCH 06/11] fix format issue --- diskann-benchmark/src/inputs/disk.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index 2223928ce0..2f21d31a50 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -155,7 +155,9 @@ impl fmt::Display for DiskSearchMode { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::Flat { .. } => write!(f, "flat"), - Self::Graph { adaptive_l, .. } if adaptive_l.is_some() => write!(f, "graph + adaptive-l"), + Self::Graph { adaptive_l, .. } if adaptive_l.is_some() => { + write!(f, "graph + adaptive-l") + } Self::Graph { .. } => write!(f, "graph"), } } From 92d8caec99d07259ad205ee2191e8ac1ff074dd3 Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Tue, 4 Aug 2026 17:54:43 +0800 Subject: [PATCH 07/11] remove unused function --- diskann-benchmark/src/disk_index/search.rs | 2 +- diskann-benchmark/src/inputs/disk.rs | 6 +----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index 65aab80d20..29e1428348 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -405,7 +405,7 @@ where num_threads: search_params.num_threads, beam_width: search_params.beam_width, recall_at: search_params.recall_at, - is_flat_search: search_params.search_mode.is_flat_search(), + is_flat_search: matches!(search_params.search_mode, DiskSearchMode::Flat { .. }), distance: search_params.distance, uses_vector_filters: search_params.search_mode.vector_filters_file().is_some(), num_nodes_to_cache: search_params.num_nodes_to_cache, diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index 2f21d31a50..739d02e281 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -98,10 +98,6 @@ impl Default for DiskSearchMode { } impl DiskSearchMode { - pub(crate) fn is_flat_search(&self) -> bool { - matches!(self, Self::Flat { .. }) - } - pub(crate) fn vector_filters_file(&self) -> Option<&InputFile> { match self { Self::Flat { @@ -175,7 +171,7 @@ mod tests { fn flat_disk_search_deserializes_without_graph_only_fields() { let mode: DiskSearchMode = serde_json::from_str(r#"{ "mode": "flat" }"#).expect("flat mode must deserialize"); - assert!(mode.is_flat_search()); + assert!(matches!(mode, DiskSearchMode::Flat { .. })); assert!(mode.post_processor().is_none()); } From a624d45fd9dee89d71322a4ba9917c408800194f Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Wed, 5 Aug 2026 10:17:39 +0800 Subject: [PATCH 08/11] change benchmark baseline config file to use main branch config file --- .github/workflows/disk-benchmarks.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/disk-benchmarks.yml b/.github/workflows/disk-benchmarks.yml index 29c0a7af86..b7960cf9ea 100644 --- a/.github/workflows/disk-benchmarks.yml +++ b/.github/workflows/disk-benchmarks.yml @@ -116,7 +116,7 @@ jobs: working-directory: baseline run: | cargo run -p diskann-benchmark --features disk-index --release -- \ - run --input-file ../diskann_rust/${{ env.PERF_INPUTS }}/${{ matrix.config }} \ + run --input-file ${{ env.PERF_INPUTS }}/${{ matrix.config }} \ --output-file target/tmp/${{ matrix.dataset }}_baseline.json - name: Run current branch benchmark @@ -144,4 +144,4 @@ jobs: 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 From 88152edfa6f3d980cc837e7ec2c9f54b139299ed Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Thu, 10 Sep 2026 14:55:43 +0800 Subject: [PATCH 09/11] Change benchmark json input to use DiskSearchStrategy struct --- .../disk-index-determinant-diversity.json | 2 +- .../example/disk-index-filter.json | 8 +- diskann-benchmark/src/disk_index/search.rs | 239 +++++++++++++----- diskann-benchmark/src/inputs/disk.rs | 234 +++++++++++++---- diskann-benchmark/src/inputs/graph_index.rs | 2 +- diskann-benchmark/src/main.rs | 164 +++++++++--- 6 files changed, 503 insertions(+), 146 deletions(-) diff --git a/diskann-benchmark/example/disk-index-determinant-diversity.json b/diskann-benchmark/example/disk-index-determinant-diversity.json index 06cbc5fc0b..cab0f93fd3 100644 --- a/diskann-benchmark/example/disk-index-determinant-diversity.json +++ b/diskann-benchmark/example/disk-index-determinant-diversity.json @@ -28,7 +28,7 @@ "recall_at": 10, "num_threads": 1, "search_mode": { - "mode": "graph", + "mode": "graph-diverse", "post_processor": { "type": "determinant-diversity", "power": 2.0, diff --git a/diskann-benchmark/example/disk-index-filter.json b/diskann-benchmark/example/disk-index-filter.json index ee55aeffa2..167b4646dc 100644 --- a/diskann-benchmark/example/disk-index-filter.json +++ b/diskann-benchmark/example/disk-index-filter.json @@ -28,8 +28,12 @@ "recall_at": 10, "num_threads": 1, "search_mode": { - "mode": "graph", - "vector_filters_file": "disk_index_10pts_idx_uint32_range_res_r_100000.bin" + "mode": "graph-inline-filter", + "vector_filters_file": "disk_index_10pts_idx_uint32_range_res_r_100000.bin", + "adaptive_l": { + "sample_count": 10, + "scale_factor": 16.0 + } }, "distance": "squared_l2" } diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index 29e1428348..43fa1d4ac1 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -19,7 +19,7 @@ use diskann_disk::{ disk_provider::DiskIndexSearcher, disk_vertex_provider_factory::DiskVertexProviderFactory, }, - search_mode::SearchMode, + search_mode::{SearchMode, SearchPredicate}, }, storage::disk_index_reader::DiskIndexReader, utils::{instrumentation::PerfLogger, statistics, QueryStatistics}, @@ -33,28 +33,83 @@ use diskann_providers::{ }; use diskann_tools::utils::{search_index_utils, KRecallAtN}; use diskann_utils::views::Matrix; +use scopeguard::defer; use serde::{Deserialize, Serialize}; use crate::{ disk_index::json_spancollector::JsonSpanCollector, - inputs::disk::{DiskIndexLoad, DiskSearchMode, DiskSearchPhase}, + inputs::disk::{DiskIndexLoad, DiskSearchMode, DiskSearchPhase, DiskSearchStrategy}, inputs::post_processor::TopkPostProcessor, utils::{datafiles, SimilarityMeasure}, }; -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Debug)] pub(super) struct DiskSearchStats { pub(super) num_threads: usize, pub(super) beam_width: usize, pub(super) recall_at: u32, - pub(crate) is_flat_search: bool, + pub(crate) search_strategy: DiskSearchStrategy, pub(crate) distance: SimilarityMeasure, - pub(crate) uses_vector_filters: bool, pub(super) num_nodes_to_cache: Option, pub(super) search_results_per_l: Vec, span_metrics: serde_json::Value, } +#[derive(Deserialize)] +struct RawDiskSearchStats { + num_threads: usize, + beam_width: usize, + recall_at: u32, + #[serde(default)] + search_strategy: Option, + #[serde(default)] + is_flat_search: Option, + distance: SimilarityMeasure, + #[serde(default)] + uses_vector_filters: Option, + num_nodes_to_cache: Option, + search_results_per_l: Vec, + span_metrics: serde_json::Value, +} + +impl<'de> Deserialize<'de> for DiskSearchStats { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let raw = RawDiskSearchStats::deserialize(deserializer)?; + let search_strategy = if let Some(strategy) = raw.search_strategy { + strategy + } else if let Some(is_flat) = raw.is_flat_search { + let uses_vector_filters = raw.uses_vector_filters.unwrap_or(false); + if is_flat { + DiskSearchStrategy::Flat { + uses_vector_filters, + } + } else { + DiskSearchStrategy::Graph { + uses_vector_filters, + } + } + } else { + return Err(serde::de::Error::custom( + "missing search_strategy or legacy is_flat_search", + )); + }; + + Ok(Self { + num_threads: raw.num_threads, + beam_width: raw.beam_width, + recall_at: raw.recall_at, + search_strategy, + distance: raw.distance, + num_nodes_to_cache: raw.num_nodes_to_cache, + search_results_per_l: raw.search_results_per_l, + span_metrics: raw.span_metrics, + }) + } +} + #[derive(Serialize, Deserialize, Debug)] pub(super) struct DiskSearchResult { pub(super) search_l: u32, @@ -160,56 +215,43 @@ impl DiskSearchResult { } } -/// Construct the disk [`SearchMode`] from the JSON-driven [`DiskSearchMode`] -/// config plus the per-query filter and post-processor supplied at search time. +/// Construct the backend [`SearchMode`] from the JSON-configured strategy, +/// the per-query vector filter, and the pre-validated adaptive-L settings. fn build_search_mode<'a>( - mode: &'a DiskSearchMode, + mode: &DiskSearchMode, vector_filter: Option<&'a HashSet>, -) -> SearchMode<'a> { - match mode { - DiskSearchMode::Flat { .. } => match vector_filter { - None => SearchMode::flat(), - Some(vector_filter) => { - SearchMode::flat_filtered(move |vid: &u32| vector_filter.contains(vid)) - } - }, - DiskSearchMode::Graph { - adaptive_l, - post_processor, - .. - } => { - let adaptive_l = adaptive_l.as_ref().map(|adaptive_l| { - graph::search::AdaptiveL::new( - adaptive_l.sample_count.into(), - adaptive_l.scale_factor, - ) - .expect("validated adaptive L must construct") - }); - - match (post_processor, adaptive_l, vector_filter) { - (Some(TopkPostProcessor::DeterminantDiversity(params)), _, None) => { - SearchMode::diverse_graph(*params) - } - (Some(TopkPostProcessor::DeterminantDiversity(params)), _, Some(vector_filter)) => { - SearchMode::diverse_graph_filtered( - move |vid: &u32| vector_filter.contains(vid), - *params, - ) - } - (None, Some(adaptive_l), None) => { - SearchMode::inline_filter(|_| true, Some(adaptive_l)) - } - (None, Some(adaptive_l), Some(vector_filter)) => SearchMode::inline_filter( - move |vid: &u32| vector_filter.contains(vid), - Some(adaptive_l), - ), - (None, None, None) => SearchMode::graph(), - (None, None, Some(vector_filter)) => { - SearchMode::graph_filtered(move |vid: &u32| vector_filter.contains(vid)) - } + adaptive_l: Option<&graph::search::AdaptiveL>, +) -> anyhow::Result> { + if adaptive_l.is_some() && !matches!(mode, DiskSearchMode::GraphInlineFilter { .. }) { + anyhow::bail!("adaptive-L is only valid for inline-filter search"); + } + + let filter = vector_filter.map(|vector_filter| { + Box::new(move |vid: &u32| vector_filter.contains(vid)) as SearchPredicate<'a> + }); + + let mode = match mode { + DiskSearchMode::Flat { .. } => SearchMode::FlatScan { filter }, + DiskSearchMode::Graph { .. } => SearchMode::Graph { filter }, + DiskSearchMode::GraphInlineFilter { .. } => { + let vector_filter = vector_filter.ok_or_else(|| { + anyhow::anyhow!("inline-filter search requires a vector filter for every query") + })?; + SearchMode::inline_filter( + move |vid: &u32| vector_filter.contains(vid), + adaptive_l.cloned(), + ) + } + DiskSearchMode::GraphDiverse { post_processor, .. } => { + let TopkPostProcessor::DeterminantDiversity(params) = post_processor; + SearchMode::DiverseGraph { + filter, + params: *params, } } - } + }; + + Ok(mode) } pub(super) fn search_disk_index( @@ -230,6 +272,9 @@ where global::set_tracer_provider(provider.clone()); Some((collector, provider)) }; + defer! { + global::set_tracer_provider(previous_tracer_provider); + } // Use PerfLogger for consistent checkpoint logging let mut logger = PerfLogger::new("search_disk_index", true); @@ -296,6 +341,7 @@ where logger.log_checkpoint("index_loaded"); + let adaptive_l = search_params.search_mode.adaptive_l()?; let pool = create_thread_pool(search_params.num_threads)?; let mut search_results_per_l = Vec::with_capacity(search_params.search_list.len()); let has_any_search_failed = AtomicBool::new(false); @@ -325,14 +371,17 @@ where .zip(statistics_vec.par_iter_mut()) .zip(result_counts.par_iter_mut()); - zipped.for_each_in_pool( + zipped.try_for_each_in_pool( pool.as_ref(), |(((((query_index, q), id_chunk), dist_chunk), stats), rc)| { let vector_filter = vector_filters .as_ref() .and_then(|filters| filters.get(query_index)); - let mode: SearchMode<'_> = - build_search_mode(&search_params.search_mode, vector_filter); + let mode = build_search_mode( + &search_params.search_mode, + vector_filter, + adaptive_l.as_ref(), + )?; match searcher.search( q, @@ -366,8 +415,10 @@ where has_any_search_failed.store(true, std::sync::atomic::Ordering::Release); } } + + Ok::<(), anyhow::Error>(()) }, - ); + )?; let total_time = start.elapsed(); if has_any_search_failed.load(std::sync::atomic::Ordering::Acquire) { @@ -399,15 +450,12 @@ where serde_json::json!({ "span_data": [] }) }; - global::set_tracer_provider(previous_tracer_provider); - Ok(DiskSearchStats { num_threads: search_params.num_threads, beam_width: search_params.beam_width, recall_at: search_params.recall_at, - is_flat_search: matches!(search_params.search_mode, DiskSearchMode::Flat { .. }), + search_strategy: search_params.search_mode.strategy(), distance: search_params.distance, - uses_vector_filters: search_params.search_mode.vector_filters_file().is_some(), num_nodes_to_cache: search_params.num_nodes_to_cache, search_results_per_l, span_metrics, @@ -489,9 +537,8 @@ 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, "Flat search, : {}", self.is_flat_search)?; + writeln!(f, "Search strategy, : {}", self.search_strategy)?; writeln!(f, "Distance, : {}", self.distance)?; - writeln!(f, "Vector filters, : {}", self.uses_vector_filters)?; writeln!( f, "Nodes to cache, : {}", @@ -538,3 +585,73 @@ impl fmt::Display for DiskSearchStats { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn legacy_search_stats_deserialize_to_strategy() { + let stats: DiskSearchStats = serde_json::from_value(serde_json::json!({ + "num_threads": 1, + "beam_width": 4, + "recall_at": 10, + "is_flat_search": false, + "distance": "squared_l2", + "num_nodes_to_cache": null, + "search_results_per_l": [], + "span_metrics": {} + })) + .unwrap(); + + assert!(matches!( + stats.search_strategy, + DiskSearchStrategy::Graph { + uses_vector_filters: false + } + )); + } + + #[test] + fn new_search_stats_serialize_only_complete_strategy() { + let stats: DiskSearchStats = serde_json::from_value(serde_json::json!({ + "num_threads": 1, + "beam_width": 4, + "recall_at": 10, + "search_strategy": { + "mode": "graph-inline-filter", + "adaptive_l": { + "sample_count": 10, + "scale_factor": 16.0 + } + }, + "is_flat_search": false, + "uses_vector_filters": true, + "distance": "squared_l2", + "num_nodes_to_cache": null, + "search_results_per_l": [], + "span_metrics": {} + })) + .unwrap(); + + let value = serde_json::to_value(stats).unwrap(); + assert_eq!(value["search_strategy"]["mode"], "graph-inline-filter"); + assert_eq!(value["search_strategy"]["adaptive_l"]["sample_count"], 10); + assert!(value.get("is_flat_search").is_none()); + assert!(value.get("uses_vector_filters").is_none()); + } + + #[test] + fn inline_search_mode_requires_query_filter() { + let mode: DiskSearchMode = serde_json::from_value(serde_json::json!({ + "mode": "graph-inline-filter", + "vector_filters_file": "filters.bin" + })) + .unwrap(); + + let error = build_search_mode(&mode, None, None) + .err() + .expect("inline-filter mode must reject a missing query filter"); + assert!(error.to_string().contains("requires a vector filter")); + } +} diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index 739d02e281..040ad90bb1 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -63,11 +63,9 @@ pub(crate) struct DiskIndexBuild { pub(crate) save_path: String, } -/// Disk search mode. The `flat` / `graph` split is encoded as an enum so that -/// combinations invalid for flat scan (adaptive L, post-processor) are -/// unrepresentable rather than rejected at validation time. Mirrors the spirit -/// of `diskann_disk::search::search_mode::SearchMode` at the config level. -#[derive(Debug, Serialize, Deserialize)] +/// Disk search mode, modeled after the four backend search strategies. +/// Strategy-specific settings live only on the variants that use them. +#[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "mode", rename_all = "kebab-case")] pub(crate) enum DiskSearchMode { /// Brute-force flat scan, optionally restricted by a per-query vector filter. @@ -75,24 +73,49 @@ pub(crate) enum DiskSearchMode { #[serde(default)] vector_filters_file: Option, }, - /// Greedy graph search, optionally with inline adaptive-L, a per-query - /// vector filter, and/or a top-k post-processor. + /// Greedy graph search, optionally post-filtered by a per-query vector filter. Graph { + #[serde(default)] + vector_filters_file: Option, + }, + /// Graph search that checks a required vector filter during traversal. + GraphInlineFilter { + vector_filters_file: InputFile, #[serde(default)] adaptive_l: Option, + }, + /// Graph search followed by determinant-diversity selection. + GraphDiverse { #[serde(default)] vector_filters_file: Option, + post_processor: TopkPostProcessor, + }, +} + +/// Path-independent search metadata written to benchmark results. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "mode", rename_all = "kebab-case")] +pub(crate) enum DiskSearchStrategy { + Flat { + uses_vector_filters: bool, + }, + Graph { + uses_vector_filters: bool, + }, + GraphInlineFilter { #[serde(default)] - post_processor: Option, + adaptive_l: Option, + }, + GraphDiverse { + uses_vector_filters: bool, + post_processor: TopkPostProcessor, }, } impl Default for DiskSearchMode { fn default() -> Self { Self::Graph { - adaptive_l: None, vector_filters_file: None, - post_processor: None, } } } @@ -105,15 +128,64 @@ impl DiskSearchMode { } | Self::Graph { vector_filters_file, + } + | Self::GraphDiverse { + vector_filters_file, .. } => vector_filters_file.as_ref(), + Self::GraphInlineFilter { + vector_filters_file, + .. + } => Some(vector_filters_file), } } pub(crate) fn post_processor(&self) -> Option<&TopkPostProcessor> { match self { - Self::Flat { .. } => None, - Self::Graph { post_processor, .. } => post_processor.as_ref(), + Self::GraphDiverse { post_processor, .. } => Some(post_processor), + _ => None, + } + } + + pub(crate) fn adaptive_l(&self) -> anyhow::Result> { + let Self::GraphInlineFilter { adaptive_l, .. } = self else { + return Ok(None); + }; + + adaptive_l + .as_ref() + .map(|adaptive_l| { + diskann::graph::search::AdaptiveL::new( + adaptive_l.sample_count.into(), + adaptive_l.scale_factor, + ) + .map_err(Into::into) + }) + .transpose() + } + + pub(crate) fn strategy(&self) -> DiskSearchStrategy { + match self { + Self::Flat { + vector_filters_file, + } => DiskSearchStrategy::Flat { + uses_vector_filters: vector_filters_file.is_some(), + }, + Self::Graph { + vector_filters_file, + } => DiskSearchStrategy::Graph { + uses_vector_filters: vector_filters_file.is_some(), + }, + Self::GraphInlineFilter { adaptive_l, .. } => DiskSearchStrategy::GraphInlineFilter { + adaptive_l: adaptive_l.clone(), + }, + Self::GraphDiverse { + vector_filters_file, + post_processor, + } => DiskSearchStrategy::GraphDiverse { + uses_vector_filters: vector_filters_file.is_some(), + post_processor: post_processor.clone(), + }, } } @@ -121,26 +193,35 @@ impl DiskSearchMode { match self { Self::Flat { vector_filters_file, + } + | Self::Graph { + vector_filters_file, } => { if let Some(vf) = vector_filters_file.as_mut() { vf.resolve(checker).context("invalid vector_filters_file")?; } } - Self::Graph { - adaptive_l, + Self::GraphInlineFilter { vector_filters_file, - post_processor, + adaptive_l, } => { - if let Some(adaptive_l) = adaptive_l.as_mut() { + vector_filters_file + .resolve(checker) + .context("invalid vector_filters_file")?; + if let Some(adaptive_l) = adaptive_l { adaptive_l.validate(checker)?; } + } + Self::GraphDiverse { + vector_filters_file, + post_processor, + } => { if let Some(vf) = vector_filters_file.as_mut() { vf.resolve(checker).context("invalid vector_filters_file")?; } - if let Some(pp) = post_processor.as_mut() { - pp.validate(checker) - .context("invalid disk search post processor")?; - } + post_processor + .validate(checker) + .context("invalid disk search post processor")?; } } Ok(()) @@ -148,13 +229,38 @@ impl DiskSearchMode { } impl fmt::Display for DiskSearchMode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.strategy().fmt(f) + } +} + +impl fmt::Display for DiskSearchStrategy { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { - Self::Flat { .. } => write!(f, "flat"), - Self::Graph { adaptive_l, .. } if adaptive_l.is_some() => { - write!(f, "graph + adaptive-l") - } - Self::Graph { .. } => write!(f, "graph"), + Self::Flat { + uses_vector_filters: false, + } => write!(f, "flat"), + Self::Flat { + uses_vector_filters: true, + } => write!(f, "flat + vector-filter"), + Self::Graph { + uses_vector_filters: false, + } => write!(f, "graph"), + Self::Graph { + uses_vector_filters: true, + } => write!(f, "graph + vector-filter"), + Self::GraphInlineFilter { adaptive_l: None } => write!(f, "graph inline-filter"), + Self::GraphInlineFilter { + adaptive_l: Some(_), + } => write!(f, "graph inline-filter + adaptive-l"), + Self::GraphDiverse { + uses_vector_filters: false, + .. + } => write!(f, "graph diverse"), + Self::GraphDiverse { + uses_vector_filters: true, + .. + } => write!(f, "graph diverse + vector-filter"), } } } @@ -163,31 +269,75 @@ impl fmt::Display for DiskSearchMode { mod tests { use super::*; - // The flat/graph split is an enum, so combinations invalid for flat scan - // (adaptive L, post-processor) are unrepresentable by construction — there - // is no runtime compatibility check left to test. These round-trip tests - // just pin the JSON shape. #[test] - fn flat_disk_search_deserializes_without_graph_only_fields() { - let mode: DiskSearchMode = + fn disk_search_modes_deserialize() { + let flat: DiskSearchMode = serde_json::from_str(r#"{ "mode": "flat" }"#).expect("flat mode must deserialize"); - assert!(matches!(mode, DiskSearchMode::Flat { .. })); - assert!(mode.post_processor().is_none()); - } + assert!(matches!(flat, DiskSearchMode::Flat { .. })); - #[test] - fn graph_disk_search_deserializes_with_adaptive_l() { - let mode: DiskSearchMode = serde_json::from_str( - r#"{ "mode": "graph", "adaptive_l": { "sample_count": 1, "scale_factor": 2.0 } }"#, + let graph: DiskSearchMode = + serde_json::from_str(r#"{ "mode": "graph", "vector_filters_file": "filters.bin" }"#) + .expect("graph mode must deserialize"); + assert!(matches!(graph, DiskSearchMode::Graph { .. })); + + let inline: DiskSearchMode = serde_json::from_str( + r#"{ + "mode": "graph-inline-filter", + "vector_filters_file": "filters.bin", + "adaptive_l": { "sample_count": 1, "scale_factor": 2.0 } + }"#, ) - .expect("graph mode with adaptive_l must deserialize"); + .expect("inline-filter mode must deserialize"); assert!(matches!( - mode, - DiskSearchMode::Graph { + inline, + DiskSearchMode::GraphInlineFilter { adaptive_l: Some(_), .. } )); + + let diverse: DiskSearchMode = serde_json::from_str( + r#"{ + "mode": "graph-diverse", + "post_processor": { + "type": "determinant-diversity", + "power": 2.0, + "eta": 1.0 + } + }"#, + ) + .expect("diverse mode must deserialize"); + assert!(matches!(diverse, DiskSearchMode::GraphDiverse { .. })); + } + + #[test] + fn strategy_specific_fields_are_required() { + let inline_error = + serde_json::from_str::(r#"{ "mode": "graph-inline-filter" }"#) + .expect_err("inline-filter mode must require a vector filter"); + assert!(inline_error.to_string().contains("vector_filters_file")); + + let diverse_error = + serde_json::from_str::(r#"{ "mode": "graph-diverse" }"#) + .expect_err("diverse mode must require a post-processor"); + assert!(diverse_error.to_string().contains("post_processor")); + } + + #[test] + fn search_strategy_omits_filter_file_paths() { + let mode: DiskSearchMode = serde_json::from_str( + r#"{ + "mode": "graph-inline-filter", + "vector_filters_file": "private/filters.bin", + "adaptive_l": { "sample_count": 1, "scale_factor": 2.0 } + }"#, + ) + .unwrap(); + + let value = serde_json::to_value(mode.strategy()).unwrap(); + assert_eq!(value["mode"], "graph-inline-filter"); + assert!(value.get("vector_filters_file").is_none()); + assert_eq!(value["adaptive_l"]["sample_count"], 1); } #[test] @@ -395,9 +545,7 @@ impl Example for DiskIndexOperation { recall_at: 10, num_threads: 8, search_mode: DiskSearchMode::Graph { - adaptive_l: None, vector_filters_file: None, - post_processor: None, }, distance: SimilarityMeasure::SquaredL2, num_nodes_to_cache: None, diff --git a/diskann-benchmark/src/inputs/graph_index.rs b/diskann-benchmark/src/inputs/graph_index.rs index 2137c6e576..21904cb587 100644 --- a/diskann-benchmark/src/inputs/graph_index.rs +++ b/diskann-benchmark/src/inputs/graph_index.rs @@ -244,7 +244,7 @@ impl MultihopFilterSearchPhase { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct AdaptiveL { pub(crate) sample_count: NonZeroUsize, pub(crate) scale_factor: f64, diff --git a/diskann-benchmark/src/main.rs b/diskann-benchmark/src/main.rs index e577eb2181..73530b80ef 100644 --- a/diskann-benchmark/src/main.rs +++ b/diskann-benchmark/src/main.rs @@ -238,6 +238,114 @@ mod tests { } } + /// Redirect disk-index build artifacts into a temporary directory. + /// + /// Only existing `save_path` fields are replaced. Jobs without a `save_path` + /// are left unchanged. + fn redirect_save_paths(raw: &mut serde_json::Value, directory: &std::path::Path) { + let Some(jobs) = raw.get_mut("jobs").and_then(Value::as_array_mut) else { + return; + }; + + for (index, job) in jobs.iter_mut().enumerate() { + let Some(save_path) = job + .get_mut("content") + .and_then(|content| content.get_mut("source")) + .and_then(Value::as_object_mut) + .and_then(|source| source.get_mut("save_path")) + else { + continue; + }; + + *save_path = Value::String( + directory + .join(format!("disk_index_job_{index}")) + .to_string_lossy() + .into_owned(), + ); + } + } + + /// Unit-test only the in-memory JSON rewrite; no index is built or loaded. + #[test] + fn redirect_save_paths_updates_builds_without_modifying_loads() { + let mut raw = serde_json::json!({ + "jobs": [ + { + "content": { + "source": { + "disk-index-source": "Build", + "save_path": "build-only-index" + } + } + }, + { + "content": { + "source": { + "disk-index-source": "Load", + "load_path": "existing-index" + } + } + }, + { + "content": { + "source": { + "disk-index-source": "Build", + "save_path": "unrelated-build" + } + } + }, + { + "content": { + "source": { + "disk-index-source": "Load", + "load_path": "unrelated-load" + } + } + } + ] + }); + let directory = std::path::Path::new("temporary-output"); + + redirect_save_paths(&mut raw, directory); + + assert_eq!( + raw["jobs"][0]["content"]["source"]["save_path"], + directory + .join("disk_index_job_0") + .to_string_lossy() + .as_ref() + ); + assert_eq!( + raw["jobs"][1]["content"]["source"]["load_path"], + "existing-index" + ); + assert_eq!( + raw["jobs"][2]["content"]["source"]["save_path"], + directory + .join("disk_index_job_2") + .to_string_lossy() + .as_ref() + ); + assert_eq!( + raw["jobs"][3]["content"]["source"]["load_path"], + "unrelated-load" + ); + for index in [1, 3] { + assert!(raw["jobs"][index]["content"]["source"] + .get("save_path") + .is_none()); + } + } + + /// Unit-test that `redirect_save_paths` safely ignores malformed JSON input. + #[test] + fn redirect_save_paths_ignores_malformed_input() { + for mut raw in [serde_json::json!(null), serde_json::json!({ "jobs": null })] { + redirect_save_paths(&mut raw, std::path::Path::new("temporary-output")); + } + } + // Retrieve the number of jobs in the raw input JSON. // // The format is @@ -264,13 +372,18 @@ mod tests { } } - fn run_integration_test(mut raw: serde_json::Value) { + fn run_integration_test(raw: serde_json::Value) { + run_integration_test_with_results(raw); + } + + fn run_integration_test_with_results(mut raw: serde_json::Value) -> Vec { // First, parse and modify the input file to establish paths relative to the // directory building the dispatcher. // let mut raw = serde_json::from_str(json_string).unwrap(); prefix_search_directories(&mut raw, &root_directory()); let tempdir = tempfile::tempdir().unwrap(); + redirect_save_paths(&mut raw, tempdir.path()); let input_path = tempdir.path().join("input.json"); save_to_file(&input_path, &raw); @@ -299,6 +412,7 @@ mod tests { let results: Vec = load_from_file(&output_path); assert_eq!(results.len(), num_jobs(&raw)); + results } //////////////////////////////// @@ -701,46 +815,20 @@ mod tests { } /// Filtered disk search end-to-end: drives the disk-index backend through - /// `disk-index-filter.json` + /// `disk-index-filter.json`. #[test] #[cfg(feature = "disk-index")] fn disk_index_filter_integration() { - let mut raw = value_from_file(&example_directory().join("disk-index-filter.json")); - prefix_search_directories(&mut raw, &root_directory()); - - let tempdir = tempfile::tempdir().unwrap(); - - // Redirect each build job's `save_path` into the tempdir so the disk index - // artifacts are not written relative to the process cwd (the repo tree). - let jobs = raw["jobs"] - .as_array_mut() - .expect("\"jobs\" should be an array"); - for (i, job) in jobs.iter_mut().enumerate() { - let save_path = tempdir.path().join(format!("disk_index_filter_job_{i}")); - job["content"]["source"]["save_path"] = - serde_json::Value::String(save_path.to_str().unwrap().to_string()); - } - - let input_path = tempdir.path().join("disk-index-filter.json"); - save_to_file(&input_path, &raw); - let output_path = tempdir.path().join("output.json"); - - let command = Commands::Run { - input_file: input_path.to_owned(), - output_file: output_path.to_owned(), - dry_run: false, - allow_debug: true, - }; - let cli = Cli::from_commands(command, true); - let mut output = Memory::new(); - let result = cli.run(&mut output); - let output_str = String::from_utf8(output.into_inner()).unwrap(); - println!("output = {}", output_str); - result.expect("disk-index-filter run failed"); - - assert!(output_path.exists()); - let results: Vec = load_from_file(&output_path); - assert_eq!(results.len(), num_jobs(&raw)); + let raw = value_from_file(&example_directory().join("disk-index-filter.json")); + let results = run_integration_test_with_results(raw); + + let strategy = &results[0]["results"]["search"]["search_strategy"]; + assert_eq!(strategy["mode"], "graph-inline-filter"); + assert_eq!(strategy["adaptive_l"]["sample_count"], 10); + assert_eq!( + results[1]["results"]["search"]["search_strategy"]["mode"], + "flat" + ); } #[test] From 546a2d5995bc776304914deff387455d2affce86 Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Thu, 10 Sep 2026 16:16:16 +0800 Subject: [PATCH 10/11] Move Adaptive_L construction from input file to search file --- diskann-benchmark/src/disk_index/search.rs | 16 +++++++++++++++- diskann-benchmark/src/inputs/disk.rs | 17 ----------------- 2 files changed, 15 insertions(+), 18 deletions(-) diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index 43fa1d4ac1..a9e63e469e 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -215,6 +215,20 @@ impl DiskSearchResult { } } +fn build_adaptive_l(mode: &DiskSearchMode) -> anyhow::Result> { + let DiskSearchMode::GraphInlineFilter { adaptive_l, .. } = mode else { + return Ok(None); + }; + + adaptive_l + .as_ref() + .map(|adaptive_l| { + graph::search::AdaptiveL::new(adaptive_l.sample_count.into(), adaptive_l.scale_factor) + .map_err(Into::into) + }) + .transpose() +} + /// Construct the backend [`SearchMode`] from the JSON-configured strategy, /// the per-query vector filter, and the pre-validated adaptive-L settings. fn build_search_mode<'a>( @@ -341,7 +355,7 @@ where logger.log_checkpoint("index_loaded"); - let adaptive_l = search_params.search_mode.adaptive_l()?; + let adaptive_l = build_adaptive_l(&search_params.search_mode)?; let pool = create_thread_pool(search_params.num_threads)?; let mut search_results_per_l = Vec::with_capacity(search_params.search_list.len()); let has_any_search_failed = AtomicBool::new(false); diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index 040ad90bb1..f0d1e03b5b 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -147,23 +147,6 @@ impl DiskSearchMode { } } - pub(crate) fn adaptive_l(&self) -> anyhow::Result> { - let Self::GraphInlineFilter { adaptive_l, .. } = self else { - return Ok(None); - }; - - adaptive_l - .as_ref() - .map(|adaptive_l| { - diskann::graph::search::AdaptiveL::new( - adaptive_l.sample_count.into(), - adaptive_l.scale_factor, - ) - .map_err(Into::into) - }) - .transpose() - } - pub(crate) fn strategy(&self) -> DiskSearchStrategy { match self { Self::Flat { From 58d857349d3fae118168f6ee232a24ff1222a693 Mon Sep 17 00:00:00 2001 From: yaohongdeng Date: Thu, 10 Sep 2026 16:35:53 +0800 Subject: [PATCH 11/11] remove deny_unknown_fields in input --- diskann-benchmark/src/inputs/disk.rs | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index f0d1e03b5b..5a33f6b811 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -322,30 +322,10 @@ mod tests { assert!(value.get("vector_filters_file").is_none()); assert_eq!(value["adaptive_l"]["sample_count"], 1); } - - #[test] - fn disk_search_phase_rejects_legacy_phase_level_search_mode_fields() { - let error = serde_json::from_str::( - r#"{ - "queries": "queries.fbin", - "groundtruth": "groundtruth.bin", - "num_threads": 1, - "beam_width": 1, - "search_list": [1], - "recall_at": 1, - "distance": "squared_l2", - "is_flat_search": true - }"#, - ) - .expect_err("legacy phase-level search settings must be rejected"); - - assert!(error.to_string().contains("is_flat_search")); - } } /// Search phase configuration #[derive(Debug, Deserialize, Serialize)] -#[serde(deny_unknown_fields)] pub(crate) struct DiskSearchPhase { pub(crate) queries: InputFile, pub(crate) groundtruth: InputFile,