Reform Range Search by magdalendobson · Pull Request #1337 · microsoft/DiskANN

Magdalen Dobson Manohar (magdalendobson) · 2026-08-14T19:39:35Z

Range search had some problems, and the new filtered range search inherited many of them.

Problem 1: Range search added some extra elements to scratch for convenience. However, given that the result size of range search can vary wildly, it doesn't make too much sense to reuse them, and it clogged up the scratch space. This PR pulls them out of scratch and manages them directly inside the search algorithm.

Problem 2: Range search has a lot of parameters, and their meaning and use was kind of opaque to people who aren't me. This PR adds a nice long docstring that actually explains them.

Problem 3: Two parameters--the initial search slack and the range search slack--weren't really tested. This PR adds some tests that ensure that they behave as intended.

Problem 4: There was a lot of weirdness in handling of max_results. With start point filtering, if start points fell within the radius, they could be excluded during post-process and the search could return fewer than max_results points even if satisfying points existed. This PR sets an "effective" max results that includes the number of start points. Furthermore, since inner radius filtering wasn't done until the end, if the inner radius ended up excluding a lot of points at the end, we could again return way fewer than max_results points even when they exist. This PR fixes that by including inner radius filtering during the second round of search.

To do this, it includes a new InRange struct which automatically handles enforcing the radius constraints, max results, and filtering if applicable. This removes the need for the DistanceFiltered buffer that was previously handed to the post-processor.

NB: this PR also removes one test from both the filtered and unfiltered range search. The test didn't make a lot of sense and it is now subsumed by the testing of initial_slack.

Copilot

Pull request overview

This PR refactors DiskANN’s range search and filtered range search implementation to reduce scratch-space bloat, clarify parameter semantics, and improve correctness around max_returned and inner-radius behavior, while updating the corresponding golden baselines.

Changes:

Moves range-search frontier/result bookkeeping out of SearchScratch and into the range-search algorithms (introducing an InRange helper).
Adds/updates tests and regenerated baselines to cover initial_slack / range_slack behavior and revised second-round triggering.
Adjusts handling of max_returned to better account for start-point filtering and apply inner-radius constraints earlier.

Reviewed changes

Copilot reviewed 6 out of 17 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
diskann/src/graph/search/range_search.rs	Introduces `InRange` and refactors unfiltered range search to manage frontier/results locally and apply radius constraints during search.
diskann/src/graph/search/filtered_range_search.rs	Adapts filtered range search to use `InRange`, adds effective max helpers, and updates second-round search flow.
diskann/src/graph/search/scratch.rs	Removes range-search-specific scratch fields (`range_frontier`, `in_range`) from shared scratch space.
diskann/src/graph/search/diverse_search.rs	Updates `SearchScratch` initialization to match removed fields.
diskann/src/graph/test/cases/range_search.rs	Reworks range-search tests to cover `initial_slack` and `range_slack`, and updates expectations for second-round behavior.
diskann/src/graph/test/cases/filtered_range_search.rs	Removes a prior max-results test and adds new slack-focused tests for filtered range search.
diskann/test/generated/graph/test/cases/range_search/two_round_search.json	Updates recorded baseline stats/results for revised two-round behavior.
diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json	Updates recorded baseline to match updated max-results/second-round behavior.
diskann/test/generated/graph/test/cases/range_search/max_results_respected_means_no_second_round.json	Removes obsolete baseline corresponding to deleted/renamed test.
diskann/test/generated/graph/test/cases/filtered_range_search/max_results_respected_means_no_second_round.json	Removes obsolete baseline corresponding to deleted test.

Suppressed comments (1)

diskann/src/graph/search/filtered_range_search.rs:284

filtered_range_search_internal stops expanding once matched_in_range.len() == search_params.max_returned(). This bypasses the earlier “effective max” (max_returned + num_start_ids) slack that the caller sets up to compensate for start points being filtered during post_process, so the second round may still return fewer than the requested results even when more exist.

    let max_returned = search_params.max_returned().unwrap_or(usize::MAX);

    'outer: while !range_frontier.is_empty() && matched_in_range.len() < max_returned {

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

        self.range_params.max_returned()
    }

