Refactor(diskann-benchmark): consolidate disk search config under DiskSearchMode - #1232
Refactor(diskann-benchmark): consolidate disk search config under DiskSearchMode#1232Yaohong Deng (dyhyfu) wants to merge 11 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR refactors the diskann-benchmark disk-index benchmark input schema to consolidate disk search configuration under DiskSearchMode, and moves diskann_disk::SearchMode construction out of the JSON schema layer into the disk search execution path.
Changes:
- Nested
vector_filters_fileandpost_processorunderDiskSearchMode, alongsideis_flat_searchandadaptive_l, and moved validation accordingly. - Added a
build_search_modehelper indisk_index/search.rsto construct backendSearchModeat execution time. - Updated benchmark JSON fixtures (examples + perf inputs) to the new nested
search_modeformat.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| diskann-benchmark/src/inputs/disk.rs | Refactors disk-index JSON schema to centralize mode-specific config/validation in DiskSearchMode. |
| diskann-benchmark/src/disk_index/search.rs | Builds backend SearchMode during execution using new helper; updates access paths to nested config. |
| diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json | Migrates perf input to nested search_mode object. |
| diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json | Migrates perf input to nested search_mode object. |
| diskann-benchmark/example/disk-index.json | Migrates example input to nested search_mode object. |
| diskann-benchmark/example/disk-index-filter.json | Migrates filter example to nested search_mode.vector_filters_file. |
| diskann-benchmark/example/disk-index-determinant-diversity.json | Migrates post-processor example to nested search_mode.post_processor. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is
❌ Your patch status has failed because the patch coverage (42.46%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #1232 +/- ##
==========================================
- Coverage 91.46% 91.44% -0.03%
==========================================
Files 516 516
Lines 98276 98340 +64
==========================================
+ Hits 89891 89928 +37
- Misses 8385 8412 +27
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Suppressed comments (2)
diskann-benchmark/src/inputs/disk.rs:72
DiskSearchModedoes not use#[serde(deny_unknown_fields)], so typos or legacy keys inside the nestedsearch_modeobject (e.g.{ "mode": "graph", "is_flat_search": true }) will be silently ignored by Serde. This undermines the intent of adding#[serde(deny_unknown_fields)]onDiskSearchPhaseto hard-fail old schemas.
Consider denying unknown fields on DiskSearchMode as well so invalid/legacy keys under search_mode are rejected deterministically.
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "mode", rename_all = "kebab-case")]
pub(crate) enum DiskSearchMode {
diskann-benchmark/src/inputs/disk.rs:138
DiskSearchMode::Graphcurrently allows specifying bothadaptive_landpost_processor, butbuild_search_modewill always pick the determinant-diversitySearchMode::{diverse_graph,_}whenpost_processoris set, effectively ignoringadaptive_l. Sincediskann_disk::search::search_mode::SearchMode::DiverseGraphhas noadaptive_lsupport, this should be rejected (or at least made explicit) rather than silently dropping part of the config.
Self::Graph {
adaptive_l,
vector_filters_file,
post_processor,
} => {
if let Some(adaptive_l) = adaptive_l.as_mut() {
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Suppressed comments (3)
diskann-benchmark/src/main.rs:721
- Using
save_path.to_str().unwrap()can panic on non-UTF-8 temp paths (e.g., if TMPDIR contains non-UTF-8). Since this is only for test JSON rewriting, preferto_string_lossy()to avoid spurious test failures on such environments.
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());
diskann-benchmark/src/inputs/disk.rs:72
DiskSearchModedoes not deny unknown fields, so typos or legacy fields nested undersearch_mode(e.g.{ "mode": "graph", "is_flat_search": true }) may be silently ignored during deserialization. Addingdeny_unknown_fieldshere would make the JSON schema stricter and align with the intent of rejecting legacy/unknown parameters.
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "mode", rename_all = "kebab-case")]
pub(crate) enum DiskSearchMode {
diskann-benchmark/src/disk_index/search.rs:165
- The doc comment for
build_search_modesays the post-processor is supplied at search time, but the post-processor now comes from the JSON-drivenDiskSearchModeconfig (only the vector filter is per-query). Updating this comment would avoid confusion about where the post-processor is sourced.
/// 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>(
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Suppressed comments (3)
diskann-benchmark/src/inputs/disk.rs:133
DiskSearchMode::Graphallows bothadaptive_landpost_processorto be set, but the search execution path currently ignoresadaptive_lwhenever apost_processoris present (seebuild_search_modeindisk_index/search.rs, which matches determinant-diversity first and discards the computedadaptive_l). This makes part of the JSON config silently ineffective. Consider rejecting this combination during validation (or otherwise making the precedence explicit).
Self::Graph {
adaptive_l,
vector_filters_file,
post_processor,
} => {
diskann-benchmark/src/disk_index/search.rs:165
- Doc comment for
build_search_modesays the post-processor is "supplied at search time", but the function signature only takesmodeandvector_filter(the post-processor comes fromDiskSearchMode::Graph { post_processor, .. }). This is misleading when reading the code and debugging configuration-driven behavior.
/// 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>(
diskann-benchmark/src/inputs/disk.rs:80
- Doc comment for
DiskSearchMode::Graphsays it can be used with adaptive-L, vector filters, and/or a post-processor. However the backendSearchModedoes not support combining determinant-diversity post-processing withadaptive_l(andbuild_search_modecurrently dropsadaptive_lwhenpost_processoris set). If you enforce mutual exclusivity in validation, this comment should be updated to avoid implying the combination is supported.
This issue also appears on line 129 of the same file.
/// Greedy graph search, optionally with inline adaptive-L, a per-query
/// vector filter, and/or a top-k post-processor.
Graph {
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.
Suppressed comments (4)
diskann-benchmark/src/inputs/disk.rs:88
DiskSearchModedoes not deny unknown fields, so serde will silently ignore unexpected keys (e.g.,{"mode":"flat","adaptive_l":...}would deserialize asFlatand dropadaptive_l). That undermines the goal of making invalid combinations unrepresentable and makes typos easy to miss; consider denying unknown fields for the enum variants.
#[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 {
diskann-benchmark/src/inputs/disk.rs:144
DiskSearchMode::Graphcurrently allowsadaptive_landpost_processorto be set together, butbuild_search_modewill always pickSearchMode::DiverseGraphwhen a post-processor is present, silently ignoringadaptive_l. Consider rejecting this combination during validation to avoid surprising config behavior.
Self::Graph {
adaptive_l,
vector_filters_file,
post_processor,
} => {
diskann-benchmark/src/main.rs:721
- This test builds a JSON string path via
save_path.to_str().unwrap(), which can panic on non-UTF8 paths. Usingto_string_lossy()avoids a hard panic and is consistent with other path-to-string conversions in the benchmark code.
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());
diskann-benchmark/src/disk_index/search.rs:165
- The doc comment for
build_search_modesays the post-processor is supplied at search time, but the implementation reads it fromDiskSearchMode::Graph { post_processor, .. }(JSON config). Updating the comment would avoid confusion about where this value comes from.
/// 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>(
Suryansh Gupta (suri-kumkaran)
left a comment
Suryansh Gupta (suri-kumkaran)
left a comment
There was a problem hiding this comment.
I like the overall direction here. Decoupling the input schema from diskann-disk, keeping backend construction in build_search_mode, and moving vector_filters_file / post_processor out of DiskSearchPhase all make the code cleaner. The Option<&HashSet> runtime filter change is a nice cleanup too.
My main concern is that the enum refactor stops a bit short. The backend already models the four actual search strategies (FlatScan, Graph, InlineFilter, DiverseGraph), but the new config still models traversal separately and keeps strategy in a few optional fields. That makes build_search_mode more complicated than it needs to be and leaves invalid combinations representable.
I also called out two pre-existing adaptive_l issues in the inline comments. They're not regressions, but this PR is already touching the relevant code, and the enum shape gives us a clean way to make both cases impossible.
The one thing I'd want settled before merge is the is_flat_search migration. This PR removes the old back-compat field, and deny_unknown_fields is currently what prevents an old config from silently switching from flat to graph. I suggest temporarily keeping it locally to verify the migration, then removing it once we've confirmed all configs are clean. See the inline comment for the details.
Nothing else blocking. Happy to talk through the enum shape live if that's easier.
| pub(crate) enum DiskSearchMode { | ||
| /// Brute-force flat scan, optionally restricted by a per-query vector filter. | ||
| Flat { | ||
| #[serde(default)] | ||
| vector_filters_file: Option<InputFile>, | ||
| }, | ||
| /// 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<AdaptiveL>, | ||
| #[serde(default)] | ||
| vector_filters_file: Option<InputFile>, | ||
| #[serde(default)] | ||
| post_processor: Option<TopkPostProcessor>, | ||
| }, | ||
| } |
There was a problem hiding this comment.
This is my main comment, and it's a "go one step further", not a "this is wrong".
Turning flat/graph into an enum was the right move. But the backend has four variants keyed on strategy, and this has two keyed on traversal with strategy in the optionals. build_search_mode still produces 8 leaf constructions for those 4 strategies, because every strategy is written twice, once for None filter and once for Some.
Mirroring the backend collapses that. Note vector_filters_file is required on the inline-filter variant, because the backend's InlineFilter.filter is a plain Box<dyn QueryLabelProvider<u32>>, not an Option:
#[serde(tag = "mode", rename_all = "kebab-case")]
pub(crate) enum DiskSearchMode {
Flat { vector_filters_file: Option<InputFile> },
Graph { vector_filters_file: Option<InputFile> },
GraphInlineFilter { vector_filters_file: InputFile, adaptive_l: Option<AdaptiveL> },
GraphDiverse { vector_filters_file: Option<InputFile>, post_processor: DeterminantDiversityParams },
}The mapping then normalizes the predicate once and matches four ways. I compiled this against diskann-disk to be sure:
let filter: Option<SearchPredicate<'a>> =
vector_filter.map(|f| Box::new(move |vid: &u32| f.contains(vid)) as SearchPredicate<'a>);
match mode {
Flat { .. } => SearchMode::FlatScan { filter },
Graph { .. } => SearchMode::Graph { filter },
GraphInlineFilter { .. } => {
let f = vector_filter.expect("inline filter requires a vector filter");
SearchMode::inline_filter(move |vid: &u32| f.contains(vid), adaptive_l)
}
GraphDiverse { post_processor, .. } =>
SearchMode::DiverseGraph { filter, params: *post_processor },
}What this buys: adaptive_l can no longer coexist with a post-processor. This also makes some comments becomes unrepresentable.
|
|
||
| /// Search phase configuration | ||
| #[derive(Debug, Deserialize, Serialize)] | ||
| #[serde(deny_unknown_fields)] |
There was a problem hiding this comment.
This is the one thing I'd want settled before merge, because it's a user-facing decision.
#[serde(deny_unknown_fields)] is the only occurrence in the repository, and it's inconsistent with the rest of the config parsing.
That said, I don't think it should just be deleted here. main deliberately carried a back-compat is_flat_search field, and this PR removes it. Without deny_unknown_fields, an old config could silently ignore is_flat_search and flip from flat scan to graph search.
I'd suggest keeping #[serde(deny_unknown_fields)] locally for this migration to verify that all existing configs are migrated correctly and no old config changes behavior silently. Once we've verified the migration, we can remove it in this PR.
If we want unknown-field strictness permanently, let's propose it repo-wide in a separate PR.
| // 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()); | ||
| } |
There was a problem hiding this comment.
Bug: the loop is unconditional, and serde_json's IndexMut creates missing keys, so a Load job gets a save_path invented for it. Not hypothetical: disk-index.json job 1 is a Load. It only works here because this fixture happens to be two Build jobs and DiskIndexLoad ignores unknown fields.
run_bftree_save_roundtrip (L405-448) already sets the precedent for extracting this. I'd go further and call it from run_integration_test (L267-302), so every example gets it:
let tempdir = tempfile::tempdir().unwrap();
+redirect_save_paths(&mut raw, tempdir.path());
let input_path = tempdir.path().join("input.json");/// Redirect each build job's `save_path` into `dir` so artifacts don't land in the repo.
/// Only rewrites keys that already exist: `Load` sources have no `save_path`.
fn redirect_save_paths(raw: &mut serde_json::Value, dir: &Path) {
let Some(jobs) = raw.get_mut("jobs").and_then(Value::as_array_mut) else { return };
for (i, job) in jobs.iter_mut().enumerate() {
let Some(source) = job.get_mut("content")
.and_then(|c| c.get_mut("source"))
.and_then(Value::as_object_mut) else { continue };
let Some(slot) = source.get_mut("save_path") else { continue };
*slot = Value::String(dir.join(format!("index_job_{i}")).to_string_lossy().into_owned());
}
}Checking for the key rather than a "Build" tag keeps it benchmark-agnostic and means it can't invent anything. It's also panic-free, unlike the current raw["jobs"] indexing.
The test then collapses to the same shape as the others, since run_integration_test already does prefix_search_directories:
let raw = value_from_file(&example_directory().join("disk-index-filter.json"));
run_integration_test(raw);Leave run_bftree_save_roundtrip as-is; it leaks its tempdir on purpose for the load-back step.
| 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), | ||
| ), |
There was a problem hiding this comment.
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.
| /// 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. |
There was a problem hiding this comment.
Both comments say combinations invalid for flat scan are "unrepresentable by construction" and that "there is no runtime compatibility check left to test."
They're representable. {"mode":"flat", "post_processor": {...}} and {"mode":"flat", "adaptive_l": {...}} both deserialize fine and the extra fields are dropped. The enum removed the check, not the state.
The main comment would make the claim true.
| 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") | ||
| }); |
There was a problem hiding this comment.
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.
| #[derive(Serialize, Deserialize, 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) distance: SimilarityMeasure, | ||
| pub(crate) uses_vector_filters: bool, | ||
| pub(super) num_nodes_to_cache: Option<usize>, | ||
| pub(super) search_results_per_l: Vec<DiskSearchResult>, | ||
| span_metrics: serde_json::Value, | ||
| } |
There was a problem hiding this comment.
pub(crate) is_flat_search: bool,
pub(crate) uses_vector_filters: bool,Two bools encode four states, and neither records whether adaptive L or a post-processor was active. Two runs with materially different search behaviour now serialize to identical stats, which makes the JSON ambiguous for regression and comparison tooling.
DiskSearchMode already implements Display and Serialize, so emitting the mode itself would be strictly more useful.
Summary
Consolidates disk-index search configuration under
DiskSearchModeand decouples the benchmark's JSON input schema from the optionaldiskann-diskbackend.Previously, the input schema both defined the config and constructed
diskann_disk::SearchMode, forcing several#[cfg(feature = "disk-index")]gates and scattering related fields (vector_filters_file,post_processor) acrossDiskSearchPhase. This PR moves the backend-specific construction into the search execution path and groups the search-mode fields together.Changes
SearchModeconstruction to the backend. The match logic that buildsdiskann_disk::SearchModenow lives in abuild_search_modehelper in the search execution module, so the input schema is pure config data with no dependency on the disk backend'sSearchModetype.DiskSearchMode.vector_filters_fileandpost_processorare nested insideDiskSearchModealongsideis_flat_searchandadaptive_l. Validation andDisplaymoved accordingly.DiskSearchModeand its fields no longer need#[cfg(feature = "disk-index")]. The only remaining gates are forQuantizationType, which is adiskann-disktype (intentionally left as-is — see below).search_modeformat.