Merge branch 'main' into u/arrayka/diskann-disk-tests3

Author: Alex Razumov (from Dev Box)

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

@@ -23,7 +23,7 @@ use crate::{
 };
 
 /// A built-in helper for benchmarking the K-nearest neighbors method
-/// [`graph::DiskANNIndex::search`].
+/// [`graph::DiskANNIndex::search`] with optional post-processing support.
 ///
 /// This is intended to be used in conjunction with [`search::search`] or
 /// [`search::search_all`] and provides some basic additional metrics for
@@ -32,21 +32,31 @@ use crate::{
 ///
 /// The provided implementation of [`Search`] accepts [`graph::search::Knn`]
 /// and returns [`Metrics`] as additional output.
+///
+/// # Type Parameters
+///
+/// - `DP`: The data provider type
+/// - `T`: The query element type
+/// - `S`: The search strategy type
+/// - `PP`: Post-processor selector. Defaults to [`Defaulted`], which uses the
+///   strategy's default post-processor. Use [`KNN::with_postprocessor`] to
+///   supply an explicit post-processor.
 #[derive(Debug)]
-pub struct KNN<DP, T, S>
+pub struct KNN<DP, T, S, PP = Defaulted>
 where
     DP: provider::DataProvider,
 {
     index: Arc<graph::DiskANNIndex<DP>>,
     queries: Arc<Matrix<T>>,
     strategy: Strategy<S>,
+    post_processor: PP,
 }
 
-impl<DP, T, S> KNN<DP, T, S>
+impl<DP, T, S> KNN<DP, T, S, Defaulted>
 where
     DP: provider::DataProvider,
 {
-    /// Construct a new [`KNN`] searcher.
+    /// Construct a new [`KNN`] searcher using the strategy's default post-processor.
     ///
     /// 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
@@ -68,10 +78,98 @@ where
             index,
             queries,
             strategy,
+            post_processor: Defaulted,
+        }))
+    }
+}
+
+impl<DP, T, S, PP> KNN<DP, T, S, Forwarded<PP>>
+where
+    DP: provider::DataProvider,
+{
+    /// Construct a new [`KNN`] searcher with an explicit post-processor.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the number of elements in `strategy` is not compatible with
+    /// the number of rows in `queries`.
+    pub fn with_postprocessor(
+        index: Arc<graph::DiskANNIndex<DP>>,
+        queries: Arc<Matrix<T>>,
+        strategy: Strategy<S>,
+        post_processor: PP,
+    ) -> anyhow::Result<Arc<Self>> {
+        strategy.length_compatible(queries.nrows())?;
+
+        Ok(Arc::new(Self {
+            index,
+            queries,
+            strategy,
+            post_processor: Forwarded(post_processor),
         }))
     }
 }
 
+impl<DP, T, S, PP> KNN<DP, T, S, PP>
+where
+    DP: provider::DataProvider,
+{
+    /// Access the index.
+    pub fn index(&self) -> &Arc<graph::DiskANNIndex<DP>> {
+        &self.index
+    }
+}
+
+/// Resolves a post-processor for [`KNN`] given a search strategy.
+///
+/// This trait lets [`KNN`] support both "use the strategy's default post-processor"
+/// ([`Defaulted`]) and "use this explicit post-processor" ([`Forwarded`]) without
+/// duplicating the search loop.
+pub trait AsPostProcessor<'a, S, DP, T>
+where
+    DP: provider::DataProvider,
+    S: glue::SearchStrategy<'a, DP, T>,
+{
+    /// The concrete post-processor used for a single search.
+    type Processor: glue::SearchPostProcess<S::SearchAccessor, T, DP::ExternalId> + Send + Sync;
+
+    /// Construct the post-processor to use for a single search.
+    fn as_post_processor(&'a self, strategy: &'a S) -> Self::Processor;
+}
+
+/// Marker indicating that [`KNN`] should use the strategy's default post-processor.
+#[derive(Debug, Clone, Copy)]
+pub struct Defaulted;
+
+impl<'a, S, DP, T> AsPostProcessor<'a, S, DP, T> for Defaulted
+where
+    DP: provider::DataProvider,
+    S: glue::DefaultPostProcessor<'a, DP, T, DP::ExternalId>,
+{
+    type Processor = S::Processor;
+
+    fn as_post_processor(&'a self, strategy: &'a S) -> Self::Processor {
+        strategy.default_post_processor()
+    }
+}
+
+/// Wraps an explicit post-processor for use with [`KNN::with_postprocessor`].
+#[derive(Debug, Clone, Copy)]
+pub struct Forwarded<PP>(PP);
+
+impl<'a, S, DP, T, PP> AsPostProcessor<'a, S, DP, T> for Forwarded<PP>
+where
+    DP: provider::DataProvider,
+    S: glue::SearchStrategy<'a, DP, T>,
+    PP: glue::SearchPostProcess<S::SearchAccessor, T, DP::ExternalId> + Clone + AsyncFriendly,
+{
+    type Processor = PP;
+
+    fn as_post_processor(&'a self, _strategy: &'a S) -> Self::Processor {
+        self.0.clone()
+    }
+}
+
 /// Additional metrics collected during [`KNN`] search.
 ///
 /// # Note
@@ -86,10 +184,11 @@ pub struct Metrics {
     pub hops: u32,
 }
 
-impl<DP, T, S> Search for KNN<DP, T, S>
+impl<DP, T, S, PP> Search for KNN<DP, T, S, PP>
 where
     DP: provider::DataProvider<Context: Default, ExternalId: search::Id>,
