implement compress and decompress

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

Author: connortsui20

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

@@ -1,16 +1,33 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+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::validity::Validity;
 use vortex::dtype::DType;
 use vortex::dtype::Nullability;
+use vortex::dtype::extension::ExtDType;
 use vortex::error::VortexResult;
 use vortex::error::vortex_ensure;
 use vortex::error::vortex_ensure_eq;
 use vortex::error::vortex_err;
+use vortex::extension::EmptyMetadata;
+use vortex::scalar_fn::EmptyOptions;
+use vortex::scalar_fn::ScalarFn;
 
+use crate::scalar_fns::l2_norm::L2Norm;
 use crate::utils::extension_element_ptype;
+use crate::utils::extension_list_size;
+use crate::utils::extension_storage;
+use crate::utils::extract_flat_elements;
 use crate::vector::Vector;
 
 /// A normalized array that stores unit-normalized vectors alongside their original L2 norms.
@@ -70,6 +87,63 @@ impl NormVectorArray {
         })
     }
 
+    /// 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> {
+        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 list_size = extension_list_size(ext)?;
+        let row_count = vector_array.len();
+
+        // 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();
+
+        // Divide each vector element by its corresponding norm.
+        let storage = extension_storage(&vector_array)?;
+        let flat = extract_flat_elements(&storage, list_size)?;
+        let norms_prim = norms.to_canonical()?.into_primitive();
+
+        match_each_float_ptype!(flat.ptype(), |T| {
+            let norms_slice = norms_prim.as_slice::<T>();
+
+            let normalized_elems: PrimitiveArray = (0..row_count)
+                .flat_map(|i| {
+                    let inv_norm = safe_inv_norm(norms_slice[i]);
+                    flat.row::<T>(i).iter().map(move |&v| v * inv_norm)
+                })
+                .collect();
+
+            let fsl = FixedSizeListArray::new(
+                normalized_elems.into_array(),
+                u32::try_from(list_size)?,
+                Validity::NonNullable,
+                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)
+        })
+    }
+
     /// Returns a reference to the backing vector array that has been unit normalized.
     pub fn vector_array(&self) -> &ArrayRef {
         &self.vector_array
@@ -80,8 +154,58 @@ impl NormVectorArray {
         &self.norms
     }
 
-    // TODO docs
-    pub(super) fn execute_into_vector(&self, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
-        todo!()
+    /// 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)?;
+        let row_count = self.vector_array.len();
+
+        let storage = extension_storage(&self.vector_array)?;
+        let flat = extract_flat_elements(&storage, list_size)?;
+
+        let norms_prim = self.norms.to_canonical()?.into_primitive();
+
+        match_each_float_ptype!(flat.ptype(), |T| {
+            let norms_slice = norms_prim.as_slice::<T>();
+
+            let result_elems: PrimitiveArray = (0..row_count)
+                .flat_map(|i| {
+                    let norm = norms_slice[i];
+                    flat.row::<T>(i).iter().map(move |&v| v * norm)
+                })
+                .collect();
+
+            let fsl = FixedSizeListArray::new(
+                result_elems.into_array(),
+                u32::try_from(list_size)?,
+                Validity::NonNullable,
+                row_count,
+            );
+
+            let ext_dtype =
+                ExtDType::<Vector>::try_new(EmptyMetadata, fsl.dtype().clone())?.erased();
+            Ok(ExtensionArray::new(ext_dtype, fsl.into_array()).into_array())
+        })
+    }
+}
+
+/// Returns `1 / norm` if the norm is non-zero, or zero otherwise.
+///
+/// This avoids division by zero for zero-length or all-zero vectors.
+fn safe_inv_norm<T: Float>(norm: T) -> T {
+    if norm == T::zero() {
+        T::zero()
+    } else {
+        T::one() / norm
     }
 }

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

@@ -9,5 +9,5 @@ pub use array::NormVectorArray;
 mod vtable;
 pub use vtable::NormVector;
 
-// #[cfg(test)]
-// mod tests;
+#[cfg(test)]
+mod tests;

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

