diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index 8bf45929163..7937051e9d1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -40,6 +40,10 @@ pub use types::OutputElement; pub use types::OutputSink; pub use types::SinkResult; pub use types::UninitElementSink; +pub use types::Utf8Column; +pub use types::Utf8Sink; +pub use types::Utf8View; +pub use types::Utf8Writer; pub use types::ViewLen; mod visitor; 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 71b0d17fe79..aea9c9c4da7 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 @@ -6,6 +6,9 @@ //! [`InputElement::Elem`] can borrow from its decoded column. Owned row computations return an //! [`OutputElement`]. Runtime-shaped outputs use an //! [`OutputSink`](crate::scalar_fn::unstable::row::OutputSink). +//! +//! Booleans and primitives implement both traits directly. UTF-8 columns decode through +//! [`Utf8Column`] and build through [`Utf8Sink`](crate::scalar_fn::unstable::row::Utf8Sink). mod bool; @@ -22,3 +25,7 @@ pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; pub use tuple::batch_const; pub(in crate::scalar_fn::unstable::row) use tuple::decoded_source; + +mod utf8; +pub use utf8::Utf8Column; +pub use utf8::Utf8View; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/utf8.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/utf8.rs new file mode 100644 index 00000000000..05eaa997697 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/utf8.rs @@ -0,0 +1,363 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! UTF-8 input columns for row functions. +//! +//! [`Utf8Column`] decodes a `Utf8` column once per batch and hands each row callback a +//! [`Utf8View`]. A view dereferences to `str`, so a callback can call `str` methods directly. It +//! also exposes the string-view metadata that a comparison kernel can use before it reads the +//! bytes. + +use std::ops::Deref; +use std::sync::Arc; + +use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::Constant; +use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::BinaryView; +use crate::arrays::varbinview::VarBinViewArrayExt as _; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::scalar::ScalarValue; +use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::ViewLen; + +/// A UTF-8 input element that yields [`Utf8View`] values. +pub struct Utf8Column; + +/// One decoded UTF-8 column. +pub struct Utf8Values { + views: Buffer, + buffers: Arc<[ByteBuffer]>, +} + +impl Utf8Values { + /// Build a single-value representation of one batch-constant string. + fn single(value: ByteBuffer) -> Self { + let view = BinaryView::make_view(value.as_slice(), 0, 0); + let buffers: Vec = if view.is_inlined() { + Vec::new() + } else { + vec![value] + }; + + Self { + views: Buffer::from(vec![view]), + buffers: Arc::from(buffers), + } + } + + fn views(&self) -> &[BinaryView] { + self.views.as_slice() + } +} + +/// A borrowed UTF-8 column prepared for a row loop. +#[derive(Clone, Copy)] +pub struct Utf8ValuesView<'a> { + views: &'a [BinaryView], + buffers: &'a [ByteBuffer], +} + +impl ViewLen for Utf8ValuesView<'_> { + fn len(&self) -> usize { + self.views.len() + } +} + +/// One UTF-8 value together with its Vortex string-view representation. +#[derive(Clone, Copy)] +pub struct Utf8View<'a> { + view: &'a BinaryView, + buffers: &'a [ByteBuffer], +} + +impl<'a> Utf8View<'a> { + /// Return the raw Vortex string view. + /// + /// The view reports the byte length without a read of the value. + pub fn raw_view(&self) -> &BinaryView { + self.view + } + + /// Return whether the complete string is stored inside the view. + pub fn is_inlined(&self) -> bool { + self.view.is_inlined() + } + + /// Return the string's first four bytes, or the complete string when it is shorter. + pub fn prefix(&self) -> &[u8] { + if self.view.is_inlined() { + let value = self.view.as_inlined().value(); + &value[..value.len().min(4)] + } else { + &self.view.as_view().prefix + } + } + + /// Return the complete UTF-8 string with the lifetime of the decoded column. + /// + /// [`Deref`] borrows from the view instead, so a value that must outlive this view uses this + /// method. + pub fn as_str(&self) -> &'a str { + let bytes = if self.view.is_inlined() { + self.view.as_inlined().value() + } else { + let view = self.view.as_view(); + &self.buffers[view.buffer_index as usize][view.as_range()] + }; + + // SAFETY: the `VarBinViewArray` Utf8 invariant requires every valid view to contain UTF-8. + // `decode_utf8` replaces each null row's view before this value can be constructed. + unsafe { std::str::from_utf8_unchecked(bytes) } + } +} + +impl Deref for Utf8View<'_> { + type Target = str; + + fn deref(&self) -> &str { + self.as_str() + } +} + +impl AsRef for Utf8View<'_> { + fn as_ref(&self) -> &str { + self.as_str() + } +} + +/// Decodes `array` into views and data buffers whose every view addresses valid UTF-8. +/// +/// The second [`VarBinViewArray::try_new`] is the sanitization step, not a round trip: it +/// validates each valid view and replaces every null row's view with an empty one. The +/// [`InputElement`] implementation relies on that, so a dense callback can read a null row's +/// payload and [`Utf8View::as_str`] can skip the UTF-8 check. Removing it makes those unsafe. +fn decode_utf8(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let array = array.execute::(ctx)?; + let views = Buffer::::from_byte_buffer(array.views_handle().try_to_host_sync()?); + let buffers: Arc<[ByteBuffer]> = Arc::from( + array + .data_buffers() + .iter() + .map(BufferHandle::try_to_host_sync) + .collect::>>()?, + ); + + let validity = array.varbinview_validity(); + let array = VarBinViewArray::try_new( + views, + Arc::clone(&buffers), + array.dtype().clone(), + validity, + ctx, + )?; + let views = Buffer::::from_byte_buffer(array.views_handle().try_to_host_sync()?); + + Ok(Utf8Values { views, buffers }) +} + +/// Extracts the one string held by a UTF-8 batch constant. +fn decode_constant_utf8(array: &ArrayRef) -> VortexResult { + let Some(constant) = array.as_opt::() else { + vortex_bail!( + "a Utf8 batch constant must use the Constant encoding, got {}", + array.encoding_id() + ); + }; + let scalar = constant.scalar(); + let Some(ScalarValue::Utf8(value)) = scalar.value() else { + vortex_bail!("a Utf8 batch constant must contain a non-null value, got {scalar}"); + }; + + Ok(Utf8Values::single(value.inner().clone())) +} + +fn utf8_view(values: &Utf8Values) -> Utf8ValuesView<'_> { + Utf8ValuesView { + views: values.views(), + buffers: &values.buffers, + } +} + +fn value_at<'a>(view: &Utf8ValuesView<'a>, index: usize) -> Utf8View<'a> { + Utf8View { + view: &view.views[index], + buffers: view.buffers, + } +} + +/// Read a UTF-8 view without checking its row index. +/// +/// # Safety +/// +/// `index` must be less than the length reported by `view`. +unsafe fn value_at_unchecked<'a>(view: &Utf8ValuesView<'a>, index: usize) -> Utf8View<'a> { + // SAFETY: forwarded from this function's contract. + let value = unsafe { view.views.get_unchecked(index) }; + + Utf8View { + view: value, + buffers: view.buffers, + } +} + +// SAFETY: `decode_utf8` validates every valid view and replaces null rows with empty views. The +// borrowed view reports that exact stable length. Therefore each view addresses valid UTF-8, and +// null rows are safe for dense callbacks whose outputs the executor masks with the input validity. +unsafe impl InputElement for Utf8Column { + type Column = Utf8Values; + type Constant = Utf8Values; + type View<'a> = Utf8ValuesView<'a>; + type Elem<'a> = Utf8View<'a>; + + const DENSE_SAFE: bool = true; + const DECODE_INFALLIBLE: bool = true; + + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + matches!(dtype, DType::Utf8(_)), + "expected a Utf8 column, got {dtype}" + ); + + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + decode_utf8(array, ctx) + } + + fn decode_constant(array: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + decode_constant_utf8(&array) + } + + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + + fn get(column: &Self::Column, index: usize) -> Self::Elem<'_> { + value_at(&utf8_view(column), index) + } + + fn get_constant(constant: &Self::Constant) -> Self::Elem<'_> { + value_at(&utf8_view(constant), 0) + } + + fn view(column: &Self::Column) -> Self::View<'_> { + utf8_view(column) + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> + where + Self: 'a, + { + value_at(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> + where + Self: 'a, + { + // SAFETY: the executor validated `index` against this view's exact view slice length. + unsafe { value_at_unchecked(view, index) } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::Buffer; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::Utf8Column; + use crate::IntoArray as _; + use crate::VortexSessionExecute as _; + use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; + use crate::arrays::VarBinViewArray; + use crate::arrays::varbinview::BinaryView; + use crate::buffer::BufferHandle; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::scalar::Scalar; + use crate::scalar_fn::unstable::row::InputElement; + use crate::validity::Validity; + + #[test] + fn input_reads_inline_referenced_and_view_metadata() -> VortexResult<()> { + let array = VarBinViewArray::from_iter_str(["short", "a referenced string"]); + let mut ctx = VortexSession::empty().create_execution_ctx(); + let column = Utf8Column::decode(array.into_array(), &mut ctx)?; + + let inline = Utf8Column::get(&column, 0); + assert_eq!(&*inline, "short"); + assert!(inline.is_inlined()); + assert_eq!(inline.prefix(), b"shor"); + + let referenced = Utf8Column::get(&column, 1); + assert_eq!(&*referenced, "a referenced string"); + assert!(!referenced.is_inlined()); + assert_eq!(referenced.prefix(), b"a re"); + assert_eq!(referenced.raw_view().len(), 19); + + Ok(()) + } + + #[test] + fn input_sanitizes_unvalidated_null_views() -> VortexResult<()> { + let invalid_view = BinaryView::from(u128::MAX); + let views = BufferHandle::new_host(Buffer::from(vec![invalid_view]).into_byte_buffer()); + let validity = BoolArray::from_iter([false]).into_array(); + let array = VarBinViewArray::new_handle( + views, + Default::default(), + DType::Utf8(Nullability::Nullable), + Validity::Array(validity), + ); + let mut ctx = VortexSession::empty().create_execution_ctx(); + let column = Utf8Column::decode(array.into_array(), &mut ctx)?; + + assert_eq!(&*Utf8Column::get(&column, 0), ""); + + Ok(()) + } + + #[test] + fn input_rejects_unvalidated_utf8() { + let invalid_view = BinaryView::make_view(&[0xff], 0, 0); + let views = BufferHandle::new_host(Buffer::from(vec![invalid_view]).into_byte_buffer()); + let array = VarBinViewArray::new_handle( + views, + Default::default(), + DType::Utf8(Nullability::NonNullable), + Validity::NonNullable, + ); + let mut ctx = VortexSession::empty().create_execution_ctx(); + + assert!(Utf8Column::decode(array.into_array(), &mut ctx).is_err()); + } + + #[rstest] + #[case("inlined")] + #[case("a referenced batch constant")] + fn constant_decodes_inlined_and_referenced_values(#[case] value: &str) -> VortexResult<()> { + let array = ConstantArray::new(Scalar::from(value), 8).into_array(); + let mut ctx = VortexSession::empty().create_execution_ctx(); + + let constant = Utf8Column::decode_constant(array, &mut ctx)?; + let view = Utf8Column::get_constant(&constant); + + assert_eq!(&*view, value); + assert_eq!(view.raw_view().len() as usize, value.len()); + + Ok(()) + } +} 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 19a55500e18..beee3c3fabb 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -13,6 +13,8 @@ pub use element::ElementTuple; pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; +pub use element::Utf8Column; +pub use element::Utf8View; pub(super) use element::batch_const; pub(in crate::scalar_fn::unstable::row) use element::decoded_source; @@ -26,6 +28,8 @@ pub use sink::InitializedElement; pub use sink::InitializedRow; pub use sink::OutputSink; pub use sink::UninitElementSink; +pub use sink::Utf8Sink; +pub use sink::Utf8Writer; mod view; pub use view::ViewLen; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs index df13ca18e88..476bf3295da 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs @@ -4,7 +4,8 @@ //! Output builders for row kernels that cannot return independent owned values. //! //! [`OutputSink`] owns the shared lifecycle and safety contract. [`UninitElementSink`] provides -//! uninitialized scalar storage, while [`FixedSizeListSink`] provides runtime-width row storage. +//! uninitialized scalar storage, [`FixedSizeListSink`] provides runtime-width row storage, and +//! [`Utf8Sink`] provides variable-length UTF-8 storage. use vortex_error::VortexResult; @@ -20,6 +21,10 @@ mod uninit_element; pub use uninit_element::InitializedElement; pub use uninit_element::UninitElementSink; +mod utf8; +pub use utf8::Utf8Sink; +pub use utf8::Utf8Writer; + /// A column allocated once per batch that a row closure writes into, one row at a time. /// /// A sink owns batch-wide state that an independent owned value cannot express, such as diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/utf8.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/utf8.rs new file mode 100644 index 00000000000..18809a1c511 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/utf8.rs @@ -0,0 +1,215 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! UTF-8 output storage for row kernels. +//! +//! [`Utf8Sink`] owns one initialized view per output row and batch-wide byte buffers. Its rows are +//! initialized to empty strings, so the sink uses `()` as its write token. A [`Utf8Writer`] +//! consumes the exact row handle when it replaces that placeholder. + +use std::sync::Arc; + +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use super::OutputSink; +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::BinaryView; +use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::ViewLen; +use crate::validity::Validity; + +/// An owned UTF-8 output sink for row functions. +pub struct Utf8Sink { + views: BufferMut, + buffers: Vec, +} + +/// A borrowed view of all UTF-8 output rows. +pub struct Utf8Rows<'a> { + views: &'a mut [BinaryView], + buffers: &'a mut Vec, +} + +impl ViewLen for Utf8Rows<'_> { + fn len(&self) -> usize { + self.views.len() + } +} + +/// The handle used to write one UTF-8 output row. +pub struct Utf8Writer<'a> { + view: &'a mut BinaryView, + buffers: &'a mut Vec, +} + +impl Utf8Writer<'_> { + /// Write a string into this row and consume the row handle. + pub fn write(self, value: impl AsRef) { + let bytes = value.as_ref().as_bytes(); + if bytes.len() <= BinaryView::MAX_INLINED_SIZE { + *self.view = BinaryView::make_view(bytes, 0, 0); + return; + } + + let needs_buffer = self + .buffers + .last() + .is_none_or(|buffer| buffer.len().saturating_add(bytes.len()) > MAX_BUFFER_LEN); + if needs_buffer { + self.buffers.push(ByteBufferMut::with_capacity(bytes.len())); + } + + let buffer_index = u32::try_from(self.buffers.len() - 1) + .vortex_expect("Utf8Sink data buffer count must fit in u32"); + let buffer = self + .buffers + .last_mut() + .vortex_expect("Utf8Sink creates a data buffer before writing"); + let offset = + u32::try_from(buffer.len()).vortex_expect("Utf8Sink buffer offset must fit in u32"); + + buffer.extend_from_slice(bytes); + *self.view = BinaryView::make_view(bytes, buffer_index, offset); + } +} + +// SAFETY: `with_capacity` initializes one distinct `BinaryView` for every row. `rows` retains the +// view slice and byte-buffer vector without changing the row mapping. `row_unchecked` lends one +// exact view slot and the executor cannot request another until the consuming `Utf8Writer` is +// dropped. Skipped rows retain initialized empty views. Errors and unwinds can drop every field, +// and `finish` only publishes initialized views referencing the sink-owned frozen buffers. +unsafe impl OutputSink for Utf8Sink { + type Params = (); + type Rows<'a> = Utf8Rows<'a>; + type Row<'a> = Utf8Writer<'a>; + type WriteToken = (); + + fn skipped_rows_initializer() -> Option)> { + Some(|_| {}) + } + + fn storage_dtype(_params: &Self::Params) -> DType { + DType::Utf8(Nullability::NonNullable) + } + + fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult { + let mut views = BufferMut::with_capacity(rows); + views.push_n(BinaryView::empty_view(), rows); + + Ok(Self { + views, + buffers: Vec::new(), + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + Utf8Rows { + views: self.views.as_mut_slice(), + buffers: &mut self.buffers, + } + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + let view = unsafe { rows.views.get_unchecked_mut(index) }; + + Utf8Writer { + view, + buffers: rows.buffers, + } + } + + unsafe fn finish(self) -> VortexResult { + let views = BufferHandle::new_host(self.views.freeze().into_byte_buffer()); + let buffers = self + .buffers + .into_iter() + .map(|buffer| BufferHandle::new_host(buffer.freeze())) + .collect::>(); + + Ok(VarBinViewArray::new_handle( + views, + Arc::from(buffers), + DType::Utf8(Nullability::NonNullable), + Validity::NonNullable, + ) + .into_array()) + } +} + +#[cfg(test)] +mod tests { + use std::borrow::Cow; + + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::Utf8Sink; + use crate::VortexSessionExecute as _; + use crate::arrays::VarBinViewArray; + use crate::scalar_fn::unstable::row::OutputSink; + + #[test] + fn sink_writes_owned_borrowed_and_cow_values() -> VortexResult<()> { + let expected = ["short", "a referenced string", "owned", "borrowed cow"]; + let referenced = String::from("a referenced string"); + let mut sink = ::with_capacity(expected.len(), &())?; + + { + let mut rows = ::rows(&mut sink); + + // SAFETY: each index is within the four-row sink and is written exactly once. + unsafe { ::row_unchecked(&mut rows, 0) }.write("short"); + // SAFETY: each index is within the four-row sink and is written exactly once. + unsafe { ::row_unchecked(&mut rows, 1) }.write(referenced); + // SAFETY: each index is within the four-row sink and is written exactly once. + unsafe { ::row_unchecked(&mut rows, 2) } + .write(Cow::Owned("owned".to_owned())); + // SAFETY: each index is within the four-row sink and is written exactly once. + unsafe { ::row_unchecked(&mut rows, 3) } + .write(Cow::Borrowed("borrowed cow")); + } + + // SAFETY: every row was initialized by the writes above. + let array = unsafe { ::finish(sink) }?; + let mut ctx = VortexSession::empty().create_execution_ctx(); + let array = array.execute::(&mut ctx)?; + let actual = (0..array.len()) + .map(|index| String::from_utf8_lossy(&array.bytes_at(index)).into_owned()) + .collect::>(); + + assert_eq!(actual, expected); + + Ok(()) + } + + #[test] + fn sink_finishes_empty_and_skipped_rows() -> VortexResult<()> { + let empty = ::with_capacity(0, &())?; + // SAFETY: a zero-row sink has no rows to initialize. + let empty = unsafe { ::finish(empty) }?; + assert!(empty.is_empty()); + + let mut skipped = ::with_capacity(2, &())?; + let initializer = ::skipped_rows_initializer() + .expect("the UTF-8 sink initializes skipped rows"); + initializer(&mut ::rows(&mut skipped)); + // SAFETY: the skipped-row initializer initialized every row. + let skipped = unsafe { ::finish(skipped) }?; + let mut ctx = VortexSession::empty().create_execution_ctx(); + let skipped = skipped.execute::(&mut ctx)?; + + assert_eq!(skipped.bytes_at(0).as_slice(), b""); + assert_eq!(skipped.bytes_at(1).as_slice(), b""); + + Ok(()) + } +}