Execute tensor L2 norm with RowFn

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

Author: connortsui20

vortex-tensor/Cargo.toml

@@ -17,7 +17,7 @@ version = { workspace = true }
 workspace = true
 
 [dependencies]
-vortex-array = { workspace = true }
+vortex-array = { workspace = true, features = ["unstable_row_fns"] }
 vortex-arrow = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-compressor = { workspace = true }

vortex-tensor/src/scalar_fns/l2_norm.rs

@@ -3,56 +3,41 @@
 
 //! L2 norm expression for tensor-like types.
 
-use num_traits::Float;
 use prost::Message;
 use vortex_array::ArrayRef;
-use vortex_array::ExecutionCtx;
-use vortex_array::IntoArray;
-use vortex_array::arrays::Constant;
-use vortex_array::arrays::ConstantArray;
-use vortex_array::arrays::ExtensionArray;
-use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding;
 use vortex_array::arrays::ScalarFnArray;
-use vortex_array::arrays::extension::ExtensionArrayExt;
 use vortex_array::arrays::scalar_fn::ScalarFnArrayExt;
 use vortex_array::arrays::scalar_fn::ScalarFnArrayView;
 use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts;
 use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable;
 use vortex_array::dtype::DType;
-use vortex_array::dtype::NativePType;
-use vortex_array::dtype::Nullability;
 use vortex_array::dtype::proto::dtype as pb;
-use vortex_array::expr::Expression;
-use vortex_array::expr::union_child_validities;
 use vortex_array::match_each_float_ptype;
-use vortex_array::scalar::Scalar;
-use vortex_array::scalar_fn::Arity;
-use vortex_array::scalar_fn::ChildName;
 use vortex_array::scalar_fn::EmptyOptions;
-use vortex_array::scalar_fn::ExecutionArgs;
 use vortex_array::scalar_fn::ScalarFnId;
-use vortex_array::scalar_fn::ScalarFnVTable;
 use vortex_array::scalar_fn::ScalarFnVTableExt;
+use vortex_array::scalar_fn::unstable::row::InitializedElement;
+use vortex_array::scalar_fn::unstable::row::RowFn;
+use vortex_array::scalar_fn::unstable::row::RowVisitor;
+use vortex_array::scalar_fn::unstable::row::UninitElementSink;
 use vortex_array::serde::ArrayChildren;
-use vortex_buffer::Buffer;
-use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
 use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
-use crate::matcher::AnyTensor;
-use crate::utils::extract_flat_elements;
-use crate::utils::validate_tensor_float_input;
+use crate::scalar_fns::row::TensorRow;
+use crate::scalar_fns::row::tensor_element_ptype;
+use crate::utils::l2_norm_row;
 
 /// L2 norm (Euclidean norm) of a tensor or vector column.
 ///
 /// Computes `||v|| = sqrt(sum(v_i^2))` over the flat backing buffer of each tensor-like type.
 ///
 /// The input must be a tensor-like extension array with a float element type. The output is a float
 /// column of the same float type.
-#[derive(Clone)]
+#[derive(Clone, Debug, Default)]
 pub struct L2Norm;
 
 impl L2Norm {
@@ -67,110 +52,41 @@ impl L2Norm {
     }
 }
 
-impl ScalarFnVTable for L2Norm {
+impl RowFn for L2Norm {
     type Options = EmptyOptions;
 
+    const ARG_NAMES: &'static [&'static str] = &["input"];
+    const INFALLIBLE: bool = true;
+
     fn id(&self) -> ScalarFnId {
         static ID: CachedId = CachedId::new("vortex.tensor.l2_norm");
         *ID
     }
 
-    fn arity(&self, _options: &Self::Options) -> Arity {
-        Arity::Exact(1)
-    }
-
-    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
-        match child_idx {
-            0 => ChildName::from("input"),
-            _ => unreachable!("L2Norm must have exactly one child"),
-        }
-    }
-
-    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
-        let input_dtype = &arg_dtypes[0];
-        let tensor_match = validate_tensor_float_input(input_dtype)?;
-        let ptype = tensor_match.element_ptype();
-
-        let nullability = Nullability::from(input_dtype.is_nullable());
-        Ok(DType::Primitive(ptype, nullability))
+    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
+        Ok(Some(vec![]))
     }
 