+    /// Returns the maximum number of results to return.


Codecov Comments Bot (codecov-commenter) · 2026-08-14T19:55:36Z

Codecov Report

❌ Patch coverage is 99.78858% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.57%. Comparing base (fabcb9b) to head (a10cfab).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann/src/graph/search/filtered_range_search.rs	97.82%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1337      +/-   ##
==========================================
+ Coverage   91.55%   91.57%   +0.01%     
==========================================
  Files         521      521              
  Lines      100371   100664     +293     
==========================================
+ Hits        91895    92182     +287     
- Misses       8476     8482       +6

Flag	Coverage Δ
miri	`91.57% <99.78%> (+0.01%)`	⬆️
unittests	`91.25% <99.78%> (+0.01%)`	⬆️

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

Files with missing lines	Coverage Δ
diskann-benchmark/src/index/benchmarks.rs	`75.99% <100.00%> (+0.04%)`	⬆️
diskann/src/graph/search/range_search.rs	`100.00% <100.00%> (+0.27%)`	⬆️
diskann/src/graph/search/scratch.rs	`98.14% <ø> (-0.07%)`	⬇️
...kann/src/graph/test/cases/filtered_range_search.rs	`99.42% <100.00%> (+0.48%)`	⬆️
diskann/src/graph/test/cases/range_search.rs	`98.90% <100.00%> (+0.80%)`	⬆️
diskann/src/graph/search/filtered_range_search.rs	`97.95% <97.82%> (-2.05%)`	⬇️

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:

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

Mark Hildebrand (hildebrandmw) · 2026-08-14T20:35:36Z

Sorry, accidentally hit the "approve" button.

Harsha Simhadri (harsha-simhadri) · 2026-08-17T23:42:44Z

 /// Builder for [`Range`] search parameters.
+///
+/// `max_returned`: If specified, the search will stop and return results once this number of points has been found within both the inner and outer radii.
+/// Note that due to adding some extra slack for start points, occasionally slightly more than this number of points may be returned. Since the initial


in this corner case, why cant we truncate the list of starting points and just honor the max_returned list?

I pushed an update that does this but it is a little awkward. The problem is that filtering out start points happens during post_process, so the truncation needs to happen after that point, and there isn't a built-in way to truncate an output buffer. The workaround is to introduce a new output buffer that handles the truncation after filtering, which I do here, with a corresponding integration test that checks that the desired behavior happens (e.g. if more than the maximum number of points are returned, but none of them are start points, truncation happens).

I do think this adds a lot of code for a relatively small gain, so would like opinions on whether to 1) stick with behavior in the PR as-is, 2) revert to sometimes over-returning, or 3) allow under-returning, thus going back to behavior on current main. CC Mark Hildebrand (@hildebrandmw)

We already have some precedent for over returning in the normal KNN search because we don't know precisely whether or not start points will be filtered in post processing. For KNN, this is usually solved by the user supplying a fixed size buffer to terminate post-processing. I don't think it's unreasonable to document max_returned as more of an early termination hint rather than hard guarantee.

That said, I've been operating under the assumption that start points could be either valid or invalid depending on the caller. This gives us the flexibility in the future to seed searches with valid points if we need to.

Restored to the previous functionality after the buffer proof-of-concept turned out to be messy.

…en/range_search_reform

…sers/magdalen/range_search_reform

Mark Hildebrand (hildebrandmw) · 2026-08-17T22:04:07Z

 }

 /// Builder for [`Range`] search parameters.
+///


Request: can you please wrap the comments at column 100 or so? Otherwise, they get quite difficult to read and edit in editors.

Mark Hildebrand (hildebrandmw) · 2026-08-17T22:05:24Z

