clean up vortex-tensor

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

Author: connortsui20

vortex-array/public-api.lock

@@ -21546,6 +21546,8 @@ pub fn vortex_array::Array<vortex_array::arrays::Extension>::new(ext_dtype: vort
 
 pub fn vortex_array::Array<vortex_array::arrays::Extension>::try_new(ext_dtype: vortex_array::dtype::extension::ExtDTypeRef, storage_array: vortex_array::ArrayRef) -> vortex_error::VortexResult<Self>
 
+pub fn vortex_array::Array<vortex_array::arrays::Extension>::try_new_from_vtable<V: vortex_array::dtype::extension::ExtVTable>(vtable: V, metadata: <V as vortex_array::dtype::extension::ExtVTable>::Metadata, storage_array: vortex_array::ArrayRef) -> vortex_error::VortexResult<Self>
+
 impl vortex_array::Array<vortex_array::arrays::Filter>
 
 pub fn vortex_array::Array<vortex_array::arrays::Filter>::new(array: vortex_array::ArrayRef, mask: vortex_mask::Mask) -> Self

vortex-array/src/arrays/extension/array.rs

@@ -13,7 +13,9 @@ use crate::array::ArrayParts;
 use crate::array::TypedArrayRef;
 use crate::arrays::Extension;
 use crate::dtype::DType;
+use crate::dtype::extension::ExtDType;
 use crate::dtype::extension::ExtDTypeRef;
+use crate::dtype::extension::ExtVTable;
 
 /// The backing storage array for this extension array.
 pub(super) const STORAGE_SLOT: usize = 0;
@@ -163,4 +165,17 @@ impl Array<Extension> {
             )
         })
     }
+
+    /// Creates a new [`ExtensionArray`](crate::arrays::ExtensionArray) from a vtable, metadata, and
+    /// a storage array.
+    pub fn try_new_from_vtable<V: ExtVTable>(
+        vtable: V,
+        metadata: V::Metadata,
+        storage_array: ArrayRef,
+    ) -> VortexResult<Self> {
+        let ext_dtype =
+            ExtDType::<V>::try_with_vtable(vtable, metadata, storage_array.dtype().clone())?
+                .erased();
+        Self::try_new(ext_dtype, storage_array)
+    }
 }

vortex-array/src/dtype/extension/vtable.rs

@@ -14,8 +14,8 @@ use crate::scalar::ScalarValue;
 
 /// The public API for defining new extension types.
 ///
