From 04c8e2c28567af56e8c4f9e920138be1c140aeb3 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 12:14:21 -0400 Subject: [PATCH 1/6] Add encoding-aware RowFn reductions Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 5 + .../scalar_fn/unstable/row/batch/execution.rs | 82 +++++++- .../src/scalar_fn/unstable/row/batch/tests.rs | 197 +++++++++++++++++- .../src/scalar_fn/unstable/row/row_fn.rs | 39 +++- .../src/scalar_fn/unstable/row/vtable.rs | 1 + 5 files changed, 306 insertions(+), 18 deletions(-) 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 781d15711be..922e3ee5095 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> BorrowedExecutionArgs<'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/execution.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs index 93aad88810e..b14b1c71452 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs @@ -3,11 +3,12 @@ //! Applies columnar semantics around one typed row kernel invocation. //! -//! [`Batch`] owns strict null propagation, constant broadcasting, execution strategy selection, and -//! output validation. The row kernel therefore handles only decoded values and its selected output -//! capability. +//! [`Batch`] owns strict null propagation, encoded reductions, constant broadcasting, execution +//! strategy selection, and output validation. The row kernel therefore handles only decoded values +//! and its selected output capability. use smallvec::SmallVec; +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -121,12 +122,17 @@ impl Batch { }) } - /// Apply constant folding and null handling around `kernel`. + /// Apply encoded reductions, constant folding, and null handling around `kernel`. /// - /// For a mixed validity mask, `try_unfiltered` may avoid filtering; `Ok(None)` selects - /// filter-and-scatter. Every kernel result is checked against the planned shape and dtype. + /// `reduce` receives the original inputs before constant broadcasting. For a mixed validity + /// mask, `try_unfiltered` may avoid filtering; `Ok(None)` selects filter-and-scatter. Every + /// kernel result is checked against the planned shape and dtype. pub fn execute( &self, + reduce: impl FnOnce( + BorrowedExecutionArgs<'_>, + &mut ExecutionCtx, + ) -> VortexResult>, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, try_unfiltered: impl FnOnce( BorrowedExecutionArgs<'_>, @@ -147,6 +153,20 @@ impl Batch { return Ok(self.all_null()); } + // An empty mask is both all-true and all-false, so deferred encoded evidence cannot be + // attributed to an observable row. Let the ordinary policy construct the typed empty + // output instead. + if self.row_count > 0 + && let Some(execution) = reduce(self.execution_args(&self.inputs, self.row_count), ctx)? + { + match execution { + RowExecution::Output(values) => return self.finalize_reduced(values, ctx), + RowExecution::DeferredError(error) => { + return self.resolve_reduced_error(error, kernel, try_unfiltered, ctx); + } + } + } + // All inputs constant, and their conjoined validity proves every row non-null. This sees // through extension and masked wrappers just like argument decoding does. if self.row_count > 0 @@ -331,6 +351,56 @@ impl Batch { ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() } + /// Reconcile an encoding-aware result and apply the batch's strict input validity. + 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) + } + + /// Resolve deferred evidence from the encoded path by executing only observable rows. + fn resolve_reduced_error( + &self, + error: VortexError, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + /// Pair an input view with this batch's planning metadata. fn execution_args<'b>( &'b 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 6c92fd52c40..e8e34082990 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -31,7 +31,6 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::InputElement; @@ -59,7 +58,13 @@ struct AddShort; struct ShortDecodeI64; #[derive(Clone)] -struct Identity; +struct OriginalInputReducer; + +#[derive(Clone)] +struct InvalidEncodedReduction; + +#[derive(Clone)] +struct DeferredOriginalReducer; #[derive(Clone)] struct SinkOptions; @@ -308,16 +313,103 @@ impl RowFn for RetryConstantAdd { }, ) } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 1 { + return Ok(Some(RowExecution::Output( + ConstantArray::new(0u8, args[0].len()).into_array(), + ))); + } + + Ok(None) + } } -impl RowFn for Identity { +impl RowFn for OriginalInputReducer { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["value"]; const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("test.identity"); + 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(RowExecution::Output( + ConstantArray::new(42_i64, 3).into_array(), + ))); + } + + Ok(None) + } +} + +impl RowFn for InvalidEncodedReduction { + type Options = usize; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; + + 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(RowExecution::Output( + PrimitiveArray::new( + vec![10_i64, 20], + Validity::from_iter((0..2).map(|index| index != *null_index)), + ) + .into_array(), + ))) + } +} + +impl RowFn for DeferredOriginalReducer { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.deferred_original_reducer"); *ID } @@ -329,6 +421,17 @@ impl RowFn for Identity { ) -> VortexResult { visitor.visit::<(i64,), i64>(|(value,)| value) } + + fn reduce_encoded( + &self, + _options: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(RowExecution::DeferredError(vortex_err!( + InvalidArgument: "encoded payload failed" + )))) + } } impl RowFn for SinkOptions { @@ -461,7 +564,7 @@ fn test_short_decode_beside_constant_is_rejected() -> VortexResult<()> { } #[test] -fn test_dense_retry_preserves_valid_row_failure() -> VortexResult<()> { +fn test_dense_retry_does_not_reduce_filtered_inputs() -> VortexResult<()> { let lhs = PrimitiveArray::new(vec![u8::MAX, 1], Validity::from_iter([true, false])).into_array(); let rhs = ConstantArray::new(1u8, 2).into_array(); @@ -489,13 +592,89 @@ fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> { Ok(()) } +#[test] +fn test_reduce_encoded_defers_errors_behind_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(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_empty_batch_skips_deferred_encoded_error() -> VortexResult<()> { + let input = PrimitiveArray::from_iter(Vec::::new()).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 0); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid(Validity::AllValid)] +#[case::mixed(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_constant_input_broadcasts_one_row() -> VortexResult<()> { let input = ConstantArray::new(7_i64, 2).into_array(); let args = VecExecutionArgs::new(vec![input.clone()], 2); let mut ctx = array_session().create_execution_ctx(); - let actual = execute_rows(&Identity, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; assert_arrays_eq!(&actual, &input, &mut ctx); Ok(()) @@ -519,7 +698,8 @@ fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResul let mut ctx = array_session().create_execution_ctx(); let actual = batch.execute( - |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), |_args, _valid, _ctx| Ok(None), &mut ctx, )?; @@ -547,7 +727,8 @@ fn test_valid_only_filters_and_scatters() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let actual = batch.execute( - |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), |_args, _valid, _ctx| Ok(None), &mut ctx, )?; 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 8a982c4fb37..0b813e2e039 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,8 +16,11 @@ 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; +use crate::scalar_fn::unstable::row::RowExecution; /// A scalar function computed one row at a time. /// @@ -35,12 +38,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 any dispatch can raise a semantic error. + /// Whether any dispatch or encoded reduction can raise a semantic error. /// /// See [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible) for a /// more detailed explanation of semantic errors. /// - /// The framework checks dispatched element and result types. A conservative `true` is allowed. + /// The framework checks dispatched element and result types, but cannot inspect + /// [`reduce_encoded`](Self::reduce_encoded). Set this to `true` when that hook can return a + /// semantic error or [`RowExecution::DeferredError`]. A conservative `true` is allowed. const FALLIBLE: bool; /// Returns the ID of the scalar function. @@ -71,4 +76,30 @@ 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. [`Output`](RowExecution::Output) may remain encoded or + /// lazy. [`DeferredError`](RowExecution::DeferredError) retries only valid rows. 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 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 9675219d547..98e43c3dd9c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -132,6 +132,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_rows_unfiltered(function, options, args, valid, ctx), ctx, From 5c6cc3e2d4acd3693994d13090410f9e62848439 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:58:24 -0400 Subject: [PATCH 2/6] Execute tensor L2 norm with RowFn Signed-off-by: Connor Tsui --- vortex-tensor/benches/l2_norm.rs | 36 +- .../src/scalar_fns/cosine_similarity.rs | 7 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 401 ++++-------------- vortex-tensor/src/scalar_fns/mod.rs | 4 + vortex-tensor/src/scalar_fns/row.rs | 165 +++++++ vortex-tensor/src/scalar_fns/tests/l2_norm.rs | 318 ++++++++++++++ vortex-tensor/src/scalar_fns/tests/mod.rs | 7 + vortex-tensor/src/scalar_fns/tests/row.rs | 98 +++++ vortex-tensor/src/utils.rs | 59 +++ 9 files changed, 773 insertions(+), 322 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-tensor/benches/l2_norm.rs b/vortex-tensor/benches/l2_norm.rs index 6f597084113..bf8832f2520 100644 --- a/vortex-tensor/benches/l2_norm.rs +++ b/vortex-tensor/benches/l2_norm.rs @@ -15,11 +15,19 @@ 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::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::l2_norm::L2Norm; @@ -54,13 +62,25 @@ 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 .counter(ItemsCount::new(input.len())) .with_inputs(|| { ( - L2Norm::try_new_array(input.clone()).unwrap().into_array(), + L2Norm + .try_new_array(input.len(), EmptyOptions, [input.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -80,3 +100,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/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index ca8fcf0efd4..76d266a2721 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -10,6 +10,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; @@ -144,8 +145,8 @@ impl ScalarFnVTable for CosineSimilarity { let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; // Compute inner product and norms as columnar operations, and propagate the options. - let norm_lhs_arr = L2Norm::try_new_array(lhs_ref.clone())?; - let norm_rhs_arr = L2Norm::try_new_array(rhs_ref.clone())?; + let norm_lhs_arr = L2Norm.try_new_array(len, EmptyOptions, [lhs_ref.clone()])?; + let norm_rhs_arr = L2Norm.try_new_array(len, EmptyOptions, [rhs_ref.clone()])?; let dot_arr = InnerProduct::try_new_array(lhs_ref, rhs_ref)?; // Execute to get the inner product and norms of the arrays. We only fully decompress @@ -288,7 +289,7 @@ impl CosineSimilarity { let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; - let norm_arr = L2Norm::try_new_array(plain_ref.clone())?; + let norm_arr = L2Norm.try_new_array(len, EmptyOptions, [plain_ref.clone()])?; let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index b7e9060ed3f..f8349d3431b 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -3,50 +3,44 @@ //! 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::InitializedElement; +use vortex_array::scalar_fn::RowExecution; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::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::l2_norm_row; use crate::utils::validate_tensor_float_input; /// L2 norm (Euclidean norm) of a tensor or vector column. @@ -62,142 +56,116 @@ 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 { /// Creates a new [`TypedScalarFnInstance`] wrapping the L2 norm operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(L2Norm, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the L2 norm over `child`. /// /// # 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_array(child: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(L2Norm::new().erased(), vec![child]) + ScalarFnArray::try_new(Self::new().erased(), vec![child]) } } -impl ScalarFnVTable for L2Norm { +impl RowFn for L2Norm { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["input"]; + 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()); - - // 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. - if input_ref.is::() { - let (_, norms) = extract_normalized_children(&input_ref); - 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(RowExecution::Output(norms))); + } - fn is_fallible(&self, _options: &Self::Options) -> bool { - false + 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(RowExecution::Output(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(RowExecution::Output(output.into_array()))) } } +vortex_array::impl_row_fn_vtable!(L2Norm); + /// Metadata for a serialized [`L2Norm`] array: the single `input` child's [`DType`], which carries /// the extension type (`FixedShapeTensor` vs `Vector`), dimension, and nullability that are not /// recoverable from the parent's primitive-float output. @@ -240,206 +208,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::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::validity::Validity; - use vortex_error::VortexResult; - - 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 scalar_fn = L2Norm::new().erased(); - 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()) - } - - #[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::new().erased(); - 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::new().erased(); - 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(()) - } - - /// 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::new().erased(); - 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(()) - } - - #[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_array(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..6f7e4d64d9a --- /dev/null +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -0,0 +1,165 @@ +// 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::InputElement; +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, +} + +// 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 Varying<'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_FALLIBLE: bool = false; + + 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, + "varying tensor row 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 get(column: &Self::Column, index: usize) -> &[T] { + let start = index * column.stride; + &column.elements.as_slice()[start..start + column.list_size] + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.rows + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + Self::get(column, index) + } + + unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + let start = index * column.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( + column.elements.as_slice().as_ptr().add(start), + column.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..930ef423d47 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs @@ -0,0 +1,318 @@ +// 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::ScalarFnFactoryExt; +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 _scalar_fn = L2Norm::new(); + let array = L2Norm::try_new_array(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_array(arr.len(), EmptyOptions, [arr])? + .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) }.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_option_iter([Some(5.0f64), None]).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) }.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 nulls carried by the `norms` child. +#[test] +fn normalized_readthrough_propagates_null_norms() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let denorm = Normalized::try_new(normalized, norms, &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_array(child.len(), EmptyOptions, [child.clone()])?; + + 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..77248329ae5 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -0,0 +1,98 @@ +// 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::scalar_fn::ScalarFnFactoryExt; +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::InitializedElement; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::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"]; + + 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)) } + }) + }) + } +} + +vortex_array::impl_row_fn_vtable!(L1Norm); + +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 = L1Norm + .try_new_array(arr.len(), EmptyOptions, [arr])? + .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 = L1Norm + .try_new_array(1, EmptyOptions, [tensor_array(&[2], &[3.0f32, -4.0])?])? + .execute(&mut ctx)?; + assert_eq!(f32_result.ptype(), PType::F32); + assert_eq!(f32_result.as_slice::(), &[7.0f32]); + + let f64_result: PrimitiveArray = L1Norm + .try_new_array(1, EmptyOptions, [tensor_array(&[2], &[3.0f64, -4.0])?])? + .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 488694bd47f..df40bdea44a 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; @@ -20,6 +23,7 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::dtype::proto::dtype as pb; use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -58,6 +62,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 @@ -110,6 +124,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 @@ -138,6 +168,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). @@ -343,6 +390,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>( From 9709f836c222ff508cc50fc31fb8df029ff1b18c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:08:55 -0400 Subject: [PATCH 3/6] Update tensor L2 for the RowFn API Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/l2_norm.rs | 2 -- vortex-tensor/src/scalar_fns/row.rs | 23 ++++++++++------------- vortex-tensor/src/scalar_fns/tests/row.rs | 2 -- 3 files changed, 10 insertions(+), 17 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index f8349d3431b..2a4fb0368cf 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -164,8 +164,6 @@ impl RowFn for L2Norm { } } -vortex_array::impl_row_fn_vtable!(L2Norm); - /// Metadata for a serialized [`L2Norm`] array: the single `input` child's [`DType`], which carries /// the extension type (`FixedShapeTensor` vs `Vector`), dimension, and nullability that are not /// recoverable from the parent's primitive-float output. diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 6f7e4d64d9a..8340e97160a 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -60,7 +60,7 @@ pub struct TensorRows { // unchecked access use the same stride and row width. unsafe impl InputElement for TensorRow { type Column = TensorRows; - type Varying<'a> = &'a 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 @@ -102,7 +102,7 @@ unsafe impl InputElement for TensorRow { vortex_ensure_eq!( stride, list_size, - "varying tensor row stride must equal its width, got {stride}", + "per-row tensor stride must equal its width, got {stride}", ); let Some(expected_elements) = rows.checked_mul(stride) else { vortex_bail!( @@ -132,34 +132,31 @@ unsafe impl InputElement for TensorRow { &column.elements.as_slice()[start..start + column.list_size] } - fn varying(column: &Self::Column) -> Self::Varying<'_> { + fn view(column: &Self::Column) -> Self::View<'_> { column } - fn varying_len(column: &Self::Varying<'_>) -> usize { - column.rows + fn view_len(view: &Self::View<'_>) -> usize { + view.rows } - fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] where Self: 'a, { - Self::get(column, index) + Self::get(view, index) } - unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] where Self: 'a, { - let start = index * column.stride; + 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( - column.elements.as_slice().as_ptr().add(start), - column.list_size, - ) + 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/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index 77248329ae5..b88a22838aa 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -57,8 +57,6 @@ impl RowFn for L1Norm { } } -vortex_array::impl_row_fn_vtable!(L1Norm); - fn l1_norm_row(row: &[T]) -> T { row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) } From e62db421e6d08cb2a73ceb136b9d1678f47ee92d Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 13:37:56 -0400 Subject: [PATCH 4/6] Opt tensor L2 into unstable RowFn Signed-off-by: Connor Tsui --- vortex-tensor/Cargo.toml | 2 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 10 +++++----- vortex-tensor/src/scalar_fns/row.rs | 2 +- vortex-tensor/src/scalar_fns/tests/row.rs | 8 ++++---- 4 files changed, 11 insertions(+), 11 deletions(-) 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/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 2a4fb0368cf..d0e99576f83 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -20,13 +20,13 @@ use vortex_array::dtype::proto::dtype as pb; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::InitializedElement; -use vortex_array::scalar_fn::RowExecution; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +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_error::VortexExpect; use vortex_error::VortexResult; diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 8340e97160a..9a9600bd6f9 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -16,7 +16,7 @@ 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::InputElement; +use vortex_array::scalar_fn::unstable::row::InputElement; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index b88a22838aa..aff3c2d0bb1 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -12,11 +12,11 @@ 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::InitializedElement; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::UninitElementSink; +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; From 154af6e6520a49ce0c4d1de0f3791863bc718e9e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 10:21:25 -0400 Subject: [PATCH 5/6] Declare tensor norm fallibility Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/l2_norm.rs | 1 + vortex-tensor/src/scalar_fns/tests/row.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index d0e99576f83..78cabc961b0 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -80,6 +80,7 @@ impl RowFn for L2Norm { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.l2_norm"); diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index aff3c2d0bb1..8a7fa23edd4 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -36,6 +36,7 @@ impl RowFn for L1Norm { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.test.l1_norm"); From 0a6bd79b538021b2a8111b312c0a6e195dfe79eb Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:25:53 -0400 Subject: [PATCH 6/6] Opt tensor rows into null-tolerant decoding Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/row.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 9a9600bd6f9..a99ad76ea6e 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -127,6 +127,10 @@ unsafe impl InputElement for TensorRow { }) } + 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]