Add algorithm and benchmark for filtered range search

diskann-benchmark-core/src/search/graph/filtered_range.rs

@@ -0,0 +1,216 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::{num::NonZeroUsize, sync::Arc};
+
+use diskann::{
+    ANNResult,
+    graph::{self, ext::labeled, glue},
+    provider,
+};
+use diskann_benchmark_runner::utils::{MicroSeconds, percentiles};
+use diskann_utils::{future::AsyncFriendly, views::Matrix};
+
+use crate::{
+    recall,
+    search::{self, Search, graph::Strategy},
+};
+
+/// A built-in helper for benchmarking the filtered range search method
+/// [`graph::DiskANNIndex::search`] with [`graph::search::FilteredRange`].
+///
+/// This is intended to be used in conjunction with [`search::search`] or
+/// [`search::search_all`] and provides some basic additional metrics for the
+/// latter. Result aggregation for [`search::search_all`] is provided by the
+/// [`Aggregator`] type.
+///
+/// The provided implementation of [`Search`] accepts
+/// [`graph::search::FilteredRange`] and returns [`Metrics`] as additional output.
+#[derive(Debug)]
+pub struct FilteredRange<DP, T, S>
+where
+    DP: provider::DataProvider,
+{
+    index: Arc<graph::DiskANNIndex<DP>>,
+    queries: Arc<Matrix<T>>,
+    strategy: Strategy<S>,
+    labels: Arc<[Arc<dyn labeled::QueryLabelProvider<DP::InternalId>>]>,
+}
+
+impl<DP, T, S> FilteredRange<DP, T, S>
+where
+    DP: provider::DataProvider,
+{
+    /// Construct a new [`FilteredRange`] searcher.
+    ///
+    /// If `strategy` is one of the container variants of [`Strategy`], its length
+    /// must match the number of rows in `queries`. If this is the case, then the
+    /// strategies will have a querywise correspondence (see [`search::SearchResults`])
+    /// with the query matrix.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the number of elements in `strategy` is not compatible with
+    /// the number of rows in `queries`.
+    pub fn new(
+        index: Arc<graph::DiskANNIndex<DP>>,
+        queries: Arc<Matrix<T>>,
+        strategy: Strategy<S>,
+        labels: Arc<[Arc<dyn labeled::QueryLabelProvider<DP::InternalId>>]>,
+    ) -> anyhow::Result<Arc<Self>> {
+        strategy.length_compatible(queries.nrows())?;
+
+        if labels.len() != queries.nrows() {
+            Err(anyhow::anyhow!(
+                "Number of label providers ({}) must be equal to the number of queries ({})",
+                labels.len(),
+                queries.nrows()
+            ))
+        } else {
+            Ok(Arc::new(Self {
+                index,
+                queries,
+                strategy,
+                labels,
+            }))
+        }
+    }
+}
+
+/// Placeholder for filtered range search metrics.
+#[derive(Debug, Clone, Copy)]
+#[non_exhaustive]
+pub struct Metrics {}
+
+impl<DP, T, S> Search for FilteredRange<DP, T, S>
+where
+    DP: provider::DataProvider<Context: Default, ExternalId: search::Id>,
+    S: for<'a> glue::DefaultSearchStrategy<
+            'a,
+            DP,
+            &'a [T],
+            DP::ExternalId,
+            SearchAccessor: glue::SearchAccessor,
+        > + Clone
+        + AsyncFriendly,
+    T: AsyncFriendly + Clone,
+{
+    type Id = DP::ExternalId;
+    type Parameters = graph::search::FilteredRange;
+    type Output = Metrics;
+
+    fn num_queries(&self) -> usize {
+        self.queries.nrows()
+    }
+
+    fn id_count(&self, parameters: &Self::Parameters) -> search::IdCount {
+        search::IdCount::Dynamic(NonZeroUsize::new(parameters.starting_l()))
+    }
+
+    async fn search<O>(
+        &self,
+        parameters: &Self::Parameters,
+        buffer: &mut O,
+        index: usize,
+    ) -> ANNResult<Self::Output>
+    where
+        O: graph::SearchOutputBuffer<DP::ExternalId> + Send,
+    {
+        let context = DP::Context::default();
+        let filtered_range_search = *parameters;
+        let strategy =
+            labeled::Filtered::new(self.strategy.get(index)?.clone(), &*self.labels[index]);
+        let _ = self
+            .index
+            .search(
+                filtered_range_search,
+                &strategy,
+                &context,
+                self.queries.row(index),
+                buffer,
+            )
+            .await?;
+
+        Ok(Metrics {})
+    }
+}
+
+/// An [`search::Aggregate`]d summary of multiple [`FilteredRange`] search runs returned by
+/// the provided [`Aggregator`].
+#[derive(Debug, Clone)]
+#[non_exhaustive]
+pub struct Summary {
+    pub setup: search::Setup,
+    pub parameters: graph::search::FilteredRange,
+    pub end_to_end_latencies: Vec<MicroSeconds>,
+    pub mean_latencies: Vec<f64>,
+    pub p90_latencies: Vec<MicroSeconds>,
+    pub p99_latencies: Vec<MicroSeconds>,
+    pub average_precision: recall::AveragePrecisionMetrics,
+}
+
+/// A [`search::Aggregate`] for collecting the results of multiple [`FilteredRange`] search
+/// runs.
+pub struct Aggregator<'a, I> {
+    groundtruth: &'a dyn crate::recall::Rows<I>,
+}
+
+impl<'a, I> Aggregator<'a, I> {
+    pub fn new(groundtruth: &'a dyn crate::recall::Rows<I>) -> Self {
+        Self { groundtruth }
+    }
+}
+
+impl<I> search::Aggregate<graph::search::FilteredRange, I, Metrics> for Aggregator<'_, I>
+where
+    I: crate::recall::RecallCompatible,
+{
+    type Output = Summary;
+
+    #[inline(never)]
+    fn aggregate(
+        &mut self,
+        run: search::Run<graph::search::FilteredRange>,
+        mut results: Vec<search::SearchResults<I, Metrics>>,
+    ) -> anyhow::Result<Summary> {
+        let average_precision = match results.first() {
+            Some(first) => {
+                crate::recall::average_precision(first.ids().as_rows(), self.groundtruth)?
+            }
+            None => anyhow::bail!("Results must be non-empty"),
+        };
+
+        let mut mean_latencies = Vec::with_capacity(results.len());
+        let mut p90_latencies = Vec::with_capacity(results.len());
+        let mut p99_latencies = Vec::with_capacity(results.len());
+
+        results.iter_mut().for_each(|r| {
+            match percentiles::compute_percentiles(r.latencies_mut()) {
+                Ok(values) => {
+                    let percentiles::Percentiles { mean, p90, p99, .. } = values;
+                    mean_latencies.push(mean);
+                    p90_latencies.push(p90);
+                    p99_latencies.push(p99);
+                }
+                Err(_) => {
+                    let zero = MicroSeconds::new(0);
+                    mean_latencies.push(0.0);
+                    p90_latencies.push(zero);
+                    p99_latencies.push(zero);
+                }
+            }
+        });
+
+        Ok(Summary {
+            setup: run.setup().clone(),
+            parameters: *run.parameters(),
+            end_to_end_latencies: results.iter().map(|r| r.end_to_end_latency()).collect(),
+            mean_latencies,
+            p90_latencies,
+            p99_latencies,
+            average_precision,
+        })
+    }
+}

