fix[turboquant]: second-round review fixes and merge conflict resolution

- Add TurboQuantMSE and TurboQuantQJL to ALLOWED_ENCODINGS in
  vortex-file so TurboQuant-encoded files can be deserialized
- Fix as_ptype() panic: use primitive.ptype() after to_canonical()
  instead of calling the panicking as_ptype() on the raw dtype
- Move rand_distr to dev-dependencies (only used in tests)
- Remove unused vortex-mask dependency
- Handle nullable storage in compress_turboquant: return None to fall
  through to default compression instead of failing
- Remove apply_inverse_srht_from_bits (dead code, only used in its own
  test) and apply_signs_from_bits helper
- Fix function-scoped import in gen_random_signs
- Add TODO for double f32 extraction in QJL encode
- Fix execute() signature after merge with develop (Arc<Array<Self>>)
- Collapse nested if-let per clippy

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

Cargo.lock

@@ -19,16 +19,15 @@ workspace = true
 [dependencies]
 prost = { workspace = true }
 rand = { workspace = true }
-rand_distr = { workspace = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-error = { workspace = true }
 vortex-fastlanes = { workspace = true }
-vortex-mask = { workspace = true }
 vortex-session = { workspace = true }
 vortex-utils = { workspace = true }
 parking_lot = { workspace = true }
 
 [dev-dependencies]
+rand_distr = { workspace = true }
 rstest = { workspace = true }
 vortex-array = { workspace = true, features = ["_test-harness"] }

encodings/turboquant/Cargo.toml

@@ -28,8 +28,6 @@ pub fn vortex_turboquant::rotation::RotationMatrix::rotate(&self, input: &[f32],
 
 pub fn vortex_turboquant::rotation::RotationMatrix::try_new(seed: u64, dimension: usize) -> vortex_error::VortexResult<Self>
 
-pub fn vortex_turboquant::rotation::apply_inverse_srht_from_bits(buf: &mut [f32], signs_bytes: &[u8], padded_dim: usize, norm_factor: f32)
-
 pub struct vortex_turboquant::TurboQuantConfig
 
 pub vortex_turboquant::TurboQuantConfig::bit_width: u8
@@ -86,7 +84,7 @@ pub fn vortex_turboquant::TurboQuantMSE::deserialize(bytes: &[u8], _dtype: &vort
 
 pub fn vortex_turboquant::TurboQuantMSE::dtype(array: &vortex_turboquant::TurboQuantMSEArray) -> &vortex_array::dtype::DType
 
-pub fn vortex_turboquant::TurboQuantMSE::execute(array: alloc::sync::Arc<Self::Array>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::executor::ExecutionResult>
+pub fn vortex_turboquant::TurboQuantMSE::execute(array: alloc::sync::Arc<vortex_array::vtable::typed::Array<Self>>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::executor::ExecutionResult>
 
 pub fn vortex_turboquant::TurboQuantMSE::id(&self) -> vortex_array::vtable::dyn_::ArrayId
 
@@ -232,7 +230,7 @@ pub fn vortex_turboquant::TurboQuantQJL::deserialize(bytes: &[u8], _dtype: &vort
 
 pub fn vortex_turboquant::TurboQuantQJL::dtype(array: &vortex_turboquant::TurboQuantQJLArray) -> &vortex_array::dtype::DType
 
-pub fn vortex_turboquant::TurboQuantQJL::execute(array: alloc::sync::Arc<Self::Array>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::executor::ExecutionResult>
+pub fn vortex_turboquant::TurboQuantQJL::execute(array: alloc::sync::Arc<vortex_array::vtable::typed::Array<Self>>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::executor::ExecutionResult>
 
 pub fn vortex_turboquant::TurboQuantQJL::id(&self) -> vortex_array::vtable::dyn_::ArrayId
 

encodings/turboquant/public-api.lock

@@ -39,8 +39,8 @@ pub struct TurboQuantConfig {
 #[allow(clippy::cast_possible_truncation)]
 fn extract_f32_elements(fsl: &FixedSizeListArray) -> VortexResult<Vec<f32>> {
     let elements = fsl.elements();
-    let ptype = elements.dtype().as_ptype();
     let primitive = elements.to_canonical()?.into_primitive();
+    let ptype = primitive.ptype();
 
     match ptype {
         PType::F32 => Ok(primitive.as_slice::<f32>().to_vec()),
@@ -196,6 +196,8 @@ pub fn turboquant_encode_qjl(
         return build_empty_qjl_array(fsl, config.bit_width, padded_dim, seed);
     }
 
+    // TODO(perf): `turboquant_encode_mse` above already extracts f32 elements
+    // internally. Refactor to share the buffer to avoid double materialization.
     let f32_elements = extract_f32_elements(fsl)?;
     #[allow(clippy::cast_possible_truncation)]
     let centroids = get_centroids(padded_dim as u32, mse_bit_width)?;

encodings/turboquant/src/compress.rs

@@ -4,6 +4,7 @@
 //! VTable implementation for TurboQuant MSE encoding.
 
 use std::hash::Hash;
+use std::ops::Deref;
 use std::sync::Arc;
 
 use vortex_array::ArrayEq;
@@ -22,6 +23,7 @@ use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::stats::StatsSetRef;
+use vortex_array::vtable::Array;
 use vortex_array::vtable::ArrayId;
 use vortex_array::vtable::NotSupported;
 use vortex_array::vtable::VTable;
@@ -209,9 +211,11 @@ impl VTable for TurboQuantMSE {
         Ok(())
     }
 
-    fn execute(array: Arc<Self::Array>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
-        let array = Arc::try_unwrap(array).unwrap_or_else(|arc| (*arc).clone());
-        Ok(ExecutionResult::done(execute_decompress_mse(array, ctx)?))
+    fn execute(array: Arc<Array<Self>>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
+        let inner = Arc::try_unwrap(array)
+            .map(|a| a.into_inner())
+            .unwrap_or_else(|arc| arc.as_ref().deref().clone());
+        Ok(ExecutionResult::done(execute_decompress_mse(inner, ctx)?))
     }
 }
 

encodings/turboquant/src/mse/vtable/mod.rs

@@ -4,6 +4,7 @@
 //! VTable implementation for TurboQuant QJL encoding.
 
 use std::hash::Hash;
+use std::ops::Deref;
 use std::sync::Arc;
 
 use vortex_array::ArrayEq;
@@ -22,6 +23,7 @@ use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::stats::StatsSetRef;
+use vortex_array::vtable::Array;
 use vortex_array::vtable::ArrayId;
 use vortex_array::vtable::NotSupported;
 use vortex_array::vtable::VTable;
@@ -204,9 +206,11 @@ impl VTable for TurboQuantQJL {
         Ok(())
     }
 
-    fn execute(array: Arc<Self::Array>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
-        let array = Arc::try_unwrap(array).unwrap_or_else(|arc| (*arc).clone());
-        Ok(ExecutionResult::done(execute_decompress_qjl(array, ctx)?))
+    fn execute(array: Arc<Array<Self>>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
+        let inner = Arc::try_unwrap(array)
+            .map(|a| a.into_inner())
+            .unwrap_or_else(|arc| arc.as_ref().deref().clone());
+        Ok(ExecutionResult::done(execute_decompress_qjl(inner, ctx)?))
     }
 }
 

encodings/turboquant/src/qjl/vtable/mod.rs

@@ -12,6 +12,7 @@
 //! For dimensions that are not powers of 2, the input is zero-padded to the
 //! next power of 2 before the transform and truncated afterward.
 
+use rand::RngExt;
 use rand::SeedableRng;
 use rand::rngs::StdRng;
 use vortex_array::arrays::BoolArray;
@@ -212,42 +213,8 @@ impl RotationMatrix {
 /// contains `3 * padded_dim` bits in inverse-application order `[D₃ | D₂ | D₁]`.
 /// Convention: bit set (1) = +1, bit unset (0) = -1 (negate).
 ///
-/// Applies: H → D₃ → H → D₂ → H → D₁ → scale
-#[inline]
-pub fn apply_inverse_srht_from_bits(
-    buf: &mut [f32],
-    signs_bytes: &[u8],
-    padded_dim: usize,
-    norm_factor: f32,
-) {
-    debug_assert!(padded_dim.is_power_of_two());
-    debug_assert_eq!(buf.len(), padded_dim);
-
-    for round in 0..3 {
-        walsh_hadamard_transform(buf);
-        apply_signs_from_bits(buf, signs_bytes, round * padded_dim);
-    }
-
-    for val in buf.iter_mut() {
-        *val *= norm_factor;
-    }
-}
-
-/// Element-wise negate coordinates where the sign bit is unset (0 = -1).
-#[inline]
-fn apply_signs_from_bits(buf: &mut [f32], signs_bytes: &[u8], bit_offset: usize) {
-    for (j, val) in buf.iter_mut().enumerate() {
-        let idx = bit_offset + j;
-        let is_positive = (signs_bytes[idx / 8] >> (idx % 8)) & 1 == 1;
-        if !is_positive {
-            *val = -*val;
-        }
-    }
-}
-
 /// Generate a vector of random ±1 signs.
 fn gen_random_signs(rng: &mut StdRng, len: usize) -> Vec<f32> {
-    use rand::RngExt;
     (0..len)
         .map(|_| {
             if rng.random_bool(0.5) {
@@ -416,48 +383,6 @@ mod tests {
         Ok(())
     }
 
-    /// Verify that the hot-path `apply_inverse_srht_from_bits` matches `inverse_rotate`.
-    #[rstest]
-    #[case(64)]
-    #[case(128)]
-    #[case(768)]
-    fn hot_path_matches_inverse_rotate(#[case] dim: usize) -> VortexResult<()> {
-        let rot = RotationMatrix::try_new(99, dim)?;
-        let padded_dim = rot.padded_dim();
-        let norm_factor = rot.norm_factor();
-
-        let signs_array = rot.export_inverse_signs_bool_array();
-        let bit_buf = signs_array.to_bit_buffer();
-        let (_, _, raw_buf) = bit_buf.into_inner();
-
-        // Create some rotated input.
-        let mut input = vec![0.0f32; padded_dim];
-        for i in 0..dim {
-            input[i] = (i as f32 + 1.0) * 0.01;
-        }
-        let mut rotated = vec![0.0f32; padded_dim];
-        rot.rotate(&input, &mut rotated);
-
-        // Inverse via the struct method.
-        let mut recovered1 = vec![0.0f32; padded_dim];
-        rot.inverse_rotate(&rotated, &mut recovered1);
-
-        // Inverse via the hot-path function.
-        let mut recovered2 = rotated.clone();
-        apply_inverse_srht_from_bits(&mut recovered2, raw_buf.as_ref(), padded_dim, norm_factor);
-
-        for i in 0..padded_dim {
-            assert!(
-                (recovered1[i] - recovered2[i]).abs() < 1e-10,
-                "Hot-path mismatch at {i}: {} vs {}",
-                recovered1[i],
-                recovered2[i]
-            );
-        }
-
-        Ok(())
-    }
-
     #[test]
     fn wht_basic() {
         // WHT of [1, 0, 0, 0] should be [1, 1, 1, 1]

encodings/turboquant/src/rotation.rs

@@ -296,10 +296,12 @@ impl CanonicalCompressor for BtrBlocksCompressor {
                 }
 
                 // Compress tensor extension types with TurboQuant if configured.
+                // Falls through to default compression for nullable storage.
                 if let Some(tq_config) = &self.turboquant_config
                     && is_tensor_extension(&ext_array)
+                    && let Some(compressed) = compress_turboquant(&ext_array, tq_config)?
                 {
-                    return compress_turboquant(&ext_array, tq_config);
+                    return Ok(compressed);
                 }
 
                 // Compress the underlying storage array.

vortex-btrblocks/src/canonical_compressor.rs

@@ -18,7 +18,11 @@ pub(crate) fn is_tensor_extension(ext_array: &ExtensionArray) -> bool {
     ext_id.as_ref() == VECTOR_EXT_ID || ext_id.as_ref() == FIXED_SHAPE_TENSOR_EXT_ID
 }
 
-/// Compress a tensor extension array using TurboQuant.
+/// Try to compress a tensor extension array using TurboQuant.
+///
+/// Returns `Ok(Some(...))` on success, or `Ok(None)` if the storage is nullable
+/// (TurboQuant requires non-nullable input). The caller should fall through to
+/// default compression when `None` is returned.
 ///
 /// Produces a `TurboQuantQJLArray` wrapping a `TurboQuantMSEArray`, stored inside
 /// the Extension wrapper. All children (codes, norms, centroids, rotation signs,
@@ -27,13 +31,19 @@ pub(crate) fn is_tensor_extension(ext_array: &ExtensionArray) -> bool {
 pub(crate) fn compress_turboquant(
     ext_array: &ExtensionArray,
     config: &TurboQuantConfig,
-) -> VortexResult<ArrayRef> {
+) -> VortexResult<Option<ArrayRef>> {
     let storage = ext_array.storage_array();
     let fsl = storage.to_canonical()?.into_fixed_size_list();
 
+    if fsl.dtype().is_nullable() {
+        return Ok(None);
+    }
+
     // Produce the cascaded QJL(MSE) structure. The layout writer will
     // recursively descend into children and compress each one.
     let qjl_array = turboquant_encode_qjl(&fsl, config)?;
 
-    Ok(ExtensionArray::new(ext_array.ext_dtype().clone(), qjl_array.into_array()).into_array())
+    Ok(Some(
+        ExtensionArray::new(ext_array.ext_dtype().clone(), qjl_array.into_array()).into_array(),
+    ))
 }

vortex-btrblocks/src/compressor/turboquant.rs

@@ -60,6 +60,8 @@ use vortex_pco::Pco;
 use vortex_runend::RunEnd;
 use vortex_sequence::Sequence;
 use vortex_sparse::Sparse;
+use vortex_turboquant::TurboQuantMSE;
+use vortex_turboquant::TurboQuantQJL;
 use vortex_utils::aliases::hash_map::HashMap;
 use vortex_zigzag::ZigZag;
 #[cfg(feature = "zstd")]
@@ -109,6 +111,8 @@ pub static ALLOWED_ENCODINGS: LazyLock<ArrayRegistry> = LazyLock::new(|| {
     session.register(Sequence);
     session.register(Sparse);
     session.register(ZigZag);
+    session.register(TurboQuantMSE);
+    session.register(TurboQuantQJL);
 
     #[cfg(feature = "zstd")]
     session.register(Zstd);