refactor[turboquant]: simplify code from review findings

- Consolidate encode_decode_mse and encode_decode_qjl test helpers into
  a single closure-parameterized encode_decode function
- Replace 14 copy-pasted benchmark functions (~200 lines) with a
  turboquant_bench! macro (~40 lines)
- Extract QJL correction scale factor to a named function with doc
  comment explaining the derivation
- Precompute centroid decision boundaries (midpoints) once before the
  row loop, replacing per-coordinate distance comparisons with a single
  partition_point lookup. This removes two abs() calls and a branch
  from the innermost quantization loop.

Net: -150 lines.

Signed-off-by: Will Manning <will@spiraldb.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Will Manning <will@willmanning.io>

Author: lwwmanning

encodings/turboquant/public-api.lock

@@ -2,7 +2,9 @@ pub mod vortex_turboquant
 
 pub mod vortex_turboquant::centroids
 
-pub fn vortex_turboquant::centroids::find_nearest_centroid(value: f32, centroids: &[f32]) -> u8
+pub fn vortex_turboquant::centroids::compute_boundaries(centroids: &[f32]) -> alloc::vec::Vec<f32>
+
+pub fn vortex_turboquant::centroids::find_nearest_centroid(value: f32, boundaries: &[f32]) -> u8
 
 pub fn vortex_turboquant::centroids::get_centroids(dimension: u32, bit_width: u8) -> vortex_error::VortexResult<alloc::vec::Vec<f32>>
 

encodings/turboquant/src/centroids.rs

@@ -147,32 +147,24 @@ fn pdf_unnormalized(x_val: f64, exponent: f64) -> f64 {
     base.powf(exponent)
 }
 
-/// Find the index of the nearest centroid to the given value.
+/// Precompute decision boundaries (midpoints between adjacent centroids).
 ///
