Clean up range search

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
+  }
+}