perf(diskann): reduce search scratch allocations

[huanglune]

Summary

• Replace per-query DiskANN unordered_set / unordered_map state with reusable flat scratch buffers for visited bits, exact distances, cached neighbors, and in-flight I/O slots.
• Reuse PQ query residual scratch and integrate neighbor-cache release with bounded frontier eviction to avoid hot-path allocations.
• Size and resize scratch state through DiskANNLoadParams, with focused DiskANN tests covering the new scratch contract.

Test plan

• git diff --check
• cmake -B build/Release -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=ON
• cmake --build build/Release --target test_diskann_pq test_diskann_beam_search test_diskann_tombstone_search -j$(nproc)
• timeout 60s ctest --test-dir build/Release -R 'test_diskann_(pq|beam_search|tombstone_search)' --output-on-failure -j$(nproc)
• uvx pre-commit run --files include/index/graph/diskann/beam_search.hpp include/index/graph/diskann/diskann_index.hpp include/index/graph/diskann/pq_table.hpp include/index/graph/diskann/search_scratch.hpp include/index/graph/diskann/visited_bitset.hpp include/index/graph/vamana/robust_prune.hpp tests/diskann/test_diskann_beam_search.cpp tests/diskann/test_diskann_pq.cpp tests/diskann/test_diskann_tombstone_search.cpp

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[gemini-code-assist]

Code Review

This pull request optimizes the DiskANN search path by replacing dynamic structures (such as std::unordered_set and std::unordered_map) with flat, pre-allocated structures like VisitedBitset and contiguous arrays in ThreadData. It also updates NeighborPriorityQueue to return insertion results and allows PQTable::preprocess_query to use a pre-allocated scratch buffer, reducing heap allocations. The review feedback focuses on further optimizing and cleaning up the new ThreadData implementation: specifically, using the standard std::isnan instead of manual bit-manipulation, tracking the in-flight count to make has_inflight() an O(1) operation, and optimizing the linear scan in remove_inflight() to avoid hot-path performance bottlenecks.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

This manual implementation to check for NaN is correct, but std::isnan() from the <cmath> header is more idiomatic, readable, and portable. Consider using it for better maintainability.

You would need to add #include <cmath> at the top of the file.

  [[nodiscard]] static bool is_missing_exact(float value) {
    return std::isnan(value);
  }

[gemini-code-assist]

The has_inflight() function performs a linear scan over the inflight vector, which can have up to 1024 elements. This can be optimized to an O(1) operation by maintaining a counter for the number of in-flight requests.

To implement this, you could:

1. Add uint32_t inflight_count = 0; to the ThreadData struct.
2. Increment this counter in set_inflight when a slot becomes occupied.
3. Decrement it in remove_inflight when a request is successfully removed.
4. Reset it to 0 in clear_inflight.
5. Change has_inflight to use the counter as suggested.

  [[nodiscard]] bool has_inflight() const {
    return inflight_count > 0;
  }

[gemini-code-assist]

remove_inflight performs a linear scan over the inflight vector. Since this vector can have up to 1024 elements, this O(N) operation (where N is the number of page slots) could be a performance bottleneck in a hot path. The previous std::unordered_map implementation was O(1) on average.

To optimize this while avoiding heap allocations, you could maintain a separate small data structure that maps an id to a page_slot for in-flight requests. For example, a std::vector<std::pair<uint32_t, uint64_t>> kept sorted by id would allow for O(log k) lookups via binary search, where k is the number of in-flight requests. This would be a significant improvement over the current linear scan.