min dimension 3

Signed-off-by: Will Manning <will@willmanning.io>

Author: lwwmanning

vortex-tensor/src/encodings/turboquant/centroids.rs

@@ -36,8 +36,8 @@ pub fn get_centroids(dimension: u32, bit_width: u8) -> VortexResult<Vec<f32>> {
     if !(1..=8).contains(&bit_width) {
         vortex_bail!("TurboQuant bit_width must be 1-8, got {bit_width}");
     }
-    if dimension < 2 {
-        vortex_bail!("TurboQuant dimension must be >= 2, got {dimension}");
+    if dimension < 3 {
+        vortex_bail!("TurboQuant dimension must be >= 3, got {dimension}");
     }
 
     if let Some(centroids) = CENTROID_CACHE.get(&(dimension, bit_width)) {
@@ -306,5 +306,6 @@ mod tests {
         assert!(get_centroids(128, 0).is_err());
         assert!(get_centroids(128, 9).is_err());
         assert!(get_centroids(1, 2).is_err());
+        assert!(get_centroids(2, 2).is_err());
     }
 }

vortex-tensor/src/encodings/turboquant/compress.rs

@@ -198,8 +198,8 @@ pub fn turboquant_encode_mse(
     );
     let dimension = fsl.list_size();
     vortex_ensure!(
-        dimension >= 2,
-        "TurboQuant requires dimension >= 2, got {dimension}"
+        dimension >= 3,
+        "TurboQuant requires dimension >= 3, got {dimension}"
     );
 
     if fsl.is_empty() {
@@ -233,8 +233,8 @@ pub fn turboquant_encode_qjl(
     );
     let dimension = fsl.list_size();
     vortex_ensure!(
-        dimension >= 2,
-        "TurboQuant requires dimension >= 2, got {dimension}"
+        dimension >= 3,
+        "TurboQuant requires dimension >= 3, got {dimension}"
     );
 
     if fsl.is_empty() {

vortex-tensor/src/encodings/turboquant/mod.rs

@@ -468,31 +468,38 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn mse_rejects_dimension_below_2() {
-        let fsl = make_fsl_dim1();
+    #[rstest]
+    #[case(1)]
+    #[case(2)]
+    fn mse_rejects_dimension_below_3(#[case] dim: usize) {
+        let fsl = make_fsl_small(dim);
         let config = TurboQuantConfig {
             bit_width: 2,
             seed: Some(0),
         };
         assert!(turboquant_encode_mse(&fsl, &config).is_err());
     }
 
-    #[test]
-    fn qjl_rejects_dimension_below_2() {
-        let fsl = make_fsl_dim1();
+    #[rstest]
+    #[case(1)]
+    #[case(2)]
+    fn qjl_rejects_dimension_below_3(#[case] dim: usize) {
+        let fsl = make_fsl_small(dim);
         let config = TurboQuantConfig {
             bit_width: 3,
             seed: Some(0),
         };
         assert!(turboquant_encode_qjl(&fsl, &config).is_err());
     }
 
-    fn make_fsl_dim1() -> FixedSizeListArray {
-        let mut buf = BufferMut::<f32>::with_capacity(1);
-        buf.push(1.0);
+    fn make_fsl_small(dim: usize) -> FixedSizeListArray {
+        let mut buf = BufferMut::<f32>::with_capacity(dim);
+        for i in 0..dim {
+            buf.push(i as f32 + 1.0);
+        }
         let elements = PrimitiveArray::new::<f32>(buf.freeze(), Validity::NonNullable);
-        FixedSizeListArray::try_new(elements.into_array(), 1, Validity::NonNullable, 1).unwrap()
+        FixedSizeListArray::try_new(elements.into_array(), dim as u32, Validity::NonNullable, 1)
+            .unwrap()
     }
 
     // -----------------------------------------------------------------------

vortex-tensor/src/encodings/turboquant/scheme.rs

@@ -130,6 +130,15 @@ fn get_tensor_element_ptype_and_length(dtype: &DType) -> VortexResult<(PType, u3
         ),
     };
 
+    // TurboQuant requires dimension >= 3: the marginal coordinate distribution
+    // (1 - x^2)^((d-3)/2) has a singularity at d=2 (arcsine distribution) that
+    // causes NaN in the Max-Lloyd centroid computation.
+    vortex_ensure!(
+        *fsl_len >= 3,
+        "TurboQuant requires dimension >= 3, got {}",
+        fsl_len
+    );
+
     if let &DType::Primitive(ptype, Nullability::NonNullable) = element_dtype.as_ref() {
         Ok((ptype, *fsl_len))
     } else {