Skip to content

Switch SearchOutputBuffer to accepting Neighbors instead of tuples - #1284

Open
hildebrandmw wants to merge 6 commits into
mainfrom
mhildebr/neighbor-2
Open

Switch SearchOutputBuffer to accepting Neighbors instead of tuples#1284
hildebrandmw wants to merge 6 commits into
mainfrom
mhildebr/neighbor-2

Conversation

@hildebrandmw

Copy link
Copy Markdown
Contributor

With #1273, the structural Eq bound on Neighbor has been removed. This PR is a follow-up that changes over SearchOutputBuffer to accept Neighbor instead of raw tuples. Outside of being symmetric, this has the nice property of allowing full-precision reranking methods to use diskann::neighbor::ord::fast_distance as their sorting callback instead of hand-rolling the implementation every time.

This PR does make the distance member of Neighbor generic, but defaults it to f32 to keep syntactic compatibility.

Key files to review:

  • diskann/src/neighbors.rs: Additional (defaulted) generic distance parameter to Neighbor.
  • diskann/src/graph/search_output_buffer.rs: Updated trait definition.

let mut process = |n: u32| {
if let Some(entry) = accessor.scratch.distance_cache.get(&n) {
Some(Ok::<((u32, _), f32), ANNError>(((n, entry.1), entry.0)))
Some(Neighbor::new((n, entry.1), entry.0))

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: The original code was unnecessarily wrapping the intermediate value in a Result but didn't call any fallible methods.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR follows up on #1273 by completing the migration from “(id, distance) tuples” to the unified Neighbor type for search result emission. It also makes Neighbor’s distance type generic (defaulting to f32) so reranking/postprocessing code can sort with neighbor::ord::fast_distance directly rather than re-implementing distance ordering.

Changes:

  • Generalize Neighbor to Neighbor<I, D = f32> and update downstream uses to accommodate distance() returning a reference.
  • Update SearchOutputBuffer to accept Neighbor<I, D> in push/extend, and adjust implementations and call sites across the workspace.
  • Switch multiple reranking implementations to sort with neighbor::ord::fast_distance instead of hand-rolled partial_cmp logic.

Reviewed changes

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

Show a summary per file
File Description
diskann/src/neighbor/mod.rs Makes Neighbor distance generic and updates BackInserter/buffer impls to accept Neighbor directly.
diskann/src/graph/search_output_buffer.rs Changes SearchOutputBuffer trait to accept Neighbor in push/extend and updates internal buffer implementations/tests.
diskann/src/graph/search/range_search.rs Updates range search logic and DistanceFiltered buffer wrapper to work with Neighbor and ref-returning distance().
diskann/src/graph/test/cases/range_search.rs Adjusts range-search test invariants for ref-returning Neighbor::distance().
diskann/src/graph/index.rs Updates distance extraction to dereference Neighbor::distance().
diskann/src/graph/glue.rs Removes tuple mapping and passes Neighbor iterators directly into output buffers.
diskann/src/flat/test/harness.rs Adjusts test harness to collect copied distances from Neighbor::distance().
diskann-tools/src/utils/ground_truth.rs Updates ground-truth writing to dereference Neighbor::distance().
diskann-providers/src/test_utils/search_utils.rs Updates distance comparisons for ref-returning distance() (but one arithmetic use still needs updating).
diskann-providers/src/model/graph/provider/async_/postprocess.rs Removes tuple conversion in postprocess output and forwards Neighbor directly.
diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs Builds reranked results as Neighbor and sorts via neighbor::ord::fast_distance.
diskann-providers/src/index/wrapped_async.rs Updates assertions to compare dereferenced distances.
diskann-providers/src/index/diskann_async.rs Updates distance comparisons/mutations to dereference Neighbor::distance().
diskann-garnet/src/provider.rs Updates reranking paths to use Neighbor and sort via fast_distance.
diskann-garnet/src/lib.rs Updates SearchOutputBuffer implementation and tests to accept/push Neighbor.
diskann-disk/src/search/provider/disk_provider.rs Migrates reranking to store Neighbor (including cached results) and sort via fast_distance.
diskann-bftree/src/provider.rs Updates reranking to produce/sort Neighbor instead of (id, f32) tuples.
diskann-benchmark-core/src/search/api.rs Updates benchmark test plumbing to emit Neighbor into output buffers.
diskann-benchmark-core/src/internal/buffer.rs Updates internal benchmark output buffer implementation to accept Neighbor items.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines 24 to +28
/// Returns a [`BufferState`] to indicate whether future insertions will succeed.
///
/// Unlike the iterator interface, implementations should return [`BufferState::Full`]
/// if **future** insertions will fail to prevent unnecessary work.
fn push(&mut self, id: I, distance: D) -> BufferState;
fn push(&mut self, neighbor: Neighbor<I, D>) -> BufferState;
Comment thread diskann/src/neighbor/mod.rs Outdated
Comment on lines +212 to +216
/// A [`SearchOutputBuffer`] wrapper around `&mut [Neighbor<I>]`. This can be used to
/// populate such a mutable slice as the result of [`crate::graph::DiskANNIndex::search`].
#[derive(Debug)]
pub struct BackInserter<'a, I> {
buffer: &'a mut [Neighbor<I>],
pub struct BackInserter<'a, I, D = f32> {
buffer: &'a mut [Neighbor<I, D>],
Comment on lines 86 to 90
/// Return the distance.
#[inline]
pub fn distance(&self) -> f32 {
self.distance
pub fn distance(&self) -> &D {
&self.distance
}
Copilot AI review requested due to automatic review settings July 28, 2026 22:07

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

diskann/src/neighbor/mod.rs:150

  • Neighbor now has a generic distance parameter (Neighbor<I, D>), but ord::fast_distance is still hard-coded to Neighbor<I> (i.e., D = f32). That prevents using fast_distance as a sorting callback for Neighbor<I, D> when D is not f32 (even though partial_cmp would work for any D: PartialOrd). Consider making this function generic over D (and updating the doc comment that currently mentions f32::NAN).
    pub fn fast_distance<I>(x: &Neighbor<I>, y: &Neighbor<I>) -> std::cmp::Ordering {

@codecov-commenter

codecov-commenter commented Jul 28, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.64151% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.62%. Comparing base (3e7c1f3) to head (7c938f5).

Files with missing lines Patch % Lines
...odel/graph/provider/async_/inmem/full_precision.rs 42.85% 4 Missing ⚠️
diskann-providers/src/index/diskann_async.rs 66.66% 1 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@           Coverage Diff           @@
##             main    #1284   +/-   ##
=======================================
  Coverage   90.62%   90.62%           
=======================================
  Files         512      512           
  Lines       98807    98817   +10     
=======================================
+ Hits        89541    89553   +12     
+ Misses       9266     9264    -2     
Flag Coverage Δ
miri 90.62% <97.64%> (+<0.01%) ⬆️
unittests 90.30% <97.64%> (+<0.01%) ⬆️

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

Files with missing lines Coverage Δ
diskann-benchmark-core/src/internal/buffer.rs 97.24% <100.00%> (+0.06%) ⬆️
diskann-benchmark-core/src/search/api.rs 97.85% <100.00%> (+<0.01%) ⬆️
diskann-bftree/src/provider.rs 90.42% <100.00%> (-0.01%) ⬇️
diskann-disk/src/search/provider/disk_provider.rs 93.92% <100.00%> (-0.01%) ⬇️
diskann-garnet/src/lib.rs 92.06% <100.00%> (+0.02%) ⬆️
diskann-garnet/src/provider.rs 80.27% <100.00%> (-0.08%) ⬇️
diskann-inmem/src/provider.rs 86.94% <100.00%> (+0.06%) ⬆️
diskann-providers/src/index/wrapped_async.rs 60.12% <100.00%> (ø)
...ers/src/model/graph/provider/async_/postprocess.rs 100.00% <100.00%> (ø)
diskann-providers/src/test_utils/search_utils.rs 88.05% <100.00%> (ø)
... and 10 more

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copilot AI review requested due to automatic review settings July 29, 2026 21:06

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 21 out of 21 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

diskann/src/neighbor/mod.rs:154

  • Neighbor now has a generic distance type parameter (D), but ord::fast_distance is still specialized to Neighbor<I> (i.e., D = f32). This prevents reusing the sorting callback for Neighbor<I, D> with other distance types (e.g., f64 or NotNan<f32>), which is one of the main benefits of making distance generic.
    pub fn fast_distance<I>(x: &Neighbor<I>, y: &Neighbor<I>) -> std::cmp::Ordering {
        x.distance()
            .partial_cmp(y.distance())
            .unwrap_or(std::cmp::Ordering::Equal)
    }

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.

4 participants