From fe7f00bd81f627ccfc579961b8911e5e747f88e6 Mon Sep 17 00:00:00 2001 From: "Xinyu Wen (from Dev Box)" Date: Thu, 29 Oct 2026 16:27:53 +0800 Subject: [PATCH 1/2] [diskann-quantization] Add per-input hints for spherical quantization Separate normalization and query rescaling hints from the distance metric. Preserve default preprocessing and serialization, propagate hints through in-memory and Garnet adapters, and retain cosine normalization for converted integer inputs. Cover mixed normalized/raw inputs, integer fallback, query layouts, rescaling, serialization, and benchmark dispatch. --- diskann-benchmark/src/exhaustive/spherical.rs | 7 +- .../src/index/bftree/spherical.rs | 2 +- .../src/index/bftree/spherical_streaming.rs | 2 +- .../src/index/inmem/spherical.rs | 2 +- .../src/index/inmem2/quantized/spherical.rs | 63 +++- diskann-disk/src/build/builder/quantizer.rs | 3 +- diskann-garnet/src/ffi_tests.rs | 34 +- diskann-garnet/src/provider.rs | 94 ++++- diskann-garnet/src/quantization.rs | 66 +++- diskann-inmem/src/repr/quantized/spherical.rs | 63 +++- diskann-providers/src/index/diskann_async.rs | 4 +- .../graph/provider/async_/inmem/spherical.rs | 7 +- diskann-quantization/src/spherical/hints.rs | 51 +++ .../src/spherical/hints_tests.rs | 342 ++++++++++++++++++ diskann-quantization/src/spherical/iface.rs | 165 +++++---- diskann-quantization/src/spherical/mod.rs | 8 +- .../src/spherical/quantizer.rs | 250 +++++++++++-- diskann-quantization/src/spherical/vectors.rs | 55 +-- diskann-quantization/src/utils.rs | 57 ++- 19 files changed, 1062 insertions(+), 213 deletions(-) create mode 100644 diskann-quantization/src/spherical/hints.rs create mode 100644 diskann-quantization/src/spherical/hints_tests.rs diff --git a/diskann-benchmark/src/exhaustive/spherical.rs b/diskann-benchmark/src/exhaustive/spherical.rs index 93e8a00fce..ca2d9cab66 100644 --- a/diskann-benchmark/src/exhaustive/spherical.rs +++ b/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(), ) diff --git a/diskann-benchmark/src/index/bftree/spherical.rs b/diskann-benchmark/src/index/bftree/spherical.rs index a5d24c15cb..c26141a289 100644 --- a/diskann-benchmark/src/index/bftree/spherical.rs +++ b/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, diff --git a/diskann-benchmark/src/index/bftree/spherical_streaming.rs b/diskann-benchmark/src/index/bftree/spherical_streaming.rs index dd9d008910..c6d7828a07 100644 --- a/diskann-benchmark/src/index/bftree/spherical_streaming.rs +++ b/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, diff --git a/diskann-benchmark/src/index/inmem/spherical.rs b/diskann-benchmark/src/index/inmem/spherical.rs index abddca6dc5..b664b9c6c9 100644 --- a/diskann-benchmark/src/index/inmem/spherical.rs +++ b/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, diff --git a/diskann-benchmark/src/index/inmem2/quantized/spherical.rs b/diskann-benchmark/src/index/inmem2/quantized/spherical.rs index 850f7f82bf..0a0b3642d9 100644 --- a/diskann-benchmark/src/index/inmem2/quantized/spherical.rs +++ b/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 { 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::::from_raw_with(&raw, None).unwrap(); + let score = MatchContext::test(&RunBuildAndSearch::::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(); + } + } } diff --git a/diskann-disk/src/build/builder/quantizer.rs b/diskann-disk/src/build/builder/quantizer.rs index 369f3a6600..35203c7e1e 100644 --- a/diskann-disk/src/build/builder/quantizer.rs +++ b/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, ); diff --git a/diskann-garnet/src/ffi_tests.rs b/diskann-garnet/src/ffi_tests.rs index 598c33d25a..58d90c8b45 100644 --- a/diskann-garnet/src/ffi_tests.rs +++ b/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( quant_type: VectorQuantType, metric: Metric, diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 2f244293df..cd90a9b893 100644 --- a/diskann-garnet/src/provider.rs +++ b/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 { } impl GarnetProvider { + // 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::() == TypeId::of::(), + ) + } + pub(crate) fn new( dim: usize, quant_type: VectorQuantType, @@ -472,9 +481,6 @@ impl GarnetProvider { | 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::() == TypeId::of::() @@ -714,7 +720,7 @@ impl GarnetProvider { 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 GarnetProvider { 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 GarnetProvider { 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 GarnetProvider { // 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::(&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 SetElement<(&[T], &[u8])> for GarnetProvider { 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::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::::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 = ::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] { diff --git a/diskann-garnet/src/quantization.rs b/diskann-garnet/src/quantization.rs index e4948db325..9aae612d49 100644 --- a/diskann-garnet/src/quantization.rs +++ b/diskann-garnet/src/quantization.rs @@ -14,7 +14,7 @@ use diskann_quantization::{ minmax::{self, MinMaxQuantizer}, num::POSITIVE_ONE_F32, spherical::{ - self, Data, PreScale, SphericalQuantizer, SupportedMetric, + self, Data, InputHints, PreScale, SphericalQuantizer, SupportedMetric, iface::{self, Opaque, OpaqueMut, Quantizer}, }, }; @@ -66,6 +66,30 @@ pub(crate) trait GarnetQuantizer: Send + Sync { fn distance_computer(&self) -> Result; /// Returns a query computer for comparing distances to a particular query fn query_computer(&self, query: &[f32]) -> Result; + /// Input-aware entry points; quantizers without these optimizations use their defaults. + fn train_with_hints( + &self, + metric: Metric, + data: rowmajor::Ref, + _hints: InputHints, + ) -> Result<(), GarnetQuantizerError> { + self.train(metric, data) + } + fn compress_with_hints( + &self, + v: &[f32], + into: &mut [u8], + _hints: InputHints, + ) -> Result<(), GarnetQuantizerError> { + self.compress(v, into) + } + fn query_computer_with_hints( + &self, + query: &[f32], + _hints: InputHints, + ) -> Result { + self.query_computer(query) + } // Serialize the quantizer state. fn serialize(&self) -> Result, GarnetQuantizerError>; // Deserialize the quantizer state. @@ -122,20 +146,25 @@ where self.inner.read().is_some() } - fn train( + fn train(&self, metric: Metric, data: rowmajor::Ref) -> Result<(), GarnetQuantizerError> { + self.train_with_hints(metric, data, InputHints::new()) + } + + fn train_with_hints( &self, metric_type: Metric, data: rowmajor::Ref, + hints: InputHints, ) -> Result<(), GarnetQuantizerError> { let mut rng = rand::rng(); - let quantizer = SphericalQuantizer::train( + let quantizer = SphericalQuantizer::train_with_hints( data.as_view(), TransformKind::DoubleHadamard { target_dim: diskann_quantization::algorithms::transforms::TargetDim::Same, }, - SupportedMetric::try_from(metric_type) - .map_err(|e| GarnetQuantizerError::Training(Box::new(e)))?, + SupportedMetric::from(metric_type), PreScale::ReciprocalMeanNorm, + hints, &mut rng, GlobalAllocator, ) @@ -151,11 +180,21 @@ where } fn compress(&self, v: &[f32], into: &mut [u8]) -> Result<(), GarnetQuantizerError> { + self.compress_with_hints(v, into, InputHints::new()) + } + + fn compress_with_hints( + &self, + v: &[f32], + into: &mut [u8], + hints: InputHints, + ) -> Result<(), GarnetQuantizerError> { let guard = self.inner.read(); if let Some(quantizer) = &*guard { - spherical::iface::Quantizer::::compress( + spherical::iface::Quantizer::::compress_with_hints( quantizer, v, + hints, OpaqueMut::new(into), ScopedAllocator::global(), ) @@ -179,13 +218,24 @@ where } fn query_computer(&self, query: &[f32]) -> Result { + self.query_computer_with_hints( + query, + InputHints::new().with_maybe_out_of_distribution(true), + ) + } + + fn query_computer_with_hints( + &self, + query: &[f32], + hints: InputHints, + ) -> Result { let guard = self.inner.read(); if let Some(quantizer) = &*guard { let computer = quantizer - .fused_query_computer( + .fused_query_computer_with_hints( query, iface::QueryLayout::FullPrecision, - true, + hints, GlobalAllocator, ScopedAllocator::global(), ) diff --git a/diskann-inmem/src/repr/quantized/spherical.rs b/diskann-inmem/src/repr/quantized/spherical.rs index 19670b60b0..990d8b7d4d 100644 --- a/diskann-inmem/src/repr/quantized/spherical.rs +++ b/diskann-inmem/src/repr/quantized/spherical.rs @@ -6,37 +6,53 @@ use diskann::{ANNError, ANNResult}; use diskann_quantization::{ alloc::{Allocator, GlobalAllocator, ScopedAllocator}, - spherical::iface, + spherical::{InputHints, iface}, }; use crate::{num::Bytes, repr::quantized}; /// A new-type wrapper around [`iface::Quantizer`] for implementing [`quantized::Quantizer`]. #[derive(Debug)] -pub struct Spherical(T); +pub struct Spherical { + inner: T, + data_hints: InputHints, + query_hints: InputHints, +} impl Spherical { /// Construct a new [`Spherical`] new-type wrapper around `quantizer`. pub fn new(quantizer: T) -> Self { - Self(quantizer) + Self { + inner: quantizer, + data_hints: InputHints::new(), + query_hints: InputHints::new().with_maybe_out_of_distribution(true), + } + } + + /// Configure properties of stored data and queries independently. Callers must ensure + /// normalized hints remain valid after any conversion into the f32 input representation. + pub fn with_input_hints(mut self, data: InputHints, query: InputHints) -> Self { + self.data_hints = data; + self.query_hints = query; + self } /// Return the wrapped quantizer. pub fn into_inner(self) -> T { - self.0 + self.inner } } impl std::ops::Deref for Spherical { type Target = T; fn deref(&self) -> &T { - &self.0 + &self.inner } } impl std::ops::DerefMut for Spherical { fn deref_mut(&mut self) -> &mut T { - &mut self.0 + &mut self.inner } } @@ -53,9 +69,10 @@ where } fn quantize(&self, quantized: &mut [u8], data: &[f32]) -> ANNResult<()> { - iface::Quantizer::compress( + iface::Quantizer::compress_with_hints( &**self, data, + self.data_hints, iface::OpaqueMut::new(quantized), ScopedAllocator::global(), ) @@ -73,11 +90,11 @@ where where V: quantized::QueryDistanceVisitor<'a>, { - let query_computer = iface::Quantizer::fused_query_computer( + let query_computer = iface::Quantizer::fused_query_computer_with_hints( &**self, query, iface::QueryLayout::SameAsData, - false, + self.data_hints.with_maybe_out_of_distribution(false), GlobalAllocator, ScopedAllocator::global(), ) @@ -90,11 +107,11 @@ where where V: quantized::QueryDistanceVisitor<'a>, { - let query_computer = iface::Quantizer::fused_query_computer( + let query_computer = iface::Quantizer::fused_query_computer_with_hints( &**self, query, iface::QueryLayout::FullPrecision, - true, + self.query_hints, GlobalAllocator, ScopedAllocator::global(), ) @@ -147,6 +164,30 @@ mod tests { const DIM: usize = 5; const TRIALS: usize = 10; + #[test] + fn normalized_data_with_independent_query_hints() { + let data = rowmajor::Owned::from_fn(2, 4, |rc| if rc.row == rc.col { 1.0 } else { 0.0 }); + let normalized = InputHints::new().with_normalized(true); + let make = || train_quantizer::<4>(data.as_view(), SupportedMetric::Cosine); + let raw_query = Spherical::new(make()).with_input_hints(normalized, InputHints::new()); + let normalized_query = Spherical::new(make()).with_input_hints(normalized, normalized); + let mut code = vec![0; raw_query.bytes().value()]; + raw_query + .quantize(&mut code, &[1.0, 0.0, 0.0, 0.0]) + .unwrap(); + let raw = raw_query.query_distance(&[3.0, 4.0, 0.0, 0.0]).unwrap(); + let unit = normalized_query + .query_distance(&[0.6, 0.8, 0.0, 0.0]) + .unwrap(); + let raw = raw.eval(&code).unwrap(); + assert!((raw - 0.4).abs() < 1e-4); + assert!((raw - unit.eval(&code).unwrap()).abs() < 1e-4); + let insert = normalized_query + .insert_distance(&[0.6, 0.8, 0.0, 0.0]) + .unwrap(); + assert!((insert.eval(&code).unwrap() - 0.4).abs() < 0.1); + } + fn train_quantizer( data: rowmajor::Ref<'_, f32>, metric: SupportedMetric, diff --git a/diskann-providers/src/index/diskann_async.rs b/diskann-providers/src/index/diskann_async.rs index 7a50f5c562..039665db5f 100644 --- a/diskann-providers/src/index/diskann_async.rs +++ b/diskann-providers/src/index/diskann_async.rs @@ -1865,7 +1865,7 @@ pub(crate) mod tests { diskann_quantization::algorithms::transforms::TransformKind::PaddingHadamard { target_dim: diskann_quantization::algorithms::transforms::TargetDim::Natural, }, - Metric::L2.try_into().unwrap(), + Metric::L2.into(), diskann_quantization::spherical::PreScale::ReciprocalMeanNorm, rng, diskann_quantization::alloc::GlobalAllocator, @@ -1978,7 +1978,7 @@ pub(crate) mod tests { diskann_quantization::algorithms::transforms::TransformKind::PaddingHadamard { target_dim: diskann_quantization::algorithms::transforms::TargetDim::Natural, }, - Metric::L2.try_into().unwrap(), + Metric::L2.into(), diskann_quantization::spherical::PreScale::ReciprocalMeanNorm, rng, diskann_quantization::alloc::GlobalAllocator, diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs b/diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs index d7164b7174..2edce23620 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs +++ b/diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs @@ -53,7 +53,6 @@ use crate::{ convert_error!(Bridge); convert_error!(Bridge); convert_error!(Bridge); -convert_error!(Bridge); /// An allocator error scoped to the spherical store. #[derive(Debug, Clone, Copy, Error)] @@ -192,10 +191,10 @@ impl SphericalStore { let qf32 = T::as_f32(query) .map_err(|e| QueryComputerError::FullPrecisionConversionErr(Box::new(e)))?; - let computer = self.plan.fused_query_computer( + let computer = self.plan.fused_query_computer_with_hints( qf32.as_ref(), layout, - allow_rescale, + spherical::InputHints::new().with_maybe_out_of_distribution(allow_rescale), GlobalAllocator, ScopedAllocator::global(), )?; @@ -753,7 +752,7 @@ mod tests { fn test_dim() { let mut rng = StdRng::seed_from_u64(0x721e3de995bc908c); let data = test_dataset(); - let store = make_store::<1>(data.as_view(), Metric::L2.try_into().unwrap(), &mut rng); + let store = make_store::<1>(data.as_view(), Metric::L2.into(), &mut rng); assert_eq!(store.input_dim(), data.ncols()); } diff --git a/diskann-quantization/src/spherical/hints.rs b/diskann-quantization/src/spherical/hints.rs new file mode 100644 index 0000000000..f29945194b --- /dev/null +++ b/diskann-quantization/src/spherical/hints.rs @@ -0,0 +1,51 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +/// Properties of the input to a single training or compression operation. +/// +/// Defaults make no assumptions about normalization or distribution. These hints are not +/// part of the quantizer's serialized state and do not change its distance metric. +#[derive(Debug, Default, Clone, Copy)] +pub struct InputHints { + normalized: bool, + maybe_out_of_distribution: bool, +} + +impl InputHints { + /// Make no assumptions about the input. + pub const fn new() -> Self { + Self { + normalized: false, + maybe_out_of_distribution: false, + } + } + + /// Declare that the input has unit L2 norm (zero vectors may remain zero). + /// + /// Cosine can skip input normalization. Incorrectly declaring non-unit inputs can + /// produce incorrect distances. Converting integer values to floats does not normalize + /// them. Training hints describe every row; compression hints describe only that input. + pub const fn with_normalized(mut self, normalized: bool) -> Self { + self.normalized = normalized; + self + } + + /// Allow an inner-product query to be rescaled to the training distribution's mean norm. + /// + /// This preserves ideal rankings for a fixed query, but changes the distance scale and + /// may change quantization error. Training and stored-data compression ignore this hint: + /// independently rescaling database vectors would change inner-product rankings. + pub const fn with_maybe_out_of_distribution(mut self, value: bool) -> Self { + self.maybe_out_of_distribution = value; + self + } + + pub(super) const fn normalized(self) -> bool { + self.normalized + } + pub(super) const fn maybe_out_of_distribution(self) -> bool { + self.maybe_out_of_distribution + } +} diff --git a/diskann-quantization/src/spherical/hints_tests.rs b/diskann-quantization/src/spherical/hints_tests.rs new file mode 100644 index 0000000000..68299b9037 --- /dev/null +++ b/diskann-quantization/src/spherical/hints_tests.rs @@ -0,0 +1,342 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_utils::views::rowmajor::{self, Matrix}; +use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction}; +use rand::{SeedableRng, rngs::StdRng}; + +use super::{ + InputHints, PreScale, SphericalQuantizer, SupportedMetric, + iface::{self, Opaque, OpaqueMut, Quantizer, QueryLayout}, +}; +use crate::{ + algorithms::TransformKind, + alloc::{AlignedAllocator, GlobalAllocator, Poly, ScopedAllocator}, +}; + +fn train(metric: SupportedMetric, hints: InputHints) -> SphericalQuantizer { + let data = rowmajor::Owned::from_fn(2, 4, |rc| if rc.row == rc.col { 1.0f32 } else { 0.0 }); + SphericalQuantizer::train_with_hints( + data.as_view(), + TransformKind::Null, + metric, + PreScale::ReciprocalMeanNorm, + hints, + &mut StdRng::seed_from_u64(7), + GlobalAllocator, + ) + .unwrap() +} + +fn close(a: f32, b: f32) { + assert!((a - b).abs() < 2.0e-4 * (1.0 + b.abs()), "{a} != {b}"); +} + +fn check_layouts() +where + iface::Impl: iface::Constructible + Quantizer, +{ + let normalized = InputHints::new().with_normalized(true); + let quantizer = train(SupportedMetric::Cosine, normalized); + let reference = train(SupportedMetric::Cosine, InputHints::new()); + assert_eq!(quantizer.shift(), reference.shift()); + assert_eq!(quantizer.mean_norm(), reference.mean_norm()); + let plan = quantizer.as_quantizer_impl::().unwrap(); + let plan: &dyn Quantizer = &plan; + let mut code = vec![0; plan.bytes()]; + let mut ordinary = vec![0; plan.bytes()]; + plan.compress_with_hints( + &[1.0, 0.0, 0.0, 0.0], + normalized, + OpaqueMut::new(&mut code), + ScopedAllocator::global(), + ) + .unwrap(); + plan.compress( + &[3.0, 0.0, 0.0, 0.0], + OpaqueMut::new(&mut ordinary), + ScopedAllocator::global(), + ) + .unwrap(); + assert_eq!(code, ordinary); + + for layout in [ + QueryLayout::SameAsData, + QueryLayout::FullPrecision, + QueryLayout::FourBitTransposed, + QueryLayout::ScalarQuantized, + ] { + if !plan.is_supported(layout) { + continue; + } + let unit = [0.6, 0.8, 0.0, 0.0]; + let raw = [3.0, 4.0, 0.0, 0.0]; + let hinted = plan + .fused_query_computer_with_hints( + &unit, + layout, + normalized, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + let ordinary = plan + .fused_query_computer( + &raw, + layout, + false, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + let expected = ordinary.evaluate_similarity(Opaque::new(&code)).unwrap(); + close( + hinted.evaluate_similarity(Opaque::new(&code)).unwrap(), + expected, + ); + if layout == QueryLayout::FullPrecision { + close(expected, 0.4); + } + + let desc = plan.query_buffer_description(layout).unwrap(); + let alloc = AlignedAllocator::new(desc.align()); + let mut query = Poly::broadcast(0u8, desc.bytes(), alloc).unwrap(); + plan.compress_query_with_hints( + &unit, + layout, + normalized, + OpaqueMut::new(&mut query), + ScopedAllocator::global(), + ) + .unwrap(); + let computer = plan.query_computer(layout, GlobalAllocator).unwrap(); + close( + computer + .evaluate_similarity(Opaque::new(&query), Opaque::new(&code)) + .unwrap(), + expected, + ); + + // Zero inputs remain valid with the normalized hint, including the full-query layout. + let zero = plan + .fused_query_computer_with_hints( + &[0.0; 4], + layout, + normalized, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + let ordinary_zero = plan + .fused_query_computer( + &[0.0; 4], + layout, + false, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + close( + zero.evaluate_similarity(Opaque::new(&code)).unwrap(), + ordinary_zero + .evaluate_similarity(Opaque::new(&code)) + .unwrap(), + ); + } +} + +#[test] +fn normalized_data_and_raw_queries_all_layouts() { + check_layouts::<1>(); + check_layouts::<2>(); + check_layouts::<3>(); + check_layouts::<4>(); + check_layouts::<8>(); +} + +fn check_rescale() +where + iface::Impl: iface::Constructible + Quantizer, +{ + let distribution = InputHints::new().with_maybe_out_of_distribution(true); + for metric in [ + SupportedMetric::InnerProduct, + SupportedMetric::SquaredL2, + SupportedMetric::Cosine, + ] { + let quantizer = train(metric, InputHints::new()); + let mut rescaled = [30.0, 40.0, 0.0, 0.0]; + if metric == SupportedMetric::InnerProduct { + quantizer.rescale(&mut rescaled); + } + let plan = quantizer.as_quantizer_impl::().unwrap(); + let plan: &dyn Quantizer = &plan; + let mut code = vec![0; plan.bytes()]; + let mut ordinary = vec![0; plan.bytes()]; + plan.compress_with_hints( + &[3.0, 4.0, 0.0, 0.0], + distribution, + OpaqueMut::new(&mut code), + ScopedAllocator::global(), + ) + .unwrap(); + plan.compress( + &[3.0, 4.0, 0.0, 0.0], + OpaqueMut::new(&mut ordinary), + ScopedAllocator::global(), + ) + .unwrap(); + assert_eq!(code, ordinary, "stored data must not be rescaled"); + + for layout in [ + QueryLayout::SameAsData, + QueryLayout::FullPrecision, + QueryLayout::FourBitTransposed, + QueryLayout::ScalarQuantized, + ] { + if !plan.is_supported(layout) { + continue; + } + let hinted = plan + .fused_query_computer_with_hints( + &[30.0, 40.0, 0.0, 0.0], + layout, + distribution, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + let manual = plan + .fused_query_computer( + &rescaled, + layout, + false, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + close( + hinted.evaluate_similarity(Opaque::new(&code)).unwrap(), + manual.evaluate_similarity(Opaque::new(&code)).unwrap(), + ); + if metric == SupportedMetric::InnerProduct { + let unit = plan + .fused_query_computer_with_hints( + &[0.6, 0.8, 0.0, 0.0], + layout, + distribution.with_normalized(true), + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + close( + unit.evaluate_similarity(Opaque::new(&code)).unwrap(), + manual.evaluate_similarity(Opaque::new(&code)).unwrap(), + ); + let zero = plan + .fused_query_computer_with_hints( + &[0.0; 4], + layout, + distribution, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + let ordinary_zero = plan + .fused_query_computer( + &[0.0; 4], + layout, + false, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + close( + zero.evaluate_similarity(Opaque::new(&code)).unwrap(), + ordinary_zero + .evaluate_similarity(Opaque::new(&code)) + .unwrap(), + ); + } + } + } +} + +#[test] +fn query_rescale_matches_manual_rescale() { + check_rescale::<1>(); + check_rescale::<2>(); + check_rescale::<3>(); + check_rescale::<4>(); + check_rescale::<8>(); +} + +#[test] +fn converted_integer_inputs_keep_cosine_semantics() { + use diskann_vector::distance::Metric; + let integers = rowmajor::Owned::from_fn(2, 8, |rc| if rc.row == 0 { 0u8 } else { 254 }); + let quantizer = SphericalQuantizer::train( + integers.as_view(), + TransformKind::Null, + Metric::CosineNormalized.into(), + PreScale::ReciprocalMeanNorm, + &mut StdRng::seed_from_u64(7), + GlobalAllocator, + ) + .unwrap() + .as_quantizer::<4>() + .unwrap(); + let mut code = vec![0; quantizer.bytes()]; + quantizer + .compress( + &[254.0; 8], + OpaqueMut::new(&mut code), + ScopedAllocator::global(), + ) + .unwrap(); + let query = quantizer + .fused_query_computer( + &[127.0; 8], + QueryLayout::FullPrecision, + false, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + close(query.evaluate_similarity(Opaque::new(&code)).unwrap(), 0.0); +} + +#[cfg(feature = "flatbuffers")] +#[test] +fn hints_are_not_serialized() { + let normalized = InputHints::new().with_normalized(true); + let hinted = train(SupportedMetric::Cosine, normalized); + let ordinary = train(SupportedMetric::Cosine, InputHints::new()); + let plan = hinted.as_quantizer_impl::<4>().unwrap(); + let bytes = plan.serialize(GlobalAllocator).unwrap(); + let ordinary = ordinary.as_quantizer_impl::<4>().unwrap(); + assert_eq!(&*bytes, &*ordinary.serialize(GlobalAllocator).unwrap()); + let restored = iface::Impl::<4>::try_deserialize(&bytes, GlobalAllocator).unwrap(); + let restored: &dyn Quantizer = &restored; + assert_eq!(restored.metric(), SupportedMetric::Cosine); + let mut code = vec![0; restored.bytes()]; + restored + .compress( + &[3.0, 0.0, 0.0, 0.0], + OpaqueMut::new(&mut code), + ScopedAllocator::global(), + ) + .unwrap(); + let query = restored + .fused_query_computer( + &[3.0, 4.0, 0.0, 0.0], + QueryLayout::FullPrecision, + false, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap(); + close(query.evaluate_similarity(Opaque::new(&code)).unwrap(), 0.4); +} diff --git a/diskann-quantization/src/spherical/iface.rs b/diskann-quantization/src/spherical/iface.rs index addeb70bef..892cc93749 100644 --- a/diskann-quantization/src/spherical/iface.rs +++ b/diskann-quantization/src/spherical/iface.rs @@ -153,8 +153,8 @@ use diskann_wide::arch::aarch64::Neon; use super::{ CompensatedCosine, CompensatedIP, CompensatedSquaredL2, Data, DataMut, DataRef, FullQuery, - FullQueryMut, FullQueryRef, Query, QueryMut, QueryRef, SphericalQuantizer, SupportedMetric, - quantizer, + FullQueryMut, FullQueryRef, InputHints, Query, QueryMut, QueryRef, SphericalQuantizer, + SupportedMetric, quantizer, }; use crate::{ AsFunctor, CompressIntoWith, @@ -291,14 +291,28 @@ where allow_rescale: bool, buffer: OpaqueMut<'_>, scratch: ScopedAllocator<'_>, + ) -> Result<(), QueryCompressionError> { + self.compress_query_with_hints( + x, + layout, + InputHints::new().with_maybe_out_of_distribution(allow_rescale), + buffer, + scratch, + ) + } + + /// Compress a query with properties of this input. Buffer and layout requirements are + /// the same as for [Self::compress_query]. + fn compress_query_with_hints( + &self, + x: &[f32], + layout: QueryLayout, + hints: InputHints, + buffer: OpaqueMut<'_>, + scratch: ScopedAllocator<'_>, ) -> Result<(), QueryCompressionError>; - /// Return a query for the argument `x` capable on operating on validly initialized - /// [`Opaque`] slices of length [`Self::bytes`]. - /// - /// These slices should be initialized by [`Self::compress`]. - /// - /// Note: Only layouts for which [`Self::is_supported`] returns `true` are supported. + /// Create a query computer. Only layouts accepted by [Self::is_supported] are supported. fn fused_query_computer( &self, x: &[f32], @@ -306,6 +320,25 @@ where allow_rescale: bool, allocator: A, scratch: ScopedAllocator<'_>, + ) -> Result, QueryComputerError> { + self.fused_query_computer_with_hints( + x, + layout, + InputHints::new().with_maybe_out_of_distribution(allow_rescale), + allocator, + scratch, + ) + } + + /// Create a query computer using properties of this query, independently of the hints + /// used to train the quantizer or compress database vectors. + fn fused_query_computer_with_hints( + &self, + x: &[f32], + layout: QueryLayout, + hints: InputHints, + allocator: A, + scratch: ScopedAllocator<'_>, ) -> Result, QueryComputerError>; /// Return whether or not this plan supports the given [`QueryLayout`]. @@ -322,6 +355,18 @@ where x: &[f32], into: OpaqueMut<'_>, scratch: ScopedAllocator<'_>, + ) -> Result<(), CompressionError> { + self.compress_with_hints(x, InputHints::new(), into, scratch) + } + + /// Compress stored data with input properties. The distribution hint is ignored: + /// independently rescaling stored vectors would change inner-product rankings. + fn compress_with_hints( + &self, + x: &[f32], + hints: InputHints, + into: OpaqueMut<'_>, + scratch: ScopedAllocator<'_>, ) -> Result<(), CompressionError>; /// Return the metric this plan was created with. @@ -1081,14 +1126,20 @@ impl Impl { fn compress_query<'a, T>( &self, query: &'a [f32], + hints: InputHints, storage: T, scratch: ScopedAllocator<'a>, ) -> Result<(), QueryCompressionError> where - SphericalQuantizer: CompressIntoWith<&'a [f32], T, ScopedAllocator<'a>, Error = quantizer::CompressionError>, + SphericalQuantizer: CompressIntoWith< + &'a [f32], + T, + (ScopedAllocator<'a>, InputHints), + Error = quantizer::CompressionError, + >, { self.quantizer - .compress_into_with(query, storage, scratch) + .compress_into_with_hints(query, storage, hints, scratch) .map_err(|err| CompressionError::CompressionError(err).into()) } @@ -1097,6 +1148,7 @@ impl Impl { fn fused_query_computer( &self, query: &[f32], + hints: InputHints, mut storage: T, allocator: B, scratch: ScopedAllocator<'_>, @@ -1114,14 +1166,14 @@ impl Impl { SphericalQuantizer: for<'a> CompressIntoWith< &'a [f32], >::Target, - ScopedAllocator<'a>, + (ScopedAllocator<'a>, InputHints), Error = quantizer::CompressionError, >, SphericalQuantizer: Dispatchable, { - if let Err(err) = self - .quantizer - .compress_into_with(query, storage.reborrow_mut(), scratch) + if let Err(err) = + self.quantizer + .compress_into_with_hints(query, storage.reborrow_mut(), hints, scratch) { return Err(CompressionError::CompressionError(err).into()); } @@ -1649,11 +1701,11 @@ where } } - fn compress_query( + fn compress_query_with_hints( &self, x: &[f32], layout: QueryLayout, - allow_rescale: bool, + hints: InputHints, mut buffer: OpaqueMut<'_>, scratch: ScopedAllocator<'_>, ) -> Result<(), QueryCompressionError> { @@ -1662,12 +1714,14 @@ where match layout { QueryLayout::SameAsData => self.compress_query( v, + hints, DataMut::<1>::from_canonical_back_mut(&mut buffer, dim) .map_err(NotCanonical::new)?, scratch, ), QueryLayout::FourBitTransposed => self.compress_query( v, + hints, QueryMut::<4, bits::BitTranspose>::from_canonical_back_mut(&mut buffer, dim) .map_err(NotCanonical::new)?, scratch, @@ -1677,6 +1731,7 @@ where } QueryLayout::FullPrecision => self.compress_query( v, + hints, FullQueryMut::from_canonical_mut(&mut buffer, dim) .map_err(NotCanonical::new)?, scratch, @@ -1684,20 +1739,14 @@ where } }; - if allow_rescale && self.quantizer.metric() == SupportedMetric::InnerProduct { - let mut copy = x.to_owned(); - self.quantizer.rescale(&mut copy); - finish(©) - } else { - finish(x) - } + finish(x) } - fn fused_query_computer( + fn fused_query_computer_with_hints( &self, x: &[f32], layout: QueryLayout, - allow_rescale: bool, + hints: InputHints, allocator: B, scratch: ScopedAllocator<'_>, ) -> Result, QueryComputerError> { @@ -1705,14 +1754,14 @@ where let finish = |v: &[f32], allocator: B| -> Result, QueryComputerError> { match layout { QueryLayout::SameAsData => self.fused_query_computer::, Data<1, _>, _>( - v, + v, hints, Data::new_in(dim, allocator.clone())?, allocator, scratch, ), QueryLayout::FourBitTransposed => self .fused_query_computer::, Query<4, bits::BitTranspose, _>, _>( - v, + v, hints, Query::new_in(dim, allocator.clone())?, allocator, scratch, @@ -1721,7 +1770,7 @@ where Err(UnsupportedQueryLayout::new(layout, "1-bit compression").into()) } QueryLayout::FullPrecision => self.fused_query_computer::, _>( - v, + v, hints, FullQuery::empty(dim, allocator.clone())?, allocator, scratch, @@ -1729,22 +1778,17 @@ where } }; - if allow_rescale && self.quantizer.metric() == SupportedMetric::InnerProduct { - let mut copy = x.to_owned(); - self.quantizer.rescale(&mut copy); - finish(©, allocator) - } else { - finish(x, allocator) - } + finish(x, allocator) } fn is_supported(&self, layout: QueryLayout) -> bool { Self::supports(layout) } - fn compress( + fn compress_with_hints( &self, x: &[f32], + hints: InputHints, mut into: OpaqueMut<'_>, scratch: ScopedAllocator<'_>, ) -> Result<(), CompressionError> { @@ -1752,7 +1796,12 @@ where let into = DataMut::<1>::from_canonical_back_mut(into.inspect(), dim) .map_err(CompressionError::not_canonical)?; self.quantizer - .compress_into_with(x, into, scratch) + .compress_into_with_hints( + x, + into, + hints.with_maybe_out_of_distribution(false), + scratch, + ) .map_err(CompressionError::CompressionError) } @@ -1858,11 +1907,11 @@ macro_rules! plan { } } - fn compress_query( + fn compress_query_with_hints( &self, x: &[f32], layout: QueryLayout, - allow_rescale: bool, + hints: InputHints, mut buffer: OpaqueMut<'_>, scratch: ScopedAllocator<'_>, ) -> Result<(), QueryCompressionError> { @@ -1870,7 +1919,7 @@ macro_rules! plan { let mut finish = |v: &[f32]| -> Result<(), QueryCompressionError> { match layout { QueryLayout::SameAsData => self.compress_query( - v, + v, hints, DataMut::<$N>::from_canonical_back_mut( &mut buffer, dim, @@ -1884,7 +1933,7 @@ macro_rules! plan { ).into()) }, QueryLayout::ScalarQuantized => self.compress_query( - v, + v, hints, QueryMut::<$N, bits::Dense>::from_canonical_back_mut( &mut buffer, dim, @@ -1892,7 +1941,7 @@ macro_rules! plan { scratch, ), QueryLayout::FullPrecision => self.compress_query( - v, + v, hints, FullQueryMut::from_canonical_mut( &mut buffer, dim, @@ -1902,20 +1951,14 @@ macro_rules! plan { } }; - if allow_rescale && self.quantizer.metric() == SupportedMetric::InnerProduct { - let mut copy = x.to_owned(); - self.quantizer.rescale(&mut copy); - finish(©) - } else { - finish(x) - } + finish(x) } - fn fused_query_computer( + fn fused_query_computer_with_hints( &self, x: &[f32], layout: QueryLayout, - allow_rescale: bool, + hints: InputHints, allocator: B, scratch: ScopedAllocator<'_>, ) -> Result, QueryComputerError> @@ -1925,7 +1968,7 @@ macro_rules! plan { match layout { QueryLayout::SameAsData => { self.fused_query_computer::, Data<$N, _>, B>( - v, + v, hints, Data::new_in(dim, allocator.clone())?, allocator, scratch, @@ -1939,7 +1982,7 @@ macro_rules! plan { }, QueryLayout::ScalarQuantized => { self.fused_query_computer::, Query<$N, bits::Dense, _>, B>( - v, + v, hints, Query::new_in(dim, allocator.clone())?, allocator, scratch, @@ -1947,7 +1990,7 @@ macro_rules! plan { }, QueryLayout::FullPrecision => { self.fused_query_computer::, B>( - v, + v, hints, FullQuery::empty(dim, allocator.clone())?, allocator, scratch, @@ -1956,23 +1999,17 @@ macro_rules! plan { } }; - let metric = >::metric(self); - if allow_rescale && metric == SupportedMetric::InnerProduct { - let mut copy = x.to_owned(); - self.quantizer.rescale(&mut copy); - finish(©) - } else { - finish(x) - } + finish(x) } fn is_supported(&self, layout: QueryLayout) -> bool { Self::supports(layout) } - fn compress( + fn compress_with_hints( &self, x: &[f32], + hints: InputHints, mut into: OpaqueMut<'_>, scratch: ScopedAllocator<'_>, ) -> Result<(), CompressionError> { @@ -1980,7 +2017,7 @@ macro_rules! plan { let into = DataMut::<$N>::from_canonical_back_mut(into.inspect(), dim) .map_err(CompressionError::not_canonical)?; - self.quantizer.compress_into_with(x, into, scratch) + self.quantizer.compress_into_with_hints(x, into, hints.with_maybe_out_of_distribution(false), scratch) .map_err(CompressionError::CompressionError) } diff --git a/diskann-quantization/src/spherical/mod.rs b/diskann-quantization/src/spherical/mod.rs index 58bd8e3f6f..549d48e98c 100644 --- a/diskann-quantization/src/spherical/mod.rs +++ b/diskann-quantization/src/spherical/mod.rs @@ -3,10 +3,16 @@ * Licensed under the MIT license. */ +mod hints; +pub use hints::InputHints; + pub mod iface; pub(crate) mod quantizer; mod vectors; +#[cfg(test)] +mod hints_tests; + #[doc(hidden)] #[cfg(feature = "codegen")] pub mod __codegen; @@ -23,5 +29,5 @@ pub use vectors::InvalidMetric; pub use vectors::{ CompensatedCosine, CompensatedIP, CompensatedSquaredL2, Data, DataMeta, DataMetaError, DataMetaF32, DataMut, DataRef, FullQuery, FullQueryMeta, FullQueryMut, FullQueryRef, Query, - QueryMeta, QueryMut, QueryRef, SupportedMetric, UnsupportedMetric, + QueryMeta, QueryMut, QueryRef, SupportedMetric, }; diff --git a/diskann-quantization/src/spherical/quantizer.rs b/diskann-quantization/src/spherical/quantizer.rs index 397b20da01..91e346e26e 100644 --- a/diskann-quantization/src/spherical/quantizer.rs +++ b/diskann-quantization/src/spherical/quantizer.rs @@ -19,7 +19,7 @@ use thiserror::Error; use super::{ CompensatedCosine, CompensatedIP, CompensatedSquaredL2, DataMeta, DataMetaError, DataMut, - FullQueryMeta, FullQueryMut, QueryMeta, QueryMut, SupportedMetric, iface, + FullQueryMeta, FullQueryMut, InputHints, QueryMeta, QueryMut, SupportedMetric, iface, }; use crate::{ AsFunctor, CompressIntoWith, @@ -30,7 +30,9 @@ use crate::{ alloc::{Allocator, AllocatorError, GlobalAllocator, Poly, ScopedAllocator, TryClone}, bits::{PermutationStrategy, Representation, Unsigned}, num::Positive, - utils::{CannotBeEmpty, compute_means_and_average_norm, compute_normalized_means}, + utils::{ + CannotBeEmpty, compute_means, compute_means_and_average_norm, compute_normalized_means, + }, }; #[cfg(feature = "flatbuffers")] use crate::{ @@ -296,6 +298,33 @@ where rng: &mut R, allocator: A, ) -> Result + where + T: Copy + Into + Into, + R: Rng, + { + Self::train_with_hints( + data, + transform, + metric, + pre_scale, + InputHints::new(), + rng, + allocator, + ) + } + + /// Train with properties of the supplied rows. Only cosine uses the normalization hint. + /// + /// See [`Self::train`] for errors and the input/metric contract. + pub fn train_with_hints( + data: rowmajor::Ref, + transform: TransformKind, + metric: SupportedMetric, + pre_scale: PreScale, + hints: InputHints, + rng: &mut R, + allocator: A, + ) -> Result where T: Copy + Into + Into, R: Rng, @@ -308,6 +337,7 @@ where transform: TransformKind, metric: SupportedMetric, pre_scale: PreScale, + hints: InputHints, rng: &mut dyn RngCore, allocator: A, ) -> Result, TrainError> @@ -325,6 +355,11 @@ where SupportedMetric::SquaredL2 | SupportedMetric::InnerProduct => { compute_means_and_average_norm(data) } + SupportedMetric::Cosine if hints.normalized() => ( + compute_means(data) + .map_err(|_: CannotBeEmpty| TrainError::DataCannotBeEmpty)?, + 1.0, + ), SupportedMetric::Cosine => ( compute_normalized_means(data) .map_err(|_: CannotBeEmpty| TrainError::DataCannotBeEmpty)?, @@ -381,7 +416,7 @@ where ) } - train(data, transform, metric, pre_scale, rng, allocator) + train(data, transform, metric, pre_scale, hints, rng, allocator) } /// Rescale the argument `v` to be in the rough dynamic range of the training dataset. @@ -401,18 +436,33 @@ where fn preprocess<'a>( &self, data: &[f32], + hints: InputHints, allocator: ScopedAllocator<'a>, ) -> Result, CompressionError> { assert_eq!(data.len(), self.input_dim(), "Data dimension is incorrect."); - // Fold in pre-scaling with the potential norm corretion for cosine. - // - // NOTE: When we're computing Cosine Similarity, we normalize the vector. As such, - // the `pre_scale` parameter become irrelvant since it just gets normalized away. + // Keep the operation order of rescale followed by pre-scale, without allocating + // a separate copy of the query. Combining the factors can overflow unnecessarily. + let rescale = + if self.metric == SupportedMetric::InnerProduct && hints.maybe_out_of_distribution() { + let norm = if hints.normalized() { + 1.0 + } else { + FastL2Norm.evaluate(data) + }; + if norm == 0.0 { + None + } else { + Some(self.mean_norm.into_inner() / norm) + } + } else { + None + }; let scale = self.pre_scale.into_inner(); - let mul: f32 = match self.metric { + let mul = match self.metric { + SupportedMetric::Cosine if hints.normalized() => 1.0, SupportedMetric::Cosine => { - let norm: f32 = (FastL2Norm).evaluate(data); + let norm = FastL2Norm.evaluate(data); if norm == 0.0 { 1.0 } else { 1.0 / norm } } SupportedMetric::SquaredL2 | SupportedMetric::InnerProduct => scale, @@ -420,7 +470,13 @@ where // Center the vector and compute the squared norm of the shifted vector. let shifted = Poly::from_iter( - std::iter::zip(data.iter(), self.shift.iter()).map(|(&f, &s)| mul * f - s), + std::iter::zip(data.iter(), self.shift.iter()).map(|(&f, &s)| { + let f = match rescale { + Some(r) => f * r, + None => f, + }; + mul * f - s + }), allocator, )?; @@ -443,6 +499,28 @@ where }) } + /// Compress with per-input properties. For stored data, leave the distribution hint unset. + /// Set it only when encoding a query (including queries encoded with a data layout). + /// + /// Errors are the same as [`CompressIntoWith::compress_into_with`]. + pub fn compress_into_with_hints<'a, T>( + &self, + from: &'a [f32], + into: T, + hints: InputHints, + scratch: ScopedAllocator<'a>, + ) -> Result<(), CompressionError> + where + Self: CompressIntoWith< + &'a [f32], + T, + (ScopedAllocator<'a>, InputHints), + Error = CompressionError, + >, + { + self.compress_into_with(from, into, (scratch, hints)) + } + /// Construct an [`iface::Impl`] `self`. pub fn as_quantizer_impl( self, @@ -808,7 +886,25 @@ impl FinishCompressing for DataMut<'_, 8> { // Data Compression // ////////////////////// -impl CompressIntoWith<&[f32], FullQueryMut<'_>, ScopedAllocator<'_>> for SphericalQuantizer +impl<'a, A, T> CompressIntoWith<&'a [f32], T, ScopedAllocator<'a>> for SphericalQuantizer +where + A: Allocator, + Self: + CompressIntoWith<&'a [f32], T, (ScopedAllocator<'a>, InputHints), Error = CompressionError>, +{ + type Error = CompressionError; + fn compress_into_with( + &self, + from: &'a [f32], + into: T, + scratch: ScopedAllocator<'a>, + ) -> Result<(), Self::Error> { + self.compress_into_with_hints(from, into, InputHints::new(), scratch) + } +} + +impl CompressIntoWith<&[f32], FullQueryMut<'_>, (ScopedAllocator<'_>, InputHints)> + for SphericalQuantizer where A: Allocator, { @@ -830,13 +926,13 @@ where &self, from: &[f32], mut into: FullQueryMut<'_>, - allocator: ScopedAllocator<'_>, + (allocator, hints): (ScopedAllocator<'_>, InputHints), ) -> Result<(), Self::Error> { let input_dim = self.shift.len(); let output_dim = self.output_dim(); check_dims(input_dim, output_dim, from.len(), into.len())?; - let mut preprocessed = self.preprocess(from, allocator)?; + let mut preprocessed = self.preprocess(from, hints, allocator)?; // If the preprocessed norm is zero, then we tried to compress the center directly. // In this case, we can get the correct behavior by setting `into` to all zeros. @@ -885,7 +981,8 @@ where } } -impl CompressIntoWith<&[f32], DataMut<'_, NBITS>, ScopedAllocator<'_>> +impl + CompressIntoWith<&[f32], DataMut<'_, NBITS>, (ScopedAllocator<'_>, InputHints)> for SphericalQuantizer where A: Allocator, @@ -910,13 +1007,13 @@ where &self, from: &[f32], mut into: DataMut<'_, NBITS>, - allocator: ScopedAllocator<'_>, + (allocator, hints): (ScopedAllocator<'_>, InputHints), ) -> Result<(), Self::Error> { let input_dim = self.shift.len(); let output_dim = self.output_dim(); check_dims(input_dim, output_dim, from.len(), into.len())?; - let mut preprocessed = self.preprocess(from, allocator)?; + let mut preprocessed = self.preprocess(from, hints, allocator)?; if preprocessed.shifted_norm == 0.0 { into.set_meta(DataMeta::default()); @@ -1210,7 +1307,7 @@ fn maximize_cosine_similarity( /////////////////////// impl - CompressIntoWith<&[f32], QueryMut<'_, NBITS, Perm>, ScopedAllocator<'_>> + CompressIntoWith<&[f32], QueryMut<'_, NBITS, Perm>, (ScopedAllocator<'_>, InputHints)> for SphericalQuantizer where Unsigned: Representation, @@ -1235,13 +1332,13 @@ where &self, from: &[f32], mut into: QueryMut<'_, NBITS, Perm>, - allocator: ScopedAllocator<'_>, + (allocator, hints): (ScopedAllocator<'_>, InputHints), ) -> Result<(), Self::Error> { let input_dim = self.shift.len(); let output_dim = self.output_dim(); check_dims(input_dim, output_dim, from.len(), into.len())?; - let mut preprocessed = self.preprocess(from, allocator)?; + let mut preprocessed = self.preprocess(from, hints, allocator)?; if preprocessed.shifted_norm == 0.0 { into.set_meta(QueryMeta::default()); @@ -1325,6 +1422,109 @@ where #[cfg(not(miri))] #[cfg(test)] mod tests { + #[test] + fn normalized_hints_skip_input_norm_only() { + use super::*; + let normalized = InputHints::new().with_normalized(true); + // Deliberately non-unit inputs distinguish the fast path from normalization. + let data = rowmajor::Owned::from_fn(2, 2, |rc| [[2.0f32, 0.0], [0.0, 4.0]][rc.row][rc.col]); + let quantizer = SphericalQuantizer::train_with_hints( + data.as_view(), + TransformKind::Null, + SupportedMetric::Cosine, + PreScale::Some(Positive::new(2.0).unwrap()), + normalized, + &mut StdRng::seed_from_u64(10), + GlobalAllocator, + ) + .unwrap(); + assert_eq!(quantizer.shift(), &[1.0, 2.0]); + assert_eq!(quantizer.mean_norm().into_inner(), 1.0); + let result = quantizer + .preprocess(&[3.0, 4.0], normalized, ScopedAllocator::global()) + .unwrap(); + assert_eq!(&*result.shifted, &[2.0, 2.0]); + assert!((result.shifted_norm - 8.0f32.sqrt()).abs() < 1e-6); + + let empty = rowmajor::Owned::from_element(0, 2, 0.0f32); + for hints in [InputHints::new(), normalized] { + let err = SphericalQuantizer::train_with_hints( + empty.as_view(), + TransformKind::Null, + SupportedMetric::Cosine, + PreScale::None, + hints, + &mut StdRng::seed_from_u64(10), + GlobalAllocator, + ) + .unwrap_err(); + assert!(matches!(err, TrainError::DataCannotBeEmpty)); + } + } + + #[test] + fn normalized_ip_query_uses_training_mean_and_prescale() { + use super::*; + let data = + rowmajor::Owned::from_fn(2, 2, |rc| [[3.0f32, 4.0], [-3.0, -4.0]][rc.row][rc.col]); + let quantizer = SphericalQuantizer::train( + data.as_view(), + TransformKind::Null, + SupportedMetric::InnerProduct, + PreScale::Some(Positive::new(0.5).unwrap()), + &mut StdRng::seed_from_u64(10), + GlobalAllocator, + ) + .unwrap(); + let hints = InputHints::new() + .with_normalized(true) + .with_maybe_out_of_distribution(true); + let result = quantizer + .preprocess(&[0.6, 0.8], hints, ScopedAllocator::global()) + .unwrap(); + assert_eq!(&*result.shifted, &[1.5, 2.0]); + let zero = quantizer + .preprocess(&[0.0; 2], hints, ScopedAllocator::global()) + .unwrap(); + assert_eq!(&*zero.shifted, &[0.0; 2]); + } + + #[test] + fn query_rescale_does_not_overflow_combined_scale() { + let data = rowmajor::Owned::from_fn(2, 2, |rc| { + if rc.col == 1 { + 0.0f32 + } else if rc.row == 0 { + 1.0e-10 + } else { + -1.0e-10 + } + }); + let quantizer = SphericalQuantizer::train( + data.as_view(), + TransformKind::Null, + SupportedMetric::InnerProduct, + PreScale::Some(Positive::new(1.0e29).unwrap()), + &mut StdRng::seed_from_u64(10), + GlobalAllocator, + ) + .unwrap(); + let raw = [1.0e-20, 0.0]; + let mut manual = raw; + quantizer.rescale(&mut manual); + let expected = quantizer + .preprocess(&manual, InputHints::new(), ScopedAllocator::global()) + .unwrap(); + let actual = quantizer + .preprocess( + &raw, + InputHints::new().with_maybe_out_of_distribution(true), + ScopedAllocator::global(), + ) + .unwrap(); + assert_eq!(&*actual.shifted, &*expected.shifted); + } + use super::*; use std::fmt::Display; @@ -1650,7 +1850,9 @@ mod tests { } = q.meta(); let shifted_square_norm = metric_specific; - let mut preprocessed = quantizer.preprocess(v, scoped_global).unwrap(); + let mut preprocessed = quantizer + .preprocess(v, InputHints::new(), scoped_global) + .unwrap(); preprocessed .shifted .iter_mut() @@ -1899,7 +2101,9 @@ mod tests { } = q.meta(); let inner_product_with_centroid = metric_specific; - let mut preprocessed = quantizer.preprocess(v, scoped_global).unwrap(); + let mut preprocessed = quantizer + .preprocess(v, InputHints::new(), scoped_global) + .unwrap(); preprocessed .shifted .iter_mut() @@ -2140,7 +2344,9 @@ mod tests { } = q.meta(); let inner_product_with_centroid = metric_specific; - let mut preprocessed = quantizer.preprocess(v, scoped_global).unwrap(); + let mut preprocessed = quantizer + .preprocess(v, InputHints::new(), scoped_global) + .unwrap(); preprocessed .shifted .iter_mut() diff --git a/diskann-quantization/src/spherical/vectors.rs b/diskann-quantization/src/spherical/vectors.rs index 3244d107dc..5d7903a452 100644 --- a/diskann-quantization/src/spherical/vectors.rs +++ b/diskann-quantization/src/spherical/vectors.rs @@ -153,32 +153,23 @@ impl SupportedMetric { } } -impl TryFrom for SupportedMetric { - type Error = UnsupportedMetric; - fn try_from(metric: diskann_vector::distance::Metric) -> Result { +impl From for SupportedMetric { + fn from(metric: diskann_vector::distance::Metric) -> Self { use diskann_vector::distance::Metric; match metric { - Metric::L2 => Ok(Self::SquaredL2), - Metric::InnerProduct => Ok(Self::InnerProduct), - Metric::Cosine => Ok(Self::Cosine), - unsupported => Err(UnsupportedMetric(unsupported)), + Metric::L2 => Self::SquaredL2, + Metric::InnerProduct => Self::InnerProduct, + Metric::Cosine | Metric::CosineNormalized => Self::Cosine, } } } impl PartialEq for SupportedMetric { fn eq(&self, metric: &diskann_vector::distance::Metric) -> bool { - match Self::try_from(*metric) { - Ok(m) => *self == m, - Err(_) => false, - } + *self == Self::from(*metric) } } -#[derive(Debug, Clone, Copy, Error)] -#[error("metric {0:?} is not supported for spherical quantization")] -pub struct UnsupportedMetric(pub(crate) diskann_vector::distance::Metric); - #[cfg(feature = "flatbuffers")] #[cfg_attr(docsrs, doc(cfg(feature = "flatbuffers")))] #[derive(Debug, Clone, Copy, PartialEq, Error)] @@ -1003,31 +994,15 @@ mod tests { #[test] fn supported_metric() { - assert_eq!( - SupportedMetric::try_from(Metric::L2).unwrap(), - SupportedMetric::SquaredL2 - ); - assert_eq!( - SupportedMetric::try_from(Metric::InnerProduct).unwrap(), - SupportedMetric::InnerProduct - ); - assert_eq!( - SupportedMetric::try_from(Metric::Cosine).unwrap(), - SupportedMetric::Cosine - ); - assert!(matches!( - SupportedMetric::try_from(Metric::CosineNormalized), - Err(UnsupportedMetric(Metric::CosineNormalized)) - )); - - assert_eq!(SupportedMetric::SquaredL2, Metric::L2); - assert_ne!(SupportedMetric::SquaredL2, Metric::InnerProduct); - assert_ne!(SupportedMetric::SquaredL2, Metric::Cosine); - assert_ne!(SupportedMetric::SquaredL2, Metric::CosineNormalized); - - assert_ne!(SupportedMetric::InnerProduct, Metric::L2); - assert_eq!(SupportedMetric::InnerProduct, Metric::InnerProduct); - assert_ne!(SupportedMetric::SquaredL2, Metric::Cosine); + for (metric, expected) in [ + (Metric::L2, SupportedMetric::SquaredL2), + (Metric::InnerProduct, SupportedMetric::InnerProduct), + (Metric::Cosine, SupportedMetric::Cosine), + (Metric::CosineNormalized, SupportedMetric::Cosine), + ] { + assert_eq!(SupportedMetric::from(metric), expected); + assert_eq!(expected, metric); + } assert_ne!(SupportedMetric::SquaredL2, Metric::CosineNormalized); } diff --git a/diskann-quantization/src/utils.rs b/diskann-quantization/src/utils.rs index 4f4535acb6..e4b243de75 100644 --- a/diskann-quantization/src/utils.rs +++ b/diskann-quantization/src/utils.rs @@ -105,6 +105,16 @@ where #[error("data cannot be empty")] pub(crate) struct CannotBeEmpty; +/// Compute the mean of each column without normalizing the input rows. +/// +/// Returns an error if `data` has no rows. +pub(crate) fn compute_means(data: rowmajor::Ref) -> Result, CannotBeEmpty> +where + T: Into + Copy, +{ + compute_means_impl::(data) +} + /// Compute the mean of each column in `data` as well as the average norm. pub(crate) fn compute_means_and_average_norm(data: rowmajor::Ref) -> (Vec, f64) where @@ -140,8 +150,19 @@ where (means, mean_norm) } -/// Compute the mean of each column in `data` as well as the average norm. +/// Compute the mean of each column after normalizing each input row. +/// +/// Returns an error if `data` has no rows. pub(crate) fn compute_normalized_means(data: rowmajor::Ref) -> Result, CannotBeEmpty> +where + T: Into + Copy, +{ + compute_means_impl::(data) +} + +fn compute_means_impl( + data: rowmajor::Ref, +) -> Result, CannotBeEmpty> where T: Into + Copy, { @@ -152,8 +173,6 @@ where return Ok(Vec::new()); } - // Compute the centroid of the dataset as well as the sums of the norms of every - // element in the dataset. let mut means: Vec = vec![0.0; data.ncols()]; let square = |x: &T| -> f64 { @@ -162,8 +181,12 @@ where }; data.rows().for_each(|row| { - let norm = row.iter().map(square).sum::().sqrt(); - let inv_norm = if norm == 0.0 { 1.0 } else { 1.0 / norm }; + let inv_norm = if NORMALIZE { + let norm = row.iter().map(square).sum::().sqrt(); + if norm == 0.0 { 1.0 } else { 1.0 / norm } + } else { + 1.0 + }; // Accumulate this row into the means. std::iter::zip(means.iter_mut(), row.iter()).for_each(|(m, r)| { @@ -224,6 +247,7 @@ mod tests { let mut rng = StdRng::seed_from_u64(seed); let test_problem = create_test_problem(nrows, ncols, &mut rng); let (means, mean_norm) = compute_means_and_average_norm(test_problem.data.as_view()); + assert_eq!(compute_means(test_problem.data.as_view()).unwrap(), means); for (i, (&expected, &got)) in std::iter::zip(test_problem.means.iter(), means.iter()).enumerate() @@ -315,15 +339,20 @@ mod tests { } #[test] - fn test_normalized_means_corner_cases() { - // If the input data has no columns, the returned vector should be empty. - let data = rowmajor::Owned::from_element(10, 0, 1.0f32); - let means = compute_normalized_means(data.as_view()).unwrap(); - assert!(means.is_empty()); - - // If the data has no rows, an error should be returned. - let data = rowmajor::Owned::from_element(0, 10, 1.0f32); - let _: CannotBeEmpty = compute_normalized_means(data.as_view()).unwrap_err(); + fn test_means_corner_cases() { + for compute in [compute_means::, compute_normalized_means::] { + // If the input data has no columns, the returned vector should be empty. + let data = rowmajor::Owned::from_element(10, 0, 1.0f32); + assert!(compute(data.as_view()).unwrap().is_empty()); + + // If the data has no rows, an error should be returned. + let data = rowmajor::Owned::from_element(0, 10, 1.0f32); + let _: CannotBeEmpty = compute(data.as_view()).unwrap_err(); + + // Zero vectors should remain zero, including when normalization is enabled. + let data = rowmajor::Owned::from_element(2, 3, 0.0f32); + assert_eq!(compute(data.as_view()).unwrap(), vec![0.0; 3]); + } } #[test] From bd15a801f2073e73ce261f5c2b019cdf1aac413e Mon Sep 17 00:00:00 2001 From: "Xinyu Wen (from Dev Box)" Date: Thu, 29 Oct 2026 17:00:06 +0800 Subject: [PATCH 2/2] [diskann-inmem] Use infallible spherical metric conversion --- diskann-inmem/integration/index/runner.rs | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs index d66aea4c07..d4dbf8d82e 100644 --- a/diskann-inmem/integration/index/runner.rs +++ b/diskann-inmem/integration/index/runner.rs @@ -645,12 +645,7 @@ impl Test { transforms::TransformKind::DoubleHadamard { target_dim: transforms::TargetDim::Same, }, - self.data.metric.try_into().map_err(|_| { - anyhow::anyhow!( - "metric {} not supported for spherical quantization", - self.data.metric - ) - })?, + self.data.metric.into(), PreScale::ReciprocalMeanNorm, &mut rand::rngs::StdRng::seed_from_u64(0xc0ff33), GlobalAllocator,