From 34f5eddd91616dbaa778e510d55aba38da770da8 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 21 Oct 2026 11:02:41 -0400 Subject: [PATCH 1/4] Support borrowed RowFn input decoding Keep the owned decoder as the required extension point and default borrowed decoding to it. Native primitive inputs share their canonical buffers directly. Retain decoded storage after the input handle is dropped. Extracted from #9979. Signed-off-by: Robert Kruszewski Signed-off-by: Connor Tsui --- .../unstable/row/types/element/input.rs | 15 +++++++++- .../unstable/row/types/element/primitive.rs | 15 ++++++++++ .../unstable/row/types/element/tuple/tests.rs | 28 +++++++++++++++++++ 3 files changed, 57 insertions(+), 1 deletion(-) 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 c6608c99202..120dec1d9c3 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 @@ -49,7 +49,8 @@ pub unsafe trait InputElement: 'static { /// null rows to the row closure. const DENSE_SAFE: bool; - /// Whether [`decode`](Self::decode) is infallible for _legal_ input data. + /// Whether [`decode`](Self::decode) and [`decode_ref`](Self::decode_ref) are infallible for + /// _legal_ input data. /// /// This excludes infrastructural failures such as IO or allocation. /// It is independent of @@ -67,6 +68,18 @@ pub unsafe trait InputElement: 'static { /// and other invocation-invariant work into this method. fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + /// Decode a borrowed `array` into its column representation. + /// + /// The default clones `array` for [`decode`](Self::decode). Override this when a common + /// representation decodes without taking ownership of the array. + /// + /// The result must preserve the valid-row values and row count of [`decode`](Self::decode), + /// and must retain its storage after `array` is dropped. [`DENSE_SAFE`](Self::DENSE_SAFE) and + /// [`DECODE_INFALLIBLE`](Self::DECODE_INFALLIBLE) apply to this method too. + fn decode_ref(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Self::decode(array.clone(), ctx) + } + /// Decode one non-null batch-constant input. /// /// `array` retains its logical batch length. Implementations can extract one value directly diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs index f20ad1410fa..2beaab89885 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs @@ -14,6 +14,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::Constant; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::NativePType; @@ -45,9 +46,23 @@ unsafe impl InputElement for T { } fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // A canonical input shares its buffer without executing the array. + if let Some(primitive) = array.as_opt::() { + return Ok(primitive.to_buffer::()); + } + Ok(array.execute::(ctx)?.into_buffer::()) } + #[inline] + fn decode_ref(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(primitive) = array.as_opt::() { + return Ok(primitive.to_buffer::()); + } + + Self::decode(array.clone(), ctx) + } + fn decode_constant(array: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { let Some(constant) = array.as_opt::() else { vortex_bail!( diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs index d49ca3af20c..f7f72241e33 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -4,6 +4,7 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use rstest::rstest; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -16,6 +17,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::BoolArray; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; @@ -105,6 +107,32 @@ i64_test_element!( ); i64_test_element!(DenseFallible, false); +#[rstest] +#[case::canonical(PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(), vec![1, 2, 3])] +#[case::encoded(ConstantArray::new(7_i64, 3).into_array(), vec![7, 7, 7])] +fn test_borrowed_primitive_decode_retains_storage( + #[case] array: ArrayRef, + #[case] expected: Vec, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let column = ::decode_ref(&array, &mut ctx)?; + drop(array); + + assert_eq!(column.as_slice(), expected); + Ok(()) +} + +#[test] +fn test_default_borrowed_decode_retains_storage() -> VortexResult<()> { + let array = BoolArray::from_iter([true, false, true]).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let column = ::decode_ref(&array, &mut ctx)?; + drop(array); + + assert_eq!(column.iter().collect::>(), [true, false, true]); + Ok(()) +} + #[test] fn test_view_lens_match_checks_each_view() { let first: &[i64] = &[1, 2]; From 6348fd29c05dc4da8bd1478c1dd0bbfe6a822744 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 21 Oct 2026 11:04:17 -0400 Subject: [PATCH 2/4] Separate RowFn batch inspection from execution planning Inspect borrowed inputs once for lengths, dtypes, constant classification, and strict validity. Pair collected inputs with a separate execution plan while retaining the existing dispatch and retry paths. Extracted from #9979. Signed-off-by: Robert Kruszewski Signed-off-by: Connor Tsui --- .../unstable/row/batch/execute/constant.rs | 6 +- .../unstable/row/batch/execute/dense.rs | 23 +-- .../unstable/row/batch/execute/filtered.rs | 7 +- .../unstable/row/batch/execute/mod.rs | 20 +-- .../unstable/row/batch/execute/output.rs | 12 +- .../unstable/row/batch/execute/valid_only.rs | 19 ++- .../scalar_fn/unstable/row/batch/inputs.rs | 53 +++++++ .../src/scalar_fn/unstable/row/batch/mod.rs | 55 ++++++-- .../scalar_fn/unstable/row/batch/planning.rs | 131 ++++++++++++++---- .../unstable/row/types/element/mod.rs | 2 +- .../row/types/element/tuple/element_tuple.rs | 29 ++++ .../unstable/row/types/element/tuple/mod.rs | 2 +- .../src/scalar_fn/unstable/row/types/mod.rs | 2 +- .../src/scalar_fn/unstable/row/vtable.rs | 51 +++++-- vortex-session/src/registry.rs | 2 + 15 files changed, 320 insertions(+), 94 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs index bcc61068562..40d9071cb1b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs @@ -10,7 +10,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; -impl RowFnExecutionArgs { +impl RowFnExecutionArgs<'_> { /// Execute all-constant inputs by evaluating one row and broadcasting the validated result. pub(super) fn execute_all_constant( &self, @@ -18,13 +18,13 @@ impl RowFnExecutionArgs { ctx: &mut ExecutionCtx, ) -> VortexResult { let result = self.validate_kernel_output( - kernel(self.execution_args(&self.inputs, 1), ctx)?, + kernel(self.execution_args(&self.batch.inputs, 1), ctx)?, 1, ctx, )?; let result = self.finalize_output(result, 1)?; let scalar = result.execute_scalar(0, ctx)?; - Ok(ConstantArray::new(scalar, self.row_count).into_array()) + Ok(ConstantArray::new(scalar, self.batch.row_count).into_array()) } } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs index cd5ab2aeea1..1fd8d0ad4d0 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs @@ -14,14 +14,17 @@ use crate::builtins::ArrayBuiltins; use crate::scalar_fn::unstable::row::execute::DenseAttempt; use crate::validity::Validity; -impl RowFnExecutionArgs { +impl RowFnExecutionArgs<'_> { /// Run every stored payload, then attach the input validity without materializing its mask. pub(super) fn execute_dense( &self, kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?; + let values = kernel( + self.execution_args(&self.batch.inputs, self.batch.row_count), + ctx, + )?; self.finalize_dense_output(values, ctx) } @@ -49,8 +52,10 @@ impl RowFnExecutionArgs { ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - let attempt = - execute_dense_attempt(self.execution_args(&self.inputs, self.row_count), ctx)?; + let attempt = execute_dense_attempt( + self.execution_args(&self.batch.inputs, self.batch.row_count), + ctx, + )?; match attempt { DenseAttempt::Values(values) => self.finalize_dense_output(values, ctx), @@ -75,7 +80,7 @@ impl RowFnExecutionArgs { ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - let valid_rows = match self.validity.execute_mask(self.row_count, ctx)? { + let valid_rows = match self.batch.validity.execute_mask(self.batch.row_count, ctx)? { // An array-backed validity can materialize to all valid even though the cheap checks // in `RowFnExecutionArgs::execute` could not prove that. The deferred error therefore // came from an observable row and remains terminal. An empty mask is both all-valid @@ -101,13 +106,13 @@ impl RowFnExecutionArgs { values: ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { - let values = self.validate_kernel_output(values, self.row_count, ctx)?; + let values = self.validate_kernel_output(values, self.batch.row_count, ctx)?; - match self.validity.clone() { + match self.batch.validity.clone() { Validity::NonNullable | Validity::AllValid => { - self.finalize_output(values, self.row_count) + self.finalize_output(values, self.batch.row_count) } - Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.batch.row_count), Validity::AllInvalid => { unreachable!("all-invalid validity is handled before dense row execution") } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/filtered.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filtered.rs index 6ed3fab2a1a..90e1ea92e11 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/filtered.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filtered.rs @@ -31,7 +31,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::builtins::ArrayBuiltins; -impl RowFnExecutionArgs { +impl RowFnExecutionArgs<'_> { /// Filter the batch to valid rows, then execute the kernel into the original row domain. pub(super) fn execute_filtered( &self, @@ -47,6 +47,7 @@ impl RowFnExecutionArgs { let filter_mask = Mask::Values(MaskValuesRef::clone(valid)); let filtered_inputs: SmallVec<[ArrayRef; 4]> = self + .batch .inputs .iter() .map(|input| input.filter(filter_mask.clone())) @@ -54,9 +55,9 @@ impl RowFnExecutionArgs { let filtered_args = self.execution_args(&filtered_inputs, filtered_len); let values = execute_filtered_rows(filtered_args, MaskValuesRef::clone(valid), ctx)?; - let values = self.validate_kernel_output(values, self.row_count, ctx)?; + let values = self.validate_kernel_output(values, self.batch.row_count, ctx)?; let mask = valid.as_ref().into_array(); - self.finalize_output(values.mask(mask)?, self.row_count) + self.finalize_output(values.mask(mask)?, self.batch.row_count) } } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs index 7298abfc9bc..9a0e58d357b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -14,9 +14,7 @@ use super::RowPolicy; use super::args::BorrowedRowFnArgs; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::arrays::Constant; use crate::scalar_fn::unstable::row::execute::DenseAttempt; -use crate::scalar_fn::unstable::row::types::batch_const; use crate::validity::Validity; mod constant; @@ -27,7 +25,7 @@ mod valid_only; mod output; pub(crate) use output::finalize_kernel_output; -impl RowFnExecutionArgs { +impl RowFnExecutionArgs<'_> { /// Apply constant folding and null handling around `kernel`. /// /// For a partially valid batch, `try_valid_rows` can avoid filtering. `Ok(None)` filters the @@ -55,29 +53,23 @@ impl RowFnExecutionArgs { ) -> VortexResult { // Strictness: an all-null batch has no observable row work. Keep the literal-constant // check explicit alongside the conjoined validity invariant. - if matches!(self.validity, Validity::AllInvalid) - || self.inputs.iter().any(|input| { - input - .as_opt::() - .is_some_and(|constant| constant.scalar().is_null()) - }) - { + if matches!(self.batch.validity, Validity::AllInvalid) || self.batch.has_null_const { return Ok(self.all_null()); } // All inputs are constant, and their conjoined validity proves that every row is non-null. // The constant check sees through extension and masked wrappers, just like argument // decoding. - if self.row_count > 0 - && self.validity.definitely_no_nulls() - && self.inputs.iter().all(|input| batch_const(input).is_some()) + if self.batch.row_count > 0 + && self.batch.validity.definitely_no_nulls() + && self.batch.all_batch_const() { return self.execute_all_constant(kernel, ctx); } // Do not resolve array-backed validity for the uncommon all-valid or all-null cases here. // That can execute and scan the full mask; each policy resolves it only when necessary. - if self.validity.definitely_no_nulls() { + if self.batch.validity.definitely_no_nulls() { return self.execute_dense(kernel, ctx); } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs index 6afea5798dd..2a3daecc58c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -15,9 +15,13 @@ use crate::dtype::DType; use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnId; -impl RowFnExecutionArgs { +impl RowFnExecutionArgs<'_> { pub(super) fn all_null(&self) -> ArrayRef { - ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() + ConstantArray::new( + Scalar::null(self.result_dtype.clone()), + self.batch.row_count, + ) + .into_array() } /// Label the finished column, validate it, and apply the row function's logical outer @@ -32,7 +36,7 @@ impl RowFnExecutionArgs { // partially valid batch carries the same metadata as a dense one. let values = self.plan.relabel_output(values)?; - validate_output(self.id, &self.result_dtype, expected_len, &values)?; + validate_output(self.batch.id, &self.result_dtype, expected_len, &values)?; cast_output_nullability(&self.result_dtype, values) } @@ -45,7 +49,7 @@ impl RowFnExecutionArgs { ctx: &mut ExecutionCtx, ) -> VortexResult { finalize_kernel_output( - self.id, + self.batch.id, self.plan.storage_dtype(), expected_len, values, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs index 62c5afb7458..33e683db7a5 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs @@ -12,7 +12,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::builtins::ArrayBuiltins; -impl RowFnExecutionArgs { +impl RowFnExecutionArgs<'_> { /// Resolve validity and try direct valid-row execution before filtered execution. pub(super) fn execute_valid_only( &self, @@ -29,15 +29,22 @@ impl RowFnExecutionArgs { ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - let validity = self.validity.clone().execute_mask(self.row_count, ctx)?; + let validity = self + .batch + .validity + .clone() + .execute_mask(self.batch.row_count, ctx)?; let valid_rows = match validity { // An empty mask is both all-valid and all-null. Preserve the all-valid behavior. Mask::AllTrue(_) | Mask::AllFalse(0) => { - let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?; - let values = self.validate_kernel_output(values, self.row_count, ctx)?; + let values = kernel( + self.execution_args(&self.batch.inputs, self.batch.row_count), + ctx, + )?; + let values = self.validate_kernel_output(values, self.batch.row_count, ctx)?; - return self.finalize_output(values, self.row_count); + return self.finalize_output(values, self.batch.row_count); } Mask::AllFalse(_) => return Ok(self.all_null()), Mask::Values(valid_rows) => valid_rows, @@ -62,7 +69,7 @@ impl RowFnExecutionArgs { ctx: &mut ExecutionCtx, ) -> VortexResult> { let Some(values) = try_valid_rows( - self.execution_args(&self.inputs, self.row_count), + self.execution_args(&self.batch.inputs, self.batch.row_count), MaskValuesRef::clone(valid), ctx, )? diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs b/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs new file mode 100644 index 00000000000..226bc5bad31 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Borrowed access to the input columns of one row-function execution. +//! +//! [`RowInputs`] lets batch execution inspect and decode its inputs without cloning them. An +//! internal caller that holds its operands implements it directly. [`ArrayInputs`] adapts the +//! columns collected from a general [`ExecutionArgs`](crate::scalar_fn::ExecutionArgs). + +use crate::ArrayRef; + +/// The same-length input columns of one row-function execution, borrowed by index. +/// +/// Batch inspection and execution share this borrowed view of the inputs. +pub(crate) trait RowInputs { + /// Return the number of input columns. + fn num_inputs(&self) -> usize; + + /// Return the number of rows in the execution scope. + fn row_count(&self) -> usize; + + /// Return input `index`. `index` **must** be less than [`num_inputs`](Self::num_inputs). + fn input(&self, index: usize) -> &ArrayRef; +} + +/// [`RowInputs`] over a slice of collected columns. +pub(crate) struct ArrayInputs<'a> { + /// The input columns. + arrays: &'a [ArrayRef], + + /// The number of rows in the execution scope. + row_count: usize, +} + +impl<'a> ArrayInputs<'a> { + pub(crate) fn new(arrays: &'a [ArrayRef], row_count: usize) -> Self { + Self { arrays, row_count } + } +} + +impl RowInputs for ArrayInputs<'_> { + fn num_inputs(&self) -> usize { + self.arrays.len() + } + + fn row_count(&self) -> usize { + self.row_count + } + + fn input(&self, index: usize) -> &ArrayRef { + &self.arrays[index] + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs index 26e5815bb68..b2709b15dd8 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -9,9 +9,10 @@ //! propagating strict validity, selecting an execution strategy, and validating the finished //! output. //! -//! [`BatchPlan`] carries the nullable execution strategy selected by a concrete dispatch. -//! [`RowFnExecutionArgs`] applies that strategy, and [`BorrowedRowFnArgs`] pairs each kernel -//! invocation with its planning metadata. +//! [`RowInputs`] borrows the input columns, and [`BatchProbe`] classifies them before dispatch. +//! [`BatchInputs`] retains the columns and their classification. [`BatchPlan`] carries the nullable +//! execution strategy selected by a concrete dispatch. [`RowFnExecutionArgs`] applies that strategy, +//! and [`BorrowedRowFnArgs`] pairs each kernel invocation with its planning metadata. use smallvec::SmallVec; @@ -23,6 +24,10 @@ use crate::validity::Validity; mod args; pub(super) use args::BorrowedRowFnArgs; +mod inputs; +pub(crate) use inputs::ArrayInputs; +pub(crate) use inputs::RowInputs; + mod execute; pub(super) use execute::finalize_kernel_output; @@ -31,8 +36,29 @@ mod planning; pub(super) use super::visitor::BatchPlan; pub(super) use super::visitor::RowPolicy; -/// The same-length input columns and metadata for one row-function execution. -pub(crate) struct RowFnExecutionArgs { +/// What batch execution derives from its borrowed inputs before a dispatch plans them. +/// +/// A probe holds no input column. [`BatchInputs`] retains the columns and their probe when a batch +/// needs owned inputs for execution. +pub(crate) struct BatchProbe { + /// The number of rows in the original execution scope. + row_count: usize, + + /// The conjoined input validity. An output row is valid exactly when it is valid in every + /// input. Conjoining is lazy, and null handling materializes the mask only when required. + validity: Validity, + + /// Bit `i` is set when input `i` is a batch constant, seen through extension and masked + /// wrappers. + const_mask: u32, + + /// Whether any input is a null literal constant. + has_null_const: bool, +} + +/// The same-length input columns collected for one row-function execution, with the facts of +/// their [`BatchProbe`]. +pub(crate) struct BatchInputs { /// The function being executed, named in the errors this raises. id: ScalarFnId, @@ -42,13 +68,26 @@ pub(crate) struct RowFnExecutionArgs { /// The input columns, collected once for validity, constant folding, filtering, and execution. inputs: SmallVec<[ArrayRef; 4]>, - /// The input dtypes, collected with the columns and reused by both planning and execution. - arg_dtypes: SmallVec<[DType; 4]>, - /// The conjoined input validity. An output row is valid exactly when it is valid in every /// input. Conjoining is lazy, and null handling materializes the mask only when required. validity: Validity, + /// Bit `i` is set when input `i` is a batch constant, seen through extension and masked + /// wrappers. + const_mask: u32, + + /// Whether any input is a null literal constant. + has_null_const: bool, +} + +/// The same-length input columns and metadata for one row-function execution. +pub(crate) struct RowFnExecutionArgs<'a> { + /// The collected input columns and their conjoined validity. + batch: &'a BatchInputs, + + /// The input dtypes, collected with the columns and reused by both planning and execution. + arg_dtypes: &'a [DType], + /// The output dtype, widened to nullable when any input is nullable. The finished column is /// reconciled against this dtype. result_dtype: DType, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs index 1ae3d525437..72bc25d0bae 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs @@ -1,69 +1,142 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Inspects borrowed columns before pairing collected inputs with a concrete execution plan. +//! +//! Input lengths, constants, and validity are properties of the batch. Planning derives the output +//! dtype and execution policy from the row signature selected by dispatch. + use smallvec::SmallVec; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; +use super::BatchInputs; use super::BatchPlan; +use super::BatchProbe; use super::RowFnExecutionArgs; +use super::RowInputs; use super::args::BorrowedRowFnArgs; use crate::ArrayRef; +use crate::arrays::Constant; use crate::dtype::DType; -use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::unstable::row::types::is_batch_const; use crate::validity::Validity; -impl RowFnExecutionArgs { - /// Collect the inputs and derive their dtypes, validity, and execution policy. +impl BatchProbe { + /// Inspect the borrowed inputs: check their row counts, classify batch constants, and conjoin + /// their validity. Return the probe with the input dtypes. /// - /// This constructor **must not** be used for a nullary function. With no input columns, there - /// is no validity to propagate and the all-constant check would pass vacuously. - pub(crate) fn new( - id: ScalarFnId, - args: &dyn ExecutionArgs, - plan: impl FnOnce(&[DType]) -> VortexResult, - ) -> VortexResult { - let row_count = args.row_count(); - let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) - .map(|index| args.get(index)) - .collect::>()?; + /// This **must not** be used for a nullary function. With no input columns, there is no + /// validity to propagate and the all-constant check would pass vacuously. + /// + /// `id` names the function in the errors this raises. It is called only on an error. + pub(crate) fn new( + id: impl Fn() -> ScalarFnId, + inputs: &I, + ) -> VortexResult<(Self, SmallVec<[DType; 4]>)> { + let row_count = inputs.row_count(); + let num_inputs = inputs.num_inputs(); + vortex_ensure!( + num_inputs <= u32::BITS as usize, + "the {} function has {num_inputs} inputs, more than a row function supports", + id(), + ); - for (index, input) in inputs.iter().enumerate() { + let mut arg_dtypes: SmallVec<[DType; 4]> = SmallVec::with_capacity(num_inputs); + let mut validity = Validity::NonNullable; + let mut const_mask = 0u32; + let mut has_null_const = false; + for index in 0..num_inputs { + let input = inputs.input(index); vortex_ensure_eq!( input.len(), row_count, - "the {id} input {index} has the wrong row count", + "the {} input {index} has the wrong row count", + id(), ); - } - let arg_dtypes: SmallVec<[DType; 4]> = - inputs.iter().map(|input| input.dtype().clone()).collect(); - let plan = plan(&arg_dtypes)?; - let result_dtype = plan.result_dtype(&arg_dtypes); + if is_batch_const(input) { + const_mask |= 1 << index; + has_null_const |= input + .as_opt::() + .is_some_and(|constant| constant.scalar().is_null()); + } - let mut validity = Validity::NonNullable; - for input in &inputs { - validity = validity.and(input.validity()?)?; + // A non-nullable input has no invalid row, and `NonNullable` is the identity of `and`. + if input.dtype().is_nullable() { + validity = validity.and(input.validity()?)?; + } + arg_dtypes.push(input.dtype().clone()); } - Ok(Self { + let probe = Self { + row_count, + validity, + const_mask, + has_null_const, + }; + + Ok((probe, arg_dtypes)) + } +} + +impl BatchInputs { + /// Collect the probed inputs of the function `id`. + pub(crate) fn collect( + id: ScalarFnId, + probe: BatchProbe, + inputs: &I, + ) -> Self { + let BatchProbe { + row_count, + validity, + const_mask, + has_null_const, + } = probe; + + Self { id, row_count, - inputs, - arg_dtypes, + inputs: (0..inputs.num_inputs()) + .map(|index| inputs.input(index).clone()) + .collect(), validity, + const_mask, + has_null_const, + } + } + + /// Whether every input is batch-constant. + pub(super) fn all_batch_const(&self) -> bool { + self.const_mask.count_ones() as usize == self.inputs.len() + } + + /// Attach the plan a dispatch selected for these inputs. + pub(crate) fn with_plan<'a>( + &'a self, + arg_dtypes: &'a [DType], + plan: BatchPlan, + ) -> RowFnExecutionArgs<'a> { + let result_dtype = plan.result_dtype(arg_dtypes); + + RowFnExecutionArgs { + batch: self, + arg_dtypes, result_dtype, plan, - }) + } } +} +impl RowFnExecutionArgs<'_> { /// Pair an input view with this batch's planning metadata. pub(super) fn execution_args<'b>( &'b self, arrays: &'b [ArrayRef], row_count: usize, ) -> BorrowedRowFnArgs<'b> { - BorrowedRowFnArgs::new(arrays, row_count, &self.arg_dtypes, &self.plan) + BorrowedRowFnArgs::new(arrays, row_count, self.arg_dtypes, &self.plan) } } 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 59209303695..26f114a559f 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 @@ -24,8 +24,8 @@ mod primitive; 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; +pub(in crate::scalar_fn::unstable::row) use tuple::is_batch_const; mod utf8; pub use utf8::Utf8Column; 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 77a0c13574d..4f7ea96c856 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 @@ -11,6 +11,7 @@ use vortex_error::vortex_ensure_eq; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::array::VTable; use crate::arrays::Constant; use crate::arrays::Extension; use crate::arrays::Masked; @@ -133,6 +134,34 @@ pub fn batch_const(array: &ArrayRef) -> Option { None } +/// Whether [`batch_const`] would return a batch-constant array, without cloning it. +/// +/// The encoding ID is a field of the array, so comparing it rejects the common non-constant input +/// without the dynamic downcasts that [`batch_const`] performs. +#[expect(clippy::inline_always)] +#[inline(always)] +pub(in crate::scalar_fn::unstable::row) fn is_batch_const(array: &ArrayRef) -> bool { + let id = array.encoding_id(); + + if id == VTable::id(&Constant) { + return array.is::(); + } + + if id == VTable::id(&Masked) { + return array + .as_opt::() + .is_some_and(|masked| masked.child().is::()); + } + + if id == VTable::id(&Extension) { + return array + .as_opt::() + .is_some_and(|extension| extension.storage_array().is::()); + } + + false +} + /// Typed argument tuples for arities zero through twelve. /// /// This trait is sealed. Add a new row representation by implementing [`InputElement`] and placing 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 d81c607fc5a..29b5a1c5221 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 @@ -8,7 +8,7 @@ mod element_tuple; pub use element_tuple::ElementTuple; -pub use element_tuple::batch_const; +pub(in crate::scalar_fn::unstable::row) use element_tuple::is_batch_const; mod indexed; pub use indexed::IndexedElementTuple; 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 29650de0521..907a4a1579c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -19,8 +19,8 @@ pub use element::InputElement; pub use element::OutputBuffer; pub use element::OutputElement; pub use element::Utf8Column; -pub(super) use element::batch_const; pub(in crate::scalar_fn::unstable::row) use element::decoded_source; +pub(super) use element::is_batch_const; mod result; pub use result::FailureEvidence; diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 41950737f49..799bdca37d4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -8,13 +8,17 @@ //! dispatch. [`row_fn_return_dtype`] and [`execute_rows`] expose the same paths to public vtables //! that delegate to a private row kernel. +use smallvec::SmallVec; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_mask::MaskValuesRef; use vortex_session::VortexSession; +use super::batch::ArrayInputs; +use super::batch::BatchInputs; +use super::batch::BatchProbe; use super::batch::BorrowedRowFnArgs; -use super::batch::RowFnExecutionArgs; +use super::batch::RowInputs; use super::batch::finalize_kernel_output; use super::row_fn::RowFn; use super::visitor::BatchPlanner; @@ -119,12 +123,39 @@ pub fn execute_rows( ) -> VortexResult { ensure_arity(function, args.num_inputs())?; - if args.num_inputs() == 0 { - return execute_nullary_rows(function, options, args.row_count(), ctx); + let arrays = (0..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>>()?; + + execute_row_inputs( + function, + options, + &ArrayInputs::new(&arrays, args.row_count()), + ctx, + ) +} + +/// Execute a row function over borrowed input columns. +/// +/// Callers that already hold their operands avoid collecting another owned argument list. +/// Batch inspection validates the lengths before planning or decoding the inputs. +pub(crate) fn execute_row_inputs( + function: &F, + options: &F::Options, + inputs: &I, + ctx: &mut ExecutionCtx, +) -> VortexResult { + ensure_arity(function, inputs.num_inputs())?; + + if inputs.num_inputs() == 0 { + return execute_nullary_rows(function, options, inputs.row_count(), ctx); } - let batch = prepare_batch(function, options, args)?; - batch.execute( + let (probe, arg_dtypes) = BatchProbe::new(|| RowFn::id(function), inputs)?; + let plan = function.dispatch(options, &arg_dtypes, BatchPlanner::::new(&arg_dtypes))?; + let batch = BatchInputs::collect(RowFn::id(function), probe, inputs); + + batch.with_plan(&arg_dtypes, plan).execute( |args, ctx| execute_row_kernel(function, options, args, ctx), |args, ctx| execute_dense_attempt(function, options, args, ctx), |args, valid, ctx| try_execute_valid_rows(function, options, args, valid, ctx), @@ -218,16 +249,6 @@ fn execute_filtered_rows( ) } -fn prepare_batch( - function: &F, - options: &F::Options, - args: &dyn ExecutionArgs, -) -> VortexResult { - RowFnExecutionArgs::new(RowFn::id(function), args, |arg_dtypes| { - function.dispatch(options, arg_dtypes, BatchPlanner::::new(arg_dtypes)) - }) -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/vortex-session/src/registry.rs b/vortex-session/src/registry.rs index fb63fa03174..292a968c298 100644 --- a/vortex-session/src/registry.rs +++ b/vortex-session/src/registry.rs @@ -134,6 +134,8 @@ impl CachedId { impl Deref for CachedId { type Target = Id; + // Inlining keeps an interned read to one acquire load in callers in other crates. + #[inline] #[expect( clippy::disallowed_methods, reason = "CachedId interns its static id once here" From 63703e65f31ceaf723de36a97ccc20b65329810b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 21 Oct 2026 11:05:05 -0400 Subject: [PATCH 3/4] Attach RowFn validity during Boolean output construction Let the packed Boolean collectors consume the batch validity when constructing their result. Keep custom output validation before masking and attach validity directly to other canonical outputs. Extracted from #9979. Signed-off-by: Connor Tsui --- .../unstable/row/batch/execute/constant.rs | 9 ++- .../unstable/row/batch/execute/dense.rs | 40 ++++++++++++- .../unstable/row/batch/execute/mod.rs | 6 +- .../unstable/row/batch/execute/output.rs | 2 +- .../unstable/row/batch/execute/valid_only.rs | 8 ++- .../unstable/row/execute/packed_bool.rs | 11 ++-- .../scalar_fn/unstable/row/execute/tests.rs | 58 ++++++++++++++++--- .../scalar_fn/unstable/row/visitor/execute.rs | 14 ++++- .../src/scalar_fn/unstable/row/vtable.rs | 54 ++++++++++++++++- 9 files changed, 178 insertions(+), 24 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs index 40d9071cb1b..b86243e0fa5 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs @@ -9,16 +9,21 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; +use crate::validity::Validity; impl RowFnExecutionArgs<'_> { /// Execute all-constant inputs by evaluating one row and broadcasting the validated result. pub(super) fn execute_all_constant( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { let result = self.validate_kernel_output( - kernel(self.execution_args(&self.batch.inputs, 1), ctx)?, + kernel(self.execution_args(&self.batch.inputs, 1), &mut None, ctx)?, 1, ctx, )?; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs index 1fd8d0ad4d0..b23777266ad 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs @@ -9,8 +9,12 @@ use vortex_mask::MaskValuesRef; use super::super::RowFnExecutionArgs; use super::super::args::BorrowedRowFnArgs; use crate::ArrayRef; +use crate::Canonical; use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::masked::mask_validity_canonical; use crate::builtins::ArrayBuiltins; +use crate::canonical::AnyCanonical; use crate::scalar_fn::unstable::row::execute::DenseAttempt; use crate::validity::Validity; @@ -18,13 +22,26 @@ impl RowFnExecutionArgs<'_> { /// Run every stored payload, then attach the input validity without materializing its mask. pub(super) fn execute_dense( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { + // Offer a partial validity to the kernel. A collector that takes it builds the output with + // that validity attached, so the output needs neither a second construction nor an + // all-valid check. + let mut validity = + matches!(self.batch.validity, Validity::Array(_)).then(|| self.batch.validity.clone()); let values = kernel( self.execution_args(&self.batch.inputs, self.batch.row_count), + &mut validity, ctx, )?; + if validity.is_none() && matches!(self.batch.validity, Validity::Array(_)) { + return self.finalize_output(values, self.batch.row_count); + } self.finalize_dense_output(values, ctx) } @@ -80,7 +97,11 @@ impl RowFnExecutionArgs<'_> { ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - let valid_rows = match self.batch.validity.execute_mask(self.batch.row_count, ctx)? { + let valid_rows = match self + .batch + .validity + .execute_mask(self.batch.row_count, ctx)? + { // An array-backed validity can materialize to all valid even though the cheap checks // in `RowFnExecutionArgs::execute` could not prove that. The deferred error therefore // came from an observable row and remains terminal. An empty mask is both all-valid @@ -112,7 +133,20 @@ impl RowFnExecutionArgs<'_> { Validity::NonNullable | Validity::AllValid => { self.finalize_output(values, self.batch.row_count) } - Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.batch.row_count), + // A canonical kernel output takes the validity directly. Masking it lazily would add + // a scalar-function node that the caller must execute again just to attach it. + Validity::Array(valid) => { + let values = match values.as_opt::() { + Some(canonical) => mask_validity_canonical( + Canonical::from(canonical), + Validity::Array(valid), + ctx, + )? + .into_array(), + None => values.mask(valid)?, + }; + self.finalize_output(values, self.batch.row_count) + } Validity::AllInvalid => { unreachable!("all-invalid validity is handled before dense row execution") } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs index 9a0e58d357b..ad8e03fb832 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -34,7 +34,11 @@ impl RowFnExecutionArgs<'_> { /// and dtype. pub(crate) fn execute( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, execute_dense_attempt: impl FnOnce( BorrowedRowFnArgs<'_>, &mut ExecutionCtx, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs index 2a3daecc58c..76ebe4560c4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -72,7 +72,7 @@ pub(crate) fn finalize_kernel_output( ) -> VortexResult { validate_output(id, result_dtype, expected_len, &values)?; vortex_ensure!( - values.all_valid(ctx)?, + !values.dtype().is_nullable() || values.all_valid(ctx)?, "the {id} row kernel must produce only valid rows, got at least one null row", ); diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs index 33e683db7a5..80dac87e875 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs @@ -11,12 +11,17 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::builtins::ArrayBuiltins; +use crate::validity::Validity; impl RowFnExecutionArgs<'_> { /// Resolve validity and try direct valid-row execution before filtered execution. pub(super) fn execute_valid_only( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, try_valid_rows: impl FnOnce( BorrowedRowFnArgs<'_>, MaskValuesRef, @@ -40,6 +45,7 @@ impl RowFnExecutionArgs<'_> { Mask::AllTrue(_) | Mask::AllFalse(0) => { let values = kernel( self.execution_args(&self.batch.inputs, self.batch.row_count), + &mut None, ctx, )?; let values = self.validate_kernel_output(values, self.batch.row_count, ctx)?; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs index 5fb3d0cadc1..f99bc1b0369 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs @@ -28,10 +28,11 @@ use crate::scalar_fn::unstable::row::IndexedElementTuple; use crate::scalar_fn::unstable::row::types::decoded_source; use crate::validity::Validity; -/// Decode every input column, then pack infallible Boolean outputs. +/// Decode every input column, then pack infallible Boolean outputs with `validity` attached. pub(crate) fn execute_owned_infallible_bool( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, + validity: Validity, apply: impl Fn(Args::Elems<'_>) -> bool, ) -> VortexResult where @@ -55,13 +56,15 @@ where BitBuffer::collect_bool_in(row_count, collect, ctx.allocator().clone()) }; - Ok(BoolArray::new(values, Validity::NonNullable).into_array()) + Ok(BoolArray::try_new(values, validity)?.into_array()) } -/// Decode every input column, then pack Boolean outputs while combining failure evidence. +/// Decode every input column, then pack Boolean outputs with `validity` attached while combining +/// failure evidence. pub(crate) fn execute_owned_bool( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, + validity: Validity, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, apply: impl Fn(&Prepared, Args::Elems<'_>) -> (bool, Fail), finish_failure: impl FnOnce(Fail) -> VortexResult<()>, @@ -104,7 +107,7 @@ where finish_failure(failure)?; - Ok(BoolArray::new(values, Validity::NonNullable).into_array()) + Ok(BoolArray::try_new(values, validity)?.into_array()) } /// Pack a dense Boolean attempt and report rejected failure evidence to the batch executor. diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs index 40ab20e255e..3f9f460e61d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs @@ -32,6 +32,7 @@ use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::Bool; +use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; use crate::arrays::Primitive; @@ -190,15 +191,22 @@ fn packed_boolean_payload_uses_allocator( .with_allocator(allocator); let output = match (multiversioned, traversal) { - (false, Traversal::Infallible) => { - execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, |(value,)| value > 2)? - } - (true, Traversal::Infallible) => { - execute_owned_infallible_bool::<(i64,), true>(&args, &mut ctx, |(value,)| value > 2)? - } + (false, Traversal::Infallible) => execute_owned_infallible_bool::<(i64,), false>( + &args, + &mut ctx, + Validity::NonNullable, + |(value,)| value > 2, + )?, + (true, Traversal::Infallible) => execute_owned_infallible_bool::<(i64,), true>( + &args, + &mut ctx, + Validity::NonNullable, + |(value,)| value > 2, + )?, (false, Traversal::Fallible) => execute_owned_bool::<(i64,), (), bool, false>( &args, &mut ctx, + Validity::NonNullable, |_| (), |_, (value,)| (value > 2, false), |_| Ok(()), @@ -206,6 +214,7 @@ fn packed_boolean_payload_uses_allocator( (true, Traversal::Fallible) => execute_owned_bool::<(i64,), (), bool, true>( &args, &mut ctx, + Validity::NonNullable, |_| (), |_, (value,)| (value > 2, false), |_| Ok(()), @@ -347,8 +356,12 @@ fn empty_outputs_and_zero_width_rows_do_not_allocate_payloads() -> VortexResult< let primitive = execute_owned_infallible::<(i64,), i64, ()>(&args, &mut ctx, |_| (), |_, (value,)| value)?; - let boolean = - execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, |(value,)| value > 0)?; + let boolean = execute_owned_infallible_bool::<(i64,), false>( + &args, + &mut ctx, + Validity::NonNullable, + |(value,)| value > 0, + )?; let scalar = execute_sink::<(i64,), (), UninitElementSink, _>( &args, &(), @@ -495,3 +508,32 @@ fn zero_sized_output_uses_its_own_storage(#[case] traversal: Traversal) -> Vorte Ok(()) } + +#[rstest] +fn packed_boolean_attaches_validity( + #[values(false, true)] multiversioned: bool, +) -> VortexResult<()> { + let values = PrimitiveArray::from_iter([1_i64, 2, 3, 4]).into_array(); + let args = VecExecutionArgs::new(vec![values], 4); + let validity = Validity::from_iter([true, false, true, true]); + let mut ctx = array_session().create_execution_ctx(); + + let output = if multiversioned { + execute_owned_infallible_bool::<(i64,), true>(&args, &mut ctx, validity, |(value,)| { + value > 2 + })? + } else { + execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, validity, |(value,)| { + value > 2 + })? + }; + + assert!(output.is::()); + let expected = BoolArray::new( + [false, false, true, true].into_iter().collect(), + Validity::from_iter([true, false, true, true]), + ) + .into_array(); + assert_arrays_eq!(&output, &expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index b5539d5b29e..00fa4313df6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -44,6 +44,7 @@ use crate::scalar_fn::unstable::row::execute::execute_owned_valid_rows; use crate::scalar_fn::unstable::row::execute::execute_sink; use crate::scalar_fn::unstable::row::execute::execute_sink_filtered; use crate::scalar_fn::unstable::row::execute::execute_sink_valid_rows; +use crate::validity::Validity; /// The runtime visit that decodes every column once and runs the selected row loop. pub(crate) struct ExecuteRows<'args, 'ctx, F: RowFn> { @@ -59,6 +60,9 @@ pub(crate) struct ExecuteRows<'args, 'ctx, F: RowFn> { /// The output dtype declared by [`RowVisitor::with_output_dtype`], if any. output_dtype: Option, + /// A batch validity that a collector able to attach it takes for its output. + validity: &'ctx mut Option, + /// The execution context used to decode the input columns. ctx: &'ctx mut ExecutionCtx, @@ -71,6 +75,7 @@ impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { args: &'args dyn ExecutionArgs, dtypes: &'args [DType], plan: &'args BatchPlan, + validity: &'ctx mut Option, ctx: &'ctx mut ExecutionCtx, ) -> Self { Self { @@ -78,6 +83,7 @@ impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { dtypes, plan, output_dtype: None, + validity, ctx, function: PhantomData, } @@ -129,7 +135,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { )?; self.plan.ensure_reproduced_by(&visited)?; - execute_owned_infallible_bool::(self.args, self.ctx, apply) + execute_owned_infallible_bool::( + self.args, + self.ctx, + self.validity.take().unwrap_or(Validity::NonNullable), + apply, + ) } fn visit_prepared_into( @@ -205,6 +216,7 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { execute_owned_bool::( self.args, self.ctx, + self.validity.take().unwrap_or(Validity::NonNullable), prepare, apply, finish_failure, diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 799bdca37d4..20906cb031d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -38,6 +38,7 @@ use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::union_child_validities; use crate::scalar_fn::unstable::row::execute::DenseAttempt; +use crate::validity::Validity; impl ScalarFnVTable for F { type Options = F::Options; @@ -156,7 +157,7 @@ pub(crate) fn execute_row_inputs( let batch = BatchInputs::collect(RowFn::id(function), probe, inputs); batch.with_plan(&arg_dtypes, plan).execute( - |args, ctx| execute_row_kernel(function, options, args, ctx), + |args, validity, ctx| execute_row_kernel(function, options, args, validity, ctx), |args, ctx| execute_dense_attempt(function, options, args, ctx), |args, valid, ctx| try_execute_valid_rows(function, options, args, valid, ctx), |args, valid, ctx| execute_filtered_rows(function, options, args, valid, ctx), @@ -177,7 +178,8 @@ fn execute_nullary_rows( // A nullary function has no input validity to propagate, so its kernel output is the finished // column and this path labels it directly. - let values = plan.relabel_output(execute_row_kernel(function, options, args, ctx)?)?; + let values = + plan.relabel_output(execute_row_kernel(function, options, args, &mut None, ctx)?)?; finalize_kernel_output(RowFn::id(function), &result_dtype, row_count, values, ctx) } @@ -194,16 +196,21 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { Ok(()) } +/// Execute the row kernel over `args`. +/// +/// A collector that can attach the validity in `validity` takes it and builds its output with it, +/// leaving `None` behind. Otherwise the output is non-nullable and `validity` is left unchanged. fn execute_row_kernel( function: &F, options: &F::Options, args: BorrowedRowFnArgs<'_>, + validity: &mut Option, ctx: &mut ExecutionCtx, ) -> VortexResult { function.dispatch( options, args.dtypes(), - ExecuteRows::::new(&args, args.dtypes(), args.plan(), ctx), + ExecuteRows::::new(&args, args.dtypes(), args.plan(), validity, ctx), ) } @@ -265,6 +272,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -450,6 +458,46 @@ mod tests { Ok(()) } + #[derive(Clone)] + struct GreaterThanTwo; + + impl RowFn for GreaterThanTwo { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.greater_than_two"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_bool::<(i64,), false>(|(value,)| value > 2) + } + } + + #[test] + fn test_partially_valid_packed_bool_output_is_canonical() -> VortexResult<()> { + let validity = Validity::from_iter([true, false, true, true]); + let input = PrimitiveArray::new(vec![1_i64, 3, 3, 0], validity.clone()).into_array(); + let args = VecExecutionArgs::new(vec![input], 4); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&GreaterThanTwo, &EmptyOptions, &args, &mut ctx)?; + let expected = + BoolArray::new([false, true, true, false].into_iter().collect(), validity).into_array(); + + assert!(actual.is_canonical(), "{actual}"); + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } + #[track_caller] fn assert_arity_error(error: VortexError) { assert!( From 5cbf00a7836bb406e169b3c21fefe6bd9fe9c2a3 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 21 Oct 2026 11:07:15 -0400 Subject: [PATCH 4/4] Plan and execute common RowFn batches in one dispatch Finish eligible batches in the planning visit and retain the planned retry and valid-row paths for the remaining partially valid batches. Decode direct Boolean inputs through private borrowed tuple adapters and keep collector selection outside the row loop. Extracted from #9979. Enforce the Boolean validity-length invariant in release builds before unchecked output construction. Signed-off-by: Robert Kruszewski Signed-off-by: Connor Tsui --- .../unstable/row/batch/execute/dense.rs | 6 +- .../unstable/row/batch/execute/mod.rs | 72 +++-- .../scalar_fn/unstable/row/batch/inputs.rs | 5 +- .../src/scalar_fn/unstable/row/batch/mod.rs | 15 +- .../scalar_fn/unstable/row/batch/planning.rs | 40 +++ .../src/scalar_fn/unstable/row/batch/tests.rs | 8 +- .../src/scalar_fn/unstable/row/execute/mod.rs | 1 + .../unstable/row/execute/packed_bool.rs | 90 +++++- .../scalar_fn/unstable/row/execute/tests.rs | 16 + .../unstable/row/types/element/mod.rs | 1 + .../row/types/element/tuple/element_tuple.rs | 22 ++ .../row/types/element/tuple/indexed.rs | 94 ++++++ .../unstable/row/types/element/tuple/mod.rs | 1 + .../src/scalar_fn/unstable/row/types/mod.rs | 3 +- .../scalar_fn/unstable/row/visitor/fused.rs | 290 ++++++++++++++++++ .../src/scalar_fn/unstable/row/visitor/mod.rs | 4 + .../src/scalar_fn/unstable/row/vtable.rs | 54 +++- 17 files changed, 662 insertions(+), 60 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/visitor/fused.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs index b23777266ad..85e8881ad7c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs @@ -103,9 +103,9 @@ impl RowFnExecutionArgs<'_> { .execute_mask(self.batch.row_count, ctx)? { // An array-backed validity can materialize to all valid even though the cheap checks - // in `RowFnExecutionArgs::execute` could not prove that. The deferred error therefore - // came from an observable row and remains terminal. An empty mask is both all-valid - // and all-null, so preserve the all-valid behavior. + // in `RowFnExecutionArgs::try_execute_without_retry` could not prove that. The + // deferred error therefore came from an observable row and remains terminal. An empty + // mask is both all-valid and all-null, so preserve the all-valid behavior. Mask::AllTrue(_) | Mask::AllFalse(0) => return Err(deferred_error), Mask::AllFalse(_) => return Ok(self.all_null()), Mask::Values(valid_rows) => valid_rows, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs index ad8e03fb832..ed817f01695 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -3,8 +3,10 @@ //! Selects a batch execution strategy. //! -//! [`RowFnExecutionArgs::execute`] handles universal fast paths, then delegates to dense or -//! valid-only execution. +//! [`RowFnExecutionArgs::try_execute_without_retry`] handles universal fast paths and dense +//! execution, so the visit that planned a batch can finish it without dispatching again. +//! [`RowFnExecutionArgs::execute_partially_valid`] then delegates the remaining partially valid +//! batches to dense retry or valid-only execution. use vortex_error::VortexResult; use vortex_mask::MaskValuesRef; @@ -26,13 +28,53 @@ mod output; pub(crate) use output::finalize_kernel_output; impl RowFnExecutionArgs<'_> { - /// Apply constant folding and null handling around `kernel`. + /// Execute every batch that `kernel` alone can finish, or return `None` for a partially valid + /// batch whose plan needs a dense retry or valid-row execution. + /// + /// `kernel` reports a rejected deferred failure as a terminal error. That is correct for each + /// batch this method executes: no null row can be the cause of the failure. + pub(crate) fn try_execute_without_retry( + &self, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + // Strictness: an all-null batch has no observable row work. Keep the literal-constant + // check explicit alongside the conjoined validity invariant. + if matches!(self.batch.validity, Validity::AllInvalid) || self.batch.has_null_const { + return Ok(Some(self.all_null())); + } + + // All inputs are constant, and their conjoined validity proves that every row is non-null. + // The constant check sees through extension and masked wrappers, just like argument + // decoding. + if self.batch.row_count > 0 + && self.batch.validity.definitely_no_nulls() + && self.batch.all_batch_const() + { + return self.execute_all_constant(kernel, ctx).map(Some); + } + + // Do not resolve array-backed validity for the uncommon all-valid or all-null cases here. + // That can execute and scan the full mask; each policy resolves it only when necessary. + if self.batch.validity.definitely_no_nulls() || self.plan.policy() == RowPolicy::Dense { + return self.execute_dense(kernel, ctx).map(Some); + } + + Ok(None) + } + + /// Execute a partially valid batch that + /// [`try_execute_without_retry`](Self::try_execute_without_retry) declined. /// /// For a partially valid batch, `try_valid_rows` can avoid filtering. `Ok(None)` filters the /// inputs to the valid rows and `execute_filtered_rows` runs the kernel over them while /// writing into the original row domain. Every result is checked against the planned shape /// and dtype. - pub(crate) fn execute( + pub(crate) fn execute_partially_valid( &self, kernel: impl FnOnce( BorrowedRowFnArgs<'_>, @@ -55,28 +97,6 @@ impl RowFnExecutionArgs<'_> { ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - // Strictness: an all-null batch has no observable row work. Keep the literal-constant - // check explicit alongside the conjoined validity invariant. - if matches!(self.batch.validity, Validity::AllInvalid) || self.batch.has_null_const { - return Ok(self.all_null()); - } - - // All inputs are constant, and their conjoined validity proves that every row is non-null. - // The constant check sees through extension and masked wrappers, just like argument - // decoding. - if self.batch.row_count > 0 - && self.batch.validity.definitely_no_nulls() - && self.batch.all_batch_const() - { - return self.execute_all_constant(kernel, ctx); - } - - // Do not resolve array-backed validity for the uncommon all-valid or all-null cases here. - // That can execute and scan the full mask; each policy resolves it only when necessary. - if self.batch.validity.definitely_no_nulls() { - return self.execute_dense(kernel, ctx); - } - match self.plan.policy() { RowPolicy::Dense => self.execute_dense(kernel, ctx), RowPolicy::DenseWithRetry => self.execute_dense_with_retry( diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs b/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs index 226bc5bad31..ffbc9de11fc 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs @@ -11,8 +11,9 @@ use crate::ArrayRef; /// The same-length input columns of one row-function execution, borrowed by index. /// -/// Batch inspection and execution share this borrowed view of the inputs. -pub(crate) trait RowInputs { +/// The trait is `pub` only so that the sealed tuple decoding can bound on it. The private module +/// keeps it out of the public API. +pub trait RowInputs { /// Return the number of input columns. fn num_inputs(&self) -> usize; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs index b2709b15dd8..bfc02cde32d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -10,14 +10,16 @@ //! output. //! //! [`RowInputs`] borrows the input columns, and [`BatchProbe`] classifies them before dispatch. -//! [`BatchInputs`] retains the columns and their classification. [`BatchPlan`] carries the nullable -//! execution strategy selected by a concrete dispatch. [`RowFnExecutionArgs`] applies that strategy, -//! and [`BorrowedRowFnArgs`] pairs each kernel invocation with its planning metadata. +//! [`BatchInputs`] collects them when a visit cannot finish the batch directly. [`BatchPlan`] +//! carries the nullable execution strategy selected by a concrete dispatch. [`RowFnExecutionArgs`] +//! applies that strategy, and [`BorrowedRowFnArgs`] pairs each kernel invocation with its planning +//! metadata. use smallvec::SmallVec; use crate::ArrayRef; use crate::dtype::DType; +use crate::dtype::Nullability; use crate::scalar_fn::ScalarFnId; use crate::validity::Validity; @@ -38,8 +40,8 @@ pub(super) use super::visitor::RowPolicy; /// What batch execution derives from its borrowed inputs before a dispatch plans them. /// -/// A probe holds no input column. [`BatchInputs`] retains the columns and their probe when a batch -/// needs owned inputs for execution. +/// A probe holds no input column. A visit that can finish the batch directly decodes the borrowed +/// inputs. Any other visit collects them into [`BatchInputs`]. pub(crate) struct BatchProbe { /// The number of rows in the original execution scope. row_count: usize, @@ -54,6 +56,9 @@ pub(crate) struct BatchProbe { /// Whether any input is a null literal constant. has_null_const: bool, + + /// The strict result nullability: nullable when any input dtype is nullable. + nullability: Nullability, } /// The same-length input columns collected for one row-function execution, with the facts of diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs index 72bc25d0bae..f53b67d97b5 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs @@ -20,6 +20,7 @@ use super::args::BorrowedRowFnArgs; use crate::ArrayRef; use crate::arrays::Constant; use crate::dtype::DType; +use crate::dtype::Nullability; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::unstable::row::types::is_batch_const; use crate::validity::Validity; @@ -48,6 +49,7 @@ impl BatchProbe { let mut validity = Validity::NonNullable; let mut const_mask = 0u32; let mut has_null_const = false; + let mut nullability = Nullability::NonNullable; for index in 0..num_inputs { let input = inputs.input(index); vortex_ensure_eq!( @@ -66,6 +68,7 @@ impl BatchProbe { // A non-nullable input has no invalid row, and `NonNullable` is the identity of `and`. if input.dtype().is_nullable() { + nullability = Nullability::Nullable; validity = validity.and(input.validity()?)?; } arg_dtypes.push(input.dtype().clone()); @@ -76,10 +79,41 @@ impl BatchProbe { validity, const_mask, has_null_const, + nullability, }; Ok((probe, arg_dtypes)) } + + /// Return the output validity of a _direct_ batch, or return the probe unchanged for any other + /// batch. + /// + /// A visit finishes a direct batch from its kernel output alone, once that output carries + /// this validity. A direct batch has no null literal, at least one input that is not + /// batch-constant, and at least one possibly valid row. It therefore needs no all-null or + /// all-constant handling. Set `dense` when the visit plans + /// [`RowPolicy::Dense`](super::RowPolicy::Dense). Without it, only a batch that provably has + /// no null row is direct, because the kernel visits every row. + pub(crate) fn into_direct_validity( + self, + num_inputs: usize, + dense: bool, + ) -> Result { + if self.has_null_const + || self.const_mask.count_ones() as usize == num_inputs + || matches!(self.validity, Validity::AllInvalid) + || !(dense || self.validity.definitely_no_nulls()) + { + return Err(self); + } + + Ok(self.validity.union_nullability(self.nullability)) + } + + /// Return the batch-constant classification of the inputs: bit `i` marks input `i`. + pub(crate) fn const_mask(&self) -> u32 { + self.const_mask + } } impl BatchInputs { @@ -94,6 +128,7 @@ impl BatchInputs { validity, const_mask, has_null_const, + nullability: _, } = probe; Self { @@ -131,6 +166,11 @@ impl BatchInputs { } impl RowFnExecutionArgs<'_> { + /// Return the plan, discarding the batch view. + pub(crate) fn into_plan(self) -> BatchPlan { + self.plan + } + /// Pair an input view with this batch's planning metadata. pub(super) 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 1aa7ee43bbf..8a0669ce555 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -1261,7 +1261,7 @@ impl RowFn for ChangingOutputDType { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["value"]; - const INFALLIBLE: bool = true; + const INFALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.changing_output_dtype"); @@ -1280,7 +1280,8 @@ impl RowFn for ChangingOutputDType { visitor }; - visitor.visit::<(i64,), i64>(|(value,)| value) + // A deferred visit plans a dense retry, so a partially valid batch dispatches again. + visitor.visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())) } } @@ -1462,7 +1463,8 @@ fn test_execution_rejects_a_changed_output_dtype() -> VortexResult<()> { let function = ChangingOutputDType { dispatches: Arc::new(AtomicUsize::new(0)), }; - let input = PrimitiveArray::from_iter(vec![1_i64, 2]).into_array(); + let validity = Validity::from_iter([true, false]); + let input = PrimitiveArray::new(vec![1_i64, 2], validity).into_array(); let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs index c9e17da1d09..3ca76d6360c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -18,6 +18,7 @@ mod packed_bool; pub(super) use packed_bool::execute_bool_dense_attempt; pub(super) use packed_bool::execute_owned_bool; pub(super) use packed_bool::execute_owned_infallible_bool; +pub(super) use packed_bool::execute_rows_infallible_bool; mod retry; pub(super) use retry::DenseAttempt; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs index f99bc1b0369..1fd14e545c6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs @@ -25,6 +25,8 @@ use crate::arrays::BoolArray; 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::batch::RowInputs; +use crate::scalar_fn::unstable::row::types::decode_rows; use crate::scalar_fn::unstable::row::types::decoded_source; use crate::validity::Validity; @@ -39,24 +41,94 @@ where Args: IndexedElementTuple, { let columns = Args::decode(args, ctx)?; - let row_count = args.row_count(); - let Some(source) = decoded_source::(&columns, row_count) else { - vortex_bail!("a decoded row input does not address exactly {row_count} rows"); + pack_infallible_bool::(&columns, args.row_count(), ctx, validity, apply) +} + +/// Decode one borrowed array per argument, then pack infallible Boolean outputs with `validity` +/// attached. +/// +/// Bit `i` of `const_mask` is set exactly when input `i` is batch-constant. +/// +/// Decoding and packing inline into this one frame, and this frame stays out of the dispatch +/// that calls it. A short batch then pays for one call boundary, not one per step, and each arm +/// of the dispatch stays small. +#[inline(never)] +pub(crate) fn execute_rows_infallible_bool( + inputs: &I, + const_mask: u32, + ctx: &mut ExecutionCtx, + validity: Validity, + apply: impl Fn(Args::Elems<'_>) -> bool, +) -> VortexResult +where + Args: IndexedElementTuple, + I: RowInputs + ?Sized, +{ + let columns = decode_rows::(inputs, const_mask, ctx)?; + + pack_infallible_bool::(&columns, inputs.row_count(), ctx, validity, apply) +} + +#[expect(clippy::inline_always)] +#[inline(always)] +fn pack_infallible_bool( + columns: &Args::Columns, + row_count: usize, + ctx: &mut ExecutionCtx, + validity: Validity, + apply: impl Fn(Args::Elems<'_>) -> bool, +) -> VortexResult +where + Args: IndexedElementTuple, +{ + // Without batch constants, read plain views. Selecting each argument's constant or row + // source per row stays in short scalar loops, which LLVM does not unswitch. + let values = if let Some(views) = Args::views_if_no_consts(columns) { + if !Args::view_lens_match(&views, row_count) { + vortex_bail!("a decoded row input does not address exactly {row_count} rows"); + } + + // SAFETY: `view_lens_match` checked that these exact views address `row_count` rows. + let source = unsafe { Args::indexed_source(views, row_count) }; + collect_source_bits::(&source, row_count, ctx, &apply) + } else { + let Some(source) = decoded_source::(columns, row_count) else { + vortex_bail!("a decoded row input does not address exactly {row_count} rows"); + }; + collect_source_bits::(&source, row_count, ctx, &apply) }; + assert!( + validity.maybe_len().is_none_or(|len| len == row_count), + "packed Boolean validity length must match the row count", + ); + + // SAFETY: the collector packs exactly `row_count` bits from offset zero, and the assertion + // above checks the validity length in every build. + Ok(unsafe { BoolArray::new_unchecked(values, validity) }.into_array()) +} + +/// Pack `apply` over the `row_count` rows of `source`. +#[expect(clippy::inline_always)] +#[inline(always)] +fn collect_source_bits( + source: &S, + row_count: usize, + ctx: &ExecutionCtx, + apply: &impl Fn(S::Item) -> bool, +) -> BitBuffer { let collect = |index| { - // SAFETY: Both collectors only invoke this closure with `index < row_count`, and the decoded - // source was constructed with exactly `row_count` rows. + // SAFETY: Both collectors only invoke this closure with `index < row_count`, and the + // source addresses exactly `row_count` rows. apply(unsafe { source.get_unchecked(index) }) }; - let values = if MULTIVERSIONED { + + if MULTIVERSIONED { BitBuffer::collect_bool_multiversioned_in(row_count, collect, ctx.allocator().clone()) } else { BitBuffer::collect_bool_in(row_count, collect, ctx.allocator().clone()) - }; - - Ok(BoolArray::try_new(values, validity)?.into_array()) + } } /// Decode every input column, then pack Boolean outputs with `validity` attached while combining diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs index 3f9f460e61d..69cdcac1427 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs @@ -537,3 +537,19 @@ fn packed_boolean_attaches_validity( assert_arrays_eq!(&output, &expected, &mut ctx); Ok(()) } + +#[test] +#[should_panic(expected = "packed Boolean validity length must match the row count")] +fn packed_boolean_rejects_wrong_validity_length() { + let values = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let args = VecExecutionArgs::new(vec![values], 2); + let validity = Validity::from_iter([true, false, true]); + let mut ctx = array_session().create_execution_ctx(); + + drop(execute_owned_infallible_bool::<(i64,), false>( + &args, + &mut ctx, + validity, + |(value,)| value > 0, + )); +} 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 26f114a559f..ef4f516c328 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 @@ -24,6 +24,7 @@ mod primitive; mod tuple; pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; +pub(in crate::scalar_fn::unstable::row) use tuple::decode_rows; pub(in crate::scalar_fn::unstable::row) use tuple::decoded_source; pub(in crate::scalar_fn::unstable::row) use tuple::is_batch_const; 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 4f7ea96c856..35493de055d 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 @@ -51,6 +51,28 @@ impl ArgColumn { Ok(Self(ArgColumnKind::Column(T::decode(array, ctx)?))) } + /// Decode an input already classified as batch-constant or not by [`is_batch_const`]. + #[inline] + pub(super) fn decode_rows( + array: &ArrayRef, + is_const: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if is_const { + // A literal constant is its own batch constant. Wrapped constants unwrap in `decode`. + if array.is::() && !array.is_empty() { + return Ok(Self(ArgColumnKind::Const(T::decode_constant( + array.clone(), + ctx, + )?))); + } + + return Self::decode(array.clone(), ctx); + } + + Ok(Self(ArgColumnKind::Column(T::decode_ref(array, ctx)?))) + } + fn decode_null_tolerant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { // Batch execution short-circuits null constants before selecting a strategy. if let Some(const_array) = batch_const(&array) 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 73605aa7d33..808ead04d1a 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 @@ -8,12 +8,16 @@ use vortex_compute::lane_kernels::IndexedSource; use vortex_compute::lane_kernels::LaneZip; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; use super::ElementTuple; use super::element_tuple::ArgColumn; use super::element_tuple::ArgColumnKind; +use crate::ExecutionCtx; use crate::scalar_fn::unstable::row::InputElement; use crate::scalar_fn::unstable::row::ViewLen; +use crate::scalar_fn::unstable::row::batch::RowInputs; /// An argument tuple that supports a validated dense indexed traversal. /// @@ -43,6 +47,44 @@ pub(in crate::scalar_fn::unstable::row) fn decoded_source<'a, Args: IndexedEleme ::decoded_source(columns, row_count) } +/// Decode one borrowed array per argument, like [`ElementTuple::decode`]. +/// +/// Bit `i` of `const_mask` is set exactly when input `i` is batch-constant, as classified by +/// `is_batch_const`. Only those inputs repeat the batch-constant checks. +#[expect(clippy::inline_always)] +#[inline(always)] +pub(in crate::scalar_fn::unstable::row) fn decode_rows( + inputs: &Inputs, + const_mask: u32, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + Args: IndexedElementTuple, + Inputs: RowInputs + ?Sized, +{ + vortex_ensure_eq!( + inputs.num_inputs(), + Args::ARITY, + "row input count must match the argument tuple arity", + ); + + ::decode_rows(inputs, const_mask, ctx) +} + +/// Decode input `index` as an argument of element type `T`. +#[expect(clippy::inline_always)] +#[inline(always)] +fn decode_input( + inputs: &Inputs, + index: usize, + const_mask: u32, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let is_const = const_mask & (1 << index) != 0; + + ArgColumn::decode_rows(inputs.input(index), is_const, ctx) +} + /// Indexed access to one decoded [`ArgColumn`]. /// /// [`ArgColumn`] already records whether an input is row-wise or batch-constant. Keeping that choice @@ -182,6 +224,14 @@ impl IndexedElementTuple for () { } impl private::DecodedSource for () { + fn decode_rows( + _inputs: &Inputs, + _const_mask: u32, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(()) + } + fn decoded_source<'a>( _columns: &'a Self::Columns, row_count: usize, @@ -202,6 +252,16 @@ impl IndexedElementTuple for (A,) { } impl private::DecodedSource for (A,) { + #[expect(clippy::inline_always)] + #[inline(always)] + fn decode_rows( + inputs: &Inputs, + const_mask: u32, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok((decode_input(inputs, 0, const_mask, ctx)?,)) + } + fn decoded_source<'a>( columns: &'a Self::Columns, row_count: usize, @@ -222,6 +282,19 @@ impl IndexedElementTuple for (A, B) { } impl private::DecodedSource for (A, B) { + #[expect(clippy::inline_always)] + #[inline(always)] + fn decode_rows( + inputs: &Inputs, + const_mask: u32, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(( + decode_input(inputs, 0, const_mask, ctx)?, + decode_input(inputs, 1, const_mask, ctx)?, + )) + } + fn decoded_source<'a>( columns: &'a Self::Columns, row_count: usize, @@ -283,6 +356,16 @@ macro_rules! indexed_element_tuple { } impl<$($t: InputElement),+> private::DecodedSource for ($($t,)+) { + #[allow(clippy::inline_always)] + #[inline(always)] + fn decode_rows( + inputs: &Inputs, + const_mask: u32, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(($(decode_input::<$t, Inputs>(inputs, $idx, const_mask, ctx)?,)+)) + } + fn decoded_source<'a>( columns: &'a Self::Columns, row_count: usize, @@ -312,14 +395,25 @@ indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: mod private { use vortex_compute::lane_kernels::IndexedSource; + use vortex_error::VortexResult; use super::ElementTuple; + use crate::ExecutionCtx; + use crate::scalar_fn::unstable::row::batch::RowInputs; /// 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 { + /// Decode one array per argument. Bit `i` of `const_mask` marks input `i` as + /// batch-constant. + fn decode_rows( + inputs: &Inputs, + const_mask: u32, + ctx: &mut ExecutionCtx, + ) -> VortexResult; + fn decoded_source<'a>( columns: &'a Self::Columns, row_count: usize, 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 29b5a1c5221..0054c36838e 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(in crate::scalar_fn::unstable::row) use element_tuple::is_batch_const; mod indexed; pub use indexed::IndexedElementTuple; +pub(in crate::scalar_fn::unstable::row) use indexed::decode_rows; pub(in crate::scalar_fn::unstable::row) use indexed::decoded_source; #[cfg(test)] 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 907a4a1579c..ee7d87954bd 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -19,8 +19,9 @@ pub use element::InputElement; pub use element::OutputBuffer; pub use element::OutputElement; pub use element::Utf8Column; +pub(in crate::scalar_fn::unstable::row) use element::decode_rows; pub(in crate::scalar_fn::unstable::row) use element::decoded_source; -pub(super) use element::is_batch_const; +pub(in crate::scalar_fn::unstable::row) use element::is_batch_const; mod result; pub use result::FailureEvidence; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/fused.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/fused.rs new file mode 100644 index 00000000000..fc659deac61 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/fused.rs @@ -0,0 +1,290 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Plans and executes a batch in a single dispatch. +//! +//! [`ExecuteFused`] borrows the inputs and owns their [`BatchProbe`]. Each visit plans its concrete +//! signature, then runs every batch that its own kernel can finish: an all-null, all-constant, or +//! null-free batch, or a partially valid batch under [`RowPolicy::Dense`]. Such a batch therefore +//! needs no second dispatch and no check that a second dispatch reproduced the plan. A visit hands +//! any other batch back as [`Fused::Partial`], and batch execution dispatches again for dense +//! retry or valid-row execution. +//! +//! A Boolean visit decodes a direct batch straight from the borrowed inputs. Every other visit +//! collects the inputs into [`BatchInputs`]. + +use vortex_error::VortexResult; + +use super::BatchPlan; +use super::RowPolicy; +use super::RowVisitor; +use super::check::assert_deferred_visit_contract; +use super::check::assert_owned_visit_contract; +use super::check::assert_sink_visit_contract; +use super::check::validate_owned_visit; +use super::check::validate_sink_visit; +use super::row_visitor::private; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::unstable::row::ElementTuple; +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::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::SinkResult; +use crate::scalar_fn::unstable::row::batch::BatchInputs; +use crate::scalar_fn::unstable::row::batch::BatchProbe; +use crate::scalar_fn::unstable::row::batch::BorrowedRowFnArgs; +use crate::scalar_fn::unstable::row::batch::RowInputs; +use crate::scalar_fn::unstable::row::execute::execute_owned; +use crate::scalar_fn::unstable::row::execute::execute_owned_bool; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible_bool; +use crate::scalar_fn::unstable::row::execute::execute_rows_infallible_bool; +use crate::scalar_fn::unstable::row::execute::execute_sink; +use crate::validity::Validity; + +/// The result of a fused planning and execution visit. +pub(crate) enum Fused { + /// The finished output column. + Done(ArrayRef), + + /// A partially valid batch that needs dense retry or valid-row execution, with its plan. + Partial(BatchInputs, BatchPlan), +} + +/// The visit that plans a batch and executes it when no retry can be necessary. +pub(crate) struct ExecuteFused<'a, 'ctx, F: RowFn, I: ?Sized> { + /// The borrowed input columns. + inputs: &'a I, + + /// What batch execution derived from the inputs. + probe: BatchProbe, + + /// The input dtypes passed to the dispatch. + dtypes: &'a [DType], + + /// The output dtype declared by [`RowVisitor::with_output_dtype`], if any. + output_dtype: Option, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, + + /// The function being executed, named in the errors of a collected batch. + function: &'a F, +} + +impl<'a, 'ctx, F: RowFn, I: RowInputs + ?Sized> ExecuteFused<'a, 'ctx, F, I> { + pub(crate) fn new( + function: &'a F, + inputs: &'a I, + probe: BatchProbe, + dtypes: &'a [DType], + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + inputs, + probe, + dtypes, + output_dtype: None, + ctx, + function, + } + } + + /// Collect the inputs and attach `plan`, then finish the batch with `kernel` when no retry can + /// be necessary. + fn execute( + self, + storage_dtype: DType, + policy: RowPolicy, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, + ) -> VortexResult { + let plan = BatchPlan::new(storage_dtype, self.output_dtype, policy)?; + let batch = BatchInputs::collect(RowFn::id(self.function), self.probe, self.inputs); + + let plan = { + let planned = batch.with_plan(self.dtypes, plan); + match planned.try_execute_without_retry(kernel, self.ctx)? { + Some(values) => return Ok(Fused::Done(values)), + None => planned.into_plan(), + } + }; + + Ok(Fused::Partial(batch, plan)) + } +} + +impl private::Sealed for ExecuteFused<'_, '_, F, I> {} + +impl RowVisitor for ExecuteFused<'_, '_, F, I> { + type VisitResult = Fused; + + fn with_output_dtype(mut self, dtype: DType) -> Self { + self.output_dtype = Some(dtype); + self + } + + fn visit_prepared( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + let storage_dtype = validate_owned_visit::(self.dtypes)?; + + self.execute( + storage_dtype, + RowPolicy::for_owned_output::(), + |args, _validity, ctx| { + execute_owned_infallible::(&args, ctx, prepare, apply) + }, + ) + } + + fn visit_bool( + mut self, + apply: impl Fn(Args::Elems<'_>) -> bool, + ) -> VortexResult + where + Args: IndexedElementTuple, + { + const { assert_owned_visit_contract::() }; + + // A direct batch packs straight into the finished column. Its storage dtype is the + // non-nullable Boolean dtype, so only the argument dtypes need validation. + if self.output_dtype.is_none() { + Args::validate(self.dtypes)?; + + let dense = RowPolicy::for_owned_output::() == RowPolicy::Dense; + let const_mask = self.probe.const_mask(); + match self + .probe + .into_direct_validity(self.inputs.num_inputs(), dense) + { + Ok(validity) => { + // The collector builds a canonical Boolean array of exactly the batch length + // with this validity attached, so the output needs no plan, label, or + // validation. + return execute_rows_infallible_bool::( + self.inputs, + const_mask, + self.ctx, + validity, + apply, + ) + .map(Fused::Done); + } + Err(probe) => self.probe = probe, + } + } + + let storage_dtype = validate_owned_visit::(self.dtypes)?; + + self.execute( + storage_dtype, + RowPolicy::for_owned_output::(), + |args, validity, ctx| { + execute_owned_infallible_bool::( + &args, + ctx, + validity.take().unwrap_or(Validity::NonNullable), + apply, + ) + }, + ) + } + + fn visit_prepared_into( + self, + params: Sink::Params, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + const { assert_sink_visit_contract::() }; + let storage_dtype = validate_sink_visit::(self.dtypes, ¶ms)?; + + self.execute( + storage_dtype, + RowPolicy::for_sink::(), + |args, _validity, ctx| { + execute_sink::( + &args, ¶ms, ctx, prepare, apply, + ) + }, + ) + } + + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: FailureEvidence, + { + const { assert_deferred_visit_contract::() }; + let storage_dtype = validate_owned_visit::(self.dtypes)?; + + self.execute( + storage_dtype, + RowPolicy::for_deferred_output::(), + |args, _validity, ctx| { + execute_owned::( + &args, + ctx, + prepare, + apply, + finish_failure, + ) + }, + ) + } + + fn visit_prepared_deferred_bool( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (bool, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Fail: FailureEvidence, + { + const { assert_deferred_visit_contract::() }; + let storage_dtype = validate_owned_visit::(self.dtypes)?; + + self.execute( + storage_dtype, + RowPolicy::for_deferred_output::(), + |args, validity, ctx| { + execute_owned_bool::( + &args, + ctx, + validity.take().unwrap_or(Validity::NonNullable), + prepare, + apply, + finish_failure, + ) + }, + ) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs index dc8d0209374..f7a3188b45c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -13,6 +13,10 @@ pub(super) use execute::ExecuteFilteredRows; pub(super) use execute::ExecuteRows; pub(super) use execute::ExecuteValidRows; +mod fused; +pub(super) use fused::ExecuteFused; +pub(super) use fused::Fused; + mod retry; pub(super) use retry::ExecuteDenseWithRetry; diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 20906cb031d..e63b275b9b2 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -15,7 +15,6 @@ use vortex_mask::MaskValuesRef; use vortex_session::VortexSession; use super::batch::ArrayInputs; -use super::batch::BatchInputs; use super::batch::BatchProbe; use super::batch::BorrowedRowFnArgs; use super::batch::RowInputs; @@ -24,8 +23,10 @@ use super::row_fn::RowFn; use super::visitor::BatchPlanner; use super::visitor::ExecuteDenseWithRetry; use super::visitor::ExecuteFilteredRows; +use super::visitor::ExecuteFused; use super::visitor::ExecuteRows; use super::visitor::ExecuteValidRows; +use super::visitor::Fused; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; @@ -153,10 +154,19 @@ pub(crate) fn execute_row_inputs( } let (probe, arg_dtypes) = BatchProbe::new(|| RowFn::id(function), inputs)?; - let plan = function.dispatch(options, &arg_dtypes, BatchPlanner::::new(&arg_dtypes))?; - let batch = BatchInputs::collect(RowFn::id(function), probe, inputs); - batch.with_plan(&arg_dtypes, plan).execute( + // Plan and execute in one dispatch. Only a partially valid batch that needs dense retry or + // valid-row execution dispatches again. + let (batch, plan) = match function.dispatch( + options, + &arg_dtypes, + ExecuteFused::::new(function, inputs, probe, &arg_dtypes, ctx), + )? { + Fused::Done(values) => return Ok(values), + Fused::Partial(batch, plan) => (batch, plan), + }; + + batch.with_plan(&arg_dtypes, plan).execute_partially_valid( |args, validity, ctx| execute_row_kernel(function, options, args, validity, ctx), |args, ctx| execute_dense_attempt(function, options, args, ctx), |args, valid, ctx| try_execute_valid_rows(function, options, args, valid, ctx), @@ -363,13 +373,15 @@ mod tests { _args: &[DType], visitor: V, ) -> VortexResult { + // The first dispatch plans dense execution with a retry, so a partially valid batch + // dispatches again for its dense attempt. if self.dispatches.fetch_add(1, Ordering::Relaxed) == 0 { - visitor.visit::<(i64,), i64>(|(value,)| value) + visitor.visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())) } else { match self.change { - DispatchChange::Policy => visitor - .visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())), - DispatchChange::Element => visitor.visit::<(u64,), u64>(|(value,)| value), + DispatchChange::Policy => visitor.visit::<(i64,), i64>(|(value,)| value), + DispatchChange::Element => visitor + .visit_deferred::<(u64,), u64, bool>(|(value,)| (value, false), |_| Ok(())), } } } @@ -407,13 +419,32 @@ mod tests { Ok(()) } + #[test] + fn test_execute_dispatches_once_without_retry() -> VortexResult<()> { + let dispatches = Arc::new(AtomicUsize::new(0)); + let function = ChangingDispatchRowFn { + dispatches: Arc::clone(&dispatches), + change: DispatchChange::Policy, + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; + + assert_eq!(dispatches.load(Ordering::Relaxed), 1); + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) + } + #[test] fn test_execute_rejects_dispatch_that_changes_after_planning() -> VortexResult<()> { let function = ChangingDispatchRowFn { dispatches: Arc::new(AtomicUsize::new(0)), change: DispatchChange::Policy, }; - let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let validity = Validity::from_iter([true, false]); + let input = PrimitiveArray::new(vec![1_i64, 2], validity).into_array(); let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = array_session().create_execution_ctx(); @@ -428,7 +459,7 @@ mod tests { "unexpected error: {error}", ); assert!( - message.contains("planned Dense, got DenseWithRetry"), + message.contains("planned DenseWithRetry, got Dense"), "unexpected error: {error}", ); Ok(()) @@ -440,7 +471,8 @@ mod tests { dispatches: Arc::new(AtomicUsize::new(0)), change: DispatchChange::Element, }; - let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let validity = Validity::from_iter([true, false]); + let input = PrimitiveArray::new(vec![1_i64, 2], validity).into_array(); let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = array_session().create_execution_ctx();