diskann-benchmark-core/src/search/graph/mod.rs

@@ -3,6 +3,7 @@
  * Licensed under the MIT license.
  */
 
+pub mod filtered_range;
 pub mod inline;
 pub mod knn;
 pub mod multihop;

diskann-benchmark-core/src/search/mod.rs

@@ -123,6 +123,7 @@
 //!
 //! * [`graph::search::Knn`]: K-nearest neighbors search for [`diskann::graph::DiskANNIndex`].
 //! * [`graph::search::Range`]: Range search for [`diskann::graph::DiskANNIndex`].
+//! * [`graph::search::FilteredRange`]: Filtered range search for [`diskann::graph::DiskANNIndex`].
 //! * [`graph::MultiHop`]: Multi-hop filtered search for [`diskann::graph::DiskANNIndex`].
 
 pub(crate) mod ids;

diskann-benchmark/example/graph-index-filter-range.json

@@ -0,0 +1,54 @@
+{
+  "search_directories": [
+    "test_data/disk_index_search"
+  ],
+  "jobs": [
+    {
+      "type": "graph-index-build",
+      "content": {
+        "source": {
+          "index-source": "Build",
+          "data_type": "float32",
+          "data": "disk_index_siftsmall_learn_256pts_data.fbin",
+          "distance": "squared_l2",
+          "max_degree": 32,
+          "l_build": 50,
+          "alpha": 1.2,
+          "backedge_ratio": 1.0,
+          "num_threads": 1,
+          "start_point_strategy": "medoid",
+          "num_insert_attempts": 1,
+          "saturate_inserts": false
+        },
+        "search_phase": {
+          "search-type": "filtered-range",
+          "queries": "disk_index_sample_query_10pts.fbin",
+          "query_predicates": "query.10.label.jsonl",
+          "data_labels": "data.256.label.jsonl",
+          "groundtruth": "groundtruth_rad_125000_filtered.rangeres",
+          "reps": 5,
+          "num_threads": [
+            1
+          ],
+          "runs": [
+            {
+              "radius": 125000,
+              "inner_radius": null,
+              "max_returned": null,
+              "beam_width": null,
+              "initial_search_l": [
+                1,
+                5,
+                10,
+                20,
+                30
+              ],
+              "initial_search_slack": 1.0,
+              "range_search_slack": 1.0
+            }
+          ]
+        }
+      }
+    }
+  ]
+}

