From d979a9b9c08359aefc79648c2c7e139830fd43d3 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 25 Aug 2026 16:03:36 -0400 Subject: [PATCH] Map constant RowFn inputs directly Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 57 +++++ .../scalar_fn/unstable/row/execute/owned.rs | 41 +--- .../unstable/row/types/element/input.rs | 4 +- .../unstable/row/types/element/mod.rs | 1 + .../row/types/element/tuple/element_tuple.rs | 4 +- .../row/types/element/tuple/indexed.rs | 202 ++++++++++++++++-- .../unstable/row/types/element/tuple/mod.rs | 1 + .../src/scalar_fn/unstable/row/types/mod.rs | 1 + 8 files changed, 264 insertions(+), 47 deletions(-) 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 99e2c074373..3f6717ba53d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -68,6 +68,9 @@ struct InvalidKernelOutput; #[derive(Clone)] struct PackedPositive; +#[derive(Clone)] +struct PackedGreaterThan; + #[derive(Clone)] struct ValidOnlyPositive; @@ -344,6 +347,27 @@ impl RowFn for PackedPositive { } } +impl RowFn for PackedGreaterThan { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.packed_greater_than"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64, i64), bool>(|(lhs, rhs)| lhs > rhs) + } +} + impl RowFn for ValidOnlyPositive { type Options = EmptyOptions; @@ -473,6 +497,39 @@ fn test_bool_output_word_boundaries(#[case] len: usize) -> VortexResult<()> { Ok(()) } +#[test] +fn test_bool_output_handles_partial_constants() -> VortexResult<()> { + let values: Vec<_> = (0_i64..65).collect(); + let varying = PrimitiveArray::from_iter(values.iter().copied()).into_array(); + let constant = ConstantArray::new(32_i64, values.len()).into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let lhs_varying = VecExecutionArgs::new(vec![varying.clone(), constant.clone()], values.len()); + let actual = execute_rows(&PackedGreaterThan, &EmptyOptions, &lhs_varying, &mut ctx)?; + let expected = BoolArray::from_iter(values.iter().map(|value| *value > 32)).into_array(); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + let rhs_varying = VecExecutionArgs::new(vec![constant, varying], values.len()); + let actual = execute_rows(&PackedGreaterThan, &EmptyOptions, &rhs_varying, &mut ctx)?; + let expected = BoolArray::from_iter(values.iter().map(|value| 32 > *value)).into_array(); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + Ok(()) +} + +#[test] +fn test_bool_output_handles_all_constant_input() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 65).into_array(); + let args = VecExecutionArgs::new(vec![input], 65); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&PackedPositive, &EmptyOptions, &args, &mut ctx)?; + let expected = BoolArray::from_iter(std::iter::repeat_n(true, 65)).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + #[test] fn test_bool_output_handles_nullary_rows() -> VortexResult<()> { let args = VecExecutionArgs::new(vec![], 65); diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs index 0a8c250ce1f..926f5398a17 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -5,12 +5,13 @@ //! //! [`execute_owned`] writes fallible row results into spare vector capacity and reduces compact //! failure evidence outside the hot loop. [`execute_owned_infallible`] lets the output type map a -//! validated non-constant row source directly into its physical representation. +//! validated row source directly into its physical representation. use std::ops::BitOrAssign; use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_mask::MaskValuesRef; @@ -21,6 +22,7 @@ use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::unstable::row::FailureEvidence; use crate::scalar_fn::unstable::row::IndexedElementTuple; use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::types::decoded_source; use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; /// Zero-sized failure accumulator for infallible owned visits. @@ -48,38 +50,13 @@ where let prepared = prepare(Args::const_values(&columns)); let row_count = args.row_count(); - if let Some(views) = Args::views_if_no_consts(&columns) { - vortex_ensure!( - Args::view_lens_match(&views, row_count), - "a decoded row input does not address exactly {row_count} rows", - ); - - // SAFETY: `view_lens_match` checked that these exact retained views address `row_count` - // rows. - let source = unsafe { Args::indexed_source(views, row_count) }; - - return Ok(Out::build_from(source, |elements| { - apply(&prepared, elements) - })); - } - - vortex_ensure!( - Args::decoded_lens_match(&columns, row_count), - "a decoded row input does not address exactly {row_count} rows", - ); - - let mut values = Vec::::with_capacity(row_count); - let output = &mut values.spare_capacity_mut()[..row_count]; - - for (index, slot) in output.iter_mut().enumerate() { - slot.write(apply(&prepared, Args::get(&columns, index))); - } - - // SAFETY: normal completion initializes every output slot exactly once, and `values` was - // allocated with at least `row_count` capacity. - unsafe { values.set_len(row_count) }; + let Some(source) = decoded_source::(&columns, row_count) else { + vortex_bail!("a decoded row input does not address exactly {row_count} rows"); + }; - Ok(Out::build(values)) + Ok(Out::build_from(source, |elements| { + apply(&prepared, elements) + })) } /// Decode nullable inputs, then store one output for each valid row from an infallible kernel. diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index e6e89fff220..a9e1c09ed24 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -93,7 +93,9 @@ pub unsafe trait InputElement: 'static { /// Borrow the representation used when this argument varies within the batch. /// /// Called once before the hot loop. Constants do not use this view because the tuple adapter - /// keeps their one-row decoded representation separate. + /// keeps their one-row decoded representation separate. For every index below the returned + /// view's length, [`get_from_view`](Self::get_from_view) must produce the same element as + /// [`get`](Self::get) on `column`. fn view(column: &Self::Column) -> Self::View<'_>; /// Read one row from a [`View`](Self::View). 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 73a0c1d47d1..71b0d17fe79 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 @@ -21,3 +21,4 @@ mod tuple; pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; pub use tuple::batch_const; +pub(in crate::scalar_fn::unstable::row) use tuple::decoded_source; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index b251ede1575..e3c071e0f8f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -24,10 +24,10 @@ use crate::scalar_fn::unstable::row::ViewLen; /// One decoded input, collapsed to a single row when it is constant for the batch. pub struct ArgColumn( /// The decoded argument, classified by how the row loop addresses it. - ArgColumnKind, + pub(super) ArgColumnKind, ); -enum ArgColumnKind { +pub(super) enum ArgColumnKind { /// A decoded column covering the full batch; executors validate its length before traversal. Column(T::Column), diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs index 3a21176e25c..9fcf1042839 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs @@ -10,6 +10,8 @@ use vortex_compute::lane_kernels::IndexedSource; use vortex_compute::lane_kernels::LaneZip; use super::ElementTuple; +use super::element_tuple::ArgColumn; +use super::element_tuple::ArgColumnKind; use crate::scalar_fn::unstable::row::InputElement; use crate::scalar_fn::unstable::row::ViewLen; @@ -17,7 +19,7 @@ use crate::scalar_fn::unstable::row::ViewLen; /// /// Every [`ElementTuple`] implements this trait. Its source delegates each lane read to the tuple's /// unchecked view access after batch execution validates every decoded column length once. -pub trait IndexedElementTuple: ElementTuple { +pub trait IndexedElementTuple: ElementTuple + private::DecodedSource { /// The source used when no input is batch-constant. /// /// Its length must be the common view length. For every valid index it must preserve row order, @@ -34,6 +36,78 @@ pub trait IndexedElementTuple: ElementTuple { unsafe fn indexed_source<'a>(views: Self::Views<'a>, row_count: usize) -> Self::Source<'a>; } +pub(in crate::scalar_fn::unstable::row) fn decoded_source<'a, Args: IndexedElementTuple>( + columns: &'a Args::Columns, + row_count: usize, +) -> Option>> { + ::decoded_source(columns, row_count) +} + +/// Indexed access to one decoded [`ArgColumn`]. +/// +/// [`ArgColumn`] already records whether an input is row-wise or batch-constant. Keeping that choice +/// in each argument source lets LLVM unswitch it before vectorizing Boolean collection into packed +/// words. Routing every input through [`ElementTuple::get`] obscures the independent choices inside +/// the row loop. +enum ArgColumnSource<'a, T: InputElement> { + Rows(T::View<'a>), + Constant { + column: &'a T::Column, + row_count: usize, + }, +} + +impl<'a, T: InputElement> ArgColumnSource<'a, T> { + fn try_new(column: &'a ArgColumn, row_count: usize) -> Option { + match &column.0 { + ArgColumnKind::Column(column) => { + let view = T::view(column); + (view.len() == row_count).then_some(Self::Rows(view)) + } + ArgColumnKind::Const(column) => Some(Self::Constant { column, row_count }), + } + } +} + +impl<'a, T: InputElement> IndexedSource for ArgColumnSource<'a, T> { + type Item = T::Elem<'a>; + + fn len(&self) -> usize { + match self { + Self::Rows(view) => view.len(), + Self::Constant { row_count, .. } => *row_count, + } + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + match self { + Self::Rows(view) => { + // SAFETY: `try_new` checked that this retained view has `row_count` rows, and the + // caller guarantees that `index` is below the source length. + unsafe { T::get_from_view_unchecked(view, index) } + } + // `ArgColumn::try_from_const` validated row zero when it constructed this column. + Self::Constant { column, .. } => T::get(column, 0), + } + } +} + +/// Indexed access to a tuple of decoded argument sources. +struct ArgTupleSource { + sources: Sources, + row_count: usize, +} + +impl IndexedSource for ArgTupleSource<()> { + type Item = (); + + fn len(&self) -> usize { + self.row_count + } + + unsafe fn get_unchecked(&self, _index: usize) -> Self::Item {} +} + /// Indexed access to one element view. pub struct ElementSource<'a, T: InputElement> { view: T::View<'a>, @@ -103,6 +177,18 @@ impl IndexedElementTuple for () { } } +impl private::DecodedSource for () { + fn decoded_source<'a>( + _columns: &'a Self::Columns, + row_count: usize, + ) -> Option>> { + Some(ArgTupleSource { + sources: (), + row_count, + }) + } +} + impl IndexedElementTuple for (A,) { type Source<'a> = UnaryTupleSource>; @@ -111,6 +197,18 @@ impl IndexedElementTuple for (A,) { } } +impl private::DecodedSource for (A,) { + fn decoded_source<'a>( + columns: &'a Self::Columns, + row_count: usize, + ) -> Option>> { + Some(ArgTupleSource { + sources: (ArgColumnSource::try_new(&columns.0, row_count)?,), + row_count, + }) + } +} + impl IndexedElementTuple for (A, B) { type Source<'a> = LaneZip, ElementSource<'a, B>>; @@ -119,8 +217,56 @@ impl IndexedElementTuple for (A, B) { } } +impl private::DecodedSource for (A, B) { + fn decoded_source<'a>( + columns: &'a Self::Columns, + row_count: usize, + ) -> Option>> { + Some(ArgTupleSource { + sources: ( + ArgColumnSource::try_new(&columns.0, row_count)?, + ArgColumnSource::try_new(&columns.1, row_count)?, + ), + row_count, + }) + } +} + +macro_rules! arg_tuple_source { + ($($source:ident: $idx:tt),+) => { + impl<$($source: IndexedSource),+> IndexedSource + for ArgTupleSource<($($source,)+)> + { + type Item = ($($source::Item,)+); + + fn len(&self) -> usize { + self.row_count + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: forwarded from this method's contract. Every source has `row_count` + // rows by construction. + ($(unsafe { self.sources.$idx.get_unchecked(index) },)+) + } + } + }; +} + +arg_tuple_source!(A: 0); +arg_tuple_source!(A: 0, B: 1); +arg_tuple_source!(A: 0, B: 1, C: 2); +arg_tuple_source!(A: 0, B: 1, C: 2, D: 3); +arg_tuple_source!(A: 0, B: 1, C: 2, D: 3, E: 4); +arg_tuple_source!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5); +arg_tuple_source!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6); +arg_tuple_source!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7); +arg_tuple_source!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8); +arg_tuple_source!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9); +arg_tuple_source!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9, K: 10); +arg_tuple_source!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9, K: 10, L: 11); + macro_rules! indexed_element_tuple { - ($($t:ident),+) => { + ($($t:ident: $idx:tt),+) => { impl<$($t: InputElement),+> IndexedElementTuple for ($($t,)+) { type Source<'a> = ElementTupleSource<'a, ($($t,)+)>; @@ -131,16 +277,48 @@ macro_rules! indexed_element_tuple { ElementTupleSource { views, row_count } } } + + impl<$($t: InputElement),+> private::DecodedSource for ($($t,)+) { + fn decoded_source<'a>( + columns: &'a Self::Columns, + row_count: usize, + ) -> Option>> { + Some(ArgTupleSource { + sources: ($(ArgColumnSource::try_new( + &columns.$idx, + row_count, + )?,)+), + row_count, + }) + } + } }; } -indexed_element_tuple!(A, B, C); -indexed_element_tuple!(A, B, C, D); -indexed_element_tuple!(A, B, C, D, E); -indexed_element_tuple!(A, B, C, D, E, F); -indexed_element_tuple!(A, B, C, D, E, F, G); -indexed_element_tuple!(A, B, C, D, E, F, G, H); -indexed_element_tuple!(A, B, C, D, E, F, G, H, I); -indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J); -indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K); -indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K, L); +indexed_element_tuple!(A: 0, B: 1, C: 2); +indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3); +indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3, E: 4); +indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5); +indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6); +indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7); +indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8); +indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9); +indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9, K: 10); +indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9, K: 10, L: 11); + +mod private { + use vortex_compute::lane_kernels::IndexedSource; + + use super::ElementTuple; + + /// The crate-private half of [`super::IndexedElementTuple`]. + /// + /// Rust trait methods have the visibility of their trait. This companion keeps source + /// construction from decoded columns out of the public unstable API. + pub trait DecodedSource: ElementTuple { + fn decoded_source<'a>( + columns: &'a Self::Columns, + row_count: usize, + ) -> Option>>; + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs index cc312a3fb14..d81c607fc5a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -12,6 +12,7 @@ pub use element_tuple::batch_const; mod indexed; pub use indexed::IndexedElementTuple; +pub(in crate::scalar_fn::unstable::row) use indexed::decoded_source; #[cfg(test)] mod tests; 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 007d9916d26..19a55500e18 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -14,6 +14,7 @@ pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; pub(super) use element::batch_const; +pub(in crate::scalar_fn::unstable::row) use element::decoded_source; mod result; pub use result::FailureEvidence;