-/// This is the non-object-safe trait that plugin authors implement to define a new extension
-/// type. It specifies the type's identity, metadata, serialization, and validation.
+/// This is the non-object-safe trait that plugin authors implement to define a new extension type.
+/// It specifies the type's identity, metadata, serialization, and validation.
 pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash {
     /// Associated type containing the deserialized metadata for this extension type.
     type Metadata: 'static + Send + Sync + Clone + Debug + Display + Eq + Hash;
@@ -39,26 +39,27 @@ pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash {
     /// Validate that the given storage type is compatible with this extension type.
     fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()>;
 
-    /// Can a value of `other` be implicitly widened into this type?
-    /// e.g. GeographyType might accept Point, LineString, etc.
+    /// Can a value of `other` be implicitly widened into this type? (e.g. GeographyType might
+    /// accept Point, LineString, etc.)
     ///
-    /// Implementors only need to override one of `can_coerce_from` or `can_coerce_to` — both
-    /// exist so that either side of the coercion can provide the logic.
+    /// Implementors only need to override one of `can_coerce_from` or `can_coerce_to`. We have both
+    /// so that either side of the coercion can provide the logic.
     fn can_coerce_from(ext_dtype: &ExtDType<Self>, other: &DType) -> bool {
         let _ = (ext_dtype, other);
         false
     }
 
     /// Can this type be implicitly widened into `other`?
     ///
-    /// Implementors only need to override one of `can_coerce_from` or `can_coerce_to` — both
-    /// exist so that either side of the coercion can provide the logic.
+    /// Implementors only need to override one of `can_coerce_from` or `can_coerce_to`. We have both
+    /// so that either side of the coercion can provide the logic.
     fn can_coerce_to(ext_dtype: &ExtDType<Self>, other: &DType) -> bool {
         let _ = (ext_dtype, other);
         false
     }
 
     /// Given two types in a Uniform context, what is their least supertype?
+    ///
     /// Return None if no supertype exists.
     fn least_supertype(ext_dtype: &ExtDType<Self>, other: &DType) -> Option<DType> {
         let _ = (ext_dtype, other);
@@ -69,7 +70,8 @@ pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash {
 
     /// Validate the given storage value is compatible with the extension type.
     ///
-    /// By default, this calls [`unpack_native()`](ExtVTable::unpack_native) and discards the result.
+    /// By default, this calls [`unpack_native()`](ExtVTable::unpack_native) and discards the
+    /// result.
     ///
     /// # Errors
     ///

vortex-tensor/public-api.lock

@@ -80,7 +80,7 @@ pub const vortex_tensor::encodings::turboquant::MIN_DIMENSION: u32
 
 pub fn vortex_tensor::encodings::turboquant::tq_validate_vector_dtype(dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_tensor::vector::VectorMatcherMetadata>
 
-pub fn vortex_tensor::encodings::turboquant::turboquant_encode(ext: vortex_array::array::view::ArrayView<'_, vortex_array::arrays::extension::vtable::Extension>, config: &vortex_tensor::encodings::turboquant::TurboQuantConfig, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
+pub fn vortex_tensor::encodings::turboquant::turboquant_encode(input: vortex_array::array::erased::ArrayRef, config: &vortex_tensor::encodings::turboquant::TurboQuantConfig, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
 
 pub unsafe fn vortex_tensor::encodings::turboquant::turboquant_encode_unchecked(ext: vortex_array::array::view::ArrayView<'_, vortex_array::arrays::extension::vtable::Extension>, config: &vortex_tensor::encodings::turboquant::TurboQuantConfig, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
 
@@ -502,7 +502,7 @@ pub fn vortex_tensor::scalar_fns::sorf_transform::SorfTransform::child_name(&sel
 
 pub fn vortex_tensor::scalar_fns::sorf_transform::SorfTransform::execute(&self, options: &Self::Options, args: &dyn vortex_array::scalar_fn::vtable::ExecutionArgs, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
 
-pub fn vortex_tensor::scalar_fns::sorf_transform::SorfTransform::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::expression::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+pub fn vortex_tensor::scalar_fns::sorf_transform::SorfTransform::fmt_sql(&self, options: &Self::Options, expr: &vortex_array::expr::expression::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
 
 pub fn vortex_tensor::scalar_fns::sorf_transform::SorfTransform::id(&self) -> vortex_array::scalar_fn::ScalarFnId
 
@@ -570,10 +570,12 @@ pub struct vortex_tensor::vector::VectorMatcherMetadata
 
 impl vortex_tensor::vector::VectorMatcherMetadata
 
-pub fn vortex_tensor::vector::VectorMatcherMetadata::dimensions(&self) -> u32
+pub fn vortex_tensor::vector::VectorMatcherMetadata::dimensions(&self) -> usize
 
 pub fn vortex_tensor::vector::VectorMatcherMetadata::element_ptype(&self) -> vortex_array::dtype::ptype::PType
 
+pub fn vortex_tensor::vector::VectorMatcherMetadata::list_size(&self) -> u32
+
 pub fn vortex_tensor::vector::VectorMatcherMetadata::try_new(element_ptype: vortex_array::dtype::ptype::PType, dimensions: u32) -> vortex_error::VortexResult<Self>
 
 impl core::clone::Clone for vortex_tensor::vector::VectorMatcherMetadata
@@ -600,12 +602,8 @@ impl core::marker::StructuralPartialEq for vortex_tensor::vector::VectorMatcherM
 
 pub mod vortex_tensor::vector_search
 
-pub fn vortex_tensor::vector_search::build_constant_query_vector<T: vortex_array::dtype::ptype::NativePType + core::convert::Into<vortex_array::scalar::typed_view::primitive::pvalue::PValue>>(query: &[T], num_rows: usize) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
-
 pub fn vortex_tensor::vector_search::build_similarity_search_tree<T: vortex_array::dtype::ptype::NativePType + core::convert::Into<vortex_array::scalar::typed_view::primitive::pvalue::PValue>>(data: vortex_array::array::erased::ArrayRef, query: &[T], threshold: T) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
 
-pub fn vortex_tensor::vector_search::compress_turboquant(data: vortex_array::array::erased::ArrayRef, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
-
 pub const vortex_tensor::SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV: &str
 
 pub fn vortex_tensor::initialize(session: &vortex_session::VortexSession)