add most implementation

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

Author: connortsui20

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

@@ -2,6 +2,16 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex::array::ArrayRef;
+use vortex::array::ExecutionCtx;
+use vortex::dtype::DType;
+use vortex::dtype::Nullability;
+use vortex::error::VortexResult;
+use vortex::error::vortex_ensure;
+use vortex::error::vortex_ensure_eq;
+use vortex::error::vortex_err;
+
+use crate::utils::extension_element_ptype;
+use crate::vector::Vector;
 
 /// A normalized array that stores unit-normalized vectors alongside their original L2 norms.
 ///
@@ -12,15 +22,54 @@ pub struct NormVectorArray {
     /// The backing vector array that has been unit normalized.
     ///
     /// The underlying elements of the vector array must be floating-point.
-    vector_array: ArrayRef,
+    pub(crate) vector_array: ArrayRef,
 
     /// The L2 (Frobenius) norms of each vector.
     ///
     /// This must have the same dtype as the elements of the vector array.
-    norms: ArrayRef,
+    pub(crate) norms: ArrayRef,
 }
 
 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.
+    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 element_ptype = extension_element_ptype(ext)?;
+
+        let expected_norms_dtype = DType::Primitive(element_ptype, Nullability::NonNullable);
+        vortex_ensure_eq!(
+            *norms.dtype(),
+            expected_norms_dtype,
+            "norms dtype must match vector element type"
+        );
+
+        vortex_ensure_eq!(
+            vector_array.len(),
+            norms.len(),
+            "vector_array and norms must have the same length"
+        );
+
+        Ok(Self {
+            vector_array,
+            norms,
+        })
+    }
+
     /// Returns a reference to the backing vector array that has been unit normalized.
     pub fn vector_array(&self) -> &ArrayRef {
         &self.vector_array
@@ -30,4 +79,9 @@ impl NormVectorArray {
     pub fn norms(&self) -> &ArrayRef {
         &self.norms
     }
+
+    // TODO docs
+    pub(super) fn execute_into_vector(&self, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
+        todo!()
+    }
 }

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

@@ -3,6 +3,8 @@
 
 use std::hash::Hasher;
 
+use vortex::array::ArrayEq;
+use vortex::array::ArrayHash;
 use vortex::array::ArrayRef;
 use vortex::array::EmptyMetadata;
 use vortex::array::ExecutionCtx;
@@ -16,10 +18,15 @@ use vortex::array::vtable::ArrayId;
 use vortex::array::vtable::VTable;
 use vortex::array::vtable::ValidityVTableFromChild;
 use vortex::dtype::DType;
+use vortex::dtype::Nullability;
 use vortex::error::VortexResult;
+use vortex::error::vortex_ensure_eq;
+use vortex::error::vortex_err;
+use vortex::error::vortex_panic;
 use vortex::session::VortexSession;
 
 use crate::encodings::norm::array::NormVectorArray;
+use crate::utils::extension_element_ptype;
 
 mod operations;
 mod validity;
@@ -52,70 +59,109 @@ impl VTable for NormVector {
     }
 
     fn array_hash<H: Hasher>(array: &NormVectorArray, state: &mut H, precision: Precision) {
-        todo!()
+        array.vector_array().array_hash(state, precision);
+        array.norms().array_hash(state, precision);
     }
 
     fn array_eq(array: &NormVectorArray, other: &NormVectorArray, precision: Precision) -> bool {
-        todo!()
+        array.norms().array_eq(other.norms(), precision)
+            && array
+                .vector_array()
+                .array_eq(other.vector_array(), precision)
     }
 
-    fn nbuffers(array: &NormVectorArray) -> usize {
-        todo!()
+    fn nbuffers(_array: &NormVectorArray) -> usize {
+        0
     }
 
-    fn buffer(array: &NormVectorArray, idx: usize) -> BufferHandle {
-        todo!()
+    fn buffer(_array: &NormVectorArray, idx: usize) -> BufferHandle {
+        vortex_panic!("NormVectorArray has no buffers (index {idx})")
     }
 
-    fn buffer_name(array: &NormVectorArray, idx: usize) -> Option<String> {
-        todo!()
+    fn buffer_name(_array: &NormVectorArray, idx: usize) -> Option<String> {
+        vortex_panic!("NormVectorArray has no buffers (index {idx})")
     }
 
-    fn nchildren(array: &NormVectorArray) -> usize {
-        todo!()
+    fn nchildren(_array: &NormVectorArray) -> usize {
+        2
     }
 
     fn child(array: &NormVectorArray, idx: usize) -> ArrayRef {
-        todo!()
+        match idx {
+            0 => array.vector_array().clone(),
+            1 => array.norms().clone(),
+            _ => vortex_panic!("NormVectorArray child index {idx} out of bounds"),
+        }
     }
 
-    fn child_name(array: &NormVectorArray, idx: usize) -> String {
-        todo!()
+    fn child_name(_array: &NormVectorArray, idx: usize) -> String {
+        match idx {
+            0 => "vector_array".to_string(),
+            1 => "norms".to_string(),
+            _ => vortex_panic!("NormVectorArray child_name index {idx} out of bounds"),
+        }
     }
 
-    fn metadata(array: &NormVectorArray) -> VortexResult<Self::Metadata> {
-        todo!()
+    fn metadata(_array: &NormVectorArray) -> VortexResult<Self::Metadata> {
+        Ok(EmptyMetadata)
     }
 
-    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
-        todo!()
+    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
+        Ok(Some(vec![]))
     }
 
     fn deserialize(
-        bytes: &[u8],
+        _bytes: &[u8],
         _dtype: &DType,
         _len: usize,
         _buffers: &[BufferHandle],
         _session: &VortexSession,
     ) -> VortexResult<Self::Metadata> {
-        todo!()
+        Ok(EmptyMetadata)
     }
 
     fn build(
         dtype: &DType,
         len: usize,
-        metadata: &Self::Metadata,
-        buffers: &[BufferHandle],
+        _metadata: &Self::Metadata,
+        _buffers: &[BufferHandle],
         children: &dyn ArrayChildren,
     ) -> VortexResult<NormVectorArray> {
-        todo!()
+        vortex_ensure_eq!(
+            children.len(),
+            2,
+            "NormVectorArray requires exactly 2 children"
+        );
+
+        let vector_array = children.get(0, dtype, len)?;
+
+        let ext = dtype.as_extension_opt().ok_or_else(|| {
+            vortex_err!("NormVectorArray dtype must be an extension type, got {dtype}")
+        })?;
+        let element_ptype = extension_element_ptype(ext)?;
+        let norms_dtype = DType::Primitive(element_ptype, Nullability::NonNullable);
+        let norms = children.get(1, &norms_dtype, len)?;
+
+        NormVectorArray::try_new(vector_array, norms)
     }
 
     fn with_children(array: &mut NormVectorArray, children: Vec<ArrayRef>) -> VortexResult<()> {
-        todo!()
+        vortex_ensure_eq!(
+            children.len(),
+            2,
+            "NormVectorArray requires exactly 2 children"
+        );
+
+        let [vector_array, norms]: [ArrayRef; 2] = children
+            .try_into()
+            .map_err(|_| vortex_err!("NormVectorArray requires exactly 2 children"))?;
+
+        array.vector_array = vector_array;
+        array.norms = norms;
+        Ok(())
     }
 
     fn execute(array: &NormVectorArray, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
-        todo!()
+        Ok(ExecutionStep::Done(array.execute_into_vector(ctx)?))
     }
 }

vortex-tensor/src/encodings/norm/vtable/operations.rs

@@ -10,6 +10,6 @@ use crate::encodings::norm::vtable::NormVector;
 
 impl OperationsVTable<NormVector> for NormVector {
     fn scalar_at(array: &NormVectorArray, index: usize) -> VortexResult<Scalar> {
-        todo!()
+        array.vector_array().scalar_at(index)
     }
 }

vortex-tensor/src/lib.rs

@@ -12,3 +12,5 @@ pub mod fixed_shape;
 pub mod vector;
 
 pub mod encodings;
+
+mod utils;

vortex-tensor/src/scalar_fns/cosine_similarity.rs

@@ -28,10 +28,10 @@ use vortex::scalar_fn::ScalarFnId;
 use vortex::scalar_fn::ScalarFnVTable;
 
 use crate::matcher::AnyTensor;
-use crate::scalar_fns::utils::extension_element_ptype;
-use crate::scalar_fns::utils::extension_list_size;
-use crate::scalar_fns::utils::extension_storage;
-use crate::scalar_fns::utils::extract_flat_elements;
+use crate::utils::extension_element_ptype;
+use crate::utils::extension_list_size;
+use crate::utils::extension_storage;
+use crate::utils::extract_flat_elements;
 
 /// Cosine similarity between two columns.
 ///
@@ -196,11 +196,11 @@ mod tests {
     use vortex::scalar_fn::ScalarFn;
 
     use crate::scalar_fns::cosine_similarity::CosineSimilarity;
-    use crate::scalar_fns::utils::test_helpers::assert_close;
-    use crate::scalar_fns::utils::test_helpers::constant_tensor_array;
-    use crate::scalar_fns::utils::test_helpers::constant_vector_array;
-    use crate::scalar_fns::utils::test_helpers::tensor_array;
-    use crate::scalar_fns::utils::test_helpers::vector_array;
+    use crate::utils::test_helpers::assert_close;
+    use crate::utils::test_helpers::constant_tensor_array;
+    use crate::utils::test_helpers::constant_vector_array;
+    use crate::utils::test_helpers::tensor_array;
+    use crate::utils::test_helpers::vector_array;
 
     /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec<f64>`.
     fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef, len: usize) -> VortexResult<Vec<f64>> {

vortex-tensor/src/scalar_fns/l2_norm.rs

@@ -28,10 +28,10 @@ use vortex::scalar_fn::ScalarFnId;
 use vortex::scalar_fn::ScalarFnVTable;
 
 use crate::matcher::AnyTensor;
-use crate::scalar_fns::utils::extension_element_ptype;
-use crate::scalar_fns::utils::extension_list_size;
-use crate::scalar_fns::utils::extension_storage;
-use crate::scalar_fns::utils::extract_flat_elements;
+use crate::utils::extension_element_ptype;
+use crate::utils::extension_list_size;
+use crate::utils::extension_storage;
+use crate::utils::extract_flat_elements;
 
 /// L2 norm (Euclidean norm) of a tensor or vector column.
 ///
@@ -163,9 +163,9 @@ mod tests {
     use vortex::scalar_fn::ScalarFn;
 
     use crate::scalar_fns::l2_norm::L2Norm;
-    use crate::scalar_fns::utils::test_helpers::assert_close;
-    use crate::scalar_fns::utils::test_helpers::tensor_array;
-    use crate::scalar_fns::utils::test_helpers::vector_array;
+    use crate::utils::test_helpers::assert_close;
+    use crate::utils::test_helpers::tensor_array;
+    use crate::utils::test_helpers::vector_array;
 
     /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec<f64>`.
     fn eval_l2_norm(input: vortex::array::ArrayRef, len: usize) -> VortexResult<Vec<f64>> {

vortex-tensor/src/scalar_fns/mod.rs

@@ -5,5 +5,3 @@
 
 pub mod cosine_similarity;
 pub mod l2_norm;
-
-mod utils;