From 5974184b1a0071714433ad40a53a226446e4c126 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 4 Sep 2026 14:12:01 -0400 Subject: [PATCH] Execute tensor L2 norm with RowFn Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/l2_norm.rs | 214 +++++++++++------------- vortex-tensor/src/utils.rs | 12 ++ 2 files changed, 108 insertions(+), 118 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index ea7b47a721e..295738eba60 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -9,45 +9,31 @@ 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::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::scalar_fns::arithmetic::l2_norm_row; -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; /// L2 norm (Euclidean norm) of a tensor or vector column. /// @@ -75,110 +61,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 { - 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>> { + Ok(Some(vec![])) } - fn execute( + fn deserialize( &self, - _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - 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::() - .vortex_expect("L2Norm::return_dtype validated the input tensor metadata"); - 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::() { - 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 = elements - .iter() - .map(|element| { - element - .as_primitive() - .as_::() - .vortex_expect("L2Norm::return_dtype validated the float element type") - }) - .collect(); - let norm = l2_norm_row::(&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 = (0..row_count) - .map(|row_index| l2_norm_row(flat.row::(row_index))) - .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 { + Ok(EmptyOptions) } - fn validity( + fn dispatch( &self, _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - 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 { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow,), UninitElementSink, _>((), |(row,), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, l2_norm_row(row)) } + }) + }) } } @@ -186,7 +103,7 @@ impl ScalarFnVTable for L2Norm { /// the extension type (`FixedShapeTensor` vs `Vector`), dimension, and nullability that are not /// recoverable from the parent's primitive-float output. #[derive(Clone, prost::Message)] -pub(super) struct L2NormMetadata { +struct L2NormMetadata { #[prost(message, optional, tag = "1")] input_dtype: Option, } @@ -240,6 +157,7 @@ mod tests { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; use vortex_array::dtype::DType; + use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::extension::ExtDType; @@ -254,13 +172,29 @@ 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`. - fn eval_l2_norm(input: ArrayRef) -> VortexResult> { + fn evaluate_l2_norm(input: ArrayRef) -> VortexResult> { 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::().to_vec()) + let output: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + Ok(output.as_slice::().to_vec()) + } + + #[test] + fn test_zero_width_and_empty_inputs() -> VortexResult<()> { + assert_close( + &evaluate_l2_norm(zero_width_vector_array::(3)?)?, + &[0.0, 0.0, 0.0], + ); + assert!(evaluate_l2_norm::(vector_array::(2, &[])?)?.is_empty()); + assert_close( + &evaluate_l2_norm(Vector::constant_array::(&[], 3)?)?, + &[0.0, 0.0, 0.0], + ); + + Ok(()) } #[rstest] @@ -274,7 +208,7 @@ mod tests { #[case] expected: &[f64], ) -> VortexResult<()> { let arr = tensor_array(shape, elements)?; - assert_close(&eval_l2_norm(arr)?, expected); + assert_close(&evaluate_l2_norm(arr)?, expected); Ok(()) } @@ -288,7 +222,7 @@ mod tests { 1.0, 1.0, 1.0, // norm = sqrt(3) ], )?; - assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); + assert_close(&evaluate_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); Ok(()) } @@ -301,7 +235,7 @@ mod tests { 3.0, 4.0, 0.0, // norm = 5.0 ], )?; - assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); + assert_close(&evaluate_l2_norm(arr)?, &[1.0, 5.0]); Ok(()) } @@ -347,6 +281,50 @@ mod tests { Ok(()) } + #[test] + fn test_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::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("L2Norm over constant-backed extension storage must be constant"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_eq!(norm.to_bits(), 5.0f64.to_bits()); + + Ok(()) + } + + #[test] + fn test_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<_> = evaluate_l2_norm::(materialized)? + .into_iter() + .map(f32::to_bits) + .collect(); + for encoded in [storage_constant, literal_constant] { + let actual: Vec<_> = evaluate_l2_norm::(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] diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 0e8768c2d13..965e26a2ac6 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -255,6 +255,18 @@ pub mod test_helpers { Vector::try_new_vector_array(flat_fsl(elements, dim)) } + /// Builds `row_count` zero-width vectors over an empty typed element buffer. + pub fn zero_width_vector_array(row_count: usize) -> VortexResult { + let storage = FixedSizeListArray::new( + Buffer::::empty().into_array(), + 0, + Validity::NonNullable, + row_count, + ) + .into_array(); + Vector::try_new_vector_array(storage) + } + /// Builds a [`FixedShapeTensor`] extension array whose storage is a [`ConstantArray`], /// representing a single query tensor broadcast to `len` rows. pub fn constant_tensor_array>(