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. 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.

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.

Validation

• cargo fmt --all -- --check
• cargo check -p diskann-disk --all-targets --all-features
• Existing 128-dimensional disk-search coverage checks both CachingStrategy::None and static caching, unchanged IDs/distances, and returned vectors against stored rows.

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.81818% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.59%. Comparing base (6f2ff48) to head (4efb698).

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

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1345      +/-   ##
==========================================
+ Coverage   91.58%   91.59%   +0.01%     
==========================================
  Files         521      521              
  Lines       99598    99788     +190     
==========================================
+ Hits        91212    91396     +184     
- Misses       8386     8392       +6     

Flag	Coverage Δ
miri	91.59% <96.81%> (+0.01%)	⬆️
unittests	91.27% <96.81%> (+0.01%)	⬆️

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.83% <96.81%> (+0.07%)	⬆️

... and 4 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[Aditya Krishnan (arkrishn94)]

Thanks for working on this Yujie. I must admit, I am bit confused by the complexity of the approach here.

There's a few problems from what I'm seeing in the current design-

1. The collector seems to be an additional abstraction just to store full precision vectors, if needed, during the traversal but the CachedDiskVertexProvider is already doing this.
2. With the way the collector is implemented (i.e. MutexLock, hash-lookup and hash-map on insertion), this seems to be pretty expensive operations in the hot-path (i.e. in ensure_loaded).
3. We're threading the collector through the strategy and the accessor and collecting results there when the glue::SearchPostProcess architecture is explicitly designed to be generic over the shape of the output. So it feels weird to me that you need to have a new structure to collect results in the accessor.

I'm curious, have we tried doing the simple thing of just loading the vertices (by calling load_vertices on the vertex provider) from the result set after the regular search path? Why can't we just do that in search_with_indexed_vectors? It seems like that should allow you to get rid of the collector completely

Let me know if I'm missing something here

[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!.