two queue filtered search with max effort

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

@@ -6,12 +6,14 @@
 pub mod knn;
 pub mod multihop;
 pub mod range;
+pub mod two_queue;
 
 pub mod strategy;
 
 pub use knn::KNN;
 pub use multihop::MultiHop;
 pub use range::Range;
+pub use two_queue::TwoQueue;
 
 pub use strategy::Strategy;
 

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

@@ -0,0 +1,142 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::sync::Arc;
+
+use diskann::{
+    ANNResult,
+    graph::{self, glue, search::TwoQueueSearch},
+    provider,
+};
+use diskann_utils::{future::AsyncFriendly, views::Matrix};
+
+use crate::search::{self, Search, graph::Strategy};
+
+/// A built-in helper for benchmarking filtered K-nearest neighbors search
+/// using the two-queue search method.
+///
+/// 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 [`search::graph::knn::Aggregator`] type (same
+/// aggregator as [`search::graph::knn::KNN`]).
+///
+/// The provided implementation of [`Search`] accepts [`graph::search::Knn`]
+/// and returns [`search::graph::knn::Metrics`] as additional output.
+#[derive(Debug)]
+pub struct TwoQueue<DP, T, S>
+where
+    DP: provider::DataProvider,
+{
+    index: Arc<graph::DiskANNIndex<DP>>,
+    queries: Arc<Matrix<T>>,
+    strategy: Strategy<S>,
+    labels: Arc<[Arc<dyn graph::index::QueryLabelProvider<DP::InternalId>>]>,
+    max_candidates: usize,
+    result_size_factor: usize,
+}
+
+impl<DP, T, S> TwoQueue<DP, T, S>
+where
+    DP: provider::DataProvider,
+{
+    /// Construct a new [`TwoQueue`] 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.
+    ///
+    /// Additionally, the length of `labels` must match the number of rows in `queries`
+    /// and will be used in querywise correspondence with `queries`.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error under the following conditions.
+    ///
+    /// 1. The number of elements in `strategy` is not compatible with the number of rows in
+    ///    `queries`.
+    ///
+    /// 2. The number of label providers in `labels` is not equal to 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 graph::index::QueryLabelProvider<DP::InternalId>>]>,
+        max_candidates: usize,
+        result_size_factor: usize,
+    ) -> 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,
+                max_candidates,
+                result_size_factor,
+            }))
+        }
+    }
+}
+
+impl<DP, T, S> Search for TwoQueue<DP, T, S>
+where
+    DP: provider::DataProvider<Context: Default, ExternalId: search::Id>,
+    S: for<'a> glue::DefaultSearchStrategy<DP, &'a [T], DP::ExternalId> + Clone + AsyncFriendly,
+    T: AsyncFriendly + Clone,
+{
+    type Id = DP::ExternalId;
+    type Parameters = graph::search::Knn;
+    type Output = super::knn::Metrics;
+
+    fn num_queries(&self) -> usize {
+        self.queries.nrows()
+    }
+
+    fn id_count(&self, parameters: &Self::Parameters) -> search::IdCount {
+        search::IdCount::Fixed(parameters.k_value())
+    }
+
+    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 two_queue_search = TwoQueueSearch::new(
+            *parameters,
+            &*self.labels[index],
+            self.max_candidates,
+            self.result_size_factor,
+        );
+        let result = self
+            .index
+            .search(
+                two_queue_search,
+                self.strategy.get(index)?,
+                &context,
+                self.queries.row(index),
+                buffer,
+            )
+            .await?;
+
+        Ok(super::knn::Metrics {
+            comparisons: result.stats.cmps,
+            hops: result.stats.hops,
+        })
+    }
+}

diskann-benchmark/example/async-two-queue-filter-ground-truth-small.json

@@ -0,0 +1,50 @@
+{
+    "search_directories": [
+      "test_data/disk_index_search"
+    ],
+    "jobs": [
+      {
+        "type": "async-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,
+            "num_start_points": 1,
+            "num_insert_attempts": 1,
+            "saturate_inserts": false,
+            "start_point_strategy": "medoid"
+          },
+          "search_phase": {
+            "search-type": "topk-two-queue-filter",
+            "queries": "disk_index_sample_query_10pts.fbin",
+            "groundtruth": "gt_small_filter.bin",
+            "query_predicates": "query.10.label.jsonl",
+            "data_labels": "data.256.label.jsonl",
+            "max_candidates": [500, 1000, 2000],
+            "result_size_factor": 10,
+            "reps": 5,
+            "num_threads": [
+              1
+            ],
+            "runs": [
+              {
+                "search_n": 20,
+                "search_l": [
+                  100,
+                  200
+                ],
+                "recall_k": 10
+              }
+            ]
+          }
+        }
+      }
+    ]
+  }

diskann-benchmark/src/backend/index/benchmarks.rs

@@ -392,6 +392,50 @@ where
             result.append(AggregatedSearchResults::Topk(search_results));
             Ok(result)
         }
