Skip to content

perf(diskann): reuse VisitedSet across searches, add search_with and allocation regression test#683

Open
ohdearquant wants to merge 2 commits into
ruvnet:mainfrom
ohdearquant:perf/diskann-reusable-visited-set
Open

perf(diskann): reuse VisitedSet across searches, add search_with and allocation regression test#683
ohdearquant wants to merge 2 commits into
ruvnet:mainfrom
ohdearquant:perf/diskann-reusable-visited-set

Conversation

@ohdearquant

Copy link
Copy Markdown
Contributor

Summary

This change removes the per-query VisitedSet allocation from hot DiskAnnIndex::search() calls while preserving the existing API and result ordering.

  • Adds DiskAnnIndex::search_with, which accepts a caller-owned VisitedSet and routes directly through VamanaGraph::greedy_search_fast.
  • Replaces the unreachable preallocated visited field with a bounded std::sync::Mutex<Vec<VisitedSet>> pool. The normal single-threaded path checks out and returns one preallocated set. An empty or contended pool falls back to a fresh set, and the pool retains at most four sets.
  • Makes generation rollover safe by fully zeroing the visited storage when the u64 generation reaches its maximum.
  • Handles index-size mismatches by resizing and fully resetting the caller's visited storage. Repeated searches at the same size retain the O(1) generation-counter reset.
  • Adds deterministic equivalence coverage across 100 queries against both search_with and the pre-fix allocating algorithm shape.
  • Adds a counting-allocator regression test covering 100 hot pooled searches after warmup.

Fixes #677

Measurements

bench_search_hot builds seeded random unit vectors with d=128, then measures 2,000 single-threaded searches at k=10 in a release build. Build settings are R=4, L_build=8, L_search=64, and alpha=1.0. The table reports the median of three runs per configuration. Additional raw samples are retained where visible host variance prompted repeat collection.

Vectors Metric Before After
100,000 p50 latency 41.333 us 29.042 us
100,000 p99 latency 301.417 us 116.208 us
100,000 Queries/sec 18,111.622 27,236.237
100,000 Visited-state allocation/query 812,504 bytes 0 bytes
250,000 p50 latency 54.666 us 33.084 us
250,000 p99 latency 486.834 us 116.417 us
250,000 Queries/sec 5,524.658 22,139.261
250,000 Visited-state allocation/query 2,031,256 bytes 0 bytes

Allocation values are the exact backing storage requested by the old VisitedSet::new path: the n * size_of::<u64>() generation array plus the bit-vector storage. The after value is the steady-state visited-state allocation after one warmup. The regression test additionally asserts that 100 hot searches allocate zero blocks at or above the 2,000,000-byte generation-array size and less than 300,000 bytes total for all other search output and candidate storage combined.

Hardware: Apple M2 Max, 32 GB RAM, macOS 26.6. Searches were single-threaded and binaries used the Cargo release profile.

Reproduce with cargo run --release -p ruvector-diskann --example bench_search_hot. The benchmark commit predates the fix commit on this branch, so before numbers come from checking out the benchmark commit alone.

Compatibility and safety

The public search signature and behavior remain unchanged. The new search_with API allows explicit reuse where callers want to own query scratch state. Search candidates, exact re-ranking, and ordering are unchanged.

The pool uses only the standard library, adds no dependency, and does not assume threads exist. This keeps the single-thread pool path suitable for wasm32. The upstream crate baseline currently fails wasm32-unknown-unknown compilation in getrandom because its JavaScript feature is not enabled, so no new wasm target regression was introduced or masked here.

Verification

  • cargo fmt -p ruvector-diskann
  • cargo clippy -p ruvector-diskann --all-targets -- -D warnings
  • cargo test -p ruvector-diskann
  • cargo build --release -p ruvector-diskann --examples, on both the unfixed and fixed source

OceanLi added 2 commits July 12, 2026 12:52
Fixes ruvnet#677

Co-Authored-By: Leo <noreply@khive.ai>
…allocation regression test

Fixes ruvnet#677

Co-Authored-By: Leo <noreply@khive.ai>
@ohdearquant

Copy link
Copy Markdown
Contributor Author

CI note — the red checks on this PR are not caused by this diff:

  • Clippy (deny warnings): pre-existing failure in crates/ruvector-rabitq/src/persist.rs:194 (useless_borrows_in_formatting on MAGIC, &magic), present on main — unrelated to this branch.
  • dependency-review (PRs only): errors with "Dependency review is not supported on this repository. Please ensure that Dependency graph is enabled" — a repo Settings toggle, fails on every PR.
  • Security audit / cargo audit / cargo deny: these are the RustSec advisories cleared by fix(deps): clear RUSTSEC advisory failures in supply-chain CI #682. This branch was cut before those dependency bumps, so it still carries them. Merging fix(deps): clear RUSTSEC advisory failures in supply-chain CI #682 first (or rebasing this branch on top of it) turns the audit gates green.

The DiskANN build and test jobs for this change all pass.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

diskann: search() allocates an O(n) VisitedSet per query — the pre-allocated reusable set is never used

1 participant