Implement determinant-diversity search in the new post-processing pipeline

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

@@ -41,6 +41,18 @@ where
     strategy: Strategy<S>,
 }
 
+/// A [`KNN`] variant that uses explicit post-processing during search.
+#[derive(Debug)]
+pub struct KNNWithPostProcessor<DP, T, S, P>
+where
+    DP: provider::DataProvider,
+{
+    index: Arc<graph::DiskANNIndex<DP>>,
+    queries: Arc<Matrix<T>>,
+    strategy: Strategy<S>,
+    post_processor: Strategy<P>,
+}
+
 impl<DP, T, S> KNN<DP, T, S>
 where
     DP: provider::DataProvider,
@@ -71,6 +83,39 @@ where
     }
 }
 
+impl<DP, T, S, P> KNNWithPostProcessor<DP, T, S, P>
+where
+    DP: provider::DataProvider,
+{
+    /// Construct a new [`KNNWithPostProcessor`] searcher.
+    ///
+    /// If `strategy` or `post_processor` 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/processors will have a querywise correspondence (see [`search::SearchResults`])
+    /// with the query matrix.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the number of elements in `strategy` or `post_processor` 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>,
+        post_processor: Strategy<P>,
+    ) -> anyhow::Result<Arc<Self>> {
+        strategy.length_compatible(queries.nrows())?;
+        post_processor.length_compatible(queries.nrows())?;
+
+        Ok(Arc::new(Self {
+            index,
+            queries,
+            strategy,
+            post_processor,
+        }))
+    }
+}
+
 /// Additional metrics collected during [`KNN`] search.
 ///
 /// # Note
@@ -132,6 +177,59 @@ where
     }
 }
 
+impl<DP, T, S, P> Search for KNNWithPostProcessor<DP, T, S, P>
+where
+    DP: provider::DataProvider<Context: Default, ExternalId: search::Id>,
+    S: glue::DefaultSearchStrategy<DP, [T], DP::ExternalId> + Clone + AsyncFriendly,
+    T: AsyncFriendly + Clone,
+    P: for<'a> glue::SearchPostProcess<S::SearchAccessor<'a>, [T], DP::ExternalId>
+        + Clone
+        + Send
+        + Sync
+        + AsyncFriendly,
+{
+    type Id = DP::ExternalId;
+    type Parameters = graph::search::Knn;
+    type Output = 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 knn_search = *parameters;
+        let stats = self
+            .index
+            .search_with(
+                knn_search,
+                self.strategy.get(index)?,
+                self.post_processor.get(index)?.clone(),
+                &context,
+                self.queries.row(index),
+                buffer,
+            )
+            .await?;
+
+        Ok(Metrics {
+            comparisons: stats.cmps,
+            hops: stats.hops,
+        })
+    }
+}
+
 /// An [`search::Aggregate`]d summary of multiple [`KNN`] search runs
 /// returned by the provided [`Aggregator`].
 ///

diskann-benchmark-runner/src/any.rs

