Fix normalized cosine distance consistency for PQ graphs

diskann-providers/src/index/diskann_async.rs

@@ -182,7 +182,7 @@ pub(crate) mod tests {
     use diskann_utils::{test_data_root, views::Matrix};
     use diskann_vector::{
         DistanceFunction, PureDistanceFunction,
-        distance::{Metric, SquaredL2},
+        distance::{CosineNormalized, Metric, SquaredL2},
     };
     use rand::{distr::Distribution, rngs::StdRng, seq::SliceRandom};
     use rstest::rstest;
@@ -1378,29 +1378,36 @@ pub(crate) mod tests {
     }
 
     const SIFTSMALL: &str = "/sift/siftsmall_learn_256pts.fbin";
+    const SIFTSMALL_NORMALIZED: &str = "/sift/siftsmall_learn_256pts_normalized.fbin";
 
     #[rstest]
     #[tokio::test]
     async fn test_sift_build_and_search<S>(
         #[values(FullPrecision, Hybrid::new(None))] build_strategy: S,
         #[values(1, 10)] batchsize: usize,
+        #[values(
+            (Metric::L2, SIFTSMALL),
+            (Metric::CosineNormalized, SIFTSMALL_NORMALIZED),
+        )]
+        metric_and_file: (Metric, &str),
     ) where
         S: for<'a> InsertStrategy<'a, TestProvider, &'a [f32]>
             + MultiInsertStrategy<TestProvider, Matrix<f32>>
             + Clone,
     {
+        let (metric, file) = metric_and_file;
         let ctx = &DefaultContext;
         let parameters = InitParams {
             l_build: 64,
             max_degree: 16,
-            metric: Metric::L2,
+            metric,
             batchsize: NonZeroUsize::new(batchsize).unwrap(),
         };
 
         let (index, data) = init_from_file(
             build_strategy.clone(),
             parameters,
-            SIFTSMALL,
+            file,
             8,
             StartPointStrategy::RandomSamples {
                 nsamples: ONE,
@@ -1433,7 +1440,11 @@ pub(crate) mod tests {
         //
         // Because this dataset is small, we can expect exact equality.
         for (q, query) in data.row_iter().enumerate() {
-            let gt = groundtruth(data.as_view(), query, |a, b| SquaredL2::evaluate(a, b));
+            let gt = groundtruth(data.as_view(), query, |a, b| match metric {
+                Metric::L2 => SquaredL2::evaluate(a, b),
+                Metric::CosineNormalized => CosineNormalized::evaluate(a, b),
+                _ => unreachable!(),
+            });
             {
                 let mut result_output_buffer =
                     search_output_buffer::IdDistance::new(&mut ids, &mut distances);
@@ -2061,8 +2072,11 @@ pub(crate) mod tests {
     /// PQ only Build & Search ///
     //////////////////////////////
 
+    #[rstest]
+    #[case(Metric::L2, SIFTSMALL)]
+    #[case(Metric::CosineNormalized, SIFTSMALL_NORMALIZED)]
     #[tokio::test]
-    async fn test_sift_pq_only_build_and_search() {
+    async fn test_sift_pq_only_build_and_search(#[case] metric: Metric, #[case] file: &str) {
         let ctx = &DefaultContext;
         let create_fn = |data: Arc<Matrix<f32>>, start_points: &[f32]| {
             let pq_table = train_pq(
@@ -2074,8 +2088,7 @@ pub(crate) mod tests {
             .unwrap();
 
             let (config, parameters) =
-                simplified_builder(64, 16, Metric::L2, data.ncols(), data.nrows(), no_modify)
-                    .unwrap();
+                simplified_builder(64, 16, metric, data.ncols(), data.nrows(), no_modify).unwrap();
 
             let index =
                 Arc::new(new_quant_only_index(config, parameters, pq_table, NoDeletes).unwrap());
@@ -2086,7 +2099,7 @@ pub(crate) mod tests {
             index
         };
         let (index, data) =
-            init_and_build_index_from_file(SIFTSMALL, create_fn, build_using_single_insert).await;
+            init_and_build_index_from_file(file, create_fn, build_using_single_insert).await;
 
         let neighbor_accessor = &mut index.provider().neighbors();
         // There should be one more reachable node than points in the dataset to account for
@@ -2131,7 +2144,17 @@ pub(crate) mod tests {
                 .await
                 .unwrap();
 
-            assert_top_k_exactly_match(q, &gt, &ids, &distances, top_k);
+            if metric == Metric::CosineNormalized {
+                let expected: Vec<_> = gt
+                    .iter()
+                    .rev()
+                    .take(top_k)
+                    .map(|neighbor| *neighbor.id())
+                    .collect();
+                assert_eq!(expected, ids, "failed on query {q}");
+            } else {
+                assert_top_k_exactly_match(q, &gt, &ids, &distances, top_k);
+            }
         }
     }
 

diskann-providers/src/model/graph/provider/async_/inmem/product.rs

@@ -42,6 +42,20 @@ use crate::model::{
 /// The default quant provider.
 pub type DefaultQuant = FastMemoryQuantVectorProviderAsync;
 
+fn quant_pruning_distance_computer(
+    quant: &DefaultQuant,
+) -> (pq::distance::DistanceComputer<'_>, Metric) {
+    let metric = match quant.metric() {
+        // Match the squared-L2 approximation used by PQ graph-search queries.
+        Metric::CosineNormalized => Metric::L2,
+        metric => metric,
+    };
+    (
+        pq::distance::DistanceComputer::new(&quant.pq_chunk_table, metric),
+        metric,
+    )
+}
+
 impl CreateVectorStore for FixedChunkPQTable {
     type Target = DefaultQuant;
     fn create(
@@ -410,10 +424,10 @@ where
     ) -> Result<Self::PruneAccessor<'a>, Self::PruneAccessorError> {
         let full = &provider.base_vectors;
         let quant = &provider.aux_vectors;
-        let metric = quant.metric();
+        let (quant_distance_computer, metric) = quant_pruning_distance_computer(quant);
 
         let distance = distances::pq::HybridComputer::new(
-            quant.distance_computer(),
+            quant_distance_computer,
             T::distance(metric, Some(full.dim())),
         );
 
@@ -579,10 +593,12 @@ where
         _context: &'a Ctx,
         _capacity: usize,
     ) -> Result<Self::PruneAccessor<'a>, Self::PruneAccessorError> {
+        let quant = &provider.aux_vectors;
+        let (distance_computer, _) = quant_pruning_distance_computer(quant);
         let accessor = PruneAccessor {
-            provider: &provider.aux_vectors,
+            provider: quant,
             neighbors: provider.neighbors(),
-            distance: provider.aux_vectors.distance_computer(),
+            distance: distance_computer,
         };
         Ok(accessor)
     }
@@ -647,3 +663,75 @@ where
             .into_ann_result()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use diskann::{graph::glue::PruneStrategy, provider::DefaultContext, utils::VectorRepr};
+    use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::Metric};
+
+    use super::{DefaultQuant, quant_pruning_distance_computer};
+    use crate::model::{
+        graph::provider::async_::{
+            FastMemoryQuantVectorProviderAsync,
+            common::{NoDeletes, NoStore, Quantized},
+            distances::pq::{Hybrid, HybridComputer},
+            inmem::{DefaultProvider, DefaultProviderParameters},
+        },
+        pq::FixedChunkPQTable,
+    };
+
+    fn test_table() -> FixedChunkPQTable {
+        FixedChunkPQTable::new(
+            4,
+            vec![1.0, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 2.0].into(),
+            vec![0, 2, 4].into(),
+        )
+        .unwrap()
+    }
+
+    #[test]
+    fn cosine_normalized_query_and_hybrid_pruning_use_squared_l2() {
+        let provider =
+            FastMemoryQuantVectorProviderAsync::new(Metric::CosineNormalized, 2, test_table());
+        let full0 = [1u8, 0, 0, 2];
+        let full1 = [2u8, 0, 0, 1];
+        let code0 = [0u8, 1];
+        let code1 = [1u8, 0];
+
+        let query = provider.query_computer(&full0).unwrap();
+        assert_eq!(query.evaluate_similarity(&code1), 2.0);
+
+        let (quant_distance_computer, metric) = quant_pruning_distance_computer(&provider);
+        let computer =
+            HybridComputer::<u8>::new(quant_distance_computer, u8::distance(metric, Some(4)));
+        for (left, right) in [
+            (Hybrid::Full(&full0[..]), Hybrid::Full(&full1[..])),
+            (Hybrid::Full(&full0[..]), Hybrid::Quant(&code1[..])),
+            (Hybrid::Quant(&code0[..]), Hybrid::Full(&full1[..])),
+            (Hybrid::Quant(&code0[..]), Hybrid::Quant(&code1[..])),
+        ] {
+            assert_eq!(computer.evaluate_similarity(left, right), 2.0);
+        }
+    }
+
+    #[test]
+    fn cosine_normalized_quantized_pruning_uses_squared_l2() {
+        let provider: DefaultProvider<NoStore, DefaultQuant> = DefaultProvider::new_empty(
+            DefaultProviderParameters::simple(2, 4, Metric::CosineNormalized, 1),
+            NoStore,
+            test_table(),
+            NoDeletes,
+        )
+        .unwrap();
+        let accessor = Quantized
+            .prune_accessor(&provider, &DefaultContext, 2)
+            .unwrap();
+
+        assert_eq!(
+            accessor
+                .distance
+                .evaluate_similarity(&[0u8, 1][..], &[1u8, 0][..]),
+            2.0,
+        );
+    }
+}