diskann-benchmark/src/index/benchmarks.rs

@@ -75,6 +75,7 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()>
         FullPrecision::<f32>::new()
             .search(plugins::Topk)
             .search(plugins::Range)
+            .search(plugins::FilteredRange)
             .search(plugins::TopkBetaFilter)
             .search(plugins::TopkMultihopFilter)
             .search(plugins::TopkInlineFilter)
@@ -87,7 +88,10 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()>
     )?;
     registry.register(
         "graph-index-full-precision-u8",
-        FullPrecision::<u8>::new().search(plugins::Topk),
+        FullPrecision::<u8>::new()
+            .search(plugins::Topk)
+            .search(plugins::TopkInlineFilter)
+            .search(plugins::FilteredRange),
     )?;
     registry.register(
         "graph-index-full-precision-i8",
@@ -571,7 +575,84 @@ where
             benchmark_core::search::graph::Strategy::broadcast(strategy.inner()),
         )?;
 
-        let result = search::range::run(&range, &groundtruth, steps)?;
+        let result = search::range::run(&range, &groundtruth, steps, Ok)?;
+        Ok(AggregatedSearchResults::Range(result))
+    }
+}
+
+//-------------------//
+// Filtered Range    //
+//-------------------//
+
+impl<DP, S> search::Plugin<DP, SearchPhase, Strategy<S>> for plugins::FilteredRange
+where
+    DP: DataProvider<Context: Default, InternalId = u32, ExternalId = u32> + QueryType,
+    S: for<'a> glue::DefaultSearchStrategy<
+            'a,
+            DP,
+            &'a [DP::Element],
+            SearchAccessor: glue::SearchAccessor,
+        > + Clone
+        + AsyncFriendly,
+{
+    fn is_match(&self, phase: &SearchPhase) -> bool {
+        plugins::FilteredRange::is_match(phase)
+    }
+
+    fn kind(&self) -> &'static str {
+        plugins::FilteredRange::as_str()
+    }
+
+    fn run(
+        &self,
+        index: Arc<DiskANNIndex<DP>>,
+        phase: &SearchPhase,
+        strategy: &Strategy<S>,
+    ) -> anyhow::Result<AggregatedSearchResults> {
+        let filtered_range = phase.as_filtered_range()?;
+
+        let queries: Arc<Matrix<DP::Element>> = Arc::new(datafiles::load_dataset(
+            datafiles::BinFile(&filtered_range.queries),
+        )?);
+
+        let groundtruth =
+            datafiles::load_range_groundtruth(datafiles::BinFile(&filtered_range.groundtruth))?;
+
+        let steps = search::range::RangeSearchSteps::new(
+            filtered_range.reps,
+            &filtered_range.num_threads,
+            &filtered_range.runs,
+        );
+
+        let bit_maps = generate_bitmaps(
+            &filtered_range.query_predicates,
+            &filtered_range.data_labels,
+        )?;
+
+        let labels: Arc<[_]> = bit_maps
+            .into_iter()
+            .map(utils::filters::as_query_label_provider)
+            .collect();
+
+        let filtered_range = benchmark_core::search::graph::filtered_range::FilteredRange::new(
+            index,
+            queries,
+            benchmark_core::search::graph::Strategy::broadcast(strategy.inner()),
+            labels,
+        )?;
+
+        let result = search::range::run(&filtered_range, &groundtruth, steps, |range_search| {
+            diskann::graph::search::FilteredRange::with_options(
+                range_search.max_returned(),
+                range_search.starting_l(),
+                range_search.beam_width(),
+                range_search.radius(),
+                range_search.inner_radius(),
+                range_search.initial_slack(),
+                range_search.range_slack(),
+            )
+            .map_err(Into::into)
+        })?;
         Ok(AggregatedSearchResults::Range(result))
     }
 }