@@ -0,0 +1,110 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex::array::VortexSessionExecute;
+use vortex::array::arrays::Extension;
+use vortex::error::VortexResult;
+
+use crate::encodings::norm::NormVectorArray;
+use crate::utils::extension_list_size;
+use crate::utils::extension_storage;
+use crate::utils::extract_flat_elements;
+use crate::utils::test_helpers::assert_close;
+use crate::utils::test_helpers::vector_array;
+
+#[test]
+fn encode_unit_vectors() -> VortexResult<()> {
+    // Already unit-length vectors: norms should be 1.0 and vectors unchanged.
+    let arr = vector_array(
+        3,
+        &[
+            1.0, 0.0, 0.0, // norm = 1.0
+            0.0, 1.0, 0.0, // norm = 1.0
+        ],
+    )?;
+
+    let norm = NormVectorArray::compress(arr)?;
+    let norms = norm.norms().to_canonical()?.into_primitive();
+    assert_close(norms.as_slice::<f64>(), &[1.0, 1.0]);
+
+    let vectors = norm.vector_array();
+    let ext = vectors.dtype().as_extension_opt().unwrap();
+    let list_size = extension_list_size(ext)?;
+    let storage = extension_storage(vectors)?;
+    let flat = extract_flat_elements(&storage, list_size)?;
+    assert_close(flat.row::<f64>(0), &[1.0, 0.0, 0.0]);
+    assert_close(flat.row::<f64>(1), &[0.0, 1.0, 0.0]);
+
+    Ok(())
+}
+
+#[test]
+fn encode_non_unit_vectors() -> VortexResult<()> {
+    let arr = vector_array(
+        2,
+        &[
+            3.0, 4.0, // norm = 5.0
+            0.0, 0.0, // norm = 0.0 (zero vector)
+        ],
+    )?;
+
+    let norm = NormVectorArray::compress(arr)?;
+    let norms = norm.norms().to_canonical()?.into_primitive();
+    assert_close(norms.as_slice::<f64>(), &[5.0, 0.0]);
+
+    let vectors = norm.vector_array();
+    let ext = vectors.dtype().as_extension_opt().unwrap();
+    let list_size = extension_list_size(ext)?;
+    let storage = extension_storage(vectors)?;
+    let flat = extract_flat_elements(&storage, list_size)?;
+    assert_close(flat.row::<f64>(0), &[3.0 / 5.0, 4.0 / 5.0]);
+    assert_close(flat.row::<f64>(1), &[0.0, 0.0]);
+
+    Ok(())
+}
+
+#[test]
+fn execute_round_trip() -> VortexResult<()> {
+    let original_elements = &[
+        3.0, 4.0, // norm = 5.0
+        6.0, 8.0, // norm = 10.0
+    ];
+    let arr = vector_array(2, original_elements)?;
+
+    let norm = NormVectorArray::compress(arr)?;
+
+    // Execute to reconstruct the original vectors.
+    let mut ctx = vortex::array::LEGACY_SESSION.create_execution_ctx();
+    let reconstructed = norm.decompress(&mut ctx)?;
+
+    // The reconstructed array should be a Vector extension array.
+    assert!(reconstructed.as_opt::<Extension>().is_some());
+
+    let ext = reconstructed.dtype().as_extension_opt().unwrap();
+    let list_size = extension_list_size(ext)?;
+    let storage = extension_storage(&reconstructed)?;
+    let flat = extract_flat_elements(&storage, list_size)?;
+    assert_close(flat.row::<f64>(0), &[3.0, 4.0]);
+    assert_close(flat.row::<f64>(1), &[6.0, 8.0]);
+
+    Ok(())
+}
+
+#[test]
+fn execute_round_trip_zero_vector() -> VortexResult<()> {
+    let arr = vector_array(2, &[0.0, 0.0])?;
+
+    let norm = NormVectorArray::compress(arr)?;
+
+    let mut ctx = vortex::array::LEGACY_SESSION.create_execution_ctx();
+    let reconstructed = norm.decompress(&mut ctx)?;
+
+    let ext = reconstructed.dtype().as_extension_opt().unwrap();
+    let list_size = extension_list_size(ext)?;
+    let storage = extension_storage(&reconstructed)?;
+    let flat = extract_flat_elements(&storage, list_size)?;
+    // Zero vector should remain zero after round-trip.
+    assert_close(flat.row::<f64>(0), &[0.0, 0.0]);
+
+    Ok(())
+}

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

@@ -162,6 +162,6 @@ impl VTable for NormVector {
     }
 
     fn execute(array: &NormVectorArray, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
-        Ok(ExecutionStep::Done(array.execute_into_vector(ctx)?))
+        Ok(ExecutionStep::Done(array.decompress(ctx)?))
     }
 }