+        SearchPhase::TopkTwoQueueFilter(search_phase) => {
+            // Handle TwoQueue Topk search phase
+            let mut result = BuildResult::new_topk(build_stats);
+
+            // Save construction stats before running queries.
+            checkpoint.checkpoint(&result)?;
+
+            let queries: Arc<Matrix<T>> = Arc::new(datafiles::load_dataset(datafiles::BinFile(
+                &search_phase.queries,
+            ))?);
+
+            let groundtruth =
+                datafiles::load_range_groundtruth(datafiles::BinFile(&search_phase.groundtruth))?;
+
+            let steps = search::knn::SearchSteps::new(
+                search_phase.reps,
+                &search_phase.num_threads,
+                &search_phase.runs,
+            );
+
+            let bit_maps =
+                generate_bitmaps(&search_phase.query_predicates, &search_phase.data_labels)?;
+
+            let labels: Arc<[_]> = bit_maps
+                .into_iter()
+                .map(utils::filters::as_query_label_provider)
+                .collect();
+
+            for &max_candidates in &search_phase.max_candidates {
+                let two_queue = benchmark_core::search::graph::TwoQueue::new(
+                    index.clone(),
+                    queries.clone(),
+                    benchmark_core::search::graph::Strategy::broadcast(search_strategy.clone()),
+                    labels.clone(),
+                    max_candidates,
+                    search_phase.result_size_factor,
+                )?;
+
+                let search_results = search::knn::run(&two_queue, &groundtruth, steps)?;
+                result.append(AggregatedSearchResults::Topk(search_results));
+            }
+
+            Ok(result)
+        }
     }
 }
 

diskann-benchmark/src/backend/index/search/knn.rs

@@ -129,3 +129,29 @@ where
         Ok(results.into_iter().map(SearchResults::new).collect())
     }
 }
+
+impl<DP, T, S> Knn<DP::InternalId> for Arc<core_search::graph::TwoQueue<DP, T, S>>
+where
+    DP: diskann::provider::DataProvider,
+    core_search::graph::TwoQueue<DP, T, S>: core_search::Search<
+        Id = DP::InternalId,
+        Parameters = diskann::graph::search::Knn,
+        Output = core_search::graph::knn::Metrics,
+    >,
+{
+    fn search_all(
+        &self,
+        parameters: Vec<core_search::Run<diskann::graph::search::Knn>>,
+        groundtruth: &dyn benchmark_core::recall::Rows<DP::InternalId>,
+        recall_k: usize,
+        recall_n: usize,
+    ) -> anyhow::Result<Vec<SearchResults>> {
+        let results = core_search::search_all(
+            self.clone(),
+            parameters.into_iter(),
+            core_search::graph::knn::Aggregator::new(groundtruth, recall_k, recall_n),
+        )?;
+
+        Ok(results.into_iter().map(SearchResults::new).collect())
+    }
+}

diskann-benchmark/src/backend/index/spherical.rs

@@ -460,6 +460,60 @@ mod imp {
                             writeln!(output, "\n\n{}", result)?;
                             Ok(result)
                         }
+                        SearchPhase::TopkTwoQueueFilter(search_phase) => {
+                            // Handle Two-Queue Filtered Topk search phase
+
+                            // Save construction stats before running queries.
+                            _checkpoint.checkpoint(&result)?;
+
+                            let queries: Arc<Matrix<f32>> = Arc::new(datafiles::load_dataset(
+                                datafiles::BinFile(&search_phase.queries),
+                            )?);
+
+                            let groundtruth = datafiles::load_groundtruth(datafiles::BinFile(
+                                &search_phase.groundtruth,
+                            ))?;
+
+                            let steps = search::knn::SearchSteps::new(
+                                search_phase.reps,
+                                &search_phase.num_threads,
+                                &search_phase.runs,
+                            );
+
+                            let bit_maps = generate_bitmaps(
+                                &search_phase.query_predicates,
+                                &search_phase.data_labels,
+                            )?;
+
+                            let labels: Arc<[_]> = bit_maps
+                                .into_iter()
+                                .map(utils::filters::as_query_label_provider)
+                                .collect();
+
+                            for &layout in self.input.query_layouts.iter() {
+                                for &max_candidates in &search_phase.max_candidates {
+                                    let two_queue = benchmark_core::search::graph::TwoQueue::new(
+                                        index.clone(),
+                                        queries.clone(),
+                                        benchmark_core::search::graph::Strategy::broadcast(
+                                            inmem::spherical::Quantized::search(layout.into()),
+                                        ),
+                                        labels.clone(),
+                                        max_candidates,
+                                        search_phase.result_size_factor,
+                                    )?;
+
+                                    let search_results =
+                                        search::knn::run(&two_queue, &groundtruth, steps)?;
+                                    result.append(SearchRun {
+                                        layout,
+                                        results: AggregatedSearchResults::Topk(search_results),
+                                    });
+                                }
+                            }
+                            writeln!(output, "\n\n{}", result)?;
+                            Ok(result)
+                        }
                     }
                 }
             }