Clean up range search (#1253)

While working on #1228, I noticed that there were some issues with our
current implementation of range search and its testing.

The main issue with testing is that there were no tests ensuring the
`max_results` parameter was respected. I have added two tests that
ensure this now.

The main code had several issues with how `max_results` was handled:
1. The `max_results` parameter was allowed to be less than the initial
L_search. This is a conceptual issue because the user expects
`max_results` to stop the search from continuing for too long, and the
compute used in the initial search will always be controlled by
`initial_search_l`.
2. A `max_results` check was not enforced before deciding to continue to
the second round search. This meant that if the max results was reached
via the initial search, it might not be respected.
3. The second round search was not terminated when `max_results` was
reached, meaning it would continue to perform unnecessary work.

This PR fixes these issues by adding additional checks of `max_results`
at the correct points in the code.

---------

Co-authored-by: Magdalen Manohar <mmanohar@microsoft.com>

Author: magdalendobson

diskann/src/graph/search/range_search.rs

@@ -36,6 +36,8 @@ pub enum RangeSearchError {
     RangeSearchSlackValueError,
     #[error("inner_radius must be less than or equal to radius")]
     InnerRadiusValueError,
+    #[error("max_returned must be greater than or equal to starting_l")]
+    MaxReturnedLessThanInitialL,
 }
 
 impl From<RangeSearchError> for ANNError {
@@ -91,6 +93,11 @@ impl Range {
         if starting_l == 0 {
             return Err(RangeSearchError::LZero);
         }
+        if let Some(max) = max_returned
+            && max < starting_l
+        {
+            return Err(RangeSearchError::MaxReturnedLessThanInitialL);
+        }
         if !(0.0..=1.0).contains(&initial_slack) {
             return Err(RangeSearchError::StartingListSlackValueError);
         }
@@ -197,7 +204,10 @@ where
 
             let mut in_range = Vec::with_capacity(self.starting_l().into_usize());
 
-            for neighbor in scratch.best.iter().take(self.starting_l().into_usize()) {
+            let starting_l = self.starting_l().into_usize();
+            let max_returned = self.max_returned().unwrap_or(usize::MAX);
+
+            for neighbor in scratch.best.iter().take(starting_l) {
                 if neighbor.distance <= self.radius() {
                     in_range.push(neighbor);
                 }
@@ -211,7 +221,8 @@ where
             scratch.in_range = in_range;
 
             let stats = if scratch.in_range.len()
-                >= ((self.starting_l() as f32) * self.initial_slack()) as usize
+                >= ((starting_l as f32) * self.initial_slack()) as usize
+                && scratch.in_range.len() < max_returned
             {
                 // Move to range search
                 let range_stats = range_search_internal(
@@ -336,7 +347,7 @@ where
 
     let max_returned = search_params.max_returned().unwrap_or(usize::MAX);
 
-    while !scratch.range_frontier.is_empty() {
+    while !scratch.range_frontier.is_empty() && scratch.in_range.len() < max_returned {
         scratch.beam_nodes.clear();
 
         // In this loop we are going to find the beam_width number of remaining nodes within the radius
@@ -401,6 +412,9 @@ mod tests {
 
         // Invalid inner radius > radius
         assert!(Range::with_options(None, 100, None, 0.5, Some(1.0), 1.0, 1.0).is_err());
+
+        // Invalid max_results < initial_l_search
+        assert!(Range::with_options(Some(50), 100, None, 0.5, None, 1.0, 1.0).is_err());
     }
 
     #[test]

diskann/src/graph/test/cases/range_search.rs

@@ -284,3 +284,120 @@ fn empty_results() {
         "empty results shouldn't trigger a second round"
     );
 }
+
+#[test]
+fn max_results_respected_means_no_second_round() {
+    let rt = current_thread_runtime();
+    let mut test_root = root();
+    let mut path = test_root.path();
+    let name = path.push("max_results_respected_means_no_second_round");
+
+    let grid_size = 5;
+    let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three);
+    let radius = 1.0e9; // every point will be in range with this radius
+    let starting_l = 4; // small set to trigger multiple rounds
+    let max_results = 4; // max_returned = starting_l, so second round should not be triggered
+
+    let range_search =
+        Range::with_options(Some(max_results), starting_l, None, radius, None, 1.0, 1.0).unwrap();
+    let mut results: Vec<Neighbor<u32>> = Vec::new();
+
+    let stats = rt
+        .block_on(index.search(
+            range_search,
+            &test_provider::Strategy::new(),
+            &test_provider::Context::new(),
+            query.as_slice(),
+            &mut results,
+        ))
+        .unwrap();
+
+    let baseline = RangeSearchBaseline {
+        grid_size,
+        query: query.clone(),
+        radius,
+        inner_radius: None,
+        starting_l,
+        results: results.iter().map(|n| (n.id, n.distance)).collect(),
+        comparisons: stats.cmps as usize,
+        hops: stats.hops as usize,
+        result_count: results.len(),
+        range_search_second_round: stats.range_search_second_round,
+    };
+
+    let expected = get_or_save_test_results(&name, &baseline);
+    assert_eq_verbose!(expected, baseline);
+
+    assert!(
+        results.len() <= max_results,
+        "result count {} exceeds max_results {}",
+        results.len(),
+        max_results
+    );
+
+    assert!(
+        !stats.range_search_second_round,
+        "If max_results is respected, a second round should not be triggered"
+    );
+    assert_range_invariants(&results, radius, None);
+    assert_no_duplicates(&results);
+}
+
+#[test]
+fn max_results_respected_and_second_round_triggered() {
+    let rt = current_thread_runtime();
+    let mut test_root = root();
+    let mut path = test_root.path();
+    let name = path.push("max_results_respected_and_second_round_triggered");
+
+    let grid_size = 5;
+    let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three);
+    let radius = 1.0e9; // every point will be in range with this radius
+    let starting_l = 4; // small set to trigger multiple rounds
+    let max_results = 5; // max_returned greater than starting_l, so second round should be triggered
+
+    let range_search =
+        Range::with_options(Some(max_results), starting_l, None, radius, None, 1.0, 1.0).unwrap();
+    let mut results: Vec<Neighbor<u32>> = Vec::new();
+
+    let stats = rt
+        .block_on(index.search(
+            range_search,
+            &test_provider::Strategy::new(),
+            &test_provider::Context::new(),
+            query.as_slice(),
+            &mut results,
+        ))
+        .unwrap();
+
+    let baseline = RangeSearchBaseline {
+        grid_size,
+        query: query.clone(),
+        radius,
+        inner_radius: None,
+        starting_l,
+        results: results.iter().map(|n| (n.id, n.distance)).collect(),
+        comparisons: stats.cmps as usize,
+        hops: stats.hops as usize,
+        result_count: results.len(),
+        range_search_second_round: stats.range_search_second_round,
+    };
+
+    let expected = get_or_save_test_results(&name, &baseline);
+    assert_eq_verbose!(expected, baseline);
+
+    assert!(
+        results.len() <= max_results,
+        "result count {} exceeds max_results {}",
+        results.len(),
+        max_results
+    );
+
+    assert!(
+        stats.range_search_second_round,
+        "If max_results is respected, a second round should be triggered"
+    );
+
+    assert_range_invariants(&results, radius, None);
+    assert_no_duplicates(&results);
+}

diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json

@@ -0,0 +1,37 @@
+{
+  "file": "diskann/src/graph/test/cases/range_search.rs",
+  "test": "graph/test/cases/range_search/max_results_respected_and_second_round_triggered",
+  "payload": {
+    "comparisons": 11,
+    "grid_size": 5,
+    "hops": 12,
+    "inner_radius": null,
+    "query": [
+      5.0,
+      5.0,
+      5.0
+    ],
+    "radius": 1000000000.0,
+    "range_search_second_round": true,
+    "result_count": 4,
+    "results": [
+      [
+        124,
+        3.0
+      ],
+      [
+        123,
+        6.0
+      ],
+      [
+        119,
+        6.0
+      ],
+      [
+        99,
+        6.0
+      ]
+    ],
+    "starting_l": 4
+  }
+}

diskann/test/generated/graph/test/cases/range_search/max_results_respected_means_no_second_round.json

@@ -0,0 +1,33 @@
+{
+  "file": "diskann/src/graph/test/cases/range_search.rs",
+  "test": "graph/test/cases/range_search/max_results_respected_means_no_second_round",
+  "payload": {
+    "comparisons": 11,
+    "grid_size": 5,
+    "hops": 5,
+    "inner_radius": null,
+    "query": [
+      5.0,
+      5.0,
+      5.0
+    ],
+    "radius": 1000000000.0,
+    "range_search_second_round": false,
+    "result_count": 3,
+    "results": [
+      [
+        124,
+        3.0
+      ],
+      [
+        123,
+        6.0
+      ],
+      [
+        119,
+        6.0
+      ]
+    ],
+    "starting_l": 4
+  }
+}