From 54057ac92d737c3bb2f1d2e42e3960af597a2722 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 2 Oct 2026 22:07:33 -0400 Subject: [PATCH] Read RowFn UTF-8 byte lengths from native metadata Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/mod.rs | 2 + .../unstable/row/types/element/mod.rs | 4 + .../row/types/element/utf8_length/mod.rs | 202 +++++++++++++++ .../row/types/element/utf8_length/tests.rs | 244 ++++++++++++++++++ .../src/scalar_fn/unstable/row/types/mod.rs | 2 + 5 files changed, 454 insertions(+) create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/utf8_length/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/utf8_length/tests.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index d33062fa06e..643ccb510c4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -45,8 +45,10 @@ pub use types::SinkResult; pub use types::UninitElementSink; pub use types::Utf8Column; pub use types::Utf8OffsetColumn; +pub use types::Utf8OffsetLengthColumn; pub use types::Utf8Sink; pub use types::Utf8View; +pub use types::Utf8ViewLengthColumn; pub use types::Utf8Writer; pub use types::ViewLen; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs index 4e489e8f877..63685024b7c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -31,6 +31,10 @@ mod utf8; pub use utf8::Utf8Column; pub use utf8::Utf8View; +mod utf8_length; +pub use utf8_length::Utf8OffsetLengthColumn; +pub use utf8_length::Utf8ViewLengthColumn; + mod utf8_offset; pub use utf8_offset::Utf8OffsetColumn; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/utf8_length/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/utf8_length/mod.rs new file mode 100644 index 00000000000..c703909cd26 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/utf8_length/mod.rs @@ -0,0 +1,202 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! UTF-8 byte lengths read from native metadata without string preparation. +//! +//! [`Utf8OffsetLengthColumn`] and [`Utf8ViewLengthColumn`] read lengths without acquiring payload +//! buffers or constructing strings. The caller selects the physical layout before row execution. +//! Null-slot lengths are unspecified, and the executor preserves strict null propagation. + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::Utf8Column; +use super::utf8_offset::OffsetMetadata; +use super::utf8_offset::constant_string; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::VarBin; +use crate::arrays::VarBinView; +use crate::arrays::varbinview::BinaryView; +use crate::dtype::DType; +use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::ViewLen; + +/// A UTF-8 input element that yields [`u64`] byte lengths from native [`u32`] offsets. +/// +/// Non-constant, nonempty inputs **must** use [`VarBin`] with non-nullable [`u32`] offsets. Other +/// layouts are rejected even when their dtype is [`DType::Utf8`]. Empty UTF-8 arrays and _batch +/// constants_ (one value shared by every row) are supported separately. +/// +/// Decoding validates offset bounds against payload length metadata without acquiring payload +/// buffers or constructing strings. Null-slot lengths are unspecified, and the executor propagates +/// validity. Layout selection must occur before dtype-only [RowFn dispatch]. +/// +/// [RowFn dispatch]: crate::scalar_fn::unstable::row::RowFn::dispatch +pub struct Utf8OffsetLengthColumn; + +/// A UTF-8 input element that yields [`u64`] byte lengths from native string-view headers. +/// +/// Non-constant, nonempty inputs **must** use [`VarBinView`]. Other layouts are rejected even when +/// their dtype is [`DType::Utf8`]. Empty UTF-8 arrays and batch constants are supported separately. +/// Decoding retains the headers without acquiring payload buffers or constructing strings, so +/// unreadable null payloads do not prevent header access. +/// +/// See [`Utf8OffsetLengthColumn`] for the shared null-propagation and layout-selection contract. +pub struct Utf8ViewLengthColumn; + +// SAFETY: OffsetMetadata::decode retains monotonic offsets with one terminal entry beyond the +// stable row domain. Length access only subtracts those offsets, so null payloads are never read. +// The executor propagates validity because null-slot lengths are unspecified. +unsafe impl InputElement for Utf8OffsetLengthColumn { + type Column = OffsetMetadata; + type Constant = u64; + type View<'a> = &'a OffsetMetadata; + type Elem<'a> = u64; + + const DENSE_SAFE: bool = true; + + // Other legal UTF-8 layouts require a different input element. + const DECODE_INFALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + Utf8Column::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + OffsetMetadata::decode(&array, ctx) + } + + fn decode_constant(array: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(constant_string(array)?.len() as u64) + } + + fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult { + Ok(array.is_empty() || array.is::()) + } + + fn get(column: &Self::Column, index: usize) -> u64 { + u64::from(column.offsets[index + 1] - column.offsets[index]) + } + + fn get_constant(constant: &u64) -> u64 { + *constant + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> u64 + where + Self: 'a, + { + Self::get(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> u64 + where + Self: 'a, + { + let offsets = view.offsets.as_slice(); + + // SAFETY: the caller bounds index by the row count, and decode retains rows + 1 offsets. + let start = unsafe { *offsets.get_unchecked(index) }; + // SAFETY: the same row bound puts index + 1 within the retained terminal offset. + let end = unsafe { *offsets.get_unchecked(index + 1) }; + + u64::from(end - start) + } +} + +/// Owned string-view headers, without string payload buffers. +pub struct Utf8LengthViews { + /// One initialized header per row, including null rows whose lengths are unspecified. + views: Buffer, +} + +impl ViewLen for &Utf8LengthViews { + fn len(&self) -> usize { + self.views.len() + } +} + +// SAFETY: one owned view corresponds to each row. Every length header is an initialized integer, +// including null slots. Null lengths are unspecified. The executor propagates validity. No bytes +// or unchecked strings are exposed, so UTF-8 validation is unnecessary. +unsafe impl InputElement for Utf8ViewLengthColumn { + type Column = Utf8LengthViews; + type Constant = u64; + type View<'a> = &'a Utf8LengthViews; + type Elem<'a> = u64; + + const DENSE_SAFE: bool = true; + + // Other legal UTF-8 layouts require a different input element. + const DECODE_INFALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + Utf8Column::validate(dtype) + } + + fn decode(array: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + Utf8Column::validate(array.dtype())?; + + if array.is_empty() { + return Ok(Utf8LengthViews { + views: Buffer::empty(), + }); + } + + let array = array.as_opt::().ok_or_else(|| { + vortex_err!( + "expected VarBinView for a Utf8 view-length column, got {}", + array.encoding_id() + ) + })?; + let views = Buffer::from_byte_buffer(array.views_handle().try_to_host_sync()?); + + Ok(Utf8LengthViews { views }) + } + + fn decode_constant(array: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(constant_string(array)?.len() as u64) + } + + fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult { + Ok(array.is_empty() || array.is::()) + } + + fn get(column: &Self::Column, index: usize) -> u64 { + u64::from(column.views[index].len()) + } + + fn get_constant(constant: &u64) -> u64 { + *constant + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> u64 + where + Self: 'a, + { + Self::get(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> u64 + where + Self: 'a, + { + // SAFETY: the caller bounds index by the exact retained view domain. + let header = unsafe { view.views.as_slice().get_unchecked(index) }; + + u64::from(header.len()) + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/utf8_length/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/utf8_length/tests.rs new file mode 100644 index 00000000000..e0354b7b998 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/utf8_length/tests.rs @@ -0,0 +1,244 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use super::Utf8OffsetLengthColumn; +use super::Utf8ViewLengthColumn; +use crate::ArrayRef; +use crate::IntoArray as _; +use crate::VortexSessionExecute as _; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::BinaryView; +use crate::assert_arrays_eq; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_rows; +use crate::scalar_fn::unstable::row::types::element::test_support::UnreadablePayload; +use crate::scalar_fn::unstable::row::types::element::test_support::unreadable_offset_array; +use crate::validity::Validity; + +#[track_caller] +fn assert_threshold(layout: Layout, input: ArrayRef, expected: ArrayRef) -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![input.clone()], input.len()); + let mut ctx = array_session().create_execution_ctx(); + let output = execute_rows(&Threshold(layout), &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&output, &expected, &mut ctx); + + Ok(()) +} + +/// The native representation selected explicitly by each test. +#[derive(Clone, Copy)] +enum Layout { + Offsets, + Views, +} + +#[derive(Clone)] +struct Threshold(Layout); + +impl RowFn for Threshold { + type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.utf8_length_threshold"); + *ID + } + + fn dispatch( + &self, + _options: &EmptyOptions, + _args: &[DType], + visitor: V, + ) -> VortexResult { + match self.0 { + Layout::Offsets => visitor.visit::<(Utf8OffsetLengthColumn,), bool>(|(len,)| len >= 13), + Layout::Views => visitor.visit::<(Utf8ViewLengthColumn,), bool>(|(len,)| len >= 13), + } + } +} + +/// Build nullable UTF-8 input in the representation under test. +fn input(layout: Layout, values: Vec>) -> ArrayRef { + match layout { + Layout::Offsets => { + VarBinArray::from_iter(values, DType::Utf8(Nullability::Nullable)).into_array() + } + Layout::Views => VarBinViewArray::from_iter_nullable_str(values).into_array(), + } +} + +#[test] +fn test_metadata_decode_never_requests_payload() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let column = + Utf8OffsetLengthColumn::decode(unreadable_offset_array(Validity::NonNullable)?, &mut ctx)?; + assert_eq!(Utf8OffsetLengthColumn::get(&column, 0), 13); + + let views = Buffer::from(vec![BinaryView::make_view(b"1234567890123", 0, 0)]); + let array = VarBinViewArray::new_handle( + BufferHandle::new_host(views.into_byte_buffer()), + Arc::from(vec![BufferHandle::new_device(Arc::new(UnreadablePayload { + len: 13, + }))]), + DType::Utf8(Nullability::NonNullable), + Validity::NonNullable, + ) + .into_array(); + let column = Utf8ViewLengthColumn::decode(array, &mut ctx)?; + assert_eq!(Utf8ViewLengthColumn::get(&column, 0), 13); + + Ok(()) +} + +#[test] +fn test_unspecified_null_header_stays_null() -> VortexResult<()> { + let views = Buffer::from(vec![ + BinaryView::make_view(b"a longer string", 0, 0), // Valid row with unreadable payload. + BinaryView::from(u128::MAX), // Null row with an unspecified header. + ]); + let array = VarBinViewArray::new_handle( + BufferHandle::new_host(views.into_byte_buffer()), + Arc::from(vec![BufferHandle::new_device(Arc::new(UnreadablePayload { + len: 15, + }))]), + DType::Utf8(Nullability::Nullable), + Validity::from_iter([true, false]), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + let column = Utf8ViewLengthColumn::decode(array.clone(), &mut ctx)?; + assert_eq!(Utf8ViewLengthColumn::get(&column, 1), u64::from(u32::MAX)); + + assert_threshold( + Layout::Views, + array, + BoolArray::from_iter([Some(true), None]).into_array(), + ) +} + +#[test] +fn test_metadata_decode_rejects_wrong_dtype_and_layout() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let primitive = PrimitiveArray::from_iter([13u32]).into_array(); + assert!(Utf8OffsetLengthColumn::decode(primitive.clone(), &mut ctx).is_err()); + assert!(Utf8ViewLengthColumn::decode(primitive, &mut ctx).is_err()); + let offsets = unreadable_offset_array(Validity::NonNullable)?; + assert!(Utf8ViewLengthColumn::decode(offsets, &mut ctx).is_err()); + let views = VarBinViewArray::from_iter_str(["text"]).into_array(); + assert!(Utf8OffsetLengthColumn::decode(views, &mut ctx).is_err()); + + Ok(()) +} + +#[rstest] +#[case::offsets(Layout::Offsets)] +#[case::views(Layout::Views)] +fn test_access_paths_use_utf8_byte_lengths(#[case] layout: Layout) -> VortexResult<()> { + let values = input( + layout, + vec![ + Some("outside"), // Excluded prefix. + Some(""), // Empty. + Some("héllo"), // Six UTF-8 bytes. + Some("a longer string"), // Fifteen ASCII bytes. + Some("outside"), // Excluded suffix. + ], + ); + let values = values.slice(1..4)?; + let mut ctx = array_session().create_execution_ctx(); + let expected = [ + 0, // Empty. + 6, // UTF-8 byte length, rather than character count. + 15, // Long ASCII. + ]; + + match layout { + Layout::Offsets => { + let column = Utf8OffsetLengthColumn::decode(values, &mut ctx)?; + let view = Utf8OffsetLengthColumn::view(&column); + + for (index, expected) in expected.into_iter().enumerate() { + assert_eq!(Utf8OffsetLengthColumn::get(&column, index), expected); + assert_eq!( + Utf8OffsetLengthColumn::get_from_view(&view, index), + expected + ); + + // SAFETY: there is one expected length for each retained row. + let unchecked = + unsafe { Utf8OffsetLengthColumn::get_from_view_unchecked(&view, index) }; + assert_eq!(unchecked, expected); + } + } + Layout::Views => { + let column = Utf8ViewLengthColumn::decode(values, &mut ctx)?; + let view = Utf8ViewLengthColumn::view(&column); + + for (index, expected) in expected.into_iter().enumerate() { + assert_eq!(Utf8ViewLengthColumn::get(&column, index), expected); + assert_eq!(Utf8ViewLengthColumn::get_from_view(&view, index), expected); + + // SAFETY: there is one expected length for each retained row. + let unchecked = + unsafe { Utf8ViewLengthColumn::get_from_view_unchecked(&view, index) }; + assert_eq!(unchecked, expected); + } + } + } + + Ok(()) +} + +#[rstest] +#[case::empty(vec![], vec![])] +#[case::all_valid(vec![Some(""), Some("a longer string")], vec![Some(false), Some(true)])] +#[case::partially_valid(vec![Some("héllo"), None, Some("a longer string")], vec![Some(false), None, Some(true)])] +#[case::all_null(vec![None, None], vec![None, None])] +fn test_row_fn_matches_scalar_lengths( + #[values(Layout::Offsets, Layout::Views)] layout: Layout, + #[case] values: Vec>, + #[case] expected: Vec>, +) -> VortexResult<()> { + assert_threshold( + layout, + input(layout, values), + BoolArray::from_iter(expected).into_array(), + ) +} + +#[rstest] +#[case::offsets(Layout::Offsets)] +#[case::views(Layout::Views)] +fn test_metadata_constants_use_one_value(#[case] layout: Layout) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let value = ConstantArray::new("a longer string", 17).into_array(); + let length = match layout { + Layout::Offsets => Utf8OffsetLengthColumn::decode_constant(value.clone(), &mut ctx)?, + Layout::Views => Utf8ViewLengthColumn::decode_constant(value.clone(), &mut ctx)?, + }; + assert_eq!(length, 15); + + assert_threshold(layout, value, ConstantArray::new(true, 17).into_array()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs index 68561c5c11a..ab09c91b731 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -20,7 +20,9 @@ pub use element::OutputBuffer; pub use element::OutputElement; pub use element::Utf8Column; pub use element::Utf8OffsetColumn; +pub use element::Utf8OffsetLengthColumn; pub use element::Utf8View; +pub use element::Utf8ViewLengthColumn; pub(super) use element::batch_const; pub(in crate::scalar_fn::unstable::row) use element::decoded_source;