diskann-benchmark refactor pt2 - shared logic overhaul#1169
diskann-benchmark refactor pt2 - shared logic overhaul#1169JordanMaples wants to merge 3 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR refactors diskann-benchmark streaming benchmarks by introducing a shared streaming execution layer (runner + common helpers), unifying bf-tree streaming/build inputs (including quantization configuration), and updating examples/docs to the new JSON schema.
Changes:
- Introduce shared streaming infrastructure (
StreamRunner,Maintainer,build_streamer,run_streaming) and extract in-mem maintenance logic intoInmemMaintainer. - Unify/rename streaming inputs (Dynamic → Streaming) and introduce scalar
StreamingSearchParams(no parameter sweep matrix) for streaming runs. - Consolidate bf-tree quantization inputs behind a single
QuantConfigenum and add a shared spherical quantizer helper.
Reviewed changes
Copilot reviewed 28 out of 28 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| diskann-benchmark/src/main.rs | Updates integration test coverage for renamed streaming example and adds bf-tree spherical streaming example test. |
| diskann-benchmark/src/inputs/graph_index.rs | Renames dynamic→streaming types; introduces StreamingSearchParams; privatizes some input fields behind accessors. |
| diskann-benchmark/src/inputs/exhaustive.rs | Makes TransformKind clonable (used by new shared quantization config). |
| diskann-benchmark/src/inputs/bftree.rs | Unifies bf-tree build/streaming inputs around QuantConfig and streaming search/runbook types. |
| diskann-benchmark/src/index/streaming/runner.rs | Adds generic shared streaming runner with pluggable maintenance. |
| diskann-benchmark/src/index/streaming/mod.rs | Adds shared streamer construction + runbook execution helpers; re-exports shared pieces. |
| diskann-benchmark/src/index/streaming/full_precision.rs | Removes legacy full-precision streaming implementation (superseded by shared runner). |
| diskann-benchmark/src/index/result.rs | Tightens visibility of result structs/enums to the index module. |
| diskann-benchmark/src/index/inmem/streaming.rs | Introduces InmemMaintainer (DropDeleted + Release) extracted from old inline logic. |
| diskann-benchmark/src/index/inmem/spherical.rs | Updates to new accessor-based input APIs. |
| diskann-benchmark/src/index/inmem/scalar.rs | Updates to new accessor-based input APIs. |
| diskann-benchmark/src/index/inmem/product.rs | Updates to new accessor-based input APIs. |
| diskann-benchmark/src/index/inmem/mod.rs | Exposes InmemMaintainer and adds inmem streaming module. |
| diskann-benchmark/src/index/build.rs | Tightens visibility of internal build helpers/types. |
| diskann-benchmark/src/index/bftree/streaming_utils.rs | Removes duplicated streaming runbook-driving helper (replaced by shared run_streaming). |
| diskann-benchmark/src/index/bftree/spherical.rs | Switches spherical bf-tree build to unified BfTreeBuild + shared quantizer helper. |
| diskann-benchmark/src/index/bftree/spherical_streaming.rs | Switches bf-tree spherical streaming to shared streamer + StreamRunner. |
| diskann-benchmark/src/index/bftree/quantizer_util.rs | Adds shared spherical quantizer construction utility for bf-tree benchmarks. |
| diskann-benchmark/src/index/bftree/mod.rs | Updates module docs and wires new shared quantizer utility. |
| diskann-benchmark/src/index/bftree/full_precision.rs | Switches bf-tree full-precision build to unified BfTreeBuild (with QuantConfig::None). |
| diskann-benchmark/src/index/bftree/full_precision_streaming.rs | Switches bf-tree full-precision streaming to shared streamer + StreamRunner. |
| diskann-benchmark/src/index/benchmarks.rs | Refactors in-mem streaming benchmark to use shared run_streaming/build_streamer + StreamRunner. |
| diskann-benchmark/README.md | Updates documented command/example to new streaming input filename and tag. |
| diskann-benchmark/example/graph-index-stream.json | Updates streaming example JSON to new schema (type, search shape, no groundtruth). |
| diskann-benchmark/example/graph-index-bftree.json | Updates bf-tree build example to unified graph-index-bftree + quantization config. |
| diskann-benchmark/example/graph-index-bftree-stream.json | Updates bf-tree streaming example to unified tag + new search + quantization config. |
| diskann-benchmark/example/graph-index-bftree-stream-spherical.json | Adds new example for bf-tree spherical streaming with unified quantization config. |
| diskann-benchmark/example/graph-index-bftree-spherical.json | Updates spherical bf-tree build example to unified tag + nested quantization config. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1169 +/- ##
==========================================
- Coverage 90.15% 90.08% -0.07%
==========================================
Files 509 511 +2
Lines 97191 97376 +185
==========================================
+ Hits 87625 87724 +99
- Misses 9566 9652 +86
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
c33c71c to
634b25c
Compare
magdalendobson
left a comment
There was a problem hiding this comment.
Looks close, just a couple issues in comments.
hildebrandmw
left a comment
There was a problem hiding this comment.
Two big concerns here:
- Moving from
pub(crate)members topub(crate)accessors adds a lot of noise to the diff without meaningfully increasing protection. The real danger with sharing inputs is having backends ignore fields entirely, which per-field accessors doesn't really change. - This PR centers around
ManagedStream, but that is a crutch for providers not being able to manage their own ID translation. The upcoming inmem 2.0 (which can do everything except make it to a PR ready state), will not need this management layer. In which case, will a separate streaming implementation be needed anyways? Is there a way to share this infrastructure without requiringManagedStream?
hildebrandmw
left a comment
There was a problem hiding this comment.
Something was bothering me about removing the ID translation layer for bf-tree, and I think I managed to figure out what it is. For some reason, bf-tree is pre-configure with a maximum capacity. The BAB streaming runs simulate the runbook to obtain the maximum number of active points, which could be less than the maximum tag. It's this maximum that is provided to the BfTreeProvider's constructor. So, for runbooks with an active set that is indeed less, then the untranslated IDs will trigger the assert linked earlier. The runbook in test_data doesn't hit this because it doesn't have the property of the active points being less than the maximum tag.
Did this latest iteration go through the wikipedia runbook.
A small example runbook that triggers this situation:
sift-small-256:
max_pts: 128
1:
operation: "insert"
start: 0
end: 128
2:
operation: "search"
3:
operation: "delete"
start: 0
end: 64
4:
operation: "insert"
start: 128
end: 192
5:
operation: "search"
6:
operation: "replace"
tags_start: 64
tags_end: 128
ids_start: 192
ids_end: 256
7:
operation: "search"Squashed work from the benchmarks-refactor-v3 branch: - Unify bftree benchmark inputs and streaming backends - Add QuantConfig enum and stream-spherical example/config - Validation for streaming search params and Managed opt-in Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The direct (non-Managed) streaming path uses absolute runbook tag IDs directly as provider slot IDs. The provider was sized to the runbook's max concurrent point count (max_pts), but absolute tag IDs span the full dataset, so an insert with id >= capacity failed with "Vector id is out of boundary in the dataset". Size the provider to data.nrows() + num_start_points instead. Drop the redundant capacity parameter from build_direct_streamer and derive it from the loaded dataset in each bftree streaming closure. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The bf-tree direct (non-Managed) streaming path uses absolute runbook tag IDs as provider slot IDs, so the provider must be sized to the dataset ID space rather than the runbook's max concurrent active-point count. The existing test_data runbook does not exercise this because its max active count equals its max tag. Add a minimal, search-free runbook whose max tag (191) exceeds its max active-point count (128) via a delete-then-insert sequence, plus a bftree streaming config and an integration test. With correct sizing the test passes; a provider mis-sized to max_pts fails with "Vector id is out of boundary in the dataset". Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
ff5cff6 to
348a43e
Compare
Follow-up to #1168 (module reorganization). This PR contains the behavioral changes that were previously combined in #1142.
Summary
Simplifies the streaming benchmark infrastructure, unifies bf-tree input types, and makes
ManagedID translation opt-in — preparing for its removal with inmem 2.0.Changes
Managed ID translation is now opt-in
run_streamingand the streaming stack are now generic over the stream type via a newStreamingOutputtraitStream<DataArgs>directly onStreamRunner, using tags as slot IDs with noManagedwrapper (no ID translation overhead)Managedfor tag-to-slot ID translation until inmem 2.0 handles this nativelySlotReclaim::Immediateandrecycle_tags(bftree-only dead code)build_direct_streamerfor non-managed paths alongside the existingbuild_managed_streamerShared streaming execution layer
StreamRunner,Managed,build_managed_streamer, andrun_streaming— a reusable stack shared across in-mem and bf-tree streaming benchmarks, replacing duplicated per-backend logicInmemMaintainer(DropDeleted + Release) extracted from the old inline implementationStreaming search overhaul (addresses @magdalendobson feedback on #1142)
StreamingSearchParamsstruct with scalarnum_threads/search_l— streaming runs a single search config between runbook stages rather than sweeping a parameter matrixInput validation
search_l >= search_ninStreamingSearchParamsandGraphSearchrecall_k <= search_ninStreamingSearchParamsandGraphSearchto fail fast before long benchmark runsInput unification
graph-index-dynamic→graph-index-stream-run)BfTreeBuildandBfTreeStreamingRunto use aQuantConfigenum (none|spherical) instead of separate input structs per quantization modequantizer_utilfor shared spherical quantizer constructionInlineFilterSearchPhasenaming (accidentally overwritten)Housekeeping
PhantomData<fn() -> T>toPhantomData<T>onStreamRunnerstreaming_utils.rsandstreaming/full_precision.rs(subsumed by the shared layer)graph-index-bftree-stream-spherical.jsonBreaking changes to JSON inputs
Streaming input files need updating:
"type": "graph-index-dynamic"→"graph-index-stream-run""search_phase"→"search"(flat struct withqueries,reps,num_threads,search_l,search_n,recall_k)"quantization": { "kind": "none" }or the spherical config