Skip to content

Refactor(diskann-benchmark): consolidate disk search config under DiskSearchMode #1232

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Yaohong Deng (dyhyfu) wants to merge 11 commits into
base: main
Choose a base branch
from
Open

Refactor(diskann-benchmark): consolidate disk search config under DiskSearchMode #1232

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
5c86c16
consolidate disk search config into DiskSearchMode and move SearchMod…
Jul 7, 2026
3014754
Merge branch 'main' into u/yaohongdeng/fixDiskModeBenchmarkCfg
Jul 7, 2026
6fe7448
remove unexpected output files from disk-filter test
Jul 8, 2026
29907bf
Fix disk benchmark search-mode configuration validation
Jul 14, 2026
521ca01
Merge branch 'main' into u/yaohongdeng/fixDiskModeBenchmarkCfg
Jul 14, 2026
ca59956
force not tolerant for old is_flat_search json config
Jul 14, 2026
2e088a2
update benchmark input for DiskSearchMode, change to enum structure
Aug 4, 2026
968b79c
Merge branch 'main' into u/yaohongdeng/fixDiskModeBenchmarkCfg
Aug 4, 2026
acd030b
fix format issue
Aug 4, 2026
92d8cae
remove unused function
Aug 4, 2026
a624d45
change benchmark baseline config file to use main branch config file
Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .github/workflows/disk-benchmarks.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -144,4 +144,4 @@ jobs:
path: |
diskann_rust/target/tmp/${{ matrix.dataset }}_target.json
baseline/target/tmp/${{ matrix.dataset }}_baseline.json
retention-days: 30
retention-days: 30
17 changes: 9 additions & 8 deletions diskann-benchmark/example/disk-index-determinant-diversity.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
"mode": "graph",
"post_processor": {
"type": "determinant-diversity",
"power": 2.0,
"eta": 1.0
}
},
"distance": "squared_l2"
}
}
}
Expand Down
16 changes: 10 additions & 6 deletions diskann-benchmark/example/disk-index-filter.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
"mode": "graph",
"vector_filters_file": "disk_index_10pts_idx_uint32_range_res_r_100000.bin"
},
"distance": "squared_l2"
}
}
},
Expand Down Expand Up @@ -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": {
"mode": "flat",
"vector_filters_file": "disk_index_10pts_idx_uint32_range_res_r_100000.bin"
},
"distance": "squared_l2"
}
}
}
Expand Down
10 changes: 4 additions & 6 deletions diskann-benchmark/example/disk-index.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": { "mode": "graph" },
"distance": "squared_l2"
}
}
},
Expand All @@ -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": { "mode": "flat" },
"distance": "squared_l2"
}
}
}
Expand Down
5 changes: 2 additions & 3 deletions diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": { "mode": "graph" },
"distance": "squared_l2"
}
}
}
Expand Down
5 changes: 2 additions & 3 deletions diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": { "mode": "graph" },
"distance": "inner_product"
}
}
}
Expand Down
94 changes: 75 additions & 19 deletions diskann-benchmark/src/disk_index/search.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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},
};

Expand Down Expand Up @@ -158,6 +160,58 @@ 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,
vector_filter: Option<&'a HashSet<u32>>,
) -> 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")
});
Comment on lines +181 to +187

@suri-kumkaran Suryansh Gupta (suri-kumkaran) Aug 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: .expect("validated adaptive L must construct")

graph_index.rs:291-298 handles the same conversion with ?. Here, build_search_mode returns a bare SearchMode, so there’s no way to propagate the error and we end up panicking instead.

The bigger issue is that this happens inside for_each_in_pool, so a failure would panic from the per-query path, and the safety of the expect depends on validate having run earlier - which isn't obvious from the function signature.

Since adaptive_l is fixed config, we can build it once before the for &l in search_list loop and use ? there. This also avoids doing the same conversion for every query. That's what the approach the main comment assumes.


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),
),
Comment on lines +189 to +205

@suri-kumkaran Suryansh Gupta (suri-kumkaran) Aug 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

adaptive_l is mishandled on both paths

Neither is a regression; both cases already exist on main. I'm calling them out because this PR is already reshaping this code, so it's a good opportunity to fix them.

With a post-processor, adaptive_l gets dropped. The _ in L190 and L193 ignores it, so a config like {"mode":"graph", "adaptive_l": {...}, "post_processor": {...}} runs with adaptive L disabled. But Display still sees adaptive_l.is_some() and prints "graph + adaptive-l", so the output is labeled as adaptive L when it isn't.

Without a filter, adaptive_l does nothing. L200 uses inline_filter(|_| true, Some(adaptive_l)). Since the predicate is always true, the filter is effectively 100% specific, so the adaptive-L scaling never kicks in.

So today, of the four possible adaptive_l paths, only one actually has an effect: graph search with a real filter. We also don't have any disk fixture covering adaptive_l.

I think both issues come from the same mismatch: the config allows combinations that SearchMode can't actually represent. The main comment fixes this cleanly by making adaptive_l and post_processor mutually exclusive variants and requiring a filter for inline filtering.

(None, None, None) => SearchMode::graph(),
(None, None, Some(vector_filter)) => {
SearchMode::graph_filtered(move |vid: &u32| vector_filter.contains(vid))
}
}
}
}
}

pub(super) fn search_disk_index<T, StorageType>(
index_load: &DiskIndexLoad,
search_params: &DiskSearchPhase,
Expand Down Expand Up @@ -185,21 +239,27 @@ 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)?
Some(search_index_utils::load_vector_filters(
storage_provider,
&vector_filters_file,
)?)
}
None => vec![HashSet::<u32>::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");
}

// 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,
Expand Down Expand Up @@ -259,24 +319,20 @@ 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())
.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(),
);
|(((((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);

match searcher.search(
q,
Expand Down Expand Up @@ -349,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: matches!(search_params.search_mode, DiskSearchMode::Flat { .. }),
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,
Expand Down
Loading
Loading