(diversity) Integrate determinant diversity in tokp full precision search and disk search

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

@@ -22,7 +22,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
@@ -31,21 +31,29 @@ 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`: Optional post-processor type (defaults to `()` for no post-processing)
 #[derive(Debug)]
-pub struct KNN<DP, T, S>
+pub struct KNN<DP, T, S, PP = ()>
 where
     DP: provider::DataProvider,
 {
     index: Arc<graph::DiskANNIndex<DP>>,
     queries: Arc<Matrix<T>>,
     strategy: Strategy<S>,
+    post_processor: Option<PP>,
 }
 
-impl<DP, T, S> KNN<DP, T, S>
+impl<DP, T, S> KNN<DP, T, S, ()>
 where
     DP: provider::DataProvider,
 {
-    /// Construct a new [`KNN`] searcher.
+    /// Construct a new [`KNN`] searcher without post-processing.
     ///
     /// 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
@@ -67,10 +75,58 @@ where
             index,
             queries,
             strategy,
+            post_processor: None,
         }))
     }
 }
 
+impl<DP, T, S, PP> KNN<DP, T, S, PP>
+where
+    DP: provider::DataProvider,
+{
+    /// Construct a new [`KNN`] searcher with post-processing.
+    ///
+    /// # 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: Some(post_processor),
+        }))
+    }
+
+    /// Access the index.
+    pub fn index(&self) -> &Arc<graph::DiskANNIndex<DP>> {
+        &self.index
+    }
+
+    /// Access the queries.
+    pub fn queries(&self) -> &Arc<Matrix<T>> {
+        &self.queries
+    }
+
+    /// Access the strategy.
+    pub fn strategy(&self) -> &Strategy<S> {
+        &self.strategy
+    }
+
+    /// Access the post-processor, if present.
+    pub fn post_processor(&self) -> &Option<PP> {
+        &self.post_processor
+    }
+}
+
 /// Additional metrics collected during [`KNN`] search.
 ///
 /// # Note
@@ -85,7 +141,14 @@ pub struct Metrics {
     pub hops: u32,
 }
 
-impl<DP, T, S> Search for KNN<DP, T, S>
+impl Metrics {
+    /// Construct a new metrics value.
+    pub fn new(comparisons: u32, hops: u32) -> Self {
+        Self { comparisons, hops }
+    }
+}
+
+impl<DP, T, S> Search for KNN<DP, T, S, ()>
 where
     DP: provider::DataProvider<Context: Default, ExternalId: search::Id>,
     S: for<'a> glue::DefaultSearchStrategy<DP, &'a [T], DP::ExternalId> + Clone + AsyncFriendly,

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

@@ -0,0 +1,51 @@
+{
+  "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",
+          "queries": "disk_index_sample_query_10pts.fbin",
+          "groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin",
+          "reps": 5,
+          "num_threads": [
+            1
+          ],
+          "post_processor": {
+            "type": "determinant-diversity",
+            "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/backend/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::{
     backend::disk_index::json_spancollector::JsonSpanCollector,
-    inputs::disk::{DiskIndexLoad, DiskSearchPhase},
+    inputs::{
+        disk::{DiskIndexLoad, DiskSearchPhase},
+        post_processor::TopkPostProcessor,
+    },
     utils::{datafiles, SimilarityMeasure},
 };
 
@@ -264,6 +268,14 @@ 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(
+                    |TopkPostProcessor::DeterminantDiversity { power, eta }| {
+                        SearchPostProcessorKind::DeterminantDiversity {
+                            power: *power,
+                            eta: *eta,
+                        }
+                    },
+                );
                 let vector_filter = if search_params.vector_filters_file.is_none() {
                     None
                 } else {
@@ -277,20 +289,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;