@@ -397,7 +397,7 @@ mod tests {
         let _: Type<f32> = value.convert::<DataType, _>().unwrap();
 
         // An invalid match should return an error.
-        let value = Any::new(0usize, "random-rag");
+        let value = Any::new(0usize, "random-determinant-diversity");
         let err = value.convert::<DataType, Type<f32>>().unwrap_err();
         let msg = err.to_string();
         assert!(msg.contains("invalid dispatch"), "{}", msg);

diskann-benchmark/example/openai-disk-determinant-diversity-compare.json

@@ -0,0 +1,52 @@
+{
+  "search_directories": [
+    "C:/data/openai"
+  ],
+  "jobs": [
+    {
+      "type": "disk-index",
+      "content": {
+        "source": {
+          "disk-index-source": "Load",
+          "data_type": "float32",
+          "load_path": "C:/data/openai/openai_index_normal"
+        },
+        "search_phase": {
+          "queries": "openai_query.bin",
+          "groundtruth": "openai_gt_50.bin",
+          "search_list": [100, 200, 400],
+          "beam_width": 4,
+          "recall_at": 10,
+          "num_threads": 8,
+          "is_flat_search": false,
+          "distance": "squared_l2",
+          "vector_filters_file": null
+        }
+      }
+    },
+    {
+      "type": "disk-index",
+      "content": {
+        "source": {
+          "disk-index-source": "Load",
+          "data_type": "float32",
+          "load_path": "C:/data/openai/openai_index_normal"
+        },
+        "search_phase": {
+          "queries": "openai_query.bin",
+          "groundtruth": "openai_gt_50.bin",
+          "search_list": [100, 200, 400],
+          "beam_width": 4,
+          "recall_at": 10,
+          "num_threads": 8,
+          "is_flat_search": false,
+          "distance": "squared_l2",
+          "vector_filters_file": null,
+          "is_determinant_diversity_search": true,
+          "determinant_diversity_eta": 0.01,
+          "determinant_diversity_power": 2.0
+        }
+      }
+    }
+  ]
+}

diskann-benchmark/src/backend/disk_index/search.rs

@@ -42,6 +42,9 @@ pub(super) struct DiskSearchStats {
     pub(super) beam_width: usize,
     pub(super) recall_at: u32,
     pub(crate) is_flat_search: bool,
+    pub(crate) is_determinant_diversity_search: bool,
+    pub(crate) determinant_diversity_eta: Option<f64>,
+    pub(crate) determinant_diversity_power: Option<f64>,
     pub(crate) distance: SimilarityMeasure,
     pub(crate) uses_vector_filters: bool,
     pub(super) num_nodes_to_cache: Option<usize>,
@@ -276,6 +279,9 @@ where
                 Some(search_params.beam_width),
                 vector_filter,
                 search_params.is_flat_search,
+                search_params.is_determinant_diversity_search,
+                search_params.determinant_diversity_eta,
+                search_params.determinant_diversity_power,
             ) {
                 Ok(search_result) => {
                     *stats = search_result.stats.query_statistics;
@@ -341,6 +347,9 @@ where
         beam_width: search_params.beam_width,
         recall_at: search_params.recall_at,
         is_flat_search: search_params.is_flat_search,
+        is_determinant_diversity_search: search_params.is_determinant_diversity_search,
+        determinant_diversity_eta: search_params.determinant_diversity_eta,
+        determinant_diversity_power: search_params.determinant_diversity_power,
         distance: search_params.distance,
         uses_vector_filters: search_params.vector_filters_file.is_some(),
         num_nodes_to_cache: search_params.num_nodes_to_cache,
@@ -425,6 +434,22 @@ impl fmt::Display for DiskSearchStats {
         writeln!(f, "Beam width,       : {}", self.beam_width)?;
         writeln!(f, "Recall at,        : {}", self.recall_at)?;
         writeln!(f, "Flat search,      : {}", self.is_flat_search)?;
+        writeln!(
+            f,
+            "Det-div search,   : {}",
+            self.is_determinant_diversity_search
+        )?;
+        writeln!(
+            f,
+            "Det-div params,   : {}",
+            match (
+                self.determinant_diversity_eta,
+                self.determinant_diversity_power,
+            ) {
+                (Some(eta), Some(power)) => format!("eta={eta}, power={power}"),
+                _ => "None".to_string(),
+            }
+        )?;
         writeln!(f, "Distance,         : {}", self.distance)?;
         writeln!(f, "Vector filters,   : {}", self.uses_vector_filters)?;
         writeln!(

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

@@ -24,7 +24,10 @@ use diskann_benchmark_runner::{
 };
 use diskann_providers::{
     index::diskann_async,
-    model::{configuration::IndexConfiguration, graph::provider::async_::common},
+    model::{
+        configuration::IndexConfiguration,
+        graph::provider::async_::{common, DeterminantDiversitySearchParams},
+    },
 };
 use diskann_utils::{
     future::AsyncFriendly,
@@ -350,6 +353,8 @@ where
         + provider::SetElement<[T]>,
     T: SampleableForStart + std::fmt::Debug + Copy + AsyncFriendly + bytemuck::Pod,
     S: glue::DefaultSearchStrategy<DP, [T]> + Clone + AsyncFriendly,
+    DeterminantDiversitySearchParams:
+        for<'a> glue::SearchPostProcess<S::SearchAccessor<'a>, [T], DP::ExternalId> + Send + Sync,
 {
     match &input {
         SearchPhase::Topk(search_phase) => {
@@ -366,19 +371,51 @@ where
             let groundtruth =
                 datafiles::load_groundtruth(datafiles::BinFile(&search_phase.groundtruth))?;
 
-            let knn = benchmark_core::search::graph::KNN::new(
-                index,
-                queries,
-                benchmark_core::search::graph::Strategy::broadcast(search_strategy),
-            )?;
-
             let steps = search::knn::SearchSteps::new(
                 search_phase.reps,
                 &search_phase.num_threads,
                 &search_phase.runs,
             );
 
-            let search_results = search::knn::run(&knn, &groundtruth, steps)?;
+            let search_results = if let (Some(eta), Some(power)) = (
+                search_phase.determinant_diversity_eta,
+                search_phase.determinant_diversity_power,
+            ) {
+                let processor = DeterminantDiversitySearchParams::new(
+                    search_phase
+                        .determinant_diversity_results_k
+                        .unwrap_or_else(|| {
+                            search_phase
+                                .runs
+                                .iter()
+                                .map(|run| run.search_n)
+                                .max()
+                                .unwrap_or(1)
+                        }),
+                    eta,
+                    power,
+                )
+                .map_err(|err| {
+                    anyhow::anyhow!("Invalid determinant-diversity parameters: {err}")
+                })?;
+
+                let knn = benchmark_core::search::graph::knn::KNNWithPostProcessor::new(
+                    index,
+                    queries,
+                    benchmark_core::search::graph::Strategy::broadcast(search_strategy),
+                    benchmark_core::search::graph::Strategy::broadcast(processor),
+                )?;
+
+                search::knn::run(&knn, &groundtruth, steps)?
+            } else {
+                let knn = benchmark_core::search::graph::KNN::new(
+                    index,
+                    queries,
+                    benchmark_core::search::graph::Strategy::broadcast(search_strategy),
+                )?;
+
+                search::knn::run(&knn, &groundtruth, steps)?
+            };
             result.append(AggregatedSearchResults::Topk(search_results));
             Ok(result)
         }