Repository navigation
Refactor(diskann-benchmark): consolidate disk search config under DiskSearchMode - #1232
Yaohong Deng (dyhyfu) wants to merge 17 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 (85.05%) 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 90.95% 91.13% +0.17%
==========================================
Files 586 592 +6
Lines 118252 123745 +5493
==========================================
+ Hits 107559 112771 +5212
- Misses 10693 10974 +281
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.
| /// 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.
There was a problem hiding this comment.
Added explicit validation during DiskSearchMode deserialization so incompatible recognized options now produce an error instead of being silently ignored. Unknown fields remain ignored, and null optional settings are treated as absent.
Updated the code comments and added tests for incompatible combinations. All six parsing tests and the formatting check pass.
Suryansh Gupta (suri-kumkaran)
left a comment
Suryansh Gupta (suri-kumkaran)
left a comment
There was a problem hiding this comment.
Thanks for the updates! Everything from my last round is addressed. I left a few inline comments.
One more ask: The PR description is out of date. It still lists is_flat_search under DiskSearchMode and points to a "see below" section that isn't there. Please also mention the workflow change, that old configs with is_flat_search on the search phase are no longer supported, and that is_flat_search and uses_vector_filters are gone from the results JSON.
| 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 }} \ |
There was a problem hiding this comment.
The baseline change in disk-benchmarks.yml makes sense for this PR, since main can't parse the new search_mode format yet.
But I'd keep it one-time. The benchmark is meant to run the same input file on both sides, so only the code differs. If this stays, a future PR that edits a perf input runs different settings on each side. Changing search_list would fail the check with a search_l mismatch, and changing num_threads would quietly skew QPS.
Could we keep it here and revert it in a small follow-up once this merges? At that point main understands the new format, so pointing the baseline back at ../diskann_rust/ works again.
There was a problem hiding this comment.
Agreed. Kept the workaround for this PR and added a reminder in the workflow. After merge, I’ll restore the ../diskann_rust/ input prefix in a small follow-up so both runs use the same configuration.
| pub(super) beam_width: usize, | ||
| pub(super) recall_at: u32, | ||
| pub(crate) is_flat_search: bool, | ||
| pub(crate) search_strategy: DiskSearchStrategy, |
There was a problem hiding this comment.
My earlier comment on the stats booleans missed this, sorry. The results file already saves the full input next to each result, so the search mode, filter file, adaptive_l and post-processor are already there. The console also prints the mode twice.
Suggestion:
- Remove
search_strategyfromDiskSearchStats. - Remove
DiskSearchStrategy,strategy(),RawDiskSearchStats, the customDeserializeand their tests, and deriveDeserializeinstead. It still reads main's old stats, since extra fields are ignored. - In the filter test, check the mode through
results[i]["input"].
There was a problem hiding this comment.
Removed search_strategy, DiskSearchStrategy, strategy(), RawDiskSearchStats, and the custom stats deserializer and associated tests. DiskSearchStats now derives Deserialize, and the filter integration test checks the mode through the saved input. Also removed the duplicate console output.
| impl<'de> Deserialize<'de> for DiskSearchMode { | ||
| fn deserialize<D>(deserializer: D) -> Result<Self, D::Error> | ||
| where | ||
| D: serde::Deserializer<'de>, | ||
| { | ||
| #[derive(Deserialize)] | ||
| #[serde(rename_all = "kebab-case")] | ||
| enum Mode { | ||
| Flat, | ||
| Graph, | ||
| GraphInlineFilter, | ||
| GraphDiverse, | ||
| } | ||
| #[derive(Deserialize)] | ||
| struct RawMode { | ||
| mode: Mode, | ||
| vector_filters_file: Option<InputFile>, | ||
| adaptive_l: Option<AdaptiveL>, | ||
| post_processor: Option<TopkPostProcessor>, | ||
| } | ||
| let raw = RawMode::deserialize(deserializer)?; | ||
| if raw.adaptive_l.is_some() && !matches!(raw.mode, Mode::GraphInlineFilter) { | ||
| return Err(serde::de::Error::custom( | ||
| "adaptive_l is only supported with mode graph-inline-filter", | ||
| )); | ||
| } | ||
| if raw.post_processor.is_some() && !matches!(raw.mode, Mode::GraphDiverse) { | ||
| return Err(serde::de::Error::custom( | ||
| "post_processor is only supported with mode graph-diverse", | ||
| )); | ||
| } | ||
| Ok(match raw.mode { | ||
| Mode::Flat => Self::Flat { | ||
| vector_filters_file: raw.vector_filters_file, | ||
| }, | ||
| Mode::Graph => Self::Graph { | ||
| vector_filters_file: raw.vector_filters_file, | ||
| }, | ||
| Mode::GraphInlineFilter => Self::GraphInlineFilter { | ||
| vector_filters_file: raw | ||
| .vector_filters_file | ||
| .ok_or_else(|| serde::de::Error::missing_field("vector_filters_file"))?, | ||
| adaptive_l: raw.adaptive_l, | ||
| }, | ||
| Mode::GraphDiverse => Self::GraphDiverse { | ||
| vector_filters_file: raw.vector_filters_file, | ||
| post_processor: raw | ||
| .post_processor | ||
| .ok_or_else(|| serde::de::Error::missing_field("post_processor"))?, | ||
| }, | ||
| }) | ||
| } | ||
| } |
There was a problem hiding this comment.
About my earlier "unrepresentable" comment: I only meant the doc comment was inaccurate. We don't need to reject extra fields. A derived Deserialize already requires the right fields for each mode, and an extra adaptive_l under flat would just be ignored like any unknown key.
Suggestion:
- Replace the custom
Deserializewith#[derive(Deserialize)], and removeincompatible_search_options_are_rejected. - Update the doc comment to say fields that don't belong to the mode are ignored.
- Remove the
#[serde(default)]s on lines 74-90 (not needed) and the unusedClone.
There was a problem hiding this comment.
Replaced the custom deserializer with derived Deserialize, removed the redundant defaults and unused Clone, and removed the incompatible-options rejection test. Updated the documentation to clarify that fields outside the selected mode are ignored. Required-field and invalid-value tests remain.
Consolidates disk-index search configuration into a backend-independent, tagged DiskSearchMode enum. Backend SearchMode construction now happens in build_search_mode in the search execution module.
Validation: cargo fmt --all --check; four disk input schema tests; the inline-search missing-filter test; the disk-index filter integration test; workspace Clippy with warnings denied; and disk-index-enabled benchmark Clippy with warnings denied all pass.