+/// Note that due to adding some extra slack for start points, occasionally slightly more than this number of points may be returned. Since the initial
+/// search phase does not respect `max_returned`, this parameter may not be set lower than `starting_l`.
+///
+/// `starting_l`: the L_search parameter for the initial search phase. Must be greater than zero.


Maybe cross reference with KNN::search_l?

Mark Hildebrand (hildebrandmw) · 2026-08-17T22:06:08Z

+///
+/// `radius`: the outer radius for the range search. Points within this distance from the query are candidates for inclusion in the results.
+///
+/// `inner_radius`: the inner radius for the range search. Points closer than this distance from the query are excluded from the results. Must be less than or equal to `radius` if specified.


I believe radii must be positive and finite as well?

Radii do NOT have to be positive because of inner product search. I guess radii should be finite, but we don't explicitly check for it at the moment, is that something you want to change?

After our side conversation, I added is_finite checks for all float values and also a note on negative radius for inner product.

Mark Hildebrand (hildebrandmw) · 2026-08-19T23:39:08Z

 /// Builder for [`Range`] search parameters.
+///
+/// `max_returned`: If specified, the search will stop and return results once this number of points has been found within both the inner and outer radii.
+/// Note that due to adding some extra slack for start points, occasionally slightly more than this number of points may be returned. Since the initial


We already have some precedent for over returning in the normal KNN search because we don't know precisely whether or not start points will be filtered in post processing. For KNN, this is usually solved by the user supplying a fixed size buffer to terminate post-processing. I don't think it's unreasonable to document max_returned as more of an early termination hint rather than hard guarantee.

That said, I've been operating under the assumption that start points could be either valid or invalid depending on the caller. This gives us the flexibility in the future to seed searches with valid points if we need to.

…en/range_search_reform

Magdalen Manohar and others added 2 commits August 12, 2026 15:31

remove fields from scratch

68d3932

clean up range search a lot yay

9187f0f

Magdalen Dobson Manohar (magdalendobson) requested review from a team and a lite review from Copilot August 14, 2026 19:39

Copilot started reviewing on behalf of Magdalen Dobson Manohar (magdalendobson) August 14, 2026 19:40 View session

Magdalen Manohar added 2 commits August 14, 2026 19:40

add missing test jsons

cdcbb84

somehow another format check

79ee339

Copilot AI reviewed Aug 14, 2026

View reviewed changes

Comment thread diskann/src/graph/search/filtered_range_search.rs

self.range_params.max_returned()

}

/// Returns the maximum number of results to return.

Mark Hildebrand (hildebrandmw) approved these changes Aug 14, 2026

View reviewed changes

Mark Hildebrand (hildebrandmw) self-requested a review August 14, 2026 20:34

Proposed change sketch.

6ae9b88

Harsha Simhadri (harsha-simhadri) reviewed Aug 17, 2026

View reviewed changes

Magdalen Manohar added 4 commits August 19, 2026 13:04

Merge branch 'main' of github.com:microsoft/DiskANN into users/magdal…

d877e9e

…en/range_search_reform

Merge branch 'mhildebr/review' of github.com:microsoft/DiskANN into u…

2e6b10c

…sers/magdalen/range_search_reform

merge with main and mhildebr/review, add tests for InRange

027cb1f

added result limiting post-processor and corresponding tests

474ad46

Mark Hildebrand (hildebrandmw) reviewed Aug 31, 2026

View reviewed changes

Magdalen Manohar added 3 commits August 31, 2026 15:22

Merge branch 'main' of github.com:microsoft/DiskANN into users/magdal…

1a3e052

…en/range_search_reform

undo buffer change, add some comments and extra checks

fbf2acf

docs, fmt

a10cfab

Uh oh!

Conversation

Magdalen Dobson Manohar (magdalendobson) commented Aug 14, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Codecov Comments Bot (codecov-commenter) commented Aug 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Mark Hildebrand (hildebrandmw) commented Aug 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Codecov Comments Bot (codecov-commenter) commented Aug 14, 2026 •

edited

Loading