From 77302ad8b6613c8425e5939072cd2ad0dd377c56 Mon Sep 17 00:00:00 2001 From: "Xinyu Wen (from Dev Box)" Date: Mon, 12 Oct 2026 16:07:56 +0800 Subject: [PATCH 1/3] [diskann-quantization] Optimize spherical encoding scale search and transforms --- .../algorithms/transforms/double_hadamard.rs | 89 ++++++++-- .../src/spherical/quantizer.rs | 166 +++++++++++++++++- 2 files changed, 231 insertions(+), 24 deletions(-) diff --git a/diskann-quantization/src/algorithms/transforms/double_hadamard.rs b/diskann-quantization/src/algorithms/transforms/double_hadamard.rs index bb84e10880..75753e9ab2 100644 --- a/diskann-quantization/src/algorithms/transforms/double_hadamard.rs +++ b/diskann-quantization/src/algorithms/transforms/double_hadamard.rs @@ -243,9 +243,23 @@ where ) -> Result<(), TransformFailed> { check_dims(dst, src, self.input_dim(), self.output_dim())?; - // Copy and flip signs + if let Some(indices) = self.subsample.as_ref() { + let mut tmp = Poly::broadcast(0.0f32, self.intermediate_dim(), allocator)?; + self.transform_full(&mut tmp, src); + let rescale = ((tmp.len() as f32) / (indices.len() as f32)).sqrt(); + dst.iter_mut() + .zip(indices.iter()) + .for_each(|(d, s)| *d = tmp[*s as usize] * rescale); + } else { + // The destination already has room for the full transform. + self.transform_full(dst, src); + } + Ok(()) + } + + fn transform_full(&self, tmp: &mut [f32], src: &[f32]) { let intermediate_dim = self.intermediate_dim(); - let mut tmp = Poly::broadcast(0.0f32, intermediate_dim, allocator)?; + tmp[src.len()..].fill(0.0); std::iter::zip(tmp.iter_mut(), src.iter()) .zip(self.signs0.iter()) @@ -269,21 +283,6 @@ where #[expect(clippy::unwrap_used)] hadamard_transform(&mut tmp[intermediate_dim - split..]).unwrap(); - - match self.subsample.as_ref() { - None => { - dst.copy_from_slice(&tmp); - } - Some(indices) => { - let rescale = ((tmp.len() as f32) / (indices.len() as f32)).sqrt(); - debug_assert_eq!(dst.len(), indices.len()); - dst.iter_mut() - .zip(indices.iter()) - .for_each(|(d, s)| *d = tmp[*s as usize] * rescale); - } - } - - Ok(()) } } @@ -397,6 +396,62 @@ mod tests { test_utils::delegate_transformer!(DoubleHadamard); + // Original allocation/copy pipeline, kept independent of transform_full so + // regressions in sign application, padding, or transform order are detected. + fn reference_transform(t: &DoubleHadamard, src: &[f32]) -> Vec { + let mut tmp = vec![0.0; t.output_dim()]; + for i in 0..src.len() { + tmp[i] = f32::from_bits(src[i].to_bits() ^ t.signs0[i]); + } + let split = 1usize << (usize::BITS - tmp.len().leading_zeros() - 1); + hadamard_transform(&mut tmp[..split]).unwrap(); + for (v, sign) in tmp.iter_mut().zip(t.signs1.iter()) { + *v = f32::from_bits(v.to_bits() ^ sign); + } + let tail = tmp.len() - split; + hadamard_transform(&mut tmp[tail..]).unwrap(); + tmp + } + + #[test] + fn transform_matches_original_with_reused_destination() { + let mut rng = StdRng::seed_from_u64(742); + for (input, output) in [(7, 7), (7, 64), (64, 64)] { + let transform = DoubleHadamard::new( + NonZeroUsize::new(input).unwrap(), + TargetDim::Override(NonZeroUsize::new(output).unwrap()), + &mut rng, + GlobalAllocator, + ) + .unwrap(); + // Offset the destination and overwrite it twice, including padded values. + let mut dst = vec![f32::NAN; output + 2]; + for zero in [false, true] { + let src: Vec<_> = (0..input) + .map(|i| { + if zero { + if i % 2 == 0 { 0.0 } else { -0.0 } + } else { + (i as f32 - input as f32 / 2.0) * 0.25 + } + }) + .collect(); + let reference = reference_transform(&transform, &src); + transform + .transform_into( + &mut dst[1..=output], + &src, + ScopedAllocator::new(&crate::test_util::AlwaysFails), + ) + .unwrap(); + for (a, b) in dst[1..=output].iter().zip(reference.iter()) { + assert_eq!(a.to_bits(), b.to_bits()); + } + assert!(dst[0].is_nan() && dst[output + 1].is_nan()); + } + } + } + #[test] fn test_double_hadamard() { // Inner product computations are more susceptible to floating point error. diff --git a/diskann-quantization/src/spherical/quantizer.rs b/diskann-quantization/src/spherical/quantizer.rs index dcb1072bce..b7f4c5284b 100644 --- a/diskann-quantization/src/spherical/quantizer.rs +++ b/diskann-quantization/src/spherical/quantizer.rs @@ -944,25 +944,31 @@ where // Additionally, ordering is reverse so that `BinaryHeap` models a min-heap instead // of a max-heap. #[derive(Debug, Clone, Copy)] -struct Pair { +struct Pair { value: f32, position: u32, } -impl PartialEq for Pair { +impl PartialEq for Pair { fn eq(&self, other: &Self) -> bool { self.value.eq(&other.value) } } -impl Eq for Pair {} -impl PartialOrd for Pair { +impl Eq for Pair {} +impl PartialOrd for Pair { fn partial_cmp(&self, other: &Self) -> Option { Some(self.cmp(other)) } } -impl Ord for Pair { +impl Ord for Pair { fn cmp(&self, other: &Self) -> std::cmp::Ordering { + if INTEGER_ORDER { + // Keys are positive_numerator / abs(component), or f32::MAX. + // With non-NaN inputs they have a clear sign bit (including +0 and + // +infinity), so unsigned bit order exactly matches float order. + return other.value.to_bits().cmp(&self.value.to_bits()); + } other .value .partial_cmp(&self.value) @@ -1010,7 +1016,23 @@ fn maximize_cosine_similarity( // This means the initial inner product between `v` and the rounded term is simply // `0.5 * sum(abs.(v))`. The absolute value is used because the latice element is // always in the direction of the components in `v`. - let mut current_ip = 0.5 * v.iter().map(|i| i.abs() as f64).sum::(); + let current_ip = 0.5 * v.iter().map(|i| i.abs() as f64).sum::(); + // The sum of nonnegative magnitudes is NaN exactly when an input is NaN. + // Reuse that result to select the comparison once, outside the heap loop. + // Preserve unordered-as-equal behavior for exceptional transform outputs. + if current_ip.is_nan() { + maximize_cosine_similarity_inner::(v, num_bits, allocator, current_ip) + } else { + maximize_cosine_similarity_inner::(v, num_bits, allocator, current_ip) + } +} + +fn maximize_cosine_similarity_inner( + v: &[f32], + num_bits: NonZeroUsize, + allocator: ScopedAllocator<'_>, + mut current_ip: f64, +) -> Result { let mut current_square_norm = 0.25 * (v.len() as f64); // Book keeping for the current value of the rounded vector. @@ -1028,7 +1050,7 @@ fn maximize_cosine_similarity( let mut base = Poly::from_iter( v.iter().enumerate().map(|(position, value)| { let value = one_and_change / value.abs(); - Pair { + Pair:: { value, position: position as u32, } @@ -1255,6 +1277,136 @@ mod tests { }; // Test cosine-similarity maximizer + #[test] + fn critical_value_order_matches_float_heap() { + let mut rng = StdRng::seed_from_u64(0x6e12_9024); + // Include ties, subnormals, both sides of exponent boundaries, and the + // sentinel relative to infinity. Positions must not break value ties. + let mut values = vec![ + 0.0, + f32::from_bits(1), + f32::MIN_POSITIVE, + f32::from_bits(1.0f32.to_bits() - 1), + 1.0, + 1.0, + f32::from_bits(1.0f32.to_bits() + 1), + f32::MAX, + f32::INFINITY, + ]; + values.extend( + (0..256).map(|_| f32::from_bits(rng.random_range(0..=f32::INFINITY.to_bits()))), + ); + for &a in &values { + for &b in &values { + let actual = Pair:: { + value: a, + position: 0, + } + .cmp(&Pair:: { + value: b, + position: 1, + }); + assert_eq!(actual, b.partial_cmp(&a).unwrap()); + } + } + let mut actual: Vec<_> = values + .iter() + .enumerate() + .map(|(i, &value)| Pair:: { + value, + position: i as u32, + }) + .collect(); + let mut expected: Vec<_> = values + .iter() + .enumerate() + .map(|(i, &value)| Pair:: { + value, + position: i as u32, + }) + .collect(); + let mut actual = SliceHeap::new(&mut actual).unwrap(); + let mut expected = SliceHeap::new(&mut expected).unwrap(); + for step in 0..=512 { + for (a, b) in actual.as_slice().iter().zip(expected.as_slice()) { + assert_eq!( + (a.value.to_bits(), a.position), + (b.value.to_bits(), b.position) + ); + } + if step < 512 { + let value = values[rng.random_range(0..values.len())]; + actual.update_root(|p| p.value = value); + expected.update_root(|p| p.value = value); + } + } + } + + #[test] + fn integer_order_preserves_scale_search() { + let mut rng = StdRng::seed_from_u64(0x8ad2_9745); + for dim in [1, 2, 7, 64, 65, 128, 255, 768, 1536] { + let random: Vec = (0..dim).map(|_| StandardNormal.sample(&mut rng)).collect(); + let mut inputs = vec![ + random, + vec![0.0; dim], + vec![-0.0; dim], + vec![1.0; dim], + (0..dim) + .map(|i| if i % 2 == 0 { -0.5 } else { 0.5 }) + .collect(), + (0..dim).map(|i| f32::from_bits((i % 16) as u32)).collect(), + ]; + for special in [ + f32::MAX, + f32::INFINITY, + f32::NEG_INFINITY, + f32::NAN, + f32::from_bits(0xffc0_0001), + ] { + let mut v = inputs[0].clone(); + v[dim / 2] = special; + inputs.push(v); + } + for bits in [2, 4, 8] { + let bits = NonZeroUsize::new(bits).unwrap(); + for v in &inputs { + let initial_ip = 0.5 * v.iter().map(|i| i.abs() as f64).sum::(); + let expected = maximize_cosine_similarity_inner::( + v, + bits, + ScopedAllocator::global(), + initial_ip, + ) + .unwrap(); + let actual = + maximize_cosine_similarity(v, bits, ScopedAllocator::global()).unwrap(); + assert_eq!( + actual.to_bits(), + expected.to_bits(), + "dim={dim}, bits={bits}" + ); + } + } + } + // The comparison dispatch must preserve allocation failure at either of + // the two scratch allocations, including on the unordered path. + for v in [[1.0, 0.5], [f32::NAN, 0.5]] { + for budget in 0..=2 { + let result = maximize_cosine_similarity( + &v, + NonZeroUsize::new(4).unwrap(), + ScopedAllocator::new(&test_util::LimitedAllocator::new(budget)), + ); + if budget < 2 { + assert_eq!(result, Err(AllocatorError)); + } else { + assert!(result.is_ok()); + } + } + } + } + #[test] fn test_cosine_similarity_maximizer() { let mut rng = StdRng::seed_from_u64(0x070d9ff8cf5e0f8c); From 9bddb31cbe809905561eaf0ea0fb0232312c8333 Mon Sep 17 00:00:00 2001 From: "Xinyu Wen (from Dev Box)" Date: Tue, 13 Oct 2026 10:41:44 +0800 Subject: [PATCH 2/3] [diskann-quantization] Simplify spherical scale search and strengthen tests --- .../algorithms/transforms/double_hadamard.rs | 93 +++----- .../src/spherical/quantizer.rs | 217 +++++------------- 2 files changed, 96 insertions(+), 214 deletions(-) diff --git a/diskann-quantization/src/algorithms/transforms/double_hadamard.rs b/diskann-quantization/src/algorithms/transforms/double_hadamard.rs index 75753e9ab2..fe3b16c10e 100644 --- a/diskann-quantization/src/algorithms/transforms/double_hadamard.rs +++ b/diskann-quantization/src/algorithms/transforms/double_hadamard.rs @@ -252,6 +252,7 @@ where .for_each(|(d, s)| *d = tmp[*s as usize] * rescale); } else { // The destination already has room for the full transform. + dst[src.len()..].fill(0.0); self.transform_full(dst, src); } Ok(()) @@ -259,7 +260,6 @@ where fn transform_full(&self, tmp: &mut [f32], src: &[f32]) { let intermediate_dim = self.intermediate_dim(); - tmp[src.len()..].fill(0.0); std::iter::zip(tmp.iter_mut(), src.iter()) .zip(self.signs0.iter()) @@ -391,67 +391,11 @@ mod tests { use crate::{ algorithms::transforms::{Transform, TransformKind, test_utils}, alloc::GlobalAllocator, - test_util::Check, + test_util::{AlwaysFails, Check, LimitedAllocator}, }; test_utils::delegate_transformer!(DoubleHadamard); - // Original allocation/copy pipeline, kept independent of transform_full so - // regressions in sign application, padding, or transform order are detected. - fn reference_transform(t: &DoubleHadamard, src: &[f32]) -> Vec { - let mut tmp = vec![0.0; t.output_dim()]; - for i in 0..src.len() { - tmp[i] = f32::from_bits(src[i].to_bits() ^ t.signs0[i]); - } - let split = 1usize << (usize::BITS - tmp.len().leading_zeros() - 1); - hadamard_transform(&mut tmp[..split]).unwrap(); - for (v, sign) in tmp.iter_mut().zip(t.signs1.iter()) { - *v = f32::from_bits(v.to_bits() ^ sign); - } - let tail = tmp.len() - split; - hadamard_transform(&mut tmp[tail..]).unwrap(); - tmp - } - - #[test] - fn transform_matches_original_with_reused_destination() { - let mut rng = StdRng::seed_from_u64(742); - for (input, output) in [(7, 7), (7, 64), (64, 64)] { - let transform = DoubleHadamard::new( - NonZeroUsize::new(input).unwrap(), - TargetDim::Override(NonZeroUsize::new(output).unwrap()), - &mut rng, - GlobalAllocator, - ) - .unwrap(); - // Offset the destination and overwrite it twice, including padded values. - let mut dst = vec![f32::NAN; output + 2]; - for zero in [false, true] { - let src: Vec<_> = (0..input) - .map(|i| { - if zero { - if i % 2 == 0 { 0.0 } else { -0.0 } - } else { - (i as f32 - input as f32 / 2.0) * 0.25 - } - }) - .collect(); - let reference = reference_transform(&transform, &src); - transform - .transform_into( - &mut dst[1..=output], - &src, - ScopedAllocator::new(&crate::test_util::AlwaysFails), - ) - .unwrap(); - for (a, b) in dst[1..=output].iter().zip(reference.iter()) { - assert_eq!(a.to_bits(), b.to_bits()); - } - assert!(dst[0].is_nan() && dst[output + 1].is_nan()); - } - } - } - #[test] fn test_double_hadamard() { // Inner product computations are more susceptible to floating point error. @@ -540,7 +484,38 @@ mod tests { &mut checker, &mut rng, ctx, - ) + ); + + let mut src = vec![1.0; input]; + let mut dst = vec![f32::NAN; output + 2]; + for value in [1.0, 0.0, -0.0] { + src.fill(value); + let result = transformer.transform_into( + &mut dst[1..=output], + &src, + ScopedAllocator::new(&AlwaysFails), + ); + if preserves_norms { + result.unwrap(); + } else { + assert_eq!( + result, + Err(TransformFailed::AllocatorError(AllocatorError)) + ); + transformer + .transform_into( + &mut dst[1..=output], + &src, + ScopedAllocator::new(&LimitedAllocator::new(1)), + ) + .unwrap(); + } + assert!(dst[1..=output].iter().all(|value| value.is_finite())); + if value == 0.0 { + assert!(dst[1..=output].iter().all(|value| *value == 0.0)); + } + assert!(dst[0].is_nan() && dst[output + 1].is_nan()); + } } // Abstract Transformer diff --git a/diskann-quantization/src/spherical/quantizer.rs b/diskann-quantization/src/spherical/quantizer.rs index b7f4c5284b..a9ee2856ee 100644 --- a/diskann-quantization/src/spherical/quantizer.rs +++ b/diskann-quantization/src/spherical/quantizer.rs @@ -943,36 +943,28 @@ where // 2. Imposes a total ordering on `f32` values so it can be used in a `BinaryHeap`. // Additionally, ordering is reverse so that `BinaryHeap` models a min-heap instead // of a max-heap. +/// Critical values are nonnegative and non-NaN, so unsigned bit order matches numeric order. #[derive(Debug, Clone, Copy)] -struct Pair { +struct Pair { value: f32, position: u32, } -impl PartialEq for Pair { +impl PartialEq for Pair { fn eq(&self, other: &Self) -> bool { - self.value.eq(&other.value) + self.value.to_bits().eq(&other.value.to_bits()) } } -impl Eq for Pair {} -impl PartialOrd for Pair { +impl Eq for Pair {} +impl PartialOrd for Pair { fn partial_cmp(&self, other: &Self) -> Option { Some(self.cmp(other)) } } -impl Ord for Pair { +impl Ord for Pair { fn cmp(&self, other: &Self) -> std::cmp::Ordering { - if INTEGER_ORDER { - // Keys are positive_numerator / abs(component), or f32::MAX. - // With non-NaN inputs they have a clear sign bit (including +0 and - // +infinity), so unsigned bit order exactly matches float order. - return other.value.to_bits().cmp(&self.value.to_bits()); - } - other - .value - .partial_cmp(&self.value) - .unwrap_or(std::cmp::Ordering::Equal) + other.value.to_bits().cmp(&self.value.to_bits()) } } @@ -998,6 +990,10 @@ impl Ord for Pair { /// /// Refer to algorithm 1 in . /// +/// # Errors +/// +/// Returns an error if `v` contains NaN or scratch allocation fails. +/// /// # Panics /// /// Panics is `v.is_empty()`. @@ -1011,28 +1007,15 @@ fn maximize_cosine_similarity( v: &[f32], num_bits: NonZeroUsize, allocator: ScopedAllocator<'_>, -) -> Result { +) -> Result { // Initially, the lattice element has the value `0.5` for all dimensions. // This means the initial inner product between `v` and the rounded term is simply // `0.5 * sum(abs.(v))`. The absolute value is used because the latice element is // always in the direction of the components in `v`. - let current_ip = 0.5 * v.iter().map(|i| i.abs() as f64).sum::(); - // The sum of nonnegative magnitudes is NaN exactly when an input is NaN. - // Reuse that result to select the comparison once, outside the heap loop. - // Preserve unordered-as-equal behavior for exceptional transform outputs. + let mut current_ip = 0.5 * v.iter().map(|i| i.abs() as f64).sum::(); if current_ip.is_nan() { - maximize_cosine_similarity_inner::(v, num_bits, allocator, current_ip) - } else { - maximize_cosine_similarity_inner::(v, num_bits, allocator, current_ip) + return Err(CompressionError::InputContainsNaN); } -} - -fn maximize_cosine_similarity_inner( - v: &[f32], - num_bits: NonZeroUsize, - allocator: ScopedAllocator<'_>, - mut current_ip: f64, -) -> Result { let mut current_square_norm = 0.25 * (v.len() as f64); // Book keeping for the current value of the rounded vector. @@ -1050,7 +1033,7 @@ fn maximize_cosine_similarity_inner( let mut base = Poly::from_iter( v.iter().enumerate().map(|(position, value)| { let value = one_and_change / value.abs(); - Pair:: { + Pair { value, position: position as u32, } @@ -1276,133 +1259,36 @@ mod tests { test_util, }; - // Test cosine-similarity maximizer #[test] - fn critical_value_order_matches_float_heap() { - let mut rng = StdRng::seed_from_u64(0x6e12_9024); - // Include ties, subnormals, both sides of exponent boundaries, and the - // sentinel relative to infinity. Positions must not break value ties. - let mut values = vec![ - 0.0, - f32::from_bits(1), - f32::MIN_POSITIVE, - f32::from_bits(1.0f32.to_bits() - 1), - 1.0, - 1.0, - f32::from_bits(1.0f32.to_bits() + 1), - f32::MAX, - f32::INFINITY, - ]; - values.extend( - (0..256).map(|_| f32::from_bits(rng.random_range(0..=f32::INFINITY.to_bits()))), - ); - for &a in &values { - for &b in &values { - let actual = Pair:: { - value: a, - position: 0, - } - .cmp(&Pair:: { - value: b, - position: 1, - }); - assert_eq!(actual, b.partial_cmp(&a).unwrap()); - } - } - let mut actual: Vec<_> = values - .iter() - .enumerate() - .map(|(i, &value)| Pair:: { - value, - position: i as u32, - }) - .collect(); - let mut expected: Vec<_> = values - .iter() - .enumerate() - .map(|(i, &value)| Pair:: { - value, - position: i as u32, - }) - .collect(); - let mut actual = SliceHeap::new(&mut actual).unwrap(); - let mut expected = SliceHeap::new(&mut expected).unwrap(); - for step in 0..=512 { - for (a, b) in actual.as_slice().iter().zip(expected.as_slice()) { + fn cosine_similarity_maximizer_errors() { + let num_bits = NonZeroUsize::new(4).unwrap(); + for nan in [f32::NAN, f32::from_bits(0xffc0_0001)] { + for position in 0..4 { + let mut target = [1.0; 4]; + target[position] = nan; assert_eq!( - (a.value.to_bits(), a.position), - (b.value.to_bits(), b.position) + maximize_cosine_similarity( + &target, + num_bits, + ScopedAllocator::new(&test_util::AlwaysFails), + ), + Err(CompressionError::InputContainsNaN), ); } - if step < 512 { - let value = values[rng.random_range(0..values.len())]; - actual.update_root(|p| p.value = value); - expected.update_root(|p| p.value = value); - } } - } - - #[test] - fn integer_order_preserves_scale_search() { - let mut rng = StdRng::seed_from_u64(0x8ad2_9745); - for dim in [1, 2, 7, 64, 65, 128, 255, 768, 1536] { - let random: Vec = (0..dim).map(|_| StandardNormal.sample(&mut rng)).collect(); - let mut inputs = vec![ - random, - vec![0.0; dim], - vec![-0.0; dim], - vec![1.0; dim], - (0..dim) - .map(|i| if i % 2 == 0 { -0.5 } else { 0.5 }) - .collect(), - (0..dim).map(|i| f32::from_bits((i % 16) as u32)).collect(), - ]; - for special in [ - f32::MAX, - f32::INFINITY, - f32::NEG_INFINITY, - f32::NAN, - f32::from_bits(0xffc0_0001), - ] { - let mut v = inputs[0].clone(); - v[dim / 2] = special; - inputs.push(v); - } - for bits in [2, 4, 8] { - let bits = NonZeroUsize::new(bits).unwrap(); - for v in &inputs { - let initial_ip = 0.5 * v.iter().map(|i| i.abs() as f64).sum::(); - let expected = maximize_cosine_similarity_inner::( - v, - bits, - ScopedAllocator::global(), - initial_ip, - ) - .unwrap(); - let actual = - maximize_cosine_similarity(v, bits, ScopedAllocator::global()).unwrap(); - assert_eq!( - actual.to_bits(), - expected.to_bits(), - "dim={dim}, bits={bits}" - ); - } - } - } - // The comparison dispatch must preserve allocation failure at either of - // the two scratch allocations, including on the unordered path. - for v in [[1.0, 0.5], [f32::NAN, 0.5]] { - for budget in 0..=2 { - let result = maximize_cosine_similarity( - &v, - NonZeroUsize::new(4).unwrap(), - ScopedAllocator::new(&test_util::LimitedAllocator::new(budget)), + for budget in 0..=2 { + let result = maximize_cosine_similarity( + &[1.0, 0.5], + num_bits, + ScopedAllocator::new(&test_util::LimitedAllocator::new(budget)), + ); + if budget < 2 { + assert_eq!( + result, + Err(CompressionError::AllocatorError(AllocatorError)) ); - if budget < 2 { - assert_eq!(result, Err(AllocatorError)); - } else { - assert!(result.is_ok()); - } + } else { + assert!(result.is_ok()); } } } @@ -1457,7 +1343,7 @@ mod tests { // We expect to either get the best value found via exhaustive search, or some // scalar multiple of it (since that will have the same cosine similarity). - let passed = if best == clamped { + let passed = if best == clamped || best_similarity == clamped_cosine.into_inner() { true } else { let ratio: Vec = std::iter::zip(best, clamped) @@ -1489,6 +1375,20 @@ mod tests { } }; + for target in [ + [1.0, 0.0, -0.0, -1.0], + [0.0, 0.5, -0.5, 0.0], + [1.0, 1.0, -1.0, -1.0], + [ + f32::from_bits(1.0f32.to_bits() - 1), + 1.0, + f32::from_bits(1.0f32.to_bits() + 1), + -1.0, + ], + ] { + run_test(target); + } + // Run targeted tests. let min = -3.5; for i0 in (0..8).step_by(2) { @@ -2575,6 +2475,7 @@ mod tests { { let mut query: Vec = quantizer.shift().to_vec(); let mut d = Data::<1, _>::new_boxed(quantizer.output_dim()); + let mut multibit_data = Data::<4, _>::new_boxed(quantizer.output_dim()); let mut q = Query::<4, BitTranspose, _>::new_boxed(quantizer.output_dim()); for i in 0..query.len() { @@ -2588,6 +2489,12 @@ mod tests { assert_eq!(err.to_string(), "input contains NaN", "failed for {}", v); + let err = quantizer + .compress_into_with(&*query, multibit_data.reborrow_mut(), scoped_global) + .unwrap_err(); + + assert_eq!(err, CompressionError::InputContainsNaN); + let err = quantizer .compress_into_with(&*query, q.reborrow_mut(), scoped_global) .unwrap_err(); From 0237d2633b9c360c9bc13b8d1414ee47440eef55 Mon Sep 17 00:00:00 2001 From: "Xinyu Wen (from Dev Box)" Date: Wed, 14 Oct 2026 15:39:36 +0800 Subject: [PATCH 3/3] [diskann-quantization] Refine spherical encoding review tests --- .../algorithms/transforms/double_hadamard.rs | 64 +++++++++---------- .../src/spherical/quantizer.rs | 13 ++-- diskann-quantization/src/test_util.rs | 4 ++ 3 files changed, 44 insertions(+), 37 deletions(-) diff --git a/diskann-quantization/src/algorithms/transforms/double_hadamard.rs b/diskann-quantization/src/algorithms/transforms/double_hadamard.rs index fe3b16c10e..50002b33b6 100644 --- a/diskann-quantization/src/algorithms/transforms/double_hadamard.rs +++ b/diskann-quantization/src/algorithms/transforms/double_hadamard.rs @@ -391,7 +391,7 @@ mod tests { use crate::{ algorithms::transforms::{Transform, TransformKind, test_utils}, alloc::GlobalAllocator, - test_util::{AlwaysFails, Check, LimitedAllocator}, + test_util::{Check, LimitedAllocator}, }; test_utils::delegate_transformer!(DoubleHadamard); @@ -485,37 +485,6 @@ mod tests { &mut rng, ctx, ); - - let mut src = vec![1.0; input]; - let mut dst = vec![f32::NAN; output + 2]; - for value in [1.0, 0.0, -0.0] { - src.fill(value); - let result = transformer.transform_into( - &mut dst[1..=output], - &src, - ScopedAllocator::new(&AlwaysFails), - ); - if preserves_norms { - result.unwrap(); - } else { - assert_eq!( - result, - Err(TransformFailed::AllocatorError(AllocatorError)) - ); - transformer - .transform_into( - &mut dst[1..=output], - &src, - ScopedAllocator::new(&LimitedAllocator::new(1)), - ) - .unwrap(); - } - assert!(dst[1..=output].iter().all(|value| value.is_finite())); - if value == 0.0 { - assert!(dst[1..=output].iter().all(|value| *value == 0.0)); - } - assert!(dst[0].is_nan() && dst[output + 1].is_nan()); - } } // Abstract Transformer @@ -538,6 +507,37 @@ mod tests { ) } } + + let mut allocation_rng = rng.clone(); + let transformer = + DoubleHadamard::new(input_nz, target, &mut allocation_rng, GlobalAllocator) + .unwrap(); + let mut src = vec![1.0; input]; + let mut dst = vec![f32::NAN; output + 2]; + for value in [1.0, 0.0, -0.0] { + src.fill(value); + if !preserves_norms { + let allocator = LimitedAllocator::new(0); + assert_eq!( + transformer.transform_into( + &mut dst[1..=output], + &src, + ScopedAllocator::new(&allocator), + ), + Err(TransformFailed::AllocatorError(AllocatorError)) + ); + } + let allocator = LimitedAllocator::new(usize::from(!preserves_norms)); + transformer + .transform_into(&mut dst[1..=output], &src, ScopedAllocator::new(&allocator)) + .unwrap(); + assert_eq!(allocator.remaining(), 0); + assert!(dst[1..=output].iter().all(|value| value.is_finite())); + if value == 0.0 { + assert!(dst[1..=output].iter().all(|value| *value == 0.0)); + } + assert!(dst[0].is_nan() && dst[output + 1].is_nan()); + } } } diff --git a/diskann-quantization/src/spherical/quantizer.rs b/diskann-quantization/src/spherical/quantizer.rs index a9ee2856ee..2a21be38f4 100644 --- a/diskann-quantization/src/spherical/quantizer.rs +++ b/diskann-quantization/src/spherical/quantizer.rs @@ -940,10 +940,9 @@ where // Informally, "what's the smallest scaling factor so `v[position]` gets rounded to // the next value. // -// 2. Imposes a total ordering on `f32` values so it can be used in a `BinaryHeap`. -// Additionally, ordering is reverse so that `BinaryHeap` models a min-heap instead -// of a max-heap. -/// Critical values are nonnegative and non-NaN, so unsigned bit order matches numeric order. +// 2. Imposes a total ordering on nonnegative, non-NaN critical values using unsigned +// bit order, which matches numeric order for these values. Additionally, ordering +// is reversed so that `BinaryHeap` models a min-heap instead of a max-heap. #[derive(Debug, Clone, Copy)] struct Pair { value: f32, @@ -1343,7 +1342,11 @@ mod tests { // We expect to either get the best value found via exhaustive search, or some // scalar multiple of it (since that will have the same cosine similarity). - let passed = if best == clamped || best_similarity == clamped_cosine.into_inner() { + let passed = if best == clamped + || (best_similarity == clamped_cosine.into_inner() + && std::iter::zip(target, std::iter::zip(best, clamped)) + .all(|(input, (expected, actual))| input == 0.0 || expected == actual)) + { true } else { let ratio: Vec = std::iter::zip(best, clamped) diff --git a/diskann-quantization/src/test_util.rs b/diskann-quantization/src/test_util.rs index d7f8f95cbb..b8abaa267e 100644 --- a/diskann-quantization/src/test_util.rs +++ b/diskann-quantization/src/test_util.rs @@ -53,6 +53,10 @@ impl LimitedAllocator { remaining: Arc::new(AtomicUsize::new(allocations)), } } + + pub(crate) fn remaining(&self) -> usize { + self.remaining.load(Ordering::Relaxed) + } } #[cfg(not(miri))]