-    S: for<'a> glue::DefaultSearchStrategy<'a, DP, &'a [T], DP::ExternalId> + Clone + AsyncFriendly,
+    S: for<'a> glue::SearchStrategy<'a, DP, &'a [T]> + Clone + AsyncFriendly,
+    PP: for<'a> AsPostProcessor<'a, S, DP, &'a [T]> + AsyncFriendly,
     graph::search::Knn:
         for<'a> graph::Search<'a, DP, S, &'a [T], Output = graph::index::SearchStats>,
     T: AsyncFriendly + Clone,
@@ -117,11 +216,15 @@ where
     {
         let context = DP::Context::default();
         let knn_search = *parameters;
+        let strategy = self.strategy.get(index)?;
+        let processor = self.post_processor.as_post_processor(strategy);
+
         let stats = self
             .index
-            .search(
+            .search_with(
                 knn_search,
-                self.strategy.get(index)?,
+                strategy,
+                processor,
                 &context,
                 self.queries.row(index),
                 buffer,

diskann-benchmark/example/async-determinant-diversity.json

@@ -0,0 +1,48 @@
+{
+  "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-determinant-diversity",
+          "queries": "disk_index_sample_query_10pts.fbin",
+          "groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin",
+          "reps": 5,
+          "num_threads": [
+            1
+          ],
+          "power": 2.0,
+          "eta": 0.01,
+          "runs": [
+            {
+              "search_n": 20,
+              "search_l": [
+                20,
+                30,
+                40
+              ],
+              "recall_k": 10
+            }
+          ]
+        }
+      }
+    }
+  ]
+}

diskann-benchmark/example/disk-index-determinant-diversity.json

@@ -0,0 +1,42 @@
+{
+  "search_directories": [
+    "test_data/disk_index_search"
+  ],
+  "jobs": [
+    {
+      "type": "disk-index",
+      "content": {
+        "source": {
+          "disk-index-source": "Build",
+          "data_type": "float32",
+          "data": "disk_index_siftsmall_learn_256pts_data.fbin",
+          "distance": "squared_l2",
+          "dim": 128,
+          "max_degree": 32,
+          "l_build": 50,
+          "num_threads": 1,
+          "build_ram_limit_gb": 2.0,
+          "num_pq_chunks": 128,
+          "quantization_type": "FP",
+          "save_path": "siftsmall_index_full_det_div"
+        },
+        "search_phase": {
+          "queries": "disk_index_sample_query_10pts.fbin",
+          "groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin",
+          "search_list": [10, 20, 40],
+          "beam_width": 4,
+          "recall_at": 10,
+          "num_threads": 1,
+          "is_flat_search": false,
+          "distance": "squared_l2",
+          "vector_filters_file": null,
+          "post_processor": {
+            "type": "determinant-diversity",
+            "power": 2.0,
+            "eta": 1.0
+          }
+        }
+      }
+    }
+  ]
+}

diskann-benchmark/src/disk_index/search.rs

@@ -14,7 +14,8 @@ use diskann_benchmark_runner::{files::InputFile, utils::MicroSeconds};
 use diskann_disk::{
     data_model::{AdHoc, CachingStrategy},
     search::provider::{
-        disk_provider::DiskIndexSearcher, disk_vertex_provider_factory::DiskVertexProviderFactory,
+        disk_provider::{DiskIndexSearcher, SearchPostProcessorKind},
+        disk_vertex_provider_factory::DiskVertexProviderFactory,
     },
     storage::disk_index_reader::DiskIndexReader,
     utils::{instrumentation::PerfLogger, statistics, AlignedFileReaderFactory, QueryStatistics},
@@ -32,7 +33,10 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     disk_index::json_spancollector::JsonSpanCollector,
-    inputs::disk::{DiskIndexLoad, DiskSearchPhase},
+    inputs::{
+        disk::{DiskIndexLoad, DiskSearchPhase},
+        post_processor::TopkPostProcessor,
+    },
     utils::{datafiles, SimilarityMeasure},
 };
 
@@ -264,6 +268,12 @@ where
         zipped.for_each_in_pool(
             pool.as_ref(),
             |(((((q, vf), id_chunk), dist_chunk), stats), rc)| {
+                let post_processor = search_params.post_processor.as_ref().map_or(
+                    SearchPostProcessorKind::None,
+                    |TopkPostProcessor::DeterminantDiversity(params)| {
+                        SearchPostProcessorKind::DeterminantDiversity(*params)
+                    },
+                );
                 let vector_filter = if search_params.vector_filters_file.is_none() {
                     None
                 } else {
@@ -277,20 +287,21 @@ where
                     l,
                     Some(search_params.beam_width),
                     vector_filter,
+                    post_processor,
                     search_params.is_flat_search,
                 ) {
                     Ok(search_result) => {
                         *stats = search_result.stats.query_statistics;
-                        *rc = search_result.results.len() as u32;
-                        let actual_results = search_result
-                            .results
-                            .len()
-                            .min(search_params.recall_at as usize);
-                        for (i, result_item) in search_result
-                            .results
-                            .iter()
-                            .take(actual_results)
-                            .enumerate()
+                        let base_count = (search_result.stats.result_count as usize)
+                            .min(search_params.recall_at as usize)
+                            .min(search_result.results.len());
+
+                        *rc = base_count as u32;
+                        id_chunk.fill(0);
+                        dist_chunk.fill(0.0);
+
+                        for (i, result_item) in
+                            search_result.results.iter().take(base_count).enumerate()
                         {
                             id_chunk[i] = result_item.vertex_id;
                             dist_chunk[i] = result_item.distance;