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..ec980eaa0ba 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,11 @@ 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 +110,15 @@ 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.batch.row_count) } - Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.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"