From 6716607ea5ccc0568b7b80851803a12a2d0bc74f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 12:57:15 -0400 Subject: [PATCH] Execute tensor L2 norm with RowFn Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 5 + .../unstable/row/batch/execute/mod.rs | 19 +- .../unstable/row/batch/execute/output.rs | 27 ++ .../src/scalar_fn/unstable/row/batch/tests.rs | 129 ++++++ .../src/scalar_fn/unstable/row/row_fn.rs | 37 +- .../src/scalar_fn/unstable/row/vtable.rs | 1 + vortex-tensor/Cargo.toml | 2 +- vortex-tensor/benches/l2_norm.rs | 30 ++ vortex-tensor/src/scalar_fns/l2_norm.rs | 413 ++++-------------- vortex-tensor/src/scalar_fns/mod.rs | 4 + vortex-tensor/src/scalar_fns/row.rs | 169 +++++++ vortex-tensor/src/scalar_fns/tests/l2_norm.rs | 323 ++++++++++++++ vortex-tensor/src/scalar_fns/tests/mod.rs | 7 + vortex-tensor/src/scalar_fns/tests/row.rs | 104 +++++ vortex-tensor/src/utils.rs | 59 +++ 15 files changed, 984 insertions(+), 345 deletions(-) create mode 100644 vortex-tensor/src/scalar_fns/row.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/l2_norm.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/mod.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/row.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index d45a66529a3..c213147b9a9 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -55,6 +55,11 @@ impl<'a> BorrowedRowFnArgs<'a> { } } + /// Return the concrete arrays used by encoding-aware execution. + pub(crate) fn arrays(&self) -> &'a [ArrayRef] { + self.arrays + } + /// Return the original input dtypes used to select the row implementation. pub(crate) fn dtypes(&self) -> &'a [DType] { self.dtypes diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs index 60a23a1154d..d3296b4ef58 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -3,8 +3,8 @@ //! Selects a batch execution strategy. //! -//! [`RowFnExecutionArgs::execute`] handles universal fast paths, then delegates to dense or -//! valid-only execution. +//! [`RowFnExecutionArgs::execute`] handles universal fast paths and encoded reductions, then +//! delegates to dense or valid-only execution. use vortex_error::VortexResult; use vortex_mask::Mask; @@ -27,12 +27,14 @@ mod output; pub(crate) use output::finalize_kernel_output; impl RowFnExecutionArgs { - /// Apply constant folding and null handling around `kernel`. + /// Apply encoded reductions, constant folding, and null handling around `kernel`. /// - /// For a partially valid batch, `try_valid_rows` executes only valid rows over the original - /// inputs. Every result is checked against the planned shape and dtype. + /// `reduce` receives the original inputs before constant broadcasting. For a partially valid + /// batch, `try_valid_rows` executes only valid rows over the original inputs. Every result is + /// checked against the planned shape and dtype. pub(crate) fn execute( &self, + reduce: impl FnOnce(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult>, kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, try_valid_rows: impl FnOnce( BorrowedRowFnArgs<'_>, @@ -53,6 +55,13 @@ impl RowFnExecutionArgs { return Ok(self.all_null()); } + // Let the ordinary policy construct the typed output for an empty batch. + if self.row_count > 0 + && let Some(values) = reduce(self.execution_args(&self.inputs, self.row_count), ctx)? + { + return self.finalize_reduced(values, ctx); + } + // All inputs are constant, and their conjoined validity proves that every row is non-null. // The constant check sees through extension and masked wrappers, just like argument // decoding. diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs index 0f8170630e3..cf44403d532 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -14,6 +14,7 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnId; +use crate::validity::Validity; impl RowFnExecutionArgs { pub(super) fn all_null(&self) -> ArrayRef { @@ -31,6 +32,32 @@ impl RowFnExecutionArgs { cast_output_nullability(&self.result_dtype, values) } + /// Reconcile an encoding-aware result and apply the batch's strict input validity. + pub(super) fn finalize_reduced( + &self, + values: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + validate_output(self.id, &self.result_dtype, self.row_count, &values)?; + + let input_valid = self.validity.execute_mask(self.row_count, ctx)?; + let output_valid = values.validity()?.execute_mask(self.row_count, ctx)?; + vortex_ensure!( + input_valid.bitand_not(&output_valid).all_false(), + "the {} encoded reduction produced nulls for valid rows", + self.id, + ); + + let values = match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => values, + Validity::Array(valid) => values.mask(valid)?, + // Handled before the encoding-aware hook runs. + Validity::AllInvalid => return Ok(self.all_null()), + }; + + cast_output_nullability(&self.result_dtype, values) + } + /// Validate the output from a row function before batch validity is attached. pub(super) fn validate_kernel_output( &self, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index dfa83cc7017..4b7a1296f42 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -3,6 +3,7 @@ use std::sync::Arc; +use rstest::rstest; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -10,6 +11,7 @@ use vortex_session::registry::CachedId; use super::finalize_kernel_output; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; @@ -30,6 +32,12 @@ use crate::scalar_fn::unstable::row::RowVisitor; use crate::scalar_fn::unstable::row::execute_rows; use crate::validity::Validity; +#[derive(Clone)] +struct OriginalInputReducer; + +#[derive(Clone)] +struct InvalidEncodedReduction; + #[derive(Clone)] struct DeferredAdd; @@ -87,6 +95,76 @@ unsafe impl OutputSink for I64Sink { } } +impl RowFn for OriginalInputReducer { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.original_input_reducer"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 3 { + return Ok(Some(ConstantArray::new(42_i64, 3).into_array())); + } + + Ok(None) + } +} + +impl RowFn for InvalidEncodedReduction { + type Options = usize; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.invalid_encoded_reduction"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + null_index: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some( + PrimitiveArray::new( + vec![10_i64, 20], + Validity::from_iter((0..2).map(|index| index != *null_index)), + ) + .into_array(), + )) + } +} + impl RowFn for DeferredAdd { type Options = EmptyOptions; @@ -182,6 +260,57 @@ fn test_finalize_kernel_output_rejects_nested_dtype_mismatch() -> VortexResult<( Ok(()) } +#[rstest] +#[case::all_valid(Validity::AllValid)] +#[case::partially_valid(Validity::from_iter([true, false]))] +fn test_reduce_encoded_rejects_nulls_on_valid_rows(#[case] validity: Validity) -> VortexResult<()> { + let input = PrimitiveArray::new(vec![10_i64, 20], validity).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&InvalidEncodedReduction, &0, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("an encoded reduction introduced a null on a valid row"), + }; + let error = error.to_string(); + + assert!( + error.contains("test.invalid_encoded_reduction"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("encoded reduction produced nulls for valid rows"), + "the boundary error must identify invalid reduced output, got {error}", + ); + Ok(()) +} + +#[test] +fn test_reduce_encoded_preserves_input_nulls() -> VortexResult<()> { + let input = + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&InvalidEncodedReduction, &1, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 3).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + let expected = ConstantArray::new(42_i64, 3).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + #[test] fn test_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index e456886e618..37685beadb4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -4,8 +4,8 @@ //! The [`RowFn`] contract for scalar functions whose natural kernel computes one row at a time. //! //! Implementations declare their arity and fallibility, then use [`RowFn::dispatch`] to select the -//! typed row signature for each supported dtype combination. Optional methods provide -//! serialization without putting persistence plumbing in the row kernel. +//! typed row signature for each supported dtype combination. Optional hooks provide serialization +//! and encoding-aware execution without putting columnar plumbing in the row kernel. use std::fmt::Debug; use std::fmt::Display; @@ -16,6 +16,8 @@ use vortex_error::vortex_bail; use vortex_session::VortexSession; use super::visitor::RowVisitor; +use crate::ArrayRef; +use crate::ExecutionCtx; use crate::dtype::DType; use crate::scalar_fn::ScalarFnId; @@ -47,12 +49,14 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// The arguments in display order. Its length is the function's exact arity. const ARG_NAMES: &'static [&'static str]; - /// Whether every dispatch is infallible. + /// Whether every dispatch and encoded reduction is infallible. /// /// See [`ScalarFnVTable::is_infallible`](crate::scalar_fn::ScalarFnVTable::is_infallible) for /// a more detailed explanation of semantic errors. /// - /// The framework checks dispatched element and result types. A conservative `false` is allowed. + /// The framework checks dispatched element and result types, but cannot inspect + /// [`reduce_encoded`](Self::reduce_encoded). Set this to `false` when that hook can return a + /// semantic error. A conservative `false` is allowed. const INFALLIBLE: bool; /// Returns the ID of the scalar function. @@ -83,4 +87,29 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { args: &[DType], visitor: V, ) -> VortexResult; + + /// Try an encoding-aware implementation before decoding the inputs into row elements. + /// + /// `None` continues to the row loop. The returned array can remain encoded or lazy. Batch + /// execution calls this hook at most once with the original nonempty inputs. Nullary functions, + /// empty batches, slices, and compacted retries skip it. + /// + /// Like a dense row closure, this hook must be total over every stored payload, including + /// payloads behind null rows. An `Err` is immediately user-visible and is never suppressed or + /// retried through the row layer. + /// + /// # Requirements + /// + /// - `output.len()` **must** equal `args[0].len()`. + /// - The output dtype **must** match the planned dtype when ignoring outer nullability. + /// - The output **must not** introduce a null where every input is valid. + fn reduce_encoded( + &self, + options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + _ = (options, args, ctx); + Ok(None) + } } diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index c2b8f44f6c3..d1f978861eb 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -118,6 +118,7 @@ pub fn execute_rows( let batch = prepare_batch(function, options, args)?; batch.execute( + |args, ctx| function.reduce_encoded(options, args.arrays(), ctx), |args, ctx| execute_row_kernel(function, options, args, ctx), |args, valid, ctx| try_execute_valid_rows(function, options, args, valid, ctx), ctx, diff --git a/vortex-tensor/Cargo.toml b/vortex-tensor/Cargo.toml index abdca676775..9706102f6d6 100644 --- a/vortex-tensor/Cargo.toml +++ b/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 } diff --git a/vortex-tensor/benches/l2_norm.rs b/vortex-tensor/benches/l2_norm.rs index bb5dc0dd0ab..f9beb7506a9 100644 --- a/vortex-tensor/benches/l2_norm.rs +++ b/vortex-tensor/benches/l2_norm.rs @@ -15,11 +15,17 @@ use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::l2_norm::L2Norm; @@ -54,6 +60,16 @@ fn vectors(width: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +fn constant_vector(width: usize) -> ArrayRef { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let children = (0..width) + .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let vector = Scalar::extension::(EmptyMetadata, storage); + ConstantArray::new(vector, ELEMENTS / width).into_array() +} + fn bench_l2_norm(bencher: Bencher, input: ArrayRef) { let session = vortex_array::array_session(); bencher @@ -80,3 +96,17 @@ fn nullable(bencher: Bencher, width: usize) { .into_array(); bench_l2_norm(bencher, input); } + +#[divan::bench(args = WIDTHS)] +fn constant(bencher: Bencher, width: usize) { + bench_l2_norm(bencher, constant_vector(width)); +} + +#[divan::bench(args = WIDTHS)] +fn nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let input = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_l2_norm(bencher, input); +} diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 0b00ec95aa4..aad84138a1f 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -3,51 +3,43 @@ //! 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_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::Normalized; -use crate::matcher::AnyTensor; -use crate::utils::extract_flat_elements; +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::extract_normalized_children; -use crate::utils::reattach_validity; +use crate::utils::l2_norm_row; use crate::utils::validate_tensor_float_input; /// L2 norm (Euclidean norm) of a tensor or vector column. @@ -63,7 +55,7 @@ use crate::utils::validate_tensor_float_input; /// of the storage contract, not a separate lossy-compute mode. /// /// [`Normalized`]: crate::encodings::normalized::Normalized -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct L2Norm; impl L2Norm { @@ -71,126 +63,99 @@ impl L2Norm { /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtype is + /// unsupported. pub fn try_new(child: ArrayRef) -> VortexResult { ScalarFnArray::try_new(L2Norm.bind(EmptyOptions), vec![child]) } } -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 serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - 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 deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn execute( + fn dispatch>( &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("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()); - - // Stored norms are authoritative. Reattach the parent validity because the child is - // non-nullable. - if input_ref.is::() { - let (_, norms) = extract_normalized_children(&input_ref); - let norms = reattach_validity(norms, input_ref.validity()?)?; - vortex_ensure_eq!(norms.dtype(), &norm_dtype); - return Ok(norms); - } - - // 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(|s| { - s.as_primitive() - .as_::() - .vortex_expect("element was somehow not the correct float") - }) - .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(|i| l2_norm_row(flat.row::(i))) - .collect(); - - // SAFETY: The buffer length equals `row_count`, which matches the source validity - // length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) + 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)) } + }) }) } - fn validity( + /// `L2Norm` over a [`Normalized`]-encoded column is defined to read back the authoritative + /// stored norms. Callers of lossy encodings opt into that storage semantics instead of forcing + /// a decode-and-recompute path here. + fn reduce_encoded( &self, _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if the input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let input = &args[0]; + if input.is::() { + let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); + let (_, norms) = extract_normalized_children(input); + vortex_ensure!( + norms.dtype().is_primitive(), + "normalized norms must be primitive, got {}", + norms.dtype(), + ); + vortex_ensure_eq!(norms.dtype().as_ptype(), element_ptype); + return Ok(Some(norms)); + } - fn is_infallible(&self, _options: &Self::Options) -> bool { - true + let Some(constant) = input.as_opt::() else { + return Ok(None); + }; + let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); + let norm_dtype = + DType::Primitive(element_ptype, input.dtype().as_extension().nullability()); + let storage = constant.scalar().as_extension().to_storage_scalar(); + + let Some(elements) = storage.as_list().elements() else { + let output = ConstantArray::new(Scalar::null(norm_dtype), input.len()); + return Ok(Some(output.into_array())); + }; + + let norm = match_each_float_ptype!(element_ptype, |T| { + let values: Vec = elements + .iter() + .map(|element| { + element + .as_primitive() + .as_::() + .vortex_expect("tensor element must match its declared ptype") + }) + .collect(); + Scalar::try_new(norm_dtype, Some(l2_norm_row::(&values).into())) + })?; + let output = ConstantArray::new(norm, input.len()); + Ok(Some(output.into_array())) } } @@ -236,225 +201,3 @@ 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(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 { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::EmptyMetadata; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::Constant; - use vortex_array::arrays::ConstantArray; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::dtype::extension::ExtDType; - use vortex_array::scalar::Scalar; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::l2_norm::L2Norm; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::literal_vector_array; - 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`. - fn eval_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()) - } - - #[rstest] - #[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])] - #[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])] - #[case::single_element(&[1], &[7.0], &[7.0])] - #[case::negative_elements(&[2], &[-3.0, -4.0], &[5.0])] - fn known_norms( - #[case] shape: &[usize], - #[case] elements: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let arr = tensor_array(shape, elements)?; - assert_close(&eval_l2_norm(arr)?, expected); - Ok(()) - } - - #[test] - fn multiple_rows() -> VortexResult<()> { - let arr = tensor_array( - &[3], - &[ - 3.0, 4.0, 0.0, // norm = 5.0 - 0.0, 0.0, 0.0, // norm = 0.0 - 1.0, 1.0, 1.0, // norm = sqrt(3) - ], - )?; - assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); - Ok(()) - } - - #[test] - fn vector_multiple_rows() -> VortexResult<()> { - let arr = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // norm = 1.0 - 3.0, 4.0, 0.0, // norm = 5.0 - ], - )?; - assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 2 rows of dim-2 vectors. Row 1 is masked as null. - let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; - let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); - - let result = L2Norm::try_new(arr)?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: norm = 5.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[5.0]); - Ok(()) - } - - /// A constant input whose scalar is a non-null tensor should short-circuit to a - /// [`ConstantArray`] output whose scalar is the precomputed norm. Uses [`execute_until`] so - /// execution stops at the [`Constant`] encoding instead of canonicalizing into a - /// [`PrimitiveArray`]. - #[test] - fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { - let input = literal_vector_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 a constant input must produce a constant output"); - assert_eq!(constant.len(), 4); - let norm = constant - .scalar() - .as_primitive() - .as_::() - .expect("norm scalar must be a non-null primitive"); - assert_close(&[norm], &[5.0]); - Ok(()) - } - - /// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of - /// the correct primitive dtype and length. - #[test] - fn constant_null_input_yields_null_constant_output() -> VortexResult<()> { - let storage_dtype = DType::FixedSizeList( - DType::Primitive(PType::F64, Nullability::NonNullable).into(), - 2, - Nullability::Nullable, - ); - let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(); - let null_scalar = Scalar::null(DType::Extension(ext_dtype)); - let input = ConstantArray::new(null_scalar, 3).into_array(); - - 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("null constant input must produce a constant output"); - assert_eq!(constant.len(), 3); - assert!(constant.scalar().is_null()); - assert_eq!( - constant.dtype(), - &DType::Primitive(PType::F64, Nullability::Nullable) - ); - Ok(()) - } - - #[test] - fn reads_through_a_nullable_normalized_column() -> VortexResult<()> { - let normalized = vector_array(2, &[0.6f64, 0.8, 1.0, 0.0])?; - let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); - - let mut ctx = SESSION.create_execution_ctx(); - let validity = Validity::from_iter([true, false]); - let input = Normalized::try_new(normalized, norms, validity, &mut ctx)?.into_array(); - - let result = L2Norm::try_new(input)?.into_array(); - let prim: PrimitiveArray = result.execute(&mut ctx)?; - - assert_eq!( - prim.dtype(), - &DType::Primitive(PType::F64, Nullability::Nullable) - ); - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[5.0]); - - Ok(()) - } - - #[rstest] - #[case::fixed_shape_tensor(l2_norm_tensor_child())] - #[case::vector(l2_norm_vector_child())] - fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { - let original = L2Norm::try_new(child.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(L2Norm); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("L2Norm serialize must produce metadata"); - - let children = vec![child]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn l2_norm_tensor_child() -> ArrayRef { - tensor_array(&[3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid tensor array") - } - - fn l2_norm_vector_child() -> ArrayRef { - vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") - } -} diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index 68f10ca6b01..da9b8950e7a 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -6,3 +6,7 @@ pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; +pub(crate) mod row; + +#[cfg(test)] +mod tests; diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs new file mode 100644 index 00000000000..ab3b6043194 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What the tensor scalar functions add to the row-function machinery: an element type that reads a +//! tensor row and the width rule they share. + +use std::marker::PhantomData; + +use num_traits::Float; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::Masked; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::masked::MaskedArraySlotsExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::scalar_fn::unstable::row::InputElement; +use vortex_array::scalar_fn::unstable::row::ViewLen; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; + +use crate::utils::extract_flat_elements; +use crate::utils::validate_tensor_float_input; +use crate::utils::validate_tensor_float_inputs; + +/// The width rule the tensor scalar functions share: every argument is the same float tensor dtype, +/// and the width is its element ptype. +pub(crate) fn tensor_element_ptype(args: &[DType]) -> VortexResult { + Ok(validate_tensor_float_inputs(args)?.element_ptype()) +} + +/// Marker for tensor-valued input elements: accepts any tensor-like extension column whose +/// elements are `T`, and presents each row as its flat elements, `&[T]`. +pub struct TensorRow(PhantomData); + +/// The decoded form of a [`TensorRow`] column: one flat typed buffer plus the stride to read it at. +/// +/// Typed at decode time rather than per row. `FlatElements::row` re-derives its typed slice on every +/// call, which costs a ptype check and a buffer downcast per row; a row loop reads every row, so it +/// pays that once here instead. +pub struct TensorRows { + /// Every row's elements, back to back. + elements: Buffer, + + /// Number of logical tensor rows, stored so zero-width tensors retain their length. + rows: usize, + + /// Elements per row, the length of each row slice. + list_size: usize, + + /// `list_size` for a full column and `0` for constant-backed storage, so `index * stride` pins a + /// constant to its single materialized row without a branch in the loop. + stride: usize, +} + +impl ViewLen for TensorRows { + fn len(&self) -> usize { + self.rows + } +} + +// SAFETY: `TensorRows` records the row count validated during decode, and both checked and +// unchecked access use the same stride and row width. +unsafe impl InputElement for TensorRow { + type Column = TensorRows; + type View<'a> = &'a TensorRows; + type Elem<'a> = &'a [T]; + + // Tensor storage is a fully materialized non-nullable primitive buffer, so the elements behind + // a null row are arbitrary values rather than an unresolvable reference. + const DENSE_SAFE: bool = true; + // Tensor storage is a primitive buffer; reading it cannot fail on account of its values. + const DECODE_INFALLIBLE: bool = true; + + fn validate(dtype: &DType) -> VortexResult<()> { + let tensor_match = validate_tensor_float_input(dtype)?; + let expected = T::PTYPE; + vortex_ensure_eq!( + tensor_match.element_ptype(), + expected, + "expected a tensor of {expected} elements, got {dtype}", + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // Dense batch execution owns the mask and restores it on the result. Decode the values + // directly so a nullable tensor does not rebuild its extension storage under that mask. + let array = match array.as_opt::() { + Some(masked) => masked.child().clone(), + None => array, + }; + + let rows = array.len(); + let list_size = validate_tensor_float_input(array.dtype())?.list_size() as usize; + let ext: ExtensionArray = array.execute(ctx)?; + let flat = extract_flat_elements(ext.storage_array(), list_size, ctx)?; + let list_size = flat.list_size(); + let stride = flat.row_stride(); + let elements = flat.into_buffer::(); + + let expected_elements = if stride == 0 { + list_size + } else { + vortex_ensure_eq!( + stride, + list_size, + "per-row tensor stride must equal its width, got {stride}", + ); + let Some(expected_elements) = rows.checked_mul(stride) else { + vortex_bail!( + "tensor row storage length must fit usize, got {rows} rows of width {stride}", + ); + }; + + expected_elements + }; + vortex_ensure_eq!( + elements.len(), + expected_elements, + "tensor row storage must contain {expected_elements} elements, got {}", + elements.len(), + ); + + Ok(TensorRows { + elements, + rows, + list_size, + stride, + }) + } + + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + + fn get(column: &Self::Column, index: usize) -> &[T] { + let start = index * column.stride; + &column.elements.as_slice()[start..start + column.list_size] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + Self::get(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + let start = index * view.stride; + + // SAFETY: decode established one complete stored row for stride 0, or `rows` contiguous + // `list_size`-element rows otherwise. The caller guarantees `index < rows`. + unsafe { + std::slice::from_raw_parts(view.elements.as_slice().as_ptr().add(start), view.list_size) + } + } +} diff --git a/vortex-tensor/src/scalar_fns/tests/l2_norm.rs b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs new file mode 100644 index 00000000000..1bf2bde1dd9 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs @@ -0,0 +1,323 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::l2_norm::L2Norm; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +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> { + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let array = L2Norm::try_new(tensor_array(&[1], &[3.0])?)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + assert_close( + &eval_l2_norm(zero_width_vector_array::(3)?)?, + &[0.0, 0.0, 0.0], + ); + assert!(eval_l2_norm(vector_array(2, &[] as &[f64])?)?.is_empty()); + + let constant = Vector::constant_array::(&[], 3)?; + assert_close(&eval_l2_norm(constant)?, &[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])] +#[case::single_element(&[1], &[7.0], &[7.0])] +#[case::negative_elements(&[2], &[-3.0, -4.0], &[5.0])] +fn known_norms( + #[case] shape: &[usize], + #[case] elements: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let arr = tensor_array(shape, elements)?; + assert_close(&eval_l2_norm(arr)?, expected); + Ok(()) +} + +#[test] +fn multiple_rows() -> VortexResult<()> { + let arr = tensor_array( + &[3], + &[ + 3.0, 4.0, 0.0, // norm = 5.0 + 0.0, 0.0, 0.0, // norm = 0.0 + 1.0, 1.0, 1.0, // norm = sqrt(3) + ], + )?; + assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); + Ok(()) +} + +#[test] +fn vector_multiple_rows() -> VortexResult<()> { + let arr = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // norm = 1.0 + 3.0, 4.0, 0.0, // norm = 5.0 + ], + )?; + assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 2 rows of dim-2 vectors. Row 1 is masked as null. + let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![arr])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: norm = 5.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +/// A constant input whose scalar is a non-null tensor should short-circuit to a +/// [`ConstantArray`] output whose scalar is the precomputed norm. Uses [`execute_until`] so +/// execution stops at the [`Constant`] encoding instead of canonicalizing into a +/// [`PrimitiveArray`]. +#[test] +fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { + let input = literal_vector_array(&[3.0f64, 4.0], 4); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![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 a constant input must produce a constant output"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_close(&[norm], &[5.0]); + Ok(()) +} + +/// An extension array over constant storage is folded just like a top-level constant instead of +/// recomputing the same norm once per row. +#[test] +fn extension_backed_constant_yields_constant_output() -> VortexResult<()> { + let input = Vector::constant_array(&[3.0f64, 4.0], 4)?; + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![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 produce a constant output"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_close(&[norm], &[5.0]); + Ok(()) +} + +/// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of +/// the correct primitive dtype and length. +#[test] +fn constant_null_input_yields_null_constant_output() -> VortexResult<()> { + let storage_dtype = DType::FixedSizeList( + DType::Primitive(PType::F64, Nullability::NonNullable).into(), + 2, + Nullability::Nullable, + ); + let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(); + let null_scalar = Scalar::null(DType::Extension(ext_dtype)); + let input = ConstantArray::new(null_scalar, 3).into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("null constant input must produce a constant output"); + assert_eq!(constant.len(), 3); + assert!(constant.scalar().is_null()); + assert_eq!( + constant.dtype(), + &DType::Primitive(PType::F64, Nullability::Nullable) + ); + Ok(()) +} + +/// An `f32` column must dispatch at `f32` and produce an `f32` result, which is the property that +/// makes width polymorphism load-bearing rather than decorative. +#[rstest] +#[case::f32(&[3.0f32, 4.0], PType::F32)] +#[case::f64(&[3.0f64, 4.0], PType::F64)] +fn dispatches_at_input_width( + #[case] elements: &[T], + #[case] expected: PType, +) -> VortexResult<()> { + let arr = tensor_array(&[2], elements)?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = L2Norm::try_new(arr)?.into_array().execute(&mut ctx)?; + assert_eq!(prim.ptype(), expected); + Ok(()) +} + +/// `L2Norm(Normalized(normalized, norms))` reads back the authoritative stored norms rather than +/// recomputing over decoded coordinates. The normalized child here is deliberately *not* +/// unit-norm, mimicking lossy storage, so readthrough and recompute disagree: row 0 decodes to +/// `[6, 8]` (norm `10`) and row 1 to `[6, 0]` (norm `6`), while the stored norms are `5` and `2`. +#[test] +fn normalized_readthrough_returns_stored_norms() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[1.2, 1.6, 3.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 2.0]).into_array(); + // SAFETY: A focused test of the lossy storage contract: the stored norms are authoritative + // even though this normalized child violates the unit-norm invariant. + let denorm = + unsafe { Normalized::new_unchecked(normalized, norms, Validity::NonNullable) }.into_array(); + + assert_close(&eval_l2_norm(denorm)?, &[5.0, 2.0]); + Ok(()) +} + +/// The readthrough must survive a partially-null column. +/// +/// This pins the dense policy the row contract derives. Filtering could hand `reduce_encoded` a +/// filtered input, which is no longer an `ExactScalarFn`, silently falling back to +/// decode-and-recompute. For a lossy child that changes the answer: row 0 below would come back as +/// `10` (recomputed from `[6, 8]`) instead of the authoritative stored `5`. +#[test] +fn normalized_readthrough_survives_null_rows() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[1.2, 1.6, 3.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 2.0]).into_array(); + // SAFETY: Intentionally lossy, as in `normalized_readthrough_returns_stored_norms`, so that + // a recompute fallback is observable. + let denorm = + unsafe { Normalized::new_unchecked(normalized, norms, Validity::from_iter([true, false])) } + .into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![denorm])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +/// The readthrough must still propagate validity carried by the `Normalized` parent. +#[test] +fn normalized_readthrough_propagates_parent_validity() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let denorm = Normalized::try_new( + normalized, + norms, + Validity::from_iter([true, false]), + &mut ctx, + )? + .into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![denorm])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +#[rstest] +#[case::fixed_shape_tensor(l2_norm_tensor_child())] +#[case::vector(l2_norm_vector_child())] +fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { + let original = L2Norm::try_new(child.clone())?.into_array(); + + let plugin = ScalarFnArrayPlugin::new(L2Norm); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("L2Norm serialize must produce metadata"); + + let children = vec![child]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn l2_norm_tensor_child() -> ArrayRef { + tensor_array(&[3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid tensor array") +} + +fn l2_norm_vector_child() -> ArrayRef { + vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/mod.rs b/vortex-tensor/src/scalar_fns/tests/mod.rs new file mode 100644 index 00000000000..5447772adda --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/mod.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for the tensor scalar functions. + +mod l2_norm; +mod row; diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs new file mode 100644 index 00000000000..338b1407379 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::Float; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::match_each_float_ptype; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnId; +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::validity::Validity; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::tensor_array; + +/// The marginal cost of a new tensor scalar function is this entire definition. Everything else +/// (null propagation, constants, validity, f16/f32/f64 dispatch, dtype checks, and constructors) is +/// derived. +#[derive(Clone, Debug, Default)] +struct L1Norm; + +impl RowFn for L1Norm { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.l1_norm"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + 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, l1_norm_row(row)) } + }) + }) + } +} + +fn l1_norm_row(row: &[T]) -> T { + row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) +} + +#[test] +fn derived_fn_executes_with_nulls() -> VortexResult<()> { + let arr = tensor_array(&[2], &[3.0, -4.0, 1.0, 1.0])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + let prim: PrimitiveArray = ScalarFnArray::try_new(L1Norm.bind(EmptyOptions), vec![arr])? + .into_array() + .execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[7.0]); + Ok(()) +} + +/// A kernel written once serves every float width. +#[test] +fn derived_fn_dispatches_at_input_width() -> VortexResult<()> { + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + + let f32_result: PrimitiveArray = ScalarFnArray::try_new( + L1Norm.bind(EmptyOptions), + vec![tensor_array(&[2], &[3.0f32, -4.0])?], + )? + .into_array() + .execute(&mut ctx)?; + assert_eq!(f32_result.ptype(), PType::F32); + assert_eq!(f32_result.as_slice::(), &[7.0f32]); + + let f64_result: PrimitiveArray = ScalarFnArray::try_new( + L1Norm.bind(EmptyOptions), + vec![tensor_array(&[2], &[3.0f64, -4.0])?], + )? + .into_array() + .execute(&mut ctx)?; + assert_eq!(f64_result.ptype(), PType::F64); + Ok(()) +} diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 3e33fe20db9..c7339787b73 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Shared helpers for tensor scalar functions. + use half::f16; +use num_traits::Float; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -22,6 +25,7 @@ use vortex_array::dtype::PType; use vortex_array::dtype::proto::dtype as pb; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::validity::Validity; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -60,6 +64,16 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { SAFETY_FACTOR as f64 * machine_epsilon * dimensions_root } +/// Computes `sqrt(sum(v_i^2))` for one row. An empty or all-zero row produces `0.0`. +pub(crate) fn l2_norm_row(row: &[T]) -> T { + let mut sum_squared = T::zero(); + for &element in row { + sum_squared = sum_squared + element * element; + } + + sum_squared.sqrt() +} + /// Extracts the `(normalized, norms)` children of a [`Normalized`]-encoded array. /// /// # Panics @@ -120,6 +134,22 @@ pub fn validate_binary_tensor_float_inputs<'a>( validate_tensor_float_input(lhs) } +/// Validates that every argument has the same float tensor dtype, ignoring nullability. +pub fn validate_tensor_float_inputs(args: &[DType]) -> VortexResult> { + let (first, rest) = args + .split_first() + .ok_or_else(|| vortex_err!("tensor expression expects at least one input"))?; + + for arg in rest { + vortex_ensure!( + first.eq_ignore_nullability(arg), + "tensor expression expects inputs to have the same dtype, got {first} and {arg}" + ); + } + + validate_tensor_float_input(first) +} + /// The flat primitive elements of a tensor storage array, with typed row access. /// /// This struct hides the stride detail that arises from the [`ConstantArray`] optimization: a @@ -148,6 +178,23 @@ impl FlatElements { let slice = self.elems.as_slice::(); &slice[row_idx * self.list_size..][..self.list_size] } + + /// Returns the number of elements in each row. + #[must_use] + pub fn list_size(&self) -> usize { + self.list_size + } + + /// Returns the physical distance between rows, or zero when every row uses one stored value. + #[must_use] + pub fn row_stride(&self) -> usize { + if self.is_constant { 0 } else { self.list_size } + } + + /// Returns the elements as a typed buffer, performing the ptype check once for the batch. + pub fn into_buffer(self) -> Buffer { + self.elems.into_buffer::() + } } /// Extracts the flat primitive elements from a tensor storage array (FixedSizeList). @@ -353,6 +400,18 @@ pub mod test_helpers { Vector::try_new_vector_array(flat_fsl(elements, dim)) } + /// Builds `rows` zero-width vectors over an empty typed element buffer. + pub fn zero_width_vector_array(rows: usize) -> VortexResult { + let storage = FixedSizeListArray::new( + Buffer::::empty().into_array(), + 0, + Validity::NonNullable, + rows, + ) + .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>(