(diversity) Refactor attribute diversity to use id trait

diskann/src/graph/misc.rs

@@ -33,30 +33,18 @@ pub enum InplaceDeleteMethod {
 
 // Parameters for diverse search
 #[cfg(feature = "experimental_diversity_search")]
-#[derive(Clone, Debug)]
-pub struct DiverseSearchParams<P>
-where
-    P: crate::neighbor::AttributeValueProvider,
-{
+#[derive(Clone, Copy, Debug)]
+pub struct DiverseSearchParams {
     pub diverse_attribute_id: usize,
     pub diverse_results_k: usize,
-    pub attribute_provider: std::sync::Arc<P>,
 }
 
 #[cfg(feature = "experimental_diversity_search")]
-impl<P> DiverseSearchParams<P>
-where
-    P: crate::neighbor::AttributeValueProvider,
-{
-    pub fn new(
-        diverse_attribute_id: usize,
-        diverse_results_k: usize,
-        attribute_provider: std::sync::Arc<P>,
-    ) -> Self {
+impl DiverseSearchParams {
+    pub fn new(diverse_attribute_id: usize, diverse_results_k: usize) -> Self {
         Self {
             diverse_attribute_id,
             diverse_results_k,
-            attribute_provider,
         }
     }
 }

diskann/src/graph/search/diverse_search.rs

@@ -18,30 +18,24 @@ use crate::{
         index::{DiskANNIndex, SearchStats},
         search_output_buffer::SearchOutputBuffer,
     },
-    neighbor::{AttributeValueProvider, DiverseNeighborQueue, NeighborQueue},
+    neighbor::{DiverseId, DiverseNeighborQueue, NeighborQueue},
     provider::DataProvider,
 };
 
 /// Parameters for diversity-aware search.
 ///
 /// Returns results that are diverse across a specified attribute.
 #[derive(Debug)]
-pub struct Diverse<P>
-where
-    P: AttributeValueProvider,
-{
+pub struct Diverse {
     /// Base k-NN search parameters.
     inner: Knn,
     /// Diversity-specific parameters.
-    diverse_params: DiverseSearchParams<P>,
+    diverse_params: DiverseSearchParams,
 }
 
-impl<P> Diverse<P>
-where
-    P: AttributeValueProvider,
-{
+impl Diverse {
     /// Create new diverse search parameters.
-    pub fn new(inner: Knn, diverse_params: DiverseSearchParams<P>) -> Self {
+    pub fn new(inner: Knn, diverse_params: DiverseSearchParams) -> Self {
         Self {
             inner,
             diverse_params,
@@ -56,25 +50,23 @@ where
 
     /// Returns a reference to the diversity-specific parameters.
     #[inline]
-    pub fn diverse_params(&self) -> &DiverseSearchParams<P> {
+    pub fn diverse_params(&self) -> &DiverseSearchParams {
         &self.diverse_params
     }
 
     /// Create search scratch with DiverseNeighborQueue for this search.
     fn create_scratch<DP>(
         &self,
         index: &DiskANNIndex<DP>,
-    ) -> SearchScratch<DP::InternalId, DiverseNeighborQueue<P>>
+    ) -> SearchScratch<DP::InternalId, DiverseNeighborQueue<DP::InternalId>>
     where
         DP: DataProvider,
-        P: AttributeValueProvider<Id = DP::InternalId>,
+        DP::InternalId: DiverseId,
     {
-        let attribute_provider = self.diverse_params.attribute_provider.clone();
         let diverse_queue = DiverseNeighborQueue::new(
             self.inner.l_value().get(),
             self.inner.k_value(),
             self.diverse_params.diverse_results_k,
-            attribute_provider,
         );
 
         SearchScratch {
@@ -92,12 +84,12 @@ where
     }
 }
 
-impl<'a, DP, S, T, P> Search<'a, DP, S, T> for Diverse<P>
+impl<'a, DP, S, T> Search<'a, DP, S, T> for Diverse
 where
     DP: DataProvider,
+    DP::InternalId: DiverseId,
     T: Copy + Send + Sync,
     S: SearchStrategy<'a, DP, T, SearchAccessor: SearchAccessor>,
-    P: AttributeValueProvider<Id = DP::InternalId>,
 {
     type Output = SearchStats;