-/// Centroids must be sorted in ascending order. Uses binary search for efficiency.
-#[inline]
-pub fn find_nearest_centroid(value: f32, centroids: &[f32]) -> u8 {
-    debug_assert!(!centroids.is_empty());
-
-    let idx = centroids.partition_point(|&c_val| c_val < value);
-
-    if idx == 0 {
-        return 0;
-    }
-    if idx >= centroids.len() {
-        #[allow(clippy::cast_possible_truncation)]
-        return (centroids.len() - 1) as u8;
-    }
-
-    let dist_left = (value - centroids[idx - 1]).abs();
-    let dist_right = (value - centroids[idx]).abs();
+/// For `k` centroids, returns `k-1` boundaries. A value below `boundaries[0]` maps
+/// to centroid 0, a value in `[boundaries[i-1], boundaries[i])` maps to centroid `i`,
+/// and a value >= `boundaries[k-2]` maps to centroid `k-1`.
+pub fn compute_boundaries(centroids: &[f32]) -> Vec<f32> {
+    centroids.windows(2).map(|w| (w[0] + w[1]) * 0.5).collect()
+}
 
-    #[allow(clippy::cast_possible_truncation)]
-    if dist_left <= dist_right {
-        (idx - 1) as u8
-    } else {
-        idx as u8
-    }
+/// Find the index of the nearest centroid using precomputed decision boundaries.
+///
+/// `boundaries` must be the output of [`compute_boundaries`] for the corresponding
+/// centroids. Uses binary search on the midpoints, avoiding distance comparisons
+/// in the inner loop.
+#[inline]
+#[allow(clippy::cast_possible_truncation)]
+pub fn find_nearest_centroid(value: f32, boundaries: &[f32]) -> u8 {
+    boundaries.partition_point(|&b| b < value) as u8
 }
 
 #[cfg(test)]
@@ -263,14 +255,15 @@ mod tests {
     #[test]
     fn find_nearest_basic() -> VortexResult<()> {
         let centroids = get_centroids(128, 2)?;
-        assert_eq!(find_nearest_centroid(-1.0, &centroids), 0);
+        let boundaries = compute_boundaries(&centroids);
+        assert_eq!(find_nearest_centroid(-1.0, &boundaries), 0);
         #[allow(clippy::cast_possible_truncation)]
         let last_idx = (centroids.len() - 1) as u8;
-        assert_eq!(find_nearest_centroid(1.0, &centroids), last_idx);
+        assert_eq!(find_nearest_centroid(1.0, &boundaries), last_idx);
         for (idx, &cv) in centroids.iter().enumerate() {
             #[allow(clippy::cast_possible_truncation)]
             let expected = idx as u8;
-            assert_eq!(find_nearest_centroid(cv, &centroids), expected);
+            assert_eq!(find_nearest_centroid(cv, &boundaries), expected);
         }
         Ok(())
     }

encodings/turboquant/src/compress.rs

@@ -17,6 +17,7 @@ use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
 use vortex_fastlanes::bitpack_compress::bitpack_encode;
 
+use crate::centroids::compute_boundaries;
 use crate::centroids::find_nearest_centroid;
 use crate::centroids::get_centroids;
 use crate::mse::array::TurboQuantMSEArray;
@@ -96,6 +97,7 @@ pub fn turboquant_encode_mse(
     let f32_elements = extract_f32_elements(fsl)?;
     #[allow(clippy::cast_possible_truncation)]
     let centroids = get_centroids(padded_dim as u32, config.bit_width)?;
+    let boundaries = compute_boundaries(&centroids);
 
     let mut all_indices = BufferMut::<u8>::with_capacity(num_rows * padded_dim);
     let mut norms_buf = BufferMut::<f32>::with_capacity(num_rows);
@@ -117,7 +119,7 @@ pub fn turboquant_encode_mse(
         rotation.rotate(&padded, &mut rotated);
 
         for j in 0..padded_dim {
-            all_indices.push(find_nearest_centroid(rotated[j], &centroids));
+            all_indices.push(find_nearest_centroid(rotated[j], &boundaries));
         }
     }
 
@@ -201,6 +203,7 @@ pub fn turboquant_encode_qjl(
     let f32_elements = extract_f32_elements(fsl)?;
     #[allow(clippy::cast_possible_truncation)]
     let centroids = get_centroids(padded_dim as u32, mse_bit_width)?;
+    let boundaries = compute_boundaries(&centroids);
 
     // QJL uses a different rotation than the MSE stage to ensure statistical
     // independence between the quantization noise and the sign projection.
@@ -232,7 +235,7 @@ pub fn turboquant_encode_qjl(
         rotation.rotate(&padded, &mut rotated);
 
         for j in 0..padded_dim {
-            let idx = find_nearest_centroid(rotated[j], &centroids);
+            let idx = find_nearest_centroid(rotated[j], &boundaries);
             dequantized_rotated[j] = centroids[idx as usize];
         }
 

encodings/turboquant/src/decompress.rs

@@ -17,6 +17,16 @@ use crate::mse::array::TurboQuantMSEArray;
 use crate::qjl::array::TurboQuantQJLArray;
 use crate::rotation::RotationMatrix;
 
+/// QJL correction scale factor: `sqrt(π/2) / padded_dim`.
+///
+/// Accounts for the SRHT normalization (`1/padded_dim^{3/2}` per transform)
+/// combined with `E[|z|] = sqrt(2/π)` for half-normal sign expectations.
+/// Verified empirically via the `qjl_inner_product_bias` test suite.
+#[inline]
+fn qjl_correction_scale(padded_dim: usize) -> f32 {
+    (std::f32::consts::FRAC_PI_2).sqrt() / (padded_dim as f32)
+}
+
 /// Decompress a `TurboQuantMSEArray` into a `FixedSizeListArray` of floats.
 ///
 /// Reads stored centroids and rotation signs from the array's children,
@@ -126,11 +136,7 @@ pub fn execute_decompress_qjl(
     let qjl_rot_signs_bool = array.rotation_signs.clone().execute::<BoolArray>(ctx)?;
     let qjl_rot = RotationMatrix::from_bool_array(&qjl_rot_signs_bool, dim)?;
 
-    // QJL correction scale: sqrt(π/2) / padded_dim.
-    // This accounts for the SRHT normalization (1/padded_dim^{3/2} per transform)
-    // combined with the E[|z|] = sqrt(2/π) expectation of half-normal signs.
-    // Verified empirically via the `qjl_inner_product_bias` test suite.
-    let qjl_scale = (std::f32::consts::FRAC_PI_2).sqrt() / (padded_dim as f32);
+    let qjl_scale = qjl_correction_scale(padded_dim);
 
     let mut output = BufferMut::<f32>::with_capacity(num_rows * dim);
     let mut qjl_signs_vec = vec![0.0f32; padded_dim];

encodings/turboquant/src/lib.rs

@@ -120,6 +120,7 @@ mod tests {
     use rand_distr::Distribution;
     use rand_distr::Normal;
     use rstest::rstest;
+    use vortex_array::ArrayRef;
     use vortex_array::IntoArray;
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::FixedSizeListArray;
@@ -186,46 +187,43 @@ mod tests {
         total / num_rows as f32
     }
 
-    /// Encode via MSE and decode, returning (original, decoded) flat f32 slices.
-    fn encode_decode_mse(
+    /// Encode and decode, returning (original, decoded) flat f32 slices.
+    fn encode_decode(
         fsl: &FixedSizeListArray,
-        config: &TurboQuantConfig,
+        encode_fn: impl FnOnce(&FixedSizeListArray) -> VortexResult<ArrayRef>,
     ) -> VortexResult<(Vec<f32>, Vec<f32>)> {
         let original: Vec<f32> = {
             let prim = fsl.elements().to_canonical().unwrap().into_primitive();
             prim.as_slice::<f32>().to_vec()
         };
-        let encoded = turboquant_encode_mse(fsl, config)?;
+        let encoded = encode_fn(fsl)?;
         let mut ctx = SESSION.create_execution_ctx();
-        let decoded = encoded
-            .into_array()
-            .execute::<FixedSizeListArray>(&mut ctx)?;
+        let decoded = encoded.execute::<FixedSizeListArray>(&mut ctx)?;
         let decoded_elements: Vec<f32> = {
             let prim = decoded.elements().to_canonical().unwrap().into_primitive();
             prim.as_slice::<f32>().to_vec()
         };
         Ok((original, decoded_elements))
     }
 
-    /// Encode via QJL and decode, returning (original, decoded) flat f32 slices.
+    fn encode_decode_mse(
+        fsl: &FixedSizeListArray,
+        config: &TurboQuantConfig,
+    ) -> VortexResult<(Vec<f32>, Vec<f32>)> {
+        let config = config.clone();
+        encode_decode(fsl, |fsl| {
+            Ok(turboquant_encode_mse(fsl, &config)?.into_array())
+        })
+    }
+
     fn encode_decode_qjl(
         fsl: &FixedSizeListArray,
         config: &TurboQuantConfig,
     ) -> VortexResult<(Vec<f32>, Vec<f32>)> {
-        let original: Vec<f32> = {
-            let prim = fsl.elements().to_canonical().unwrap().into_primitive();
-            prim.as_slice::<f32>().to_vec()
-        };
-        let encoded = turboquant_encode_qjl(fsl, config)?;
-        let mut ctx = SESSION.create_execution_ctx();
-        let decoded = encoded
-            .into_array()
-            .execute::<FixedSizeListArray>(&mut ctx)?;
-        let decoded_elements: Vec<f32> = {
-            let prim = decoded.elements().to_canonical().unwrap().into_primitive();
-            prim.as_slice::<f32>().to_vec()
-        };
-        Ok((original, decoded_elements))
+        let config = config.clone();
+        encode_decode(fsl, |fsl| {
+            Ok(turboquant_encode_qjl(fsl, &config)?.into_array())
+        })
     }
 
     // -----------------------------------------------------------------------