more fixups

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

Author: connortsui20

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

@@ -11,6 +11,7 @@
 
 use std::sync::LazyLock;
 
+use vortex_buffer::Buffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
 use vortex_utils::aliases::dash_map::DashMap;
@@ -29,14 +30,14 @@ const INTEGRATION_POINTS: usize = 1000;
 
 // TODO(connor): Maybe we should just store an `ArrayRef` here?
 /// Global centroid cache keyed by (dimension, bit_width).
-static CENTROID_CACHE: LazyLock<DashMap<(u32, u8), Vec<f32>>> = LazyLock::new(DashMap::default);
+static CENTROID_CACHE: LazyLock<DashMap<(u32, u8), Buffer<f32>>> = LazyLock::new(DashMap::default);
 
 /// Get or compute cached centroids for the given dimension and bit width.
 ///
 /// Returns `2^bit_width` centroids sorted in ascending order, representing optimal scalar
 /// quantization levels for the coordinate distribution after random rotation in
 /// `dimension`-dimensional space.
-pub fn get_centroids(dimension: u32, bit_width: u8) -> VortexResult<Vec<f32>> {
+pub fn get_centroids(dimension: u32, bit_width: u8) -> VortexResult<Buffer<f32>> {
     vortex_ensure!(
         (1..=MAX_BIT_WIDTH).contains(&bit_width),
         "TurboQuant bit_width must be 1-{}, got {bit_width}",
@@ -92,7 +93,7 @@ impl HalfIntExponent {
 /// The probability distribution function is:
 ///   `f(x) = C_d * (1 - x^2)^((d-3)/2)` on `[-1, 1]`
 /// where `C_d` is the normalizing constant.
-fn max_lloyd_centroids(dimension: u32, bit_width: u8) -> Vec<f32> {
+fn max_lloyd_centroids(dimension: u32, bit_width: u8) -> Buffer<f32> {
     debug_assert!((1..=MAX_BIT_WIDTH).contains(&bit_width));
     let num_centroids = 1usize << bit_width;
 
@@ -288,7 +289,7 @@ mod tests {
     #[case(128, 4)]
     fn centroids_within_bounds(#[case] dim: u32, #[case] bits: u8) -> VortexResult<()> {
         let centroids = get_centroids(dim, bits)?;
-        for &val in &centroids {
+        for &val in centroids.iter() {
             assert!(
                 (-1.0..=1.0).contains(&val),
                 "centroid out of [-1, 1]: {val}",

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

@@ -171,7 +171,7 @@ pub unsafe fn turboquant_encode_unchecked(
 
     let core = turboquant_quantize_core(&fsl, seed, config.bit_width, config.num_rounds, ctx)?;
     let quantized_fsl =
-        build_quantized_fsl(num_rows, core.all_indices, &core.centroids, core.padded_dim)?;
+        build_quantized_fsl(num_rows, core.all_indices, core.centroids, core.padded_dim)?;
     let padded_vector = Vector::try_new_vector_array(quantized_fsl)?;
 
     let sorf_options = SorfOptions {
@@ -185,8 +185,8 @@ pub unsafe fn turboquant_encode_unchecked(
 
 /// Shared intermediate results from the quantization loop.
 struct QuantizationResult {
-    centroids: Vec<f32>,
-    all_indices: BufferMut<u8>,
+    centroids: Buffer<f32>,
+    all_indices: Buffer<u8>,
     padded_dim: usize,
 }
 
@@ -202,8 +202,7 @@ fn turboquant_quantize_core(
     num_rounds: u8,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<QuantizationResult> {
-    let dimension =
-        usize::try_from(fsl.list_size()).vortex_expect("u32 FixedSizeList dimension fits in usize");
+    let dimension = fsl.list_size() as usize;
     let num_rows = fsl.len();
 
     let rotation = SorfMatrix::try_new(seed, dimension, num_rounds as usize)?;
@@ -238,7 +237,7 @@ fn turboquant_quantize_core(
 
     Ok(QuantizationResult {
         centroids,
-        all_indices,
+        all_indices: all_indices.freeze(),
         padded_dim,
     })
 }
@@ -250,13 +249,12 @@ fn turboquant_quantize_core(
 /// without knowledge of the rotation.
 fn build_quantized_fsl(
     num_rows: usize,
-    all_indices: BufferMut<u8>,
-    centroids: &[f32],
+    all_indices: Buffer<u8>,
+    centroids: Buffer<f32>,
     padded_dim: usize,
 ) -> VortexResult<ArrayRef> {
-    let codes = PrimitiveArray::new::<u8>(all_indices.freeze(), Validity::NonNullable);
-    let centroids_array =
-        PrimitiveArray::new::<f32>(Buffer::copy_from(centroids), Validity::NonNullable);
+    let codes = PrimitiveArray::new::<u8>(all_indices, Validity::NonNullable);
+    let centroids_array = PrimitiveArray::new::<f32>(centroids, Validity::NonNullable);
 
     let dict = DictArray::try_new(codes.into_array(), centroids_array.into_array())?;
 

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

@@ -169,7 +169,7 @@ pub fn tq_validate_vector_dtype(dtype: &DType) -> VortexResult<VectorMatcherMeta
             vortex_err!("TurboQuant dtype must be a Vector extension type, got {dtype}")
         })?;
 
-    let dimensions = vector_metadata.list_size();
+    let dimensions = vector_metadata.dimensions();
     vortex_ensure!(
         dimensions >= MIN_DIMENSION,
         "TurboQuant requires dimension >= {MIN_DIMENSION}, got {dimensions}",

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

@@ -83,7 +83,7 @@ impl Scheme for TurboQuantScheme {
             .bit_width()
             .try_into()
             .vortex_expect("invalid bit width for TurboQuant");
-        let dimension = vector_metadata.list_size();
+        let dimension = vector_metadata.dimensions();
 
         CompressionEstimate::Verdict(EstimateVerdict::Ratio(estimate_compression_ratio(
             element_bit_width,

vortex-tensor/src/encodings/turboquant/tests/roundtrip.rs

@@ -311,37 +311,3 @@ fn f16_input_encodes_successfully() -> VortexResult<()> {
     assert_eq!(decoded_fsl.len(), num_rows);
     Ok(())
 }
-
-/// Verify that the checked encode accepts normalized f16 input.
-#[test]
-fn checked_encode_accepts_normalized_f16_input() -> VortexResult<()> {
-    let num_rows = 10;
-    let dim = 128;
-    let mut rng = StdRng::seed_from_u64(99);
-    let normal = Normal::new(0.0f32, 1.0).unwrap();
-
-    let mut buf = BufferMut::<half::f16>::with_capacity(num_rows * dim);
-    for _ in 0..(num_rows * dim) {
-        buf.push(half::f16::from_f32(normal.sample(&mut rng)));
-    }
-    let elements = PrimitiveArray::new::<half::f16>(buf.freeze(), Validity::NonNullable);
-    let fsl = FixedSizeListArray::try_new(
-        elements.into_array(),
-        dim.try_into().unwrap(),
-        Validity::NonNullable,
-        num_rows,
-    )?;
-
-    let ext = make_vector_ext(&fsl);
-    let config = TurboQuantConfig {
-        bit_width: 3,
-        seed: Some(42),
-        num_rounds: 3,
-    };
-
-    let mut ctx = SESSION.create_execution_ctx();
-    let normalized = normalize_as_l2_denorm(ext, &mut ctx)?.child_at(0).clone();
-    let encoded = turboquant_encode(normalized, &config, &mut ctx)?;
-    assert_eq!(encoded.len(), num_rows);
-    Ok(())
-}

vortex-tensor/src/fixed_shape/matcher.rs

@@ -31,7 +31,7 @@ pub struct FixedShapeTensorMatcherMetadata<'a> {
     ///
     /// This matches the `FixedSizeList` list size in the storage dtype, which is the product of
     /// the logical shape dimensions.
-    flat_list_size: usize,
+    flat_list_size: u32,
 }
 
 impl Matcher for AnyFixedShapeTensor {
@@ -64,7 +64,7 @@ impl Matcher for AnyFixedShapeTensor {
         Some(FixedShapeTensorMatcherMetadata {
             metadata,
             element_ptype: element_dtype.as_ptype(),
-            flat_list_size: *list_size as usize,
+            flat_list_size: *list_size,
         })
     }
 }
@@ -81,7 +81,7 @@ impl FixedShapeTensorMatcherMetadata<'_> {
     }
 
     /// Returns the flattened element count for each tensor row.
-    pub fn list_size(&self) -> usize {
+    pub fn flat_list_size(&self) -> u32 {
         self.flat_list_size
     }
 }
@@ -118,7 +118,7 @@ mod tests {
 
         let metadata = ext_dtype.metadata::<AnyFixedShapeTensor>();
         assert_eq!(metadata.element_ptype(), PType::F32);
-        assert_eq!(metadata.list_size(), 24);
+        assert_eq!(metadata.flat_list_size(), 24);
         assert_eq!(metadata.metadata().logical_shape(), &[2, 3, 4]);
         Ok(())
     }

vortex-tensor/src/matcher.rs

@@ -42,9 +42,9 @@ impl TensorMatch<'_> {
     }
 
     /// Returns the flattened element count for each logical tensor row.
-    pub fn list_size(self) -> usize {
+    pub fn list_size(self) -> u32 {
         match self {
-            Self::FixedShapeTensor(metadata) => metadata.list_size(),
+            Self::FixedShapeTensor(metadata) => metadata.flat_list_size(),
             Self::Vector(metadata) => metadata.dimensions(),
         }
     }

vortex-tensor/src/scalar_fns/inner_product.rs

@@ -182,7 +182,7 @@ impl ScalarFnVTable for InnerProduct {
         let tensor_match = ext
             .metadata_opt::<AnyTensor>()
             .vortex_expect("we already validated this in `return_dtype`");
-        let dimensions = tensor_match.list_size();
+        let dimensions = tensor_match.list_size() as usize;
 
         // Extract the storage array from each extension input. We pass the storage (FSL) rather
         // than the extension array to avoid canonicalizing the extension wrapper.

vortex-tensor/src/scalar_fns/l2_denorm.rs

@@ -240,7 +240,7 @@ impl ScalarFnVTable for L2Denorm {
             .as_extension()
             .metadata_opt::<AnyTensor>()
             .vortex_expect("we already validated this in `return_dtype`");
-        let tensor_flat_size = tensor_match.list_size();
+        let tensor_flat_size = tensor_match.list_size() as usize;
 
         let flat = extract_flat_elements(normalized.storage_array(), tensor_flat_size, ctx)?;
 
@@ -423,7 +423,7 @@ pub fn normalize_as_l2_denorm(
 ) -> VortexResult<ScalarFnArray> {
     let row_count = input.len();
     let tensor_match = validate_tensor_float_input(input.dtype())?;
-    let tensor_flat_size = tensor_match.list_size();
+    let tensor_flat_size = tensor_match.list_size() as usize;
 
     // Constant fast path: if the input is a constant-backed extension, normalize the single
     // stored row once and return an `L2Denorm` whose children are both `ConstantArray`s.
@@ -520,7 +520,7 @@ pub(crate) fn try_build_constant_l2_denorm(
         .as_extension()
         .metadata_opt::<AnyTensor>()
         .vortex_expect("caller validated input has AnyTensor metadata");
-    let list_size = tensor_match.list_size();
+    let list_size = tensor_match.list_size() as usize;
     let original_nullability = input.dtype().nullability();
     let ext_dtype = input.dtype().as_extension().clone();
     let storage_fsl_nullability = storage.dtype().nullability();
@@ -630,7 +630,7 @@ fn validate_l2_normalized_rows_against_norms(
     let tensor_match = validate_tensor_float_input(normalized.dtype())?;
     let element_ptype = tensor_match.element_ptype();
     let tolerance = unit_norm_tolerance(element_ptype);
-    let tensor_flat_size = tensor_match.list_size();
+    let tensor_flat_size = tensor_match.list_size() as usize;
 
     if let Some(norms) = norms {
         vortex_ensure_eq!(

vortex-tensor/src/scalar_fns/l2_norm.rs

@@ -132,7 +132,7 @@ impl ScalarFnVTable for L2Norm {
         let tensor_match = ext
             .metadata_opt::<AnyTensor>()
             .vortex_expect("we already validated this in `return_dtype`");
-        let tensor_flat_size = tensor_match.list_size();
+        let tensor_flat_size = tensor_match.list_size() as usize;
         let element_ptype = tensor_match.element_ptype();
 
         let norm_dtype = DType::Primitive(element_ptype, ext.nullability());