Reform Range Search

[Magdalen Dobson Manohar (magdalendobson)]

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.

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 99.77064% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.53%. Comparing base (6ffe6d3) to head (5229445).
⚠️ Report is 19 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.53%   -0.03%     
==========================================
  Files         527      527              
  Lines      103254   103352      +98     
==========================================
+ Hits        94537    94606      +69     
- Misses       8717     8746      +29     

Flag	Coverage Δ
miri	91.53% <99.77%> (-0.03%)	⬇️
unittests	91.28% <99.77%> (-0.02%)	⬇️

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.80% <100.00%> (+0.70%)	⬆️
diskann/src/graph/search/filtered_range_search.rs	97.95% <97.82%> (-2.05%)	⬇️

... and 82 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)]

Sorry, accidentally hit the "approve" button.

[Mark Hildebrand (hildebrandmw)]

Thanks Magdalen - Implementation looks good, but can you please take another look at the tests? At least in range search, the tests don't seem to break with changes in initial_slack.

[Mark Hildebrand (hildebrandmw)]

Thanks - found one small issue with one of the range search tests, otherwise looks like a good improvement.