fix problems

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

Author: connortsui20

vortex-tensor/public-api.lock

@@ -22,7 +22,7 @@ pub type vortex_tensor::encodings::norm::NormVector::Metadata = vortex_array::me
 
 pub type vortex_tensor::encodings::norm::NormVector::OperationsVTable = vortex_tensor::encodings::norm::NormVector
 
-pub type vortex_tensor::encodings::norm::NormVector::ValidityVTable = vortex_array::vtable::validity::ValidityVTableFromChild
+pub type vortex_tensor::encodings::norm::NormVector::ValidityVTable = vortex_array::vtable::validity::ValidityVTableFromValidityHelper
 
 pub fn vortex_tensor::encodings::norm::NormVector::array_eq(array: &vortex_tensor::encodings::norm::NormVectorArray, other: &vortex_tensor::encodings::norm::NormVectorArray, precision: vortex_array::hash::Precision) -> bool
 
@@ -52,7 +52,7 @@ pub fn vortex_tensor::encodings::norm::NormVector::metadata(_array: &vortex_tens
 
 pub fn vortex_tensor::encodings::norm::NormVector::nbuffers(_array: &vortex_tensor::encodings::norm::NormVectorArray) -> usize
 
-pub fn vortex_tensor::encodings::norm::NormVector::nchildren(_array: &vortex_tensor::encodings::norm::NormVectorArray) -> usize
+pub fn vortex_tensor::encodings::norm::NormVector::nchildren(array: &vortex_tensor::encodings::norm::NormVectorArray) -> usize
 
 pub fn vortex_tensor::encodings::norm::NormVector::serialize(_metadata: Self::Metadata) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
 
@@ -66,21 +66,17 @@ impl vortex_array::vtable::operations::OperationsVTable<vortex_tensor::encodings
 
 pub fn vortex_tensor::encodings::norm::NormVector::scalar_at(array: &vortex_tensor::encodings::norm::NormVectorArray, index: usize) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
 
-impl vortex_array::vtable::validity::ValidityChild<vortex_tensor::encodings::norm::NormVector> for vortex_tensor::encodings::norm::NormVector
-
-pub fn vortex_tensor::encodings::norm::NormVector::validity_child(array: &vortex_tensor::encodings::norm::NormVectorArray) -> &vortex_array::array::ArrayRef
-
 pub struct vortex_tensor::encodings::norm::NormVectorArray
 
 impl vortex_tensor::encodings::norm::NormVectorArray
 
-pub fn vortex_tensor::encodings::norm::NormVectorArray::compress(vector_array: vortex_array::array::ArrayRef) -> vortex_error::VortexResult<Self>
+pub fn vortex_tensor::encodings::norm::NormVectorArray::compress(vector_array: vortex_array::array::ArrayRef, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<Self>
 
-pub fn vortex_tensor::encodings::norm::NormVectorArray::decompress(&self, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
+pub fn vortex_tensor::encodings::norm::NormVectorArray::decompress(&self, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
 
 pub fn vortex_tensor::encodings::norm::NormVectorArray::norms(&self) -> &vortex_array::array::ArrayRef
 
-pub fn vortex_tensor::encodings::norm::NormVectorArray::try_new(vector_array: vortex_array::array::ArrayRef, norms: vortex_array::array::ArrayRef) -> vortex_error::VortexResult<Self>
+pub fn vortex_tensor::encodings::norm::NormVectorArray::try_new(vector_array: vortex_array::array::ArrayRef, norms: vortex_array::array::ArrayRef, validity: vortex_array::validity::Validity) -> vortex_error::VortexResult<Self>
 
 pub fn vortex_tensor::encodings::norm::NormVectorArray::vector_array(&self) -> &vortex_array::array::ArrayRef
 
@@ -114,6 +110,10 @@ impl vortex_array::array::IntoArray for vortex_tensor::encodings::norm::NormVect
 
 pub fn vortex_tensor::encodings::norm::NormVectorArray::into_array(self) -> vortex_array::array::ArrayRef
 
+impl vortex_array::vtable::validity::ValidityHelper for vortex_tensor::encodings::norm::NormVectorArray
+
+pub fn vortex_tensor::encodings::norm::NormVectorArray::validity(&self) -> &vortex_array::validity::Validity
+
 pub mod vortex_tensor::fixed_shape
 
 pub struct vortex_tensor::fixed_shape::FixedShapeTensor
@@ -250,7 +250,7 @@ pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::arity(&se
 
 pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::child_name(&self, _options: &Self::Options, child_idx: usize) -> vortex_array::scalar_fn::vtable::ChildName
 
-pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::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::ArrayRef>
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::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::ArrayRef>
 
 pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::expression::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
 
@@ -280,7 +280,7 @@ pub fn vortex_tensor::scalar_fns::l2_norm::L2Norm::arity(&self, _options: &Self:
 
 pub fn vortex_tensor::scalar_fns::l2_norm::L2Norm::child_name(&self, _options: &Self::Options, child_idx: usize) -> vortex_array::scalar_fn::vtable::ChildName
 
-pub fn vortex_tensor::scalar_fns::l2_norm::L2Norm::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::ArrayRef>
+pub fn vortex_tensor::scalar_fns::l2_norm::L2Norm::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::ArrayRef>
 
 pub fn vortex_tensor::scalar_fns::l2_norm::L2Norm::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::expression::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
 

vortex-tensor/src/encodings/mod.rs

@@ -2,4 +2,4 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 pub mod norm;
-// mod spherical;
+// TODO: Spherical coordinate encoding.

vortex-tensor/src/encodings/norm/array.rs

@@ -5,20 +5,22 @@ use num_traits::Float;
 use vortex::array::ArrayRef;
 use vortex::array::ExecutionCtx;
 use vortex::array::IntoArray;
-use vortex::array::ToCanonical;
 use vortex::array::arrays::ExtensionArray;
 use vortex::array::arrays::FixedSizeListArray;
 use vortex::array::arrays::PrimitiveArray;
-use vortex::array::arrays::ScalarFnArray;
 use vortex::array::match_each_float_ptype;
+use vortex::array::stats::ArrayStats;
 use vortex::array::validity::Validity;
 use vortex::dtype::DType;
 use vortex::dtype::Nullability;
 use vortex::dtype::extension::ExtDType;
+use vortex::dtype::extension::ExtDTypeRef;
 use vortex::error::VortexResult;
 use vortex::error::vortex_ensure;
 use vortex::error::vortex_ensure_eq;
 use vortex::error::vortex_err;
+use vortex::expr::Expression;
+use vortex::expr::root;
 use vortex::extension::EmptyMetadata;
 use vortex::scalar_fn::EmptyOptions;
 use vortex::scalar_fn::ScalarFn;
@@ -34,41 +36,39 @@ use crate::vector::Vector;
 ///
 /// Each vector in the array is divided by its L2 norm, producing a unit-normalized vector. The
 /// original norms are stored separately so that the original vectors can be reconstructed.
+///
+/// The `vector_array` child carries its own validity and nullability, so a nullable input vector
+/// array produces a nullable `NormVectorArray`.
 #[derive(Debug, Clone)]
 pub struct NormVectorArray {
     /// The backing vector array that has been unit normalized.
     ///
-    /// The underlying elements of the vector array must be floating-point.
+    /// The underlying elements of the vector array must be floating-point. This child may be
+    /// nullable; its validity determines the validity of the `NormVectorArray`.
     pub(crate) vector_array: ArrayRef,
 
-    /// The L2 (Frobenius) norms of each vector.
+    /// The L2 norms of each vector.
     ///
     /// This must have the same dtype as the elements of the vector array.
     pub(crate) norms: ArrayRef,
+
+    /// Stats set owned by this array.
+    pub(crate) stats_set: ArrayStats,
 }
 
 impl NormVectorArray {
     /// Creates a new [`NormVectorArray`] from a unit-normalized vector array and its L2 norms.
     ///
     /// The `vector_array` must be a [`Vector`] extension array with floating-point elements, and
-    /// `norms` must be a primitive array of the same float type with the same length.
+    /// `norms` must be a primitive array of the same float type with the same length. The
+    /// `vector_array` may be nullable.
     pub fn try_new(vector_array: ArrayRef, norms: ArrayRef) -> VortexResult<Self> {
-        let ext = vector_array.dtype().as_extension_opt().ok_or_else(|| {
-            vortex_err!(
-                "vector_array dtype must be an extension type, got {}",
-                vector_array.dtype()
-            )
-        })?;
-
-        vortex_ensure!(
-            ext.is::<Vector>(),
-            "vector_array must have the Vector extension type, got {}",
-            vector_array.dtype()
-        );
+        let ext = Self::validate(&vector_array)?;
 
-        let element_ptype = extension_element_ptype(ext)?;
+        let element_ptype = extension_element_ptype(&ext)?;
 
-        let expected_norms_dtype = DType::Primitive(element_ptype, Nullability::NonNullable);
+        let nullability = Nullability::from(vector_array.dtype().is_nullable());
+        let expected_norms_dtype = DType::Primitive(element_ptype, nullability);
         vortex_ensure_eq!(
             *norms.dtype(),
             expected_norms_dtype,
@@ -84,14 +84,13 @@ impl NormVectorArray {
         Ok(Self {
             vector_array,
             norms,
+            stats_set: ArrayStats::default(),
         })
     }
 
-    /// Encodes a [`Vector`] extension array into a [`NormVectorArray`] by computing L2 norms and
-    /// dividing each vector by its norm.
-    ///
-    /// The input must be a [`Vector`] extension array with floating-point elements.
-    pub fn compress(vector_array: ArrayRef) -> VortexResult<Self> {
+    /// Validates that the given array has the [`Vector`] extension type and returns the extension
+    /// dtype.
+    fn validate(vector_array: &ArrayRef) -> VortexResult<ExtDTypeRef> {
         let ext = vector_array.dtype().as_extension_opt().ok_or_else(|| {
             vortex_err!(
                 "vector_array dtype must be an extension type, got {}",
@@ -105,19 +104,32 @@ impl NormVectorArray {
             vector_array.dtype()
         );
 
-        let list_size = extension_list_size(ext)?;
-        let row_count = vector_array.len();
+        Ok(ext.clone())
+    }
 
-        // Compute L2 norms using the scalar function.
-        let l2_norm_fn = ScalarFn::new(L2Norm, EmptyOptions).erased();
-        let norms = ScalarFnArray::try_new(l2_norm_fn, vec![vector_array.clone()], row_count)?
-            .to_primitive()
-            .into_array();
+    /// Encodes a [`Vector`] extension array into a [`NormVectorArray`] by computing L2 norms and
+    /// dividing each vector by its norm.
+    ///
+    /// The input must be a [`Vector`] extension array with floating-point elements. Nullable inputs
+    /// are supported; the validity mask is preserved and the normalized data for null rows is
+    /// unspecified.
+    pub fn compress(vector_array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
+        let ext = Self::validate(&vector_array)?;
+
+        let list_size = extension_list_size(&ext)?;
+        let row_count = vector_array.len();
+        let nullability = Nullability::from(vector_array.dtype().is_nullable());
 
-        // Divide each vector element by its corresponding norm.
+        // Compute L2 norms using the scalar function. If the input is nullable, the norms will
+        // also be nullable (null vectors produce null norms).
         let storage = extension_storage(&vector_array)?;
-        let flat = extract_flat_elements(&storage, list_size)?;
-        let norms_prim = norms.to_canonical()?.into_primitive();
+        let l2_norm_expr =
+            Expression::try_new(ScalarFn::new(L2Norm, EmptyOptions).erased(), [root()])?;
+        let norms_prim: PrimitiveArray = vector_array.apply(&l2_norm_expr)?.execute(ctx)?;
+        let norms_array = norms_prim.clone().into_array();
+
+        // Extract flat elements from the (always non-nullable) storage for normalization.
+        let flat = extract_flat_elements(&storage, list_size, ctx)?;
 
         match_each_float_ptype!(flat.ptype(), |T| {
             let norms_slice = norms_prim.as_slice::<T>();
@@ -129,18 +141,20 @@ impl NormVectorArray {
                 })
                 .collect();
 
+            // Reconstruct the vector array with the same nullability as the input.
+            let validity = Validity::from(nullability);
             let fsl = FixedSizeListArray::new(
                 normalized_elems.into_array(),
                 u32::try_from(list_size)?,
-                Validity::NonNullable,
+                validity,
                 row_count,
             );
 
             let ext_dtype =
                 ExtDType::<Vector>::try_new(EmptyMetadata, fsl.dtype().clone())?.erased();
             let normalized_vector = ExtensionArray::new(ext_dtype, fsl.into_array()).into_array();
 
-            Self::try_new(normalized_vector, norms)
+            Self::try_new(normalized_vector, norms_array)
         })
     }
 
@@ -149,31 +163,26 @@ impl NormVectorArray {
         &self.vector_array
     }
 
-    /// Returns a reference to the L2 (Frobenius) norms of each vector.
+    /// Returns a reference to the L2 norms of each vector.
     pub fn norms(&self) -> &ArrayRef {
         &self.norms
     }
 
     /// Reconstructs the original vectors by multiplying each unit-normalized vector by its L2 norm.
-    pub fn decompress(&self, _ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
-        let ext_dtype = self
-            .vector_array
-            .dtype()
-            .as_extension_opt()
-            .ok_or_else(|| {
-                vortex_err!(
-                    "expected Vector extension dtype, got {}",
-                    self.vector_array.dtype()
-                )
-            })?;
-
-        let list_size = extension_list_size(ext_dtype)?;
+    ///
+    /// The returned array has the same dtype (including nullability) as the original
+    /// `vector_array` child.
+    pub fn decompress(&self, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
+        let ext = Self::validate(&self.vector_array)?;
+        let nullability = Nullability::from(self.vector_array.dtype().is_nullable());
+
+        let list_size = extension_list_size(&ext)?;
         let row_count = self.vector_array.len();
 
         let storage = extension_storage(&self.vector_array)?;
-        let flat = extract_flat_elements(&storage, list_size)?;
+        let flat = extract_flat_elements(&storage, list_size, ctx)?;
 
-        let norms_prim = self.norms.to_canonical()?.into_primitive();
+        let norms_prim: PrimitiveArray = self.norms.clone().execute(ctx)?;
 
         match_each_float_ptype!(flat.ptype(), |T| {
             let norms_slice = norms_prim.as_slice::<T>();
@@ -185,10 +194,11 @@ impl NormVectorArray {
                 })
                 .collect();
 
+            let validity = Validity::from(nullability);
             let fsl = FixedSizeListArray::new(
                 result_elems.into_array(),
                 u32::try_from(list_size)?,
-                Validity::NonNullable,
+                validity,
                 row_count,
             );
 

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

@@ -4,7 +4,7 @@
 mod array;
 pub use array::NormVectorArray;
 
-// pub(crate) mod compute;
+// TODO: Compute operations for NormVector.
 
 mod vtable;
 pub use vtable::NormVector;