Reform Range Search

diskann/src/graph/search/diverse_search.rs

@@ -169,8 +169,6 @@ where
             ),
             id_scratch: Vec::with_capacity(index.max_degree_with_slack()),
             beam_nodes: Vec::with_capacity(self.inner.beam_width().get()),
-            range_frontier: std::collections::VecDeque::new(),
-            in_range: Vec::new(),
             hops: 0,
             cmps: 0,
         }

diskann/src/graph/search/filtered_range_search.rs

@@ -17,8 +17,10 @@ use crate::{
         index::{DiskANNIndex, InternalSearchStats, SearchStats},
         search::inline_filter_search::{Ret, inline_filter_search_internal},
         search::{
-            Range, RangeSearchError, Search, range_search::DistanceFiltered,
-            range_search::RangeBuilder, record::NoopSearchRecord, scratch::SearchScratch,
+            Range, RangeSearchError, Search,
+            range_search::{InRange, LimitedOutputBuffer, RangeBuilder},
+            record::NoopSearchRecord,
+            scratch::SearchScratch,
         },
         search_output_buffer::SearchOutputBuffer,
     },
@@ -55,6 +57,12 @@ impl FilteredRange {
         self.range_params.max_returned()
     }
 
+    /// Returns the maximum number of results to return.
+    #[inline]
+    pub fn effective_max_returned(&self, inc: usize) -> usize {
+        self.range_params.effective_max_returned(inc)
+    }
+
     /// Returns the initial search list size.
     #[inline]
     pub fn starting_l(&self) -> NonZeroUsize {
@@ -135,12 +143,12 @@ where
                 .search_accessor(&index.data_provider, context, query)
                 .into_ann_result()?;
             let num_start_ids = accessor.num_starting_points().await?;
-            let mut scratch = index.search_scratch(self.starting_l().get(), num_start_ids);
-
-            // Perform an initial inline filtered search, store both filtered and unfiltered results
+            let starting_l = self.starting_l().get();
+            let mut scratch = index.search_scratch(starting_l, num_start_ids);
 
             let search_knn = self.range_params.to_knn();
 
+            // Perform an initial inline filtered search, store both filtered and unfiltered results
             let Ret {
                 cmps,
                 hops,
@@ -155,51 +163,58 @@ where
             )
             .await?;
 
-            let max_returned = self.max_returned().unwrap_or(usize::MAX);
+            let max_returned = self.effective_max_returned(num_start_ids);
 
-            // merge matched_results with the best results from the first round, filtering by radius
-
-            let mut in_range: Vec<_> = scratch
-                .best
-                .iter()
-                .take(self.starting_l().get())
-                .chain(matched_results.iter().copied())
-                .filter(|neighbor| *neighbor.distance() <= self.radius())
-                .collect();
+            // Filter matched results by radius.
+            //
+            // This will be used to decide if `max_results` has been reached.
+            let mut matched_in_outer_range = InRange::new(
+                self.radius(),
+                self.inner_radius(),
+                max_returned,
+                matched_results,
+            );
 
-            in_range.sort_unstable_by(crate::neighbor::ord::fast_distance_total);
-            in_range.dedup_by(|left, right| left.id() == right.id());
+            // Merge `matched_results` with the best results from the first round,
+            // filtering by radius
+            let mut in_outer_range = InRange::new(
+                self.radius(),
+                None,
+                usize::MAX,
+                scratch
+                    .best
+                    .iter()
+                    .take(starting_l)
+                    .chain(matched_in_outer_range.iter()),
+            );
 
-            // filter matched results by radius; this will be used to decide if `max_results` has been reached
-            let mut matched_within_radius = Vec::with_capacity(matched_results.len());
-            for neighbor in matched_results.iter().copied() {
-                if *neighbor.distance() <= self.radius() {
-                    matched_within_radius.push(neighbor);
-                }
-            }
+            in_outer_range.sort_and_dedup();
 
-            let stats = if in_range.len()
-                >= ((self.starting_l().get() as f32) * self.initial_slack()) as usize
-                && matched_within_radius.len() < max_returned
+            let stats = if in_outer_range.len()
+                >= ((starting_l as f32) * self.initial_slack()) as usize
+                && matched_in_outer_range.len() < max_returned
             {
                 // clear the visited set and repopulate it with all in-range points found so far, filtered and unfiltered
-                // also add these points to range_frontier for seeding the second-round search
                 scratch.visited.clear();
-                scratch.range_frontier.clear();
                 scratch
                     .visited
-                    .extend(in_range.iter().map(|neighbor| *neighbor.id()));
-                scratch
-                    .range_frontier
-                    .extend(in_range.iter().map(|neighbor| *neighbor.id()));
+                    .extend(in_outer_range.iter().map(|neighbor| *neighbor.id()));
+
+                // Create a range frontier for seeding the second-round search
+                let mut range_frontier: std::collections::VecDeque<_> = in_outer_range
+                    .take()
+                    .iter()
+                    .map(|neighbor| *neighbor.id())
+                    .collect();
 
                 // Move to filtered range search
                 let range_stats = filtered_range_search_internal(
                     index.max_degree_with_slack(),
                     &self,
                     &mut accessor,
                     &mut scratch,
-                    &mut matched_within_radius,
+                    &mut matched_in_outer_range,
+                    &mut range_frontier,
                 )
                 .await?;
 
@@ -216,24 +231,15 @@ where
                 }
             };
 
