Return native indexed vectors from disk search

[Yujie Zhang (yjiez)]

Summary

Adds an opt-in disk search API that returns each valid ANN result with its canonical native indexed vector:

search_with_indexed_vectors(...)

The existing search() API and its padded result behavior remain unchanged.

Public contract

pub struct SearchResultItemWithIndexedVector<A, V> {
    pub vertex_id: u32,
    pub data: A,
    pub distance: f32,
    pub indexed_vector: Box<[V]>,
}

• The vector is the native representation stored in the graph and used for exact scoring; it is not a PQ code and may differ from the original input vector.
• The indexed-vector API returns valid results only: results.len() == stats.result_count.
• A missing vector is returned as ANNError, not a panic.

For reviewers

Recommended review order:

1. Public API: SearchResultWithIndexedVectors, SearchResultItemWithIndexedVector, and search_with_indexed_vectors.
2. Output compatibility: SearchPayload and SearchOutput extend the existing slice-backed output without changing legacy search().
3. Traversal reuse: DiskAccessor::ensure_loaded optionally stores the native vector beside the existing exact-distance cache entry.
4. Final output: extend_output moves traversal-cached winner vectors into results or copies winners from the post-process batch already loaded by the existing rerank path.
5. Scratch cleanup: DiskAccessor::drop immediately releases vectors left in the request-local cache.
6. Coverage: the existing 128-dimensional disk-search test checks result parity and vector values with both no cache and static cache.

The intended invariants are:

• A returned vector always belongs to the same vertex_id as its result.
• Enabling vector output does not change IDs or distances.
• Legacy search() does not allocate or copy indexed vectors.
• Returning vectors does not add a final disk-read round.

Key data structures

Internal search payload

type SearchPayload<A, V> = (u32, A, Option<Box<[V]>>);

The Option is internal: legacy search writes None; indexed-vector search writes Some(vector). The public indexed-vector result contains a non-optional Box<[V]>.

Output adapter

struct SearchOutput<'a, A, V> {
    output: IdDistanceAssociatedData<'a, u32, A>,
    indexed_vectors: Option<&'a mut [Option<Box<[V]>>]>,
}

SearchOutput keeps the existing ID/distance/associated-data buffers and adds an optional vector lane. One push writes all fields at the same position, preserving ID/vector alignment without replacing the existing search pipeline.

Request-local traversal cache

distance_cache: HashMap<
    u32,
    (f32, AssociatedData, Option<Box<[VectorData]>>),
>

The map already existed for exact distance and associated data. When indexed vectors are requested, the same entry also owns the native vector loaded during traversal. There is no collector, mutex, or second hash map. The map is cleared when the query accessor is dropped so loser vectors are not retained in the scratch pool between requests.

Data flow

search_with_indexed_vectors
    -> enable vector capture for this query
    -> traversal load: distance_cache stores distance + data + vector
    -> existing rerank path loads uncached candidates as one batch
    -> extend_output:
         cached winner   -> move Box from distance_cache
         uncached winner -> copy from current in-memory post-process batch
    -> return only stats.result_count valid results

extend_output never calls load_vertices, so it does not introduce a final fallback I/O pass.

Benchmark

GitHub Actions run with immediate cleanup, using K=1000, L=2000, recall@100, and four search threads:

Dataset	Mean API latency vs legacy	QPS vs legacy	Mean I/O count	Indexed-vector peak delta
Wikipedia-100K, 768d	+1.43%	-1.41%	unchanged	+23.25 MiB (+26.78%)
OpenAI-100K, 1536d	+1.42%	-2.09%	unchanged	+46.75 MiB (+13.79%)

Expected Peak-Memory Increase

The indexed-vector API retains the full-precision vectors needed by active searches. Its expected vector-payload increase can therefore be approximated as:
Expected peak delta ≈ C × L × D × S

where:

• C is the number of concurrent searches;
• L is the search-list size;
• D is the vector dimension;
• S is the size of each vector element (4 bytes for f32).

For Wikipedia-100K:
4 × 2000 × 768 × 4 bytes ≈ 23.44 MiB
This closely matches the measured increase of 23.25 MiB.

Closes #1339

[Copilot]

Pull request overview

This PR extends diskann-disk’s disk-index search surface with a new API that returns only valid ANN hits, each bundled with its canonical (native) indexed vector as stored in the on-disk graph. This supports downstream consumers that need full-precision indexed vectors alongside IDs/distances, while keeping the existing search() API unchanged.

Changes:

• Added search_with_indexed_vectors(...) plus new result types to return owned per-hit indexed vectors without padding.
• Implemented request-local indexed-vector capture via an IndexedVectorCollector, with different capture strategies per search mode and a post-rerank “missing winner” batch fetch.
• Expanded tests (including builder tests) to validate result parity, ownership/capacity behavior, and stored-vector round-trip behavior (including MinMax).

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
diskann-disk/src/search/provider/disk_provider.rs	Adds the new public search API, request-local vector collector, capture logic across traversal/post-processing, and associated tests.
diskann-disk/src/build/builder/tests.rs	Extends disk index build tests to validate returned indexed vectors against stored rows.
diskann-disk/src/build/builder/core.rs	Adds a test helper that verifies indexed vectors returned by search match the stored dataset rows.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 96.45161% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.55%. Comparing base (6f2ff48) to head (e554e3f).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-disk/src/search/provider/disk_provider.rs	96.45%	11 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1345      +/-   ##
==========================================
- Coverage   91.58%   91.55%   -0.03%     
==========================================
  Files         521      521              
  Lines       99598   100347     +749     
==========================================
+ Hits        91212    91874     +662     
- Misses       8386     8473      +87     

Flag	Coverage Δ
miri	91.55% <96.45%> (-0.03%)	⬇️
unittests	91.23% <96.45%> (-0.03%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
diskann-disk/src/search/provider/disk_provider.rs	95.81% <96.45%> (+0.05%)	⬆️

... and 15 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[juchen-ms (partychen)]

The API and its intended behavior make sense, but I’m concerned about the complexity and performance cost of the collector design.

ensure_loaded copies every loaded vector into a request-local Mutex<HashMap>, so memory usage, allocations, locking, and hash lookups scale with all visited vertices rather than the final result count. Could we instead complete the regular search first, then batch-load only the final result vectors? That would keep this logic out of the traversal hot path and eliminate most of the additional plumbing.

I also think the expect() and assert_ne!() on public search paths should return ANNError rather than panic, and the per-query info! log should probably be lowered to debug!.

[Aditya Krishnan (arkrishn94)]

Thanks Yujie, I have left some comments but they can mostly come in follow-up since I know this is time sensitive. I am approving but would be good to beef-up the testing a bit before merging. From what I can tell there is only one test with a single query testing this new path. Would also be good to test the error cases.