Extract Vector::wrap_storage and use it at four call sites

The two-line `ExtDType::<Vector>::try_new(EmptyMetadata, storage.dtype().clone()) +
ExtensionArray::new(ext_dtype, storage).into_array()` incantation for
wrapping a storage array in a [`Vector`] extension appeared verbatim in:

- `compress::wrap_padded_as_vector` (private, used twice)
- `sorf_transform::vtable` (empty-array branch + `inverse_rotate_typed`)
- `vector_search::build_constant_query_vector`

Promote it to `Vector::wrap_storage(storage)`, an associated function on
the [`Vector`] vtable struct that is the natural home for the operation.
Each call site drops to a single line and the `ExtDType`/`EmptyMetadata`
imports go away where they were only pulled in for this pattern. The old
private helper is deleted.

Signed-off-by: Claude <noreply@anthropic.com>

Author: claude

vortex-tensor/public-api.lock

@@ -528,6 +528,10 @@ pub fn vortex_tensor::vector::AnyVector::try_match<'a>(ext_dtype: &'a vortex_arr
 
 pub struct vortex_tensor::vector::Vector
 
+impl vortex_tensor::vector::Vector
+
+pub fn vortex_tensor::vector::Vector::wrap_storage(storage: vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
+
 impl core::clone::Clone for vortex_tensor::vector::Vector
 
 pub fn vortex_tensor::vector::Vector::clone(&self) -> vortex_tensor::vector::Vector

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

@@ -15,16 +15,13 @@ use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
 use vortex_array::arrays::Extension;
-use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::dict::DictArray;
 use vortex_array::arrays::extension::ExtensionArrayExt;
 use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
 use vortex_array::arrays::scalar_fn::ScalarFnArrayExt;
 use vortex_array::dtype::Nullability;
-use vortex_array::dtype::extension::ExtDType;
-use vortex_array::extension::EmptyMetadata;
 use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
@@ -242,7 +239,7 @@ pub unsafe fn turboquant_encode_unchecked(
             Validity::NonNullable,
             0,
         )?;
-        let empty_padded_vector = wrap_padded_as_vector(empty_fsl.into_array())?;
+        let empty_padded_vector = Vector::wrap_storage(empty_fsl.into_array())?;
 
         let sorf_options = SorfOptions {
             seed,
@@ -258,7 +255,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)?;
-    let padded_vector = wrap_padded_as_vector(quantized_fsl)?;
+    let padded_vector = Vector::wrap_storage(quantized_fsl)?;
 
     let sorf_options = SorfOptions {
         seed,
@@ -269,13 +266,6 @@ pub unsafe fn turboquant_encode_unchecked(
     Ok(SorfTransform::try_new_array(&sorf_options, padded_vector, num_rows)?.into_array())
 }
 
-/// Wrap an `FSL<f32, padded_dim>` in a [`Vector`](crate::vector::Vector) extension so it can be
-/// passed as the child of [`SorfTransform`], which expects a `Vector<padded_dim>` input.
-fn wrap_padded_as_vector(fsl: ArrayRef) -> VortexResult<ArrayRef> {
-    let ext_dtype = ExtDType::<Vector>::try_new(EmptyMetadata, fsl.dtype().clone())?.erased();
-    Ok(ExtensionArray::new(ext_dtype, fsl).into_array())
-}
-
 /// Apply the full TurboQuant compression pipeline to a [`Vector`](crate::vector::Vector)
 /// extension array: normalize the rows via [`normalize_as_l2_denorm`], quantize the normalized
 /// child via [`turboquant_encode_unchecked`], and reattach the stored norms as the outer

vortex-tensor/src/scalar_fns/sorf_transform/vtable.rs

@@ -143,9 +143,7 @@ impl ScalarFnVTable for SorfTransform {
                     validity,
                     0,
                 )?;
-                let ext_dtype =
-                    ExtDType::<Vector>::try_new(EmptyMetadata, fsl.dtype().clone())?.erased();
-                Ok(ExtensionArray::new(ext_dtype, fsl.into_array()).into_array())
+                Vector::wrap_storage(fsl.into_array())
             });
         }
 
@@ -330,7 +328,5 @@ fn inverse_rotate_typed<T: NativePType + Float + FromPrimitive>(
 
     let elements = PrimitiveArray::new::<T>(output.freeze(), Validity::NonNullable);
     let fsl = FixedSizeListArray::try_new(elements.into_array(), dim_u32, validity, num_rows)?;
-
-    let ext_dtype = ExtDType::<Vector>::try_new(EmptyMetadata, fsl.dtype().clone())?.erased();
-    Ok(ExtensionArray::new(ext_dtype, fsl.into_array()).into_array())
+    Vector::wrap_storage(fsl.into_array())
 }

vortex-tensor/src/vector/mod.rs

@@ -3,10 +3,36 @@
 
 //! Vector extension type for fixed-length float vectors (e.g., embeddings).
 
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::arrays::ExtensionArray;
+use vortex_array::dtype::extension::ExtDType;
+use vortex_array::extension::EmptyMetadata;
+use vortex_error::VortexResult;
+
 /// The Vector extension type.
 #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
 pub struct Vector;
 
+impl Vector {
+    /// Wrap a `FixedSizeList`-valued `storage` array in a [`Vector`] extension array.
+    ///
+    /// The storage's dtype is reused verbatim for the extension's storage dtype, so the caller
+    /// is responsible for having already constructed an FSL with the float element ptype and
+    /// non-nullable elements that [`Vector::validate_dtype`](ExtVTable::validate_dtype) requires.
+    ///
+    /// [`ExtVTable::validate_dtype`]: vortex_array::dtype::extension::ExtVTable::validate_dtype
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if `storage` does not satisfy [`Vector`]'s storage-dtype contract (e.g.
+    /// it is not a `FixedSizeList` of non-nullable floats).
+    pub fn wrap_storage(storage: ArrayRef) -> VortexResult<ArrayRef> {
+        let ext_dtype = ExtDType::<Self>::try_new(EmptyMetadata, storage.dtype().clone())?.erased();
+        Ok(ExtensionArray::new(ext_dtype, storage).into_array())
+    }
+}
+
 mod matcher;
 
 pub use matcher::AnyVector;

vortex-tensor/src/vector_search.rs

@@ -46,13 +46,10 @@ use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
 use vortex_array::arrays::ConstantArray;
-use vortex_array::arrays::ExtensionArray;
 use vortex_array::builtins::ArrayBuiltins;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::NativePType;
 use vortex_array::dtype::Nullability;
-use vortex_array::dtype::extension::ExtDType;
-use vortex_array::extension::EmptyMetadata;
 use vortex_array::scalar::PValue;
 use vortex_array::scalar::Scalar;
 use vortex_array::scalar_fn::fns::operators::Operator;
@@ -106,11 +103,8 @@ pub fn build_constant_query_vector<T: NativePType + Into<PValue>>(
         .map(|&v| Scalar::primitive(v, Nullability::NonNullable))
         .collect();
     let storage_scalar = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable);
-
     let storage = ConstantArray::new(storage_scalar, num_rows).into_array();
-
-    let ext_dtype = ExtDType::<Vector>::try_new(EmptyMetadata, storage.dtype().clone())?.erased();
-    Ok(ExtensionArray::new(ext_dtype, storage).into_array())
+    Vector::wrap_storage(storage)
 }
 
 /// Build the lazy similarity-search expression tree for a prepared database array and a