-            // Post-process results directly into the output buffer, filtering by radius.
-            let inner_radius = self.inner_radius();
-
-            // Note matched_in_range is assumed to satisfy the radius filter, so we only apply the inner radius filter here
-            let mut filtered = DistanceFiltered::new(output, |dist| {
-                if let Some(ir) = inner_radius
-                    && dist <= ir
-                {
-                    false
-                } else {
-                    true
-                }
-            });
-
-            let truncated_matched = matched_within_radius.iter().copied().take(max_returned);
-
+            let mut limited_output =
+                LimitedOutputBuffer::new(output, self.max_returned().unwrap_or(usize::MAX));
             let result_count = processor
-                .post_process(&mut accessor, query, truncated_matched, &mut filtered)
+                .post_process(
+                    &mut accessor,
+                    query,
+                    matched_in_outer_range.iter(),
+                    &mut limited_output,
+                )
                 .await
                 .into_ann_result()?;
 
@@ -261,7 +267,8 @@ pub(crate) async fn filtered_range_search_internal<A>(
     search_params: &FilteredRange,
     accessor: &mut A,
     scratch: &mut SearchScratch<A::Id>,
-    matched_in_range: &mut Vec<Neighbor<A::Id>>,
+    matched_in_range: &mut InRange<A::Id>,
+    range_frontier: &mut std::collections::VecDeque<A::Id>,
 ) -> ANNResult<InternalSearchStats>
 where
     A: FilteredAccessor,
@@ -270,15 +277,13 @@ where
 
     let mut neighbors = Vec::with_capacity(max_degree_with_slack);
 
-    let max_returned = search_params.max_returned().unwrap_or(usize::MAX);
-
-    while !scratch.range_frontier.is_empty() && matched_in_range.len() < max_returned {
+    while !range_frontier.is_empty() && !matched_in_range.is_full() {
         scratch.beam_nodes.clear();
 
         // In this loop we are going to find the beam_width number of remaining nodes within the radius
         // Each of these nodes will be a frontier node.
-        while !scratch.range_frontier.is_empty() && scratch.beam_nodes.len() < beam_width {
-            let next = scratch.range_frontier.pop_front();
+        while !range_frontier.is_empty() && scratch.beam_nodes.len() < beam_width {
+            let next = range_frontier.pop_front();
             if let Some(next_node) = next {
                 scratch.beam_nodes.push(next_node);
             }
@@ -298,18 +303,19 @@ where
         // but only accepted IDs are added to in-range results.
         let navigation_radius = search_params.radius() * search_params.range_slack();
         for (decision, distance) in neighbors.iter().copied() {
-            if distance <= navigation_radius {
-                let id = decision.into_inner();
-                scratch.range_frontier.push_back(id);
-
-                if distance <= search_params.radius()
-                    && decision.is_accept()
-                    && matched_in_range.len() < max_returned
-                {
-                    matched_in_range.push(Neighbor::new(id, distance));
-                }
+            if matched_in_range.is_full() {
+                break;
+            }
+
+            if let glue::Decision::Accept(id) = decision
+                && matched_in_range.push(Neighbor::new(id.into_inner(), distance))
+            {
+                range_frontier.push_back(id.into_inner());
+            } else if distance <= navigation_radius {
+                range_frontier.push_back(decision.into_inner());
             }
         }
+
         scratch.cmps += neighbors.len() as u32;
         scratch.hops += scratch.beam_nodes.len() as u32;
     }