[diskann-quantization] Add input hints for spherical normalization and query rescaling

diskann-benchmark/src/exhaustive/spherical.rs

@@ -111,7 +111,7 @@ mod imp {
             let quantizer = diskann_quantization::spherical::SphericalQuantizer::train(
                 data.as_view(),
                 (&input.transform_kind).into(),
-                metric.try_into()?,
+                metric.into(),
                 input.pre_scale.try_into()?,
                 &mut rand::rngs::StdRng::seed_from_u64(input.seed),
                 GlobalAllocator,
@@ -497,10 +497,11 @@ mod imp {
             let allow_rescale = self.measure == SimilarityMeasure::InnerProduct;
             Ok(store
                 .plan
-                .fused_query_computer(
+                .fused_query_computer_with_hints(
                     query,
                     self.layout.into(),
-                    allow_rescale,
+                    diskann_quantization::spherical::InputHints::new()
+                        .with_maybe_out_of_distribution(allow_rescale),
                     GlobalAllocator,
                     ScopedAllocator::global(),
                 )

diskann-benchmark/src/index/bftree/spherical.rs

@@ -126,7 +126,7 @@ impl Benchmark for BfTreeSpherical {
         let quantizer = diskann_quantization::spherical::SphericalQuantizer::train(
             data.as_view(),
             (input.transform_kind()).into(),
-            m.try_into()?,
+            m.into(),
             pre_scale,
             &mut rand::rngs::StdRng::seed_from_u64(input.seed()),
             GlobalAllocator,

diskann-benchmark/src/index/bftree/spherical_streaming.rs

@@ -241,7 +241,7 @@ fn bftree_sq_streaming_impl(
     let quantizer = diskann_quantization::spherical::SphericalQuantizer::train(
         data.as_view(),
         (input.transform_kind()).into(),
-        m.try_into()?,
+        m.into(),
         pre_scale,
         &mut rand::rngs::StdRng::seed_from_u64(input.seed()),
         GlobalAllocator,

diskann-benchmark/src/index/inmem/spherical.rs

@@ -283,7 +283,7 @@ mod imp {
                     let quantizer = diskann_quantization::spherical::SphericalQuantizer::train(
                         data.as_view(),
                         (&input.transform_kind).into(),
-                        m.try_into()?,
+                        m.into(),
                         pre_scale,
                         &mut rand::rngs::StdRng::seed_from_u64(input.seed),
                         GlobalAllocator,

diskann-benchmark/src/index/inmem2/quantized/spherical.rs

@@ -74,6 +74,7 @@ pub(super) struct Spherical {
     bits: Bits,
     seed: u64,
     metric: dq::spherical::SupportedMetric,
+    input_hints: dq::spherical::InputHints,
 }
 
 impl inmem2::quantized::Quantizer for Spherical {
@@ -88,7 +89,9 @@ impl inmem2::quantized::Quantizer for Spherical {
         Ok(Self {
             bits: Bits::from_raw(bits),
             seed: *seed,
-            metric: metric.try_into()?,
+            metric: metric.into(),
+            input_hints: dq::spherical::InputHints::new()
+                .with_normalized(metric == dv::distance::Metric::CosineNormalized),
         })
     }
 
@@ -99,24 +102,41 @@ impl inmem2::quantized::Quantizer for Spherical {
     fn train(&self, data: rowmajor::Ref<'_, f32>) -> anyhow::Result<di::repr::quantized::Backend> {
         use di::repr::quantized::{spherical::Spherical, Backend};
 
-        let Self { bits, seed, metric } = self;
+        let Self {
+            bits,
+            seed,
+            metric,
+            input_hints,
+        } = self;
+        let hints = *input_hints;
+        let query_hints = hints.with_maybe_out_of_distribution(true);
 
         let mut rng = StdRng::seed_from_u64(*seed);
-        let quantizer = dq::spherical::SphericalQuantizer::train(
+        let quantizer = dq::spherical::SphericalQuantizer::train_with_hints(
             data,
             dq::algorithms::transforms::TransformKind::DoubleHadamard {
                 target_dim: dq::algorithms::transforms::TargetDim::Same,
             },
             *metric,
             dq::spherical::PreScale::ReciprocalMeanNorm,
+            hints,
             &mut rng,
             dq::alloc::GlobalAllocator,
         )?;
 
         let quantizer = match bits {
-            Bits::One => Backend::new(Spherical::new(quantizer.as_quantizer_impl::<1>()?)),
-            Bits::Two => Backend::new(Spherical::new(quantizer.as_quantizer_impl::<2>()?)),
-            Bits::Four => Backend::new(Spherical::new(quantizer.as_quantizer_impl::<4>()?)),
+            Bits::One => Backend::new(
+                Spherical::new(quantizer.as_quantizer_impl::<1>()?)
+                    .with_input_hints(hints, query_hints),
+            ),
+            Bits::Two => Backend::new(
+                Spherical::new(quantizer.as_quantizer_impl::<2>()?)
+                    .with_input_hints(hints, query_hints),
+            ),
+            Bits::Four => Backend::new(
+                Spherical::new(quantizer.as_quantizer_impl::<4>()?)
+                    .with_input_hints(hints, query_hints),
+            ),
         };
 
         Ok(quantizer)
@@ -125,7 +145,9 @@ impl inmem2::quantized::Quantizer for Spherical {
 
 impl std::fmt::Display for Spherical {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let Self { bits, seed, metric } = self;
+        let Self {
+            bits, seed, metric, ..
+        } = self;
 
         let mut kv = dbr::utils::fmt::KeyValue::new();
         kv.push("bits", bits);
@@ -151,4 +173,31 @@ mod tests {
         )
         .is_ok());
     }
+
+    #[test]
+    fn spherical_build_matches_cosine_normalized() {
+        use dbr::benchmark::{MatchContext, TestScore};
+        use diskann_utils::views::rowmajor::Matrix;
+        use inmem2::{
+            common::FromRawWith,
+            quantized::{BuildAndSearch, Quantizer, RunBuildAndSearch},
+        };
+
+        for bits in [dto::Bits::One, dto::Bits::Two, dto::Bits::Four] {
+            let raw = inmem2::quantized::dto::BuildAndSearch::example(dto::Spherical {
+                bits,
+                ..dto::Spherical::example()
+            });
+            let mut raw = serde_json::to_value(raw).unwrap();
+            raw["data"]["distance"] =
+                serde_json::to_value(crate::utils::SimilarityMeasure::CosineNormalized).unwrap();
+            let raw = serde_json::from_value(raw).unwrap();
+            let input = BuildAndSearch::<Spherical>::from_raw_with(&raw, None).unwrap();
+            let score = MatchContext::test(&RunBuildAndSearch::<Spherical>::new(), &input);
+            assert!(matches!(score, TestScore::Success(0)), "{score:?}");
+            let data =
+                rowmajor::Owned::from_fn(4, 4, |rc| if rc.row == rc.col { 1.0f32 } else { 0.0 });
+            let _backend = input.quantizer.train(data.as_view()).unwrap();
+        }
+    }
 }

diskann-disk/src/build/builder/quantizer.rs

@@ -130,8 +130,7 @@ impl BuildQuantizer {
                 Ok(Self::Scalar1Bit(WithBits::<1>::new(quantizer)))
             }
             QuantizationType::Spherical(SphericalBits::One) => {
-                let metric: SupportedMetric =
-                    index_configuration.dist_metric.try_into().bridge_err()?;
+                let metric: SupportedMetric = index_configuration.dist_metric.into();
                 let rng = diskann_providers::utils::create_rnd_provider_from_optional_seed(
                     index_configuration.random_seed,
                 );

diskann-garnet/src/ffi_tests.rs

@@ -1843,14 +1843,24 @@ mod tests {
 
     #[test]
     fn spherical_lifecycle_all_types() {
-        for metric in [Metric::L2, Metric::Cosine, Metric::InnerProduct] {
+        for metric in [
+            Metric::L2,
+            Metric::Cosine,
+            Metric::CosineNormalized,
+            Metric::InnerProduct,
+        ] {
             for quant_type in [
                 VectorQuantType::Bin,
                 VectorQuantType::XSpherical2,
                 VectorQuantType::XSpherical4,
             ] {
                 check_spherical_lifecycle(quant_type, metric, |id| {
-                    [(id % 127) as f32 - 63.5, ((id * 47) % 127) as f32 - 63.5]
+                    let mut vector = [(id % 127) as f32 - 63.5, ((id * 47) % 127) as f32 - 63.5];
+                    if metric == Metric::CosineNormalized {
+                        let norm = (vector[0] * vector[0] + vector[1] * vector[1]).sqrt();
+                        vector.iter_mut().for_each(|v| *v /= norm);
+                    }
+                    vector
                 });
             }
             for quant_type in [
@@ -1874,26 +1884,6 @@ mod tests {
         }
     }
 
-    #[test]
-    fn spherical_rejects_unsupported_metric() {
-        for quant_type in [
-            VectorQuantType::Bin,
-            VectorQuantType::XBinI8,
-            VectorQuantType::XBinU8,
-            VectorQuantType::XSpherical2,
-            VectorQuantType::XSpherical2I8,
-            VectorQuantType::XSpherical2U8,
-            VectorQuantType::XSpherical4,
-            VectorQuantType::XSpherical4I8,
-            VectorQuantType::XSpherical4U8,
-        ] {
-            let store = Store::new();
-            let (index_ptr, _) =
-                create_test_index_with_metric(&store, quant_type, Metric::CosineNormalized as i32);
-            assert!(index_ptr.is_null());
-        }
-    }
-
     fn check_spherical_lifecycle<T: bytemuck::Pod>(
         quant_type: VectorQuantType,
         metric: Metric,

diskann-garnet/src/provider.rs

@@ -25,7 +25,7 @@ use diskann::{
 };
 use diskann_quantization::{
     alloc::{AllocatorError, Poly},
-    spherical::SupportedMetric,
+    spherical::InputHints,
 };
 use diskann_utils::{
     object_pool::{AsPooled, ObjectPool, PooledRef, Undef},
@@ -349,6 +349,15 @@ pub(crate) struct GarnetProvider<T: VectorRepr> {
 }
 
 impl<T: VectorRepr> GarnetProvider<T> {
+    // Determine hints before as_f32 erases the source representation. Only native f32
+    // storage inherits the caller's normalized-cosine promise; converted integers do not.
+    fn input_hints(&self) -> InputHints {
+        InputHints::new().with_normalized(
+            self.metric_type == Metric::CosineNormalized
+                && TypeId::of::<T>() == TypeId::of::<f32>(),
+        )
+    }
+
     pub(crate) fn new(
         dim: usize,
         quant_type: VectorQuantType,
@@ -472,9 +481,6 @@ impl<T: VectorRepr> GarnetProvider<T> {
             | VectorQuantType::XSpherical4
             | VectorQuantType::XSpherical4I8
             | VectorQuantType::XSpherical4U8 => {
-                SupportedMetric::try_from(metric_type)
-                    .map_err(|_| GarnetProviderError::InvalidQuantizer)?;
-
                 let valid_type = match quant_type {
                     VectorQuantType::XSpherical2 | VectorQuantType::XSpherical4 => {
                         TypeId::of::<T>() == TypeId::of::<f32>()
@@ -714,7 +720,7 @@ impl<T: VectorRepr> GarnetProvider<T> {
                 let mut qpoint = vec![0u8; quantizer.bytes()];
                 let point_f32 =
                     T::as_f32(point).map_err(|e| GarnetQuantizerError::Compression(Box::new(e)))?;
-                quantizer.compress(&point_f32, &mut qpoint)?;
+                quantizer.compress_with_hints(&point_f32, &mut qpoint, self.input_hints())?;
 
                 if !self.callbacks.write_iid(
                     &context.term(Term::Quantized),
@@ -903,7 +909,7 @@ impl<T: VectorRepr> GarnetProvider<T> {
             Ok(v) => v,
             Err(_) => return false,
         };
-        match quantizer.train(self.metric_type, view) {
+        match quantizer.train_with_hints(self.metric_type, view, self.input_hints()) {
             Ok(()) => {
                 let quant_state = if let Ok(s) = quantizer.serialize() {
                     s
@@ -995,7 +1001,10 @@ impl<T: VectorRepr> GarnetProvider<T> {
             if T::as_f32_into(&v, &mut f).is_err() {
                 continue;
             }
-            if quantizer.compress(&f, &mut q).is_err() {
+            if quantizer
+                .compress_with_hints(&f, &mut q, self.input_hints())
+                .is_err()
+            {
                 continue;
             };
 
@@ -1016,7 +1025,9 @@ impl<T: VectorRepr> GarnetProvider<T> {
             // The final thread to finish backfilling will add the quantized started points.
             if let Some(v) = self.start_point_cache.get(&self.start_point_id)
                 && let Ok(v_f32) = T::as_f32(bytemuck::cast_slice::<u8, T>(&v))
-                && quantizer.compress(&v_f32, &mut q).is_ok()
+                && quantizer
+                    .compress_with_hints(&v_f32, &mut q, self.input_hints())
+                    .is_ok()
             {
                 let _ = self.callbacks.write_iid(
                     &context.term(Term::Quantized),
@@ -1859,7 +1870,7 @@ impl<T: VectorRepr> SetElement<(&[T], &[u8])> for GarnetProvider<T> {
                 let element_f32 = T::as_f32(element.0).map_err(|e| {
                     GarnetProviderError::Quantizer(GarnetQuantizerError::Compression(Box::new(e)))
                 })?;
-                quantizer.compress(&element_f32, &mut quant)?;
+                quantizer.compress_with_hints(&element_f32, &mut quant, self.input_hints())?;
                 self.callbacks
                     .write_iid(&context.term(Term::Quantized), internal_id.id(), &quant)
                     .then_some(())
@@ -2032,7 +2043,10 @@ impl<'a, T: VectorRepr> DynamicAccessor<'a, T> {
             })?;
 
             quantizer
-                .query_computer(&from_f32)
+                .query_computer_with_hints(
+                    &from_f32,
+                    provider.input_hints().with_maybe_out_of_distribution(true),
+                )
                 .map_err(|e| GarnetQuantizerError::QueryComputer(Box::new(e)))?
         } else {
             GarnetQueryComputer::new(FullPrecisionQueryDistance::<T>(T::query_distance(
@@ -4317,6 +4331,66 @@ mod tests {
         ));
     }
 
+    #[tokio::test]
+    async fn spherical_normalized_cosine_integer_fallback() {
+        use super::DynamicAccessor;
+        use diskann::utils::VectorRepr;
+        use diskann_vector::PreprocessedDistanceFunction;
+        for (kind, dim) in [
+            (VectorQuantType::XSpherical4U8, 2),
+            (VectorQuantType::XSpherical4U8, 8),
+        ] {
+            let store = Store::new();
+            let ctx = Context::new(0);
+            let provider = GarnetProvider::<u8>::new(
+                dim,
+                kind,
+                Metric::CosineNormalized,
+                10,
+                DEFAULT_START_POINT_ID,
+                store.callbacks(),
+                &ctx,
+            )
+            .unwrap();
+            let (a, b, q) = if dim == 2 {
+                (vec![3u8, 0], vec![4u8, 1], vec![1u8, 0])
+            } else {
+                (vec![0u8; 8], vec![254u8; 8], vec![127u8; 8])
+            };
+            provider.maybe_set_start_point(&ctx, &b).unwrap();
+            let quantizer = provider.quantizer.as_ref().unwrap();
+            for id in 1..=quantizer.required_vectors() as u32 {
+                let external_id = GarnetId::from(&id.to_ne_bytes()[..]);
+                let vector = if id % 2 == 0 { &a } else { &b };
+                provider
+                    .set_element(&ctx, &external_id, (vector.as_slice(), &[]))
+                    .await
+                    .unwrap()
+                    .complete()
+                    .await;
+            }
+            assert!(provider.train_quantizer(&ctx));
+            assert!(provider.backfill_quant_vectors(&ctx, 0, 1));
+            let accessor = DynamicAccessor::new(&provider, &ctx, &q, true).unwrap();
+            let mut distances = Vec::new();
+            for vector in [&a, &b] {
+                let mut code = vec![0u8; quantizer.bytes()];
+                let converted = <u8 as VectorRepr>::as_f32(vector).unwrap();
+                quantizer
+                    .compress_with_hints(&converted, &mut code, provider.input_hints())
+                    .unwrap();
+                distances.push(accessor.computer.evaluate_similarity(&code));
+            }
+            if dim == 2 {
+                assert!(distances[0] < distances[1], "{distances:?}");
+                assert!(distances[0].abs() < 0.02);
+                assert!((distances[1] - (1.0 - 4.0 / 17.0f32.sqrt())).abs() < 0.02);
+            } else {
+                assert!(distances[1].abs() < 0.02, "{distances:?}");
+            }
+        }
+    }
+
     #[test]
     fn import_start_point_terms() {
         for start_point_id in [DEFAULT_START_POINT_ID, 0, 2] {