From 5eba5b2b8e44963519059ebdd89800e61c233c47 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 4 Sep 2026 14:09:39 -0400 Subject: [PATCH] Add tensor row inputs for RowFn Signed-off-by: Connor Tsui --- vortex-tensor/Cargo.toml | 2 +- vortex-tensor/src/scalar_fns/mod.rs | 1 + vortex-tensor/src/scalar_fns/row.rs | 211 ++++++++++++++++++++++++++++ vortex-tensor/src/utils.rs | 18 +++ 4 files changed, 231 insertions(+), 1 deletion(-) create mode 100644 vortex-tensor/src/scalar_fns/row.rs diff --git a/vortex-tensor/Cargo.toml b/vortex-tensor/Cargo.toml index d76d561d3be..cd1900dffc6 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/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index eb074b6d0b8..14d6bb140bf 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -18,3 +18,4 @@ pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; pub mod l2_normalize; +pub mod row; diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs new file mode 100644 index 00000000000..00c7eee6089 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -0,0 +1,211 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tensor row inputs for scalar functions. +//! +//! [`TensorRow`] decodes one flat row from either an ordinary tensor column or either supported +//! batch-constant representation. Batch execution owns input validity, so decoding handles only +//! the non-null coordinate storage. + +use std::marker::PhantomData; + +use num_traits::Float; +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::Extension; +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; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; + +use crate::utils::extract_flat_elements; +use crate::utils::validate_tensor_float_input; + +/// Returns the common float element type of tensor row-function arguments. +pub fn tensor_element_ptype(args: &[DType]) -> VortexResult { + let (first, rest) = args + .split_first() + .ok_or_else(|| vortex_err!("tensor row function requires at least one input"))?; + + for argument in rest { + vortex_ensure!( + first.eq_ignore_nullability(argument), + "tensor row-function inputs must have the same dtype, got {first} and {argument}", + ); + } + + Ok(validate_tensor_float_input(first)?.element_ptype()) +} + +/// A tensor-valued row-function input whose element is one flat tensor row. +pub struct TensorRow(PhantomData); + +/// A decoded tensor column with constant-width rows. +pub struct TensorRows { + elements: Buffer, + row_count: usize, + row_width: usize, + stride: usize, +} + +impl ViewLen for TensorRows { + fn len(&self) -> usize { + self.row_count + } +} + +fn decode_tensor_storage( + storage: &ArrayRef, + row_count: usize, + row_width: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let flat = extract_flat_elements(storage, row_width, ctx)?; + let row_width = flat.list_size(); + let stride = flat.row_stride(); + let elements = flat.into_buffer::(); + + let expected_elements = if stride == 0 { + row_width + } else { + vortex_ensure_eq!( + stride, + row_width, + "per-row tensor stride must equal its width, got {stride}", + ); + let Some(expected_elements) = row_count.checked_mul(stride) else { + vortex_bail!( + "tensor row storage length must fit usize, got {row_count} 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, + row_count, + row_width, + stride, + }) +} + +// 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 Constant = Buffer; + type View<'a> = &'a TensorRows; + type Elem<'a> = &'a [T]; + + const DENSE_SAFE: bool = true; + const DECODE_INFALLIBLE: bool = true; + + fn validate(dtype: &DType) -> VortexResult<()> { + let tensor_match = validate_tensor_float_input(dtype)?; + let expected_element_ptype = T::PTYPE; + vortex_ensure_eq!( + tensor_match.element_ptype(), + expected_element_ptype, + "tensor row input must use {expected_element_ptype} elements, got {dtype}", + ); + + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // Batch execution owns the mask and restores it on the result. + let array = match array.as_opt::() { + Some(masked) => masked.child().clone(), + None => array, + }; + + let row_count = array.len(); + let row_width = validate_tensor_float_input(array.dtype())?.list_size() as usize; + let extension: ExtensionArray = array.execute(ctx)?; + decode_tensor_storage(extension.storage_array(), row_count, row_width, ctx) + } + + fn decode_constant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let row_width = validate_tensor_float_input(array.dtype())?.list_size() as usize; + let storage = if let Some(constant) = array.as_opt::() { + let scalar = constant.scalar().as_extension().to_storage_scalar(); + ConstantArray::new(scalar, 1).into_array() + } else if let Some(extension) = array.as_opt::() { + extension.storage_array().clone() + } else { + vortex_bail!( + "a tensor batch constant must use the Constant encoding or constant extension \ + storage, got {}", + array.encoding_id() + ); + }; + + let decoded = decode_tensor_storage::(&storage, 1, row_width, ctx)?; + vortex_ensure_eq!( + decoded.elements.len(), + row_width, + "decoded tensor constant must contain {row_width} elements, got {}", + decoded.elements.len(), + ); + + Ok(decoded.elements) + } + + 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.row_width] + } + + fn get_constant(constant: &Self::Constant) -> &[T] { + constant.as_slice() + } + + 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 `row_count` + // contiguous `row_width`-element rows otherwise. The caller guarantees + // `index < row_count`. + unsafe { + std::slice::from_raw_parts(view.elements.as_slice().as_ptr().add(start), view.row_width) + } + } +} diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 540bb625d0e..0e8768c2d13 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -19,6 +19,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::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -87,6 +88,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).