-    fn execute(
+    fn deserialize(
         &self,
-        _options: &Self::Options,
-        args: &dyn ExecutionArgs,
-        ctx: &mut ExecutionCtx,
-    ) -> VortexResult<ArrayRef> {
-        let input_ref = args.get(0)?;
-        let row_count = args.row_count();
-
-        let ext = input_ref.dtype().as_extension();
-        let tensor_match = ext
-            .metadata_opt::<AnyTensor>()
-            .vortex_expect("we already validated this in `return_dtype`");
-        let tensor_flat_size = tensor_match.list_size() as usize;
-        let element_ptype = tensor_match.element_ptype();
-
-        let norm_dtype = DType::Primitive(element_ptype, ext.nullability());
-
-        // Optimize for the constant array case.
-        if let Some(array) = input_ref.as_opt::<Constant>() {
-            let scalar = array.scalar().as_extension().to_storage_scalar();
-
-            let Some(elements) = scalar.as_list().elements() else {
-                return Ok(ConstantArray::new(Scalar::null(norm_dtype), row_count).into_array());
-            };
-
-            let norm_scalar = match_each_float_ptype!(element_ptype, |T| {
-                let values: Vec<T> = elements
-                    .iter()
-                    .map(|s| {
-                        s.as_primitive()
-                            .as_::<T>()
-                            .vortex_expect("element was somehow not the correct float")
-                    })
-                    .collect();
-                let norm = l2_norm_row::<T>(&values);
-
-                Scalar::try_new(norm_dtype, Some(norm.into()))
-            })?;
-
-            let norms = ConstantArray::new(norm_scalar, row_count).into_array();
-            return Ok(norms);
-        }
-
-        let input: ExtensionArray = input_ref.execute(ctx)?;
-        let validity = input.as_ref().validity()?;
-
-        let storage = input.storage_array();
-        let flat = extract_flat_elements(storage, tensor_flat_size, ctx)?;
-
-        match_each_float_ptype!(flat.ptype(), |T| {
-            let buffer: Buffer<T> = (0..row_count)
-                .map(|i| l2_norm_row(flat.row::<T>(i)))
-                .collect();
-
-            // SAFETY: The buffer length equals `row_count`, which matches the source validity
-            // length.
-            Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array())
-        })
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
-    fn validity(
+    fn dispatch<V: RowVisitor>(
         &self,
         _options: &Self::Options,
-        expression: &Expression,
-    ) -> VortexResult<Option<Expression>> {
-        // The result is null if the input tensor is null.
-        union_child_validities(expression)
-    }
-
-    fn is_strict(&self, _options: &Self::Options) -> bool {
-        true
-    }
-
-    fn is_infallible(&self, _options: &Self::Options) -> bool {
-        true
+        args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        match_each_float_ptype!(tensor_element_ptype(args)?, |T| {
+            visitor.visit_into::<(TensorRow<T>,), UninitElementSink<T>, _>((), |(row,), output| {
+                // SAFETY: `output` is the `UninitElementSink` row supplied for this callback.
+                unsafe { InitializedElement::write(output, l2_norm_row(row)) }
+            })
+        })
     }
 }
 
@@ -217,17 +133,6 @@ impl ScalarFnArrayVTable for L2Norm {
     }
 }
 
-/// Computes the L2 norm (Euclidean norm) of a float slice.
-///
-/// Returns `sqrt(sum(v_i^2))`. A zero-length or all-zero input produces `0.0`.
-fn l2_norm_row<T: Float + NativePType>(v: &[T]) -> T {
-    let mut sum_sq = T::zero();
-    for &x in v {
-        sum_sq = sum_sq + x * x;
-    }
-    sum_sq.sqrt()
-}
-
 #[cfg(test)]
 mod tests {
 
@@ -258,6 +163,7 @@ mod tests {
     use crate::utils::test_helpers::literal_vector_array;
     use crate::utils::test_helpers::tensor_array;
     use crate::utils::test_helpers::vector_array;
+    use crate::utils::test_helpers::zero_width_vector_array;
 
     /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec<f64>`.
     fn eval_l2_norm(input: ArrayRef) -> VortexResult<Vec<f64>> {
@@ -267,6 +173,27 @@ mod tests {
         Ok(prim.as_slice::<f64>().to_vec())
     }
 
+    fn eval_l2_norm_f32(input: ArrayRef) -> VortexResult<Vec<f32>> {
+        let result = L2Norm::try_new(input)?;
+        let mut ctx = SESSION.create_execution_ctx();
+        let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?;
+        Ok(prim.as_slice::<f32>().to_vec())
+    }
+
+    #[test]
+    fn zero_width_and_empty_inputs() -> VortexResult<()> {
+        assert_close(
+            &eval_l2_norm(zero_width_vector_array::<f64>(3)?)?,
+            &[0.0, 0.0, 0.0],
+        );
+        assert!(eval_l2_norm(vector_array(2, &[] as &[f64])?)?.is_empty());
+        assert_close(
+            &eval_l2_norm(Vector::constant_array::<f64>(&[], 3)?)?,
+            &[0.0, 0.0, 0.0],
+        );
+        Ok(())
+    }
+
     #[rstest]
     #[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])]
     #[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])]
@@ -351,6 +278,48 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn extension_backed_constant_yields_constant_output() -> VortexResult<()> {
+        let input = Vector::constant_array(&[3.0f64, 4.0], 4)?;
+
+        let result = L2Norm::try_new(input)?.into_array();
+        let mut ctx = SESSION.create_execution_ctx();
+        let output = result.execute_until::<Constant>(&mut ctx)?;
+
+        let constant = output
+            .as_opt::<Constant>()
+            .expect("L2Norm over constant-backed extension storage must be constant");
+        assert_eq!(constant.len(), 4);
+        let norm = constant
+            .scalar()
+            .as_primitive()
+            .as_::<f64>()
+            .expect("norm scalar must be a non-null primitive");
+        assert_eq!(norm.to_bits(), 5.0f64.to_bits());
+        Ok(())
+    }
+
+    #[test]
+    fn encoded_constants_match_materialized_rows_bitwise() -> VortexResult<()> {
+        let row = [f32::MAX, 1.0, -1.0];
+        let materialized = vector_array(3, &[row, row, row].concat())?;
+        let storage_constant = Vector::constant_array(&row, 3)?;
+        let literal_constant = literal_vector_array(&row, 3);
+
+        let expected: Vec<_> = eval_l2_norm_f32(materialized)?
+            .into_iter()
+            .map(f32::to_bits)
+            .collect();
+        for encoded in [storage_constant, literal_constant] {
+            let actual: Vec<_> = eval_l2_norm_f32(encoded)?
+                .into_iter()
+                .map(f32::to_bits)
+                .collect();
+            assert_eq!(actual, expected);
+        }
+        Ok(())
+    }
+
     /// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of
     /// the correct primitive dtype and length.
     #[test]

vortex-tensor/src/scalar_fns/mod.rs

@@ -7,3 +7,5 @@ pub mod cosine_similarity;
 pub mod inner_product;
 pub mod l2_norm;
 pub mod l2_normalize;
+
+mod row;