(RFC) Refactor search interface with unified SearchDispatch trait

agents.md

@@ -152,6 +152,20 @@ CI workflow is defined in [`.github/workflows/ci.yml`](.github/workflows/ci.yml)
 
 ---
 
+## Pre-commit Checklist
+
+Before committing changes, always run:
+
+```bash
+# Format all code
+cargo fmt --all
+
+# Run clippy with warnings as errors
+cargo clippy --workspace --all-targets -- -D warnings
+```
+
+---
+
 **End of Agent Onboarding Guide**
 
 *This guide should be updated when major changes occur to the repository structure or development workflows.*

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

@@ -5,7 +5,7 @@
 
 //! A built-in helper for benchmarking K-nearest neighbors.
 
-use std::{num::NonZeroUsize, sync::Arc};
+use std::sync::Arc;
 
 use diskann::{
     ANNResult,
@@ -29,7 +29,7 @@ use crate::{
 /// the latter. Result aggregation for [`search::search_all`] is provided
 /// by the [`Aggregator`] type.
 ///
-/// The provided implementation of [`Search`] accepts [`graph::SearchParams`]
+/// The provided implementation of [`Search`] accepts [`graph::search::Knn`]
 /// and returns [`Metrics`] as additional output.
 #[derive(Debug)]
 pub struct KNN<DP, T, S>
@@ -92,15 +92,15 @@ where
     T: AsyncFriendly + Clone,
 {
     type Id = DP::ExternalId;
-    type Parameters = graph::SearchParams;
+    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(NonZeroUsize::new(parameters.k_value).unwrap_or(diskann::utils::ONE))
+        search::IdCount::Fixed(parameters.k_value())
     }
 
     async fn search<O>(
@@ -113,13 +113,14 @@ where
         O: graph::SearchOutputBuffer<DP::ExternalId> + Send,
     {
         let context = DP::Context::default();
+        let knn_search = *parameters;
         let stats = self
             .index
             .search(
+                knn_search,
                 self.strategy.get(index)?,
                 &context,
                 self.queries.row(index),
-                parameters,
                 buffer,
             )
             .await?;
@@ -142,7 +143,7 @@ pub struct Summary {
     pub setup: search::Setup,
 
     /// The [`Search::Parameters`] used for the batch of runs.
-    pub parameters: graph::SearchParams,
+    pub parameters: graph::search::Knn,
 
     /// The end-to-end latency for each repetition in the batch.
     pub end_to_end_latencies: Vec<MicroSeconds>,
@@ -207,15 +208,15 @@ impl<'a, I> Aggregator<'a, I> {
     }
 }
 
-impl<I> search::Aggregate<graph::SearchParams, I, Metrics> for Aggregator<'_, I>
+impl<I> search::Aggregate<graph::search::Knn, I, Metrics> for Aggregator<'_, I>
 where
     I: crate::recall::RecallCompatible,
 {
     type Output = Summary;
 
     fn aggregate(
         &mut self,
-        run: search::Run<graph::SearchParams>,
+        run: search::Run<graph::search::Knn>,
         mut results: Vec<search::SearchResults<I, Metrics>>,
     ) -> anyhow::Result<Summary> {
         // Compute the recall using just the first result.
@@ -280,6 +281,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use std::num::NonZeroUsize;
+
     use super::*;
 
     use diskann::graph::test::provider;
@@ -310,7 +313,7 @@ mod tests {
         let rt = crate::tokio::runtime(2).unwrap();
         let results = search::search(
             knn.clone(),
-            graph::SearchParams::new(nearest_neighbors, 10, None).unwrap(),
+            graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
             NonZeroUsize::new(2).unwrap(),
             &rt,
         )
@@ -334,11 +337,11 @@ mod tests {
         // Try the aggregated strategy.
         let parameters = [
             search::Run::new(
-                graph::SearchParams::new(nearest_neighbors, 10, None).unwrap(),
+                graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
                 setup.clone(),
             ),
             search::Run::new(
-                graph::SearchParams::new(nearest_neighbors, 15, None).unwrap(),
+                graph::search::Knn::new(nearest_neighbors, 15, None).unwrap(),
                 setup.clone(),
             ),
         ];

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

@@ -3,7 +3,7 @@
  * Licensed under the MIT license.
  */
 
-use std::{num::NonZeroUsize, sync::Arc};
+use std::sync::Arc;
 
 use diskann::{
     ANNResult,
@@ -15,14 +15,14 @@ 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 [multi-hop](graph::DiskANNIndex::multihop_search) search method.
+/// using the multi-hop 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`]).
+/// aggregator as [`search::graph::knn::KNN`]).
 ///
-/// The provided implementation of [`Search`] accepts [`graph::SearchParams`]
+/// The provided implementation of [`Search`] accepts [`graph::search::Knn`]
 /// and returns [`search::graph::knn::Metrics`] as additional output.
 #[derive(Debug)]
 pub struct MultiHop<DP, T, S>
@@ -90,15 +90,15 @@ where
     T: AsyncFriendly + Clone,
 {
     type Id = DP::ExternalId;
-    type Parameters = graph::SearchParams;
+    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(NonZeroUsize::new(parameters.k_value).unwrap_or(diskann::utils::ONE))
+        search::IdCount::Fixed(parameters.k_value())
     }
 
     async fn search<O>(
@@ -111,15 +111,15 @@ where
         O: graph::SearchOutputBuffer<DP::ExternalId> + Send,
     {
         let context = DP::Context::default();
+        let multihop_search = graph::search::MultihopSearch::new(*parameters, &*self.labels[index]);
         let stats = self
             .index
-            .multihop_search(
+            .search(
+                multihop_search,
                 self.strategy.get(index)?,
                 &context,
                 self.queries.row(index),
-                parameters,
                 buffer,
-                &*self.labels[index],
             )
             .await?;
 
@@ -136,6 +136,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use std::num::NonZeroUsize;
+
     use super::*;
 
     use diskann::graph::{index::QueryLabelProvider, test::provider};
@@ -179,7 +181,7 @@ mod tests {
         let rt = crate::tokio::runtime(2).unwrap();
         let results = search::search(
             multihop.clone(),
-            graph::SearchParams::new(nearest_neighbors, 10, None).unwrap(),
+            graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
             NonZeroUsize::new(2).unwrap(),
             &rt,
         )
@@ -207,11 +209,11 @@ mod tests {
         // Try the aggregated strategy.
         let parameters = [
             search::Run::new(
-                graph::SearchParams::new(nearest_neighbors, 10, None).unwrap(),
+                graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
                 setup.clone(),
             ),
             search::Run::new(
-                graph::SearchParams::new(nearest_neighbors, 15, None).unwrap(),
+                graph::search::Knn::new(nearest_neighbors, 15, None).unwrap(),
                 setup.clone(),
             ),
         ];

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

@@ -27,7 +27,7 @@ use crate::{
 /// by the [`Aggregator`] type.
 ///
 /// The provided implementation of [`Search`] accepts
-/// [`graph::RangeSearchParams`] and returns [`Metrics`] as additional output.
+/// [`graph::search::Range`] and returns [`Metrics`] as additional output.
 #[derive(Debug)]
 pub struct Range<DP, T, S>
 where
@@ -83,15 +83,15 @@ where
     T: AsyncFriendly + Clone,
 {
     type Id = DP::ExternalId;
-    type Parameters = graph::RangeSearchParams;
+    type Parameters = graph::search::Range;
     type Output = Metrics;
 
     fn num_queries(&self) -> usize {
         self.queries.nrows()
     }
 
     fn id_count(&self, parameters: &Self::Parameters) -> search::IdCount {
-        search::IdCount::Dynamic(NonZeroUsize::new(parameters.starting_l_value))
+        search::IdCount::Dynamic(NonZeroUsize::new(parameters.starting_l()))
     }
 
     async fn search<O>(
@@ -104,16 +104,21 @@ where
         O: graph::SearchOutputBuffer<DP::ExternalId> + Send,
     {
         let context = DP::Context::default();
-        let (_, ids, distances) = self
+        let range_search = *parameters;
+        let result = self
             .index
-            .range_search(
+            .search(
+                range_search,
                 self.strategy.get(index)?,
                 &context,
                 self.queries.row(index),
-                parameters,
+                &mut (),
             )
             .await?;
-        buffer.extend(std::iter::zip(ids.into_iter(), distances.into_iter()));
+        buffer.extend(std::iter::zip(
+            result.ids.into_iter(),
+            result.distances.into_iter(),
+        ));
 
         Ok(Metrics {})
     }
@@ -129,8 +134,8 @@ pub struct Summary {
     /// The [`search::Setup`] used for the batch of runs.
     pub setup: search::Setup,
 
-    /// The [`graph::RangeSearchParams`] used for the batch of runs.
-    pub parameters: graph::RangeSearchParams,
+    /// The [`graph::search::Range`] used for the batch of runs.
+    pub parameters: graph::search::Range,
 
     /// The end-to-end latency for each repetition in the batch.
     pub end_to_end_latencies: Vec<MicroSeconds>,
@@ -174,7 +179,7 @@ impl<'a, I> Aggregator<'a, I> {
     }
 }
 
-impl<I> search::Aggregate<graph::RangeSearchParams, I, Metrics> for Aggregator<'_, I>
+impl<I> search::Aggregate<graph::search::Range, I, Metrics> for Aggregator<'_, I>
 where
     I: crate::recall::RecallCompatible,
 {
@@ -183,7 +188,7 @@ where
     #[inline(never)]
     fn aggregate(
         &mut self,
-        run: search::Run<graph::RangeSearchParams>,
+        run: search::Run<graph::search::Range>,
         mut results: Vec<search::SearchResults<I, Metrics>>,
     ) -> anyhow::Result<Summary> {
         // Compute the recall using just the first result.
@@ -261,7 +266,7 @@ mod tests {
         let rt = crate::tokio::runtime(2).unwrap();
         let results = search::search(
             range.clone(),
-            graph::RangeSearchParams::new(None, 10, None, 2.0, None, 0.8, 1.2).unwrap(),
+            graph::search::Range::with_options(None, 10, None, 2.0, None, 0.8, 1.2).unwrap(),
             NonZeroUsize::new(2).unwrap(),
             &rt,
         )
@@ -280,11 +285,11 @@ mod tests {
         // Try the aggregated strategy.
         let parameters = [
             search::Run::new(
-                graph::RangeSearchParams::new(None, 10, None, 2.0, None, 0.8, 1.2).unwrap(),
+                graph::search::Range::with_options(None, 10, None, 2.0, None, 0.8, 1.2).unwrap(),
                 setup.clone(),
             ),
             search::Run::new(
-                graph::RangeSearchParams::new(None, 15, None, 2.0, None, 0.8, 1.2).unwrap(),
+                graph::search::Range::with_options(None, 15, None, 2.0, None, 0.8, 1.2).unwrap(),
                 setup.clone(),
             ),
         ];

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

@@ -121,8 +121,8 @@
 //!
 //! ## Graph Index
 //!
-//! * [`graph::KNN`]: K-nearest neighbors search for [`diskann::graph::DiskANNIndex`].
-//! * [`graph::Range`]: Range search for [`diskann::graph::DiskANNIndex`].
+//! * [`graph::search::Knn`]: K-nearest neighbors search for [`diskann::graph::DiskANNIndex`].
+//! * [`graph::search::Range`]: Range search for [`diskann::graph::DiskANNIndex`].
 //! * [`graph::MultiHop`]: Multi-hop filtered search for [`diskann::graph::DiskANNIndex`].
 
 pub(crate) mod ids;

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

@@ -143,8 +143,8 @@ impl SearchResults {
 
         Self {
             num_tasks: setup.tasks.into(),
-            search_n: parameters.k_value,
-            search_l: parameters.l_value,
+            search_n: parameters.k_value().get(),
+            search_l: parameters.l_value().get(),
             qps,
             search_latencies: end_to_end_latencies,
             mean_latencies,
@@ -284,7 +284,7 @@ impl RangeSearchResults {
 
         Self {
             num_tasks: setup.tasks.into(),
-            initial_l: parameters.starting_l_value,
+            initial_l: parameters.starting_l(),
             qps,
             search_latencies: end_to_end_latencies,
             mean_latencies,