Inline filtered search with adaptive L

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

@@ -0,0 +1,309 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::sync::Arc;
+
+use diskann::{
+    ANNResult,
+    graph::{self, glue, search::AdaptiveL},
+    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 inline 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 InlineFilterSearch<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>>]>,
+    adaptive_l: Option<AdaptiveL>,
+}
+
+impl<DP, T, S> InlineFilterSearch<DP, T, S>
+where
+    DP: provider::DataProvider,
+{
+    /// Construct a new [`InlineFilterSearch`] 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>>]>,
+        adaptive_l: Option<AdaptiveL>,
+    ) -> 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,
+                adaptive_l,
+            }))
+        }
+    }
+}
+
+impl<DP, T, S> Search for InlineFilterSearch<DP, T, S>
+where
+    DP: provider::DataProvider<Context: Default, ExternalId: search::Id>,
+    S: for<'a> glue::DefaultSearchStrategy<'a, 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 inline_search = graph::search::InlineFilterSearch::new(
+            *parameters,
+            &*self.labels[index],
+            self.adaptive_l.clone(),
+        );
+        let stats = self
+            .index
+            .search(
+                inline_search,
+                self.strategy.get(index)?,
+                &context,
+                self.queries.row(index),
+                buffer,
+            )
+            .await?;
+
+        Ok(super::knn::Metrics {
+            comparisons: stats.cmps,
+            hops: stats.hops,
+        })
+    }
+}
+
+///////////
+// Tests //
+///////////
+
+#[cfg(test)]
+mod tests {
+    use std::num::NonZeroUsize;
+
+    use super::*;
+
+    use crate::recall::GroundTruthMode;
+    use diskann::graph::{index::QueryLabelProvider, test::provider};
+
+    // A simple [`QueryLabelProvider`] that rejects odd indices.
+    #[derive(Debug)]
+    struct NoOdds;
+
+    impl graph::index::QueryLabelProvider<u32> for NoOdds {
+        fn is_match(&self, id: u32) -> bool {
+            id.is_multiple_of(2)
+        }
+    }
+
+    #[test]
+    fn test_inline() {
+        let nearest_neighbors = 5;
+
+        let index = search::graph::test_grid_provider();
+
+        let mut queries = Matrix::new(0.0f32, 5, index.provider().dim());
+        queries.row_mut(0).copy_from_slice(&[0.0, 0.0, 0.0, 0.0]);
+        queries.row_mut(1).copy_from_slice(&[4.0, 0.0, 0.0, 0.0]);
+        queries.row_mut(2).copy_from_slice(&[0.0, 4.0, 0.0, 0.0]);
+        queries.row_mut(3).copy_from_slice(&[0.0, 0.0, 4.0, 0.0]);
+        queries.row_mut(4).copy_from_slice(&[0.0, 0.0, 0.0, 4.0]);
+
+        let queries = Arc::new(queries);
+
+        let adaptive_l = graph::search::AdaptiveL::new(10, 16.0).unwrap();
+
+        let inline = InlineFilterSearch::new(
+            index,
+            queries.clone(),
+            Strategy::broadcast(provider::Strategy::new()),
+            (0..queries.nrows())
+                .map(|_| -> Arc<dyn QueryLabelProvider<_>> { Arc::new(NoOdds {}) })
+                .collect(),
+            Some(adaptive_l),
+        )
+        .unwrap();
+
+        // Test the standard search interface.
+        let rt = crate::tokio::runtime(2).unwrap();
+        let results = search::search(
+            inline.clone(),
+            graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
+            NonZeroUsize::new(2).unwrap(),
+            &rt,
+        )
+        .unwrap();
+
+        assert_eq!(results.len(), queries.nrows());
+        let rows = results.ids().as_rows();
+        assert_eq!(*rows.row(0).first().unwrap(), 0);
+
+        // Check that only even IDs are returned.
+        for r in 0..rows.nrows() {
+            assert_eq!(rows.row(r).len(), nearest_neighbors);
+            for &id in rows.row(r) {
+                assert_eq!(id % 2, 0, "Found odd ID {} in row {}", id, r);
+            }
+        }
+
+        const TWO: NonZeroUsize = NonZeroUsize::new(2).unwrap();
+        let setup = search::Setup {
+            threads: TWO,
+            tasks: TWO,
+            reps: TWO,
+        };
+
+        // Try the aggregated strategy.
+        let parameters = [
+            search::Run::new(
+                graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
+                setup.clone(),
+            ),
+            search::Run::new(
+                graph::search::Knn::new(nearest_neighbors, 15, None).unwrap(),
+                setup.clone(),
+            ),
+        ];
+
+        let recall_k = nearest_neighbors;
+        let recall_n = nearest_neighbors;
+
+        let all = search::search_all(
+            inline,
+            parameters,
+            search::graph::knn::Aggregator::new(
+                rows,
+                recall_k,
+                recall_n,
+                GroundTruthMode::Flexible,
+            ),
+        )
+        .unwrap();
+
+        assert_eq!(all.len(), 2);
+        for summary in all {
+            assert_eq!(summary.setup, setup);
+            assert_eq!(summary.end_to_end_latencies.len(), TWO.get());
+            assert_eq!(summary.mean_latencies.len(), TWO.get());
+            assert_eq!(summary.p90_latencies.len(), TWO.get());
+            assert_eq!(summary.p99_latencies.len(), TWO.get());
+
+            assert_ne!(summary.mean_cmps, 0.0);
+            assert_ne!(summary.mean_hops, 0.0);
+
+            let recall = summary.recall;
+            assert_eq!(recall.recall_k, recall_k);
+            assert_eq!(recall.recall_n, recall_n);
+            assert_eq!(recall.num_queries, queries.nrows());
+            assert_eq!(recall.average, 1.0, "we used a search as the groundtruth");
+        }
+    }
+
+    #[test]
+    fn test_inline_error() {
+        let index = search::graph::test_grid_provider();
+        let queries = Arc::new(Matrix::new(0.0f32, 2, index.provider().dim()));
+
+        let labels: Arc<[_]> = (0..queries.nrows() + 1)
+            .map(|_| -> Arc<dyn QueryLabelProvider<_>> { Arc::new(NoOdds {}) })
+            .collect();
+
+        let strategy = provider::Strategy::new();
+
+        // Error for a mismatch between strategies and queries.
+        let err = InlineFilterSearch::new(
+            index.clone(),
+            queries.clone(),
+            Strategy::collection([strategy.clone()]),
+            labels.clone(),
+            None,
+        )
+        .unwrap_err();
+        let msg = err.to_string();
+        assert!(
+            msg.contains("1 strategy was provided when 2 were expected"),
+            "failed with {msg}"
+        );
+
+        // Error for a mismatch between label providers and queries.
+        let err = InlineFilterSearch::new(
+            index,
+            queries.clone(),
+            Strategy::broadcast(strategy.clone()),
+            labels.clone(),
+            None,
+        )
+        .unwrap_err();
+        let msg = err.to_string();
+        assert!(
+            msg.contains(
+                "Number of label providers (3) must be equal to the number of queries (2)"
+            ),
+            "failed with {msg}"
+        );
+    }
+}

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

@@ -3,12 +3,14 @@
  * Licensed under the MIT license.
  */
 
+pub mod inline;
 pub mod knn;
 pub mod multihop;
 pub mod range;
 
 pub mod strategy;
 
+pub use inline::InlineFilterSearch;
 pub use knn::KNN;
 pub use multihop::MultiHop;
 pub use range::Range;

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

@@ -0,0 +1,52 @@
+{
+  "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": "topk-inline-filter",
+          "queries": "disk_index_sample_query_10pts.fbin",
+          "groundtruth": "disk_index_10pts_idx_uint32_truth_search_filter_res.bin",
+          "query_predicates": "query.10.label.jsonl",
+          "data_labels": "data.256.label.jsonl",
+          "reps": 5,
+          "num_threads": [
+            1
+          ],
+          "runs": [
+            {
+              "search_n": 20,
+              "search_l": [
+                20,
+                30,
+                40
+              ],
+              "recall_k": 10
+            }
+          ],
+          "adaptive_l": {
+            "sample_count": 10,
+            "scale_factor": 16.0
+          }
+        }
+      }
+    }
+  ]
+}