Return native indexed vectors from disk search (#1345) ## Summary Adds an opt-in disk search API that returns each valid ANN result with its canonical native indexed vector: ```rust search_with_indexed_vectors(...) ``` The existing `search()` API and its padded result behavior remain unchanged. ## Public contract ```rust pub struct SearchResultItemWithIndexedVector { 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 ```rust type SearchPayload = (u32, A, Option>); ``` 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 ```rust struct SearchOutput<'a, A, V> { output: IdDistanceAssociatedData<'a, u32, A>, indexed_vectors: Option<&'a mut [Option>]>, } ``` `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 ```rust distance_cache: HashMap< u32, (f32, AssociatedData, Option>), > ``` 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 ```text 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](https://github.com/microsoft/DiskANN/actions/runs/32388396932), 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 --------- Co-authored-by: Claude Opus 4.8 (1M context)

Uh oh!

Add support for passing full-precision vectors along with results in disk-index search. #1339

Description

Activity

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions