[diskann-quantization] Optimize spherical encoding scale search

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.
+            dst[src.len()..].fill(0.0);
+            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)?;
 
         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(())
     }
 }
 
@@ -392,7 +391,7 @@ mod tests {
     use crate::{
         algorithms::transforms::{Transform, TransformKind, test_utils},
         alloc::GlobalAllocator,
-        test_util::Check,
+        test_util::{Check, LimitedAllocator},
     };
 
     test_utils::delegate_transformer!(DoubleHadamard<GlobalAllocator>);
@@ -485,7 +484,7 @@ mod tests {
                         &mut checker,
                         &mut rng,
                         ctx,
-                    )
+                    );
                 }
 
                 // Abstract Transformer
@@ -508,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());
+            }
         }
     }
 

diskann-quantization/src/spherical/quantizer.rs

@@ -940,9 +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.
+// 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,
@@ -951,7 +951,7 @@ struct 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())
     }
 }
 
@@ -963,10 +963,7 @@ impl PartialOrd for Pair {
 }
 impl Ord for Pair {
     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-        other
-            .value
-            .partial_cmp(&self.value)
-            .unwrap_or(std::cmp::Ordering::Equal)
+        other.value.to_bits().cmp(&self.value.to_bits())
     }
 }
 
@@ -992,6 +989,10 @@ impl Ord for Pair {
 ///
 /// Refer to algorithm 1 in <https://arxiv.org/pdf/2409.09913>.
 ///
+/// # Errors
+///
+/// Returns an error if `v` contains NaN or scratch allocation fails.
+///
 /// # Panics
 ///
 /// Panics is `v.is_empty()`.
@@ -1005,12 +1006,15 @@ fn maximize_cosine_similarity(
     v: &[f32],
     num_bits: NonZeroUsize,
     allocator: ScopedAllocator<'_>,
-) -> Result<f32, AllocatorError> {
+) -> Result<f32, CompressionError> {
     // 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 mut current_ip = 0.5 * v.iter().map(|i| i.abs() as f64).sum::<f64>();
+    if current_ip.is_nan() {
+        return Err(CompressionError::InputContainsNaN);
+    }
     let mut current_square_norm = 0.25 * (v.len() as f64);
 
     // Book keeping for the current value of the rounded vector.
@@ -1254,7 +1258,40 @@ mod tests {
         test_util,
     };
 
-    // Test cosine-similarity maximizer
+    #[test]
+    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!(
+                    maximize_cosine_similarity(
+                        &target,
+                        num_bits,
+                        ScopedAllocator::new(&test_util::AlwaysFails),
+                    ),
+                    Err(CompressionError::InputContainsNaN),
+                );
+            }
+        }
+        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))
+                );
+            } else {
+                assert!(result.is_ok());
+            }
+        }
+    }
+
     #[test]
     fn test_cosine_similarity_maximizer() {
         let mut rng = StdRng::seed_from_u64(0x070d9ff8cf5e0f8c);
@@ -1305,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 {
+            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<f32> = std::iter::zip(best, clamped)
@@ -1337,6 +1378,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) {
@@ -2423,6 +2478,7 @@ mod tests {
         {
             let mut query: Vec<f32> = 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() {
@@ -2436,6 +2492,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();

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))]