fix casting issues and other minor things

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-tensor/public-api.lock

@@ -10,6 +10,8 @@ impl vortex_tensor::encodings::turboquant::TurboQuant
 
 pub const vortex_tensor::encodings::turboquant::TurboQuant::ID: vortex_array::array::ArrayId
 
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::try_new_array(dtype: vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<vortex_tensor::encodings::turboquant::TurboQuantArray>
+
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::validate_dtype(dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<&vortex_array::dtype::extension::erased::ExtDTypeRef>
 
 impl core::clone::Clone for vortex_tensor::encodings::turboquant::TurboQuant
@@ -24,8 +26,6 @@ impl vortex_array::array::vtable::VTable for vortex_tensor::encodings::turboquan
 
 pub type vortex_tensor::encodings::turboquant::TurboQuant::ArrayData = vortex_tensor::encodings::turboquant::TurboQuantData
 
-pub type vortex_tensor::encodings::turboquant::TurboQuant::Metadata = vortex_tensor::encodings::turboquant::TurboQuantMetadata
-
 pub type vortex_tensor::encodings::turboquant::TurboQuant::OperationsVTable = vortex_tensor::encodings::turboquant::TurboQuant
 
 pub type vortex_tensor::encodings::turboquant::TurboQuant::ValidityVTable = vortex_array::array::vtable::validity::ValidityVTableFromChild
@@ -38,35 +38,25 @@ pub fn vortex_tensor::encodings::turboquant::TurboQuant::buffer(_array: vortex_a
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::buffer_name(_array: vortex_array::array::view::ArrayView<'_, Self>, _idx: usize) -> core::option::Option<alloc::string::String>
 
-pub fn vortex_tensor::encodings::turboquant::TurboQuant::build(dtype: &vortex_array::dtype::DType, len: usize, metadata: &Self::Metadata, _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren) -> vortex_error::VortexResult<vortex_tensor::encodings::turboquant::TurboQuantData>
-
-pub fn vortex_tensor::encodings::turboquant::TurboQuant::deserialize(bytes: &[u8], _dtype: &vortex_array::dtype::DType, _len: usize, _buffers: &[vortex_array::buffer::BufferHandle], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Metadata>
-
-pub fn vortex_tensor::encodings::turboquant::TurboQuant::dtype(array: &vortex_tensor::encodings::turboquant::TurboQuantData) -> &vortex_array::dtype::DType
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::ArrayData>
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::execute(array: vortex_array::array::typed::Array<Self>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::executor::ExecutionResult>
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::execute_parent(array: vortex_array::array::view::ArrayView<'_, Self>, parent: &vortex_array::array::erased::ArrayRef, child_idx: usize, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::id(&self) -> vortex_array::array::ArrayId
 
-pub fn vortex_tensor::encodings::turboquant::TurboQuant::len(array: &vortex_tensor::encodings::turboquant::TurboQuantData) -> usize
-
-pub fn vortex_tensor::encodings::turboquant::TurboQuant::metadata(array: vortex_array::array::view::ArrayView<'_, Self>) -> vortex_error::VortexResult<Self::Metadata>
-
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::nbuffers(_array: vortex_array::array::view::ArrayView<'_, Self>) -> usize
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::reduce_parent(array: vortex_array::array::view::ArrayView<'_, Self>, parent: &vortex_array::array::erased::ArrayRef, child_idx: usize) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
 
-pub fn vortex_tensor::encodings::turboquant::TurboQuant::serialize(metadata: Self::Metadata) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::serialize(array: vortex_array::array::view::ArrayView<'_, Self>) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::slot_name(_array: vortex_array::array::view::ArrayView<'_, Self>, idx: usize) -> alloc::string::String
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::slots(array: vortex_array::array::view::ArrayView<'_, Self>) -> &[core::option::Option<vortex_array::array::erased::ArrayRef>]
 
-pub fn vortex_tensor::encodings::turboquant::TurboQuant::stats(array: &vortex_tensor::encodings::turboquant::TurboQuantData) -> &vortex_array::stats::array::ArrayStats
-
-pub fn vortex_tensor::encodings::turboquant::TurboQuant::vtable(_array: &Self::ArrayData) -> &Self
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::validate(&self, data: &Self::ArrayData, dtype: &vortex_array::dtype::DType, len: usize) -> vortex_error::VortexResult<()>
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::with_slots(array: &mut vortex_tensor::encodings::turboquant::TurboQuantData, slots: alloc::vec::Vec<core::option::Option<vortex_array::array::erased::ArrayRef>>) -> vortex_error::VortexResult<()>
 
@@ -116,46 +106,26 @@ pub fn vortex_tensor::encodings::turboquant::TurboQuantData::codes(&self) -> &vo
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuantData::dimension(&self) -> u32
 
-pub unsafe fn vortex_tensor::encodings::turboquant::TurboQuantData::new_unchecked(dtype: vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> Self
+pub unsafe fn vortex_tensor::encodings::turboquant::TurboQuantData::new_unchecked(dtype: &vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> Self
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuantData::norms(&self) -> &vortex_array::array::erased::ArrayRef
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuantData::padded_dim(&self) -> u32
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuantData::rotation_signs(&self) -> &vortex_array::array::erased::ArrayRef
 
-pub fn vortex_tensor::encodings::turboquant::TurboQuantData::try_new(dtype: vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<Self>
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::try_new(dtype: &vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<Self>
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuantData::validate(dtype: &vortex_array::dtype::DType, codes: &vortex_array::array::erased::ArrayRef, norms: &vortex_array::array::erased::ArrayRef, centroids: &vortex_array::array::erased::ArrayRef, rotation_signs: &vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<()>
 
 impl core::clone::Clone for vortex_tensor::encodings::turboquant::TurboQuantData
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuantData::clone(&self) -> vortex_tensor::encodings::turboquant::TurboQuantData
 
-impl core::convert::From<vortex_tensor::encodings::turboquant::TurboQuantData> for vortex_array::array::erased::ArrayRef
-
-pub fn vortex_array::array::erased::ArrayRef::from(value: vortex_tensor::encodings::turboquant::TurboQuantData) -> vortex_array::array::erased::ArrayRef
-
 impl core::fmt::Debug for vortex_tensor::encodings::turboquant::TurboQuantData
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuantData::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
 
-impl vortex_array::array::IntoArray for vortex_tensor::encodings::turboquant::TurboQuantData
-
-pub fn vortex_tensor::encodings::turboquant::TurboQuantData::into_array(self) -> vortex_array::array::erased::ArrayRef
-
-pub struct vortex_tensor::encodings::turboquant::TurboQuantMetadata
-
-pub vortex_tensor::encodings::turboquant::TurboQuantMetadata::bit_width: u8
-
-impl core::clone::Clone for vortex_tensor::encodings::turboquant::TurboQuantMetadata
-
-pub fn vortex_tensor::encodings::turboquant::TurboQuantMetadata::clone(&self) -> vortex_tensor::encodings::turboquant::TurboQuantMetadata
-
-impl core::fmt::Debug for vortex_tensor::encodings::turboquant::TurboQuantMetadata
-
-pub fn vortex_tensor::encodings::turboquant::TurboQuantMetadata::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
-
 pub struct vortex_tensor::encodings::turboquant::TurboQuantScheme
 
 impl core::clone::Clone for vortex_tensor::encodings::turboquant::TurboQuantScheme
@@ -186,10 +156,10 @@ pub fn vortex_tensor::encodings::turboquant::TurboQuantScheme::matches(&self, ca
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuantScheme::scheme_name(&self) -> &'static str
 
-pub fn vortex_tensor::encodings::turboquant::initialize(session: &mut vortex_session::VortexSession)
-
 pub fn vortex_tensor::encodings::turboquant::turboquant_encode(ext: &vortex_array::arrays::extension::vtable::ExtensionArray, config: &vortex_tensor::encodings::turboquant::TurboQuantConfig, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
 
+pub type vortex_tensor::encodings::turboquant::TurboQuantArray = vortex_array::array::typed::Array<vortex_tensor::encodings::turboquant::TurboQuant>
+
 pub mod vortex_tensor::fixed_shape
 
 pub struct vortex_tensor::fixed_shape::FixedShapeTensor

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

@@ -11,3 +11,15 @@ pub(crate) mod centroids;
 pub(crate) mod rotation;
 
 pub(crate) mod scheme;
+
+use num_traits::Float;
+use num_traits::FromPrimitive;
+use vortex_error::VortexExpect;
+
+/// Convert an f32 value to a float type `T`.
+///
+/// `FromPrimitive::from_f32` is infallible for all Vortex float types: f16 saturates via the
+/// inherent `f16::from_f32()`, f32 is identity, f64 is lossless widening.
+pub(crate) fn float_from_f32<T: Float + FromPrimitive>(v: f32) -> T {
+    FromPrimitive::from_f32(v).vortex_expect("f32-to-float conversion is infallible")
+}

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

@@ -3,6 +3,7 @@
 
 //! TurboQuant encoding (quantization) logic.
 
+use num_traits::ToPrimitive;
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
@@ -15,6 +16,7 @@ use vortex_array::dtype::PType;
 use vortex_array::match_each_float_ptype;
 use vortex_array::validity::Validity;
 use vortex_buffer::BufferMut;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
@@ -105,14 +107,18 @@ fn turboquant_quantize_core(
     // the input, so null vectors get null norms automatically.
     let norms_sfn = L2Norm::try_new_array(&ApproxOptions::Exact, ext.as_ref().clone(), num_rows)?;
     let norms_array: ArrayRef = norms_sfn.into_array().execute(ctx)?;
-    let norms_prim: PrimitiveArray = norms_array.to_canonical()?.into_primitive();
+    let norms_prim: PrimitiveArray = norms_array.clone().execute(ctx)?;
 
     // Extract f32 norms for the internal quantization loop.
     let f32_norms: Vec<f32> = match_each_float_ptype!(norms_prim.ptype(), |T| {
         norms_prim
             .as_slice::<T>()
             .iter()
-            .map(|&v| num_traits::ToPrimitive::to_f32(&v).unwrap_or(0.0))
+            .map(|&v| {
+                // `ToPrimitive::to_f32` is infallible for all float types: f16 -> f32 is lossless,
+                // f32 is identity, and f64 -> f32 saturates to +-inf.
+                ToPrimitive::to_f32(&v).vortex_expect("float-to-f32 conversion is infallible")
+            })
             .collect()
     });
 

vortex-tensor/src/encodings/turboquant/compute/cosine_similarity.rs

@@ -30,8 +30,6 @@
 //! usually sufficient -- the relative ordering of cosine similarities is preserved
 //! even if the absolute values have bounded error.
 
-use num_traits::FromPrimitive;
-use num_traits::Zero;
 use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
@@ -44,19 +42,9 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_ensure_eq;
 
 use crate::encodings::turboquant::TurboQuant;
+use crate::encodings::turboquant::array::float_from_f32;
 use crate::utils::extension_element_ptype;
 
-/// Convert an f32 value to `T`, returning `T::zero()` if the conversion fails.
-///
-/// This helper exists because `half::f16` has an inherent `from_f32` method that shadows
-/// the [`FromPrimitive`] trait method, causing compilation errors when used inside
-/// [`match_each_float_ptype!`].
-#[inline]
-fn f32_to_t<T: FromPrimitive + Zero>(v: f32) -> T {
-    // TODO(connor): Is this actually correct? How should we handle f64 overflow?
-    FromPrimitive::from_f32(v).unwrap_or_else(T::zero)
-}
-
 /// Compute the per-row unit-norm dot products in f32 (centroids are always f32).
 ///
 /// Returns a `Vec<f32>` of length `num_rows`.
@@ -124,7 +112,7 @@ pub fn cosine_similarity_quantized_column(
         let mut result = BufferMut::<T>::with_capacity(dots.len());
         for &dot in &dots {
             // SAFETY: We allocated the correct amount.
-            unsafe { result.push_unchecked(f32_to_t(dot)) };
+            unsafe { result.push_unchecked(float_from_f32(dot)) };
         }
 
         // SAFETY: `result` has the same length as the input arrays, matching `validity`.
@@ -164,7 +152,7 @@ pub fn dot_product_quantized_column(
 
         let mut result = BufferMut::<T>::with_capacity(num_rows);
         for row in 0..num_rows {
-            let dot_t: T = f32_to_t(dots[row]);
+            let dot_t: T = float_from_f32(dots[row]);
             // SAFETY: We allocated the correct amount.
             unsafe { result.push_unchecked(na[row] * nb[row] * dot_t) };
         }

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

@@ -3,8 +3,8 @@
 
 //! TurboQuant decoding (dequantization) logic.
 
+use num_traits::Float;
 use num_traits::FromPrimitive;
-use num_traits::Zero;
 use vortex_array::Array;
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
@@ -20,6 +20,7 @@ use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 
 use crate::encodings::turboquant::TurboQuant;
+use crate::encodings::turboquant::array::float_from_f32;
 use crate::encodings::turboquant::array::rotation::RotationMatrix;
 use crate::utils::extension_element_ptype;
 
@@ -103,7 +104,7 @@ pub fn execute_decompress(
 }
 
 /// Typed decompress: reads norms as `T`, dequantizes in f32, and produces output as `T`.
-fn decompress_typed<T: NativePType + FromPrimitive + Zero>(
+fn decompress_typed<T: NativePType + Float + FromPrimitive>(
     norms_prim: &PrimitiveArray,
     centroids: &[f32],
     rotation: &RotationMatrix,
@@ -129,8 +130,7 @@ fn decompress_typed<T: NativePType + FromPrimitive + Zero>(
         rotation.inverse_rotate(&dequantized, &mut unrotated);
 
         for idx in 0..dim {
-            // Convert f32 dequantized value to T, then scale by the native-precision norm.
-            let val = T::from_f32(unrotated[idx]).unwrap_or_else(T::zero) * norm;
+            let val = float_from_f32::<T>(unrotated[idx]) * norm;
             output.push(val);
         }
     }

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

@@ -98,6 +98,7 @@ pub(crate) mod compute;
 
 mod vtable;
 pub use vtable::TurboQuant;
+pub use vtable::TurboQuantArray;
 
 mod compress;
 pub use compress::TurboQuantConfig;