From 386ce38439d1863a5cc157ad0b488ec8e45e376f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 27 Oct 2026 15:34:53 -0400 Subject: [PATCH 1/3] 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 buffer without executing the array. Signed-off-by: Connor Tsui --- .../unstable/row/types/element/input.rs | 15 +++++++++- .../unstable/row/types/element/primitive.rs | 17 +++++++++++ .../unstable/row/types/element/tuple/tests.rs | 28 +++++++++++++++++++ 3 files changed, 59 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..9cc2a158344 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,25 @@ 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::()); + } + + let executed = array.clone().execute::(ctx)?; + + Ok(executed.into_buffer::()) + } + 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 16f639c4254..a5213af0f4e 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 @@ -5,6 +5,7 @@ use std::sync::LazyLock; 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; @@ -17,6 +18,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::arrays::BoolArray; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; @@ -108,6 +110,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 = 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 = 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 7932560597e2da9111c59aa197a4f88e3f70f9eb Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 27 Oct 2026 15:34:56 -0400 Subject: [PATCH 2/3] Separate RowFn batch inspection from execution planning Inspect the borrowed inputs once for row counts, batch-constant classification, null literals, and strict validity. Batch execution reads that probe instead of scanning the inputs again, and borrows the columns and dtypes instead of collecting owned copies. Execution still plans with one dispatch and executes with a second one. Signed-off-by: Connor Tsui --- .../unstable/row/batch/execute/constant.rs | 6 +- .../unstable/row/batch/execute/dense.rs | 28 ++-- .../unstable/row/batch/execute/filtered.rs | 7 +- .../unstable/row/batch/execute/mod.rs | 20 +-- .../unstable/row/batch/execute/output.rs | 14 +- .../unstable/row/batch/execute/valid_only.rs | 21 ++- .../scalar_fn/unstable/row/batch/inputs.rs | 71 +++++++++ .../src/scalar_fn/unstable/row/batch/mod.rs | 56 +++++-- .../scalar_fn/unstable/row/batch/planning.rs | 138 ++++++++++++++---- .../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 | 62 ++++++-- vortex-session/src/registry.rs | 2 + 15 files changed, 357 insertions(+), 103 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..62411f48b1a 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.probe.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 50835c0ded5..ac125b81709 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 @@ -18,14 +18,17 @@ use crate::canonical::AnyCanonical; 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.probe.row_count), + ctx, + )?; self.finalize_dense_output(values, ctx) } @@ -53,8 +56,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.probe.row_count), + ctx, + )?; match attempt { DenseAttempt::Values(values) => self.finalize_dense_output(values, ctx), @@ -79,7 +84,12 @@ 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 + .probe + .validity + .execute_mask(self.batch.probe.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 @@ -105,11 +115,11 @@ 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.probe.row_count, ctx)?; - match self.validity.clone() { + match self.batch.probe.validity.clone() { Validity::NonNullable | Validity::AllValid => { - self.finalize_output(values, self.row_count) + self.finalize_output(values, self.batch.probe.row_count) } Validity::Array(valid) => { // `values.mask(valid)` builds and optimizes a ScalarFnArray to dispatch through @@ -126,7 +136,7 @@ impl RowFnExecutionArgs { .into_array(), None => values.mask(valid)?, }; - self.finalize_output(values, self.row_count) + self.finalize_output(values, self.batch.probe.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..f167766be1e 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.probe.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.probe.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..c1e3c58664e 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,12 +53,8 @@ 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.probe.validity, Validity::AllInvalid) + || self.batch.probe.has_null_const { return Ok(self.all_null()); } @@ -68,16 +62,16 @@ impl RowFnExecutionArgs { // 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.probe.row_count > 0 + && self.batch.probe.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.probe.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..74c75069be5 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.probe.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, @@ -68,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 62c5afb7458..b593249c4e3 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,24 @@ impl RowFnExecutionArgs { ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - let validity = self.validity.clone().execute_mask(self.row_count, ctx)?; + let validity = self + .batch + .probe + .validity + .clone() + .execute_mask(self.batch.probe.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.probe.row_count), + ctx, + )?; + let values = + self.validate_kernel_output(values, self.batch.probe.row_count, ctx)?; - return self.finalize_output(values, self.row_count); + return self.finalize_output(values, self.batch.probe.row_count); } Mask::AllFalse(_) => return Ok(self.all_null()), Mask::Values(valid_rows) => valid_rows, @@ -62,7 +71,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.probe.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..1ea379e8ff9 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs @@ -0,0 +1,71 @@ +// 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; +use crate::dtype::DType; + +/// The same-length input columns of one row-function execution, borrowed by index. +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; + + /// Return the input dtypes in order, if this already holds them contiguously. + /// + /// Batch execution otherwise clones them into a collected list. + fn dtypes(&self) -> Option<&[DType]> { + None + } + + /// Return the input columns in order, if this already holds them contiguously. + /// + /// The slice **must** hold exactly [`num_inputs`](Self::num_inputs) columns. Batch execution + /// otherwise clones the columns. + fn as_slice(&self) -> Option<&[ArrayRef]> { + None + } +} + +/// [`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] + } + + fn as_slice(&self) -> Option<&[ArrayRef]> { + Some(self.arrays) + } +} 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..f37cf0b32f4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -9,11 +9,12 @@ //! 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`] keeps the columns with their probe. [`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 std::borrow::Cow; use crate::ArrayRef; use crate::dtype::DType; @@ -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,24 +36,45 @@ 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 { - /// The function being executed, named in the errors this raises. - id: ScalarFnId, - +/// What batch execution derives from its borrowed inputs before a dispatch plans them. +/// +/// A probe holds no input column. [`BatchInputs`] pairs it with the columns. +pub(crate) struct BatchProbe { /// The number of rows in the original execution scope. row_count: usize, - /// 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 of one row-function execution, with their [`BatchProbe`]. +pub(crate) struct BatchInputs<'a> { + /// The function being executed, named in the errors this raises. + id: ScalarFnId, + + /// The input columns, borrowed when the caller holds them as one slice and cloned otherwise. + inputs: Cow<'a, [ArrayRef]>, + + /// What batch execution derived from `inputs` before planning. + probe: BatchProbe, +} + +/// The same-length input columns and metadata for one row-function execution. +pub(crate) struct RowFnExecutionArgs<'a> { + /// The input columns and their conjoined validity. + batch: &'a BatchInputs<'a>, + + /// The input dtypes, shared by 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..9f23ca1fe60 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,147 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use smallvec::SmallVec; +//! 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 std::borrow::Cow; + 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. /// - /// 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, + /// 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 { - let row_count = args.row_count(); - let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) - .map(|index| args.get(index)) - .collect::>()?; + 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 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); - let mut validity = Validity::NonNullable; - for input in &inputs { - validity = validity.and(input.validity()?)?; + if is_batch_const(input) { + const_mask |= 1 << index; + has_null_const |= input + .as_opt::() + .is_some_and(|constant| constant.scalar().is_null()); + } } + // A null literal makes every row invalid, so skip conjoining the other validities: each + // conjunction of two array-backed validities builds and optimizes an expression. A + // non-nullable input has no invalid row, and `NonNullable` is the identity of `and`. + let validity = if has_null_const { + Validity::AllInvalid + } else { + (0..num_inputs) + .map(|index| inputs.input(index)) + .filter(|input| input.dtype().is_nullable()) + .try_fold(Validity::NonNullable, |validity, input| { + validity.and(input.validity()?) + })? + }; + Ok(Self { - id, row_count, - inputs, - arg_dtypes, validity, + const_mask, + has_null_const, + }) + } +} + +impl<'a> BatchInputs<'a> { + /// Keep the probed inputs of the function `id`. + /// + /// The columns are borrowed when [`RowInputs::as_slice`] provides them, so a caller that + /// already collected them does not pay for a second copy. + pub(crate) fn new( + id: ScalarFnId, + probe: BatchProbe, + inputs: &'a I, + ) -> Self { + let columns = match inputs.as_slice() { + Some(columns) => Cow::Borrowed(columns), + None => Cow::Owned( + (0..inputs.num_inputs()) + .map(|index| inputs.input(index).clone()) + .collect(), + ), + }; + + Self { + id, + inputs: columns, + probe, + } + } + + /// Whether every input is batch-constant. + pub(super) fn all_batch_const(&self) -> bool { + self.probe.const_mask.count_ones() as usize == self.inputs.len() + } + + /// Attach the plan a dispatch selected for these inputs. + pub(crate) fn with_plan<'b>( + &'b self, + arg_dtypes: &'b [DType], + plan: BatchPlan, + ) -> RowFnExecutionArgs<'b> { + 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..2ce2e9e3e69 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(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/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index bc4d793404f..55e686d1b09 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; @@ -117,14 +121,50 @@ pub fn execute_rows( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult { - ensure_arity(function, args.num_inputs())?; + let arrays = (0..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>>()?; - if args.num_inputs() == 0 { - return execute_nullary_rows(function, options, args.row_count(), ctx); + 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 = BatchProbe::new(|| RowFn::id(function), inputs)?; + let collected: SmallVec<[DType; 4]>; + let arg_dtypes = match inputs.dtypes() { + Some(dtypes) => dtypes, + None => { + collected = (0..inputs.num_inputs()) + .map(|index| inputs.input(index).dtype().clone()) + .collect(); + &collected + } + }; + + let plan = function.dispatch(options, arg_dtypes, BatchPlanner::::new(arg_dtypes))?; + let batch = BatchInputs::new(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 +258,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 0ddfefa190804abccc6e06912e78e4736cd35e39 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 27 Oct 2026 15:34:59 -0400 Subject: [PATCH 3/3] Cover RowFn input short circuits with regression tests Pin the behavior that a later single-dispatch path must keep: valid-only decoding skips invalid payloads, a null literal skips earlier decoding, lazy division, and device transfers, and validity conjunction stays lazy for device-backed validity. Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 227 ++++++++++++++++++ 1 file changed, 227 insertions(+) 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 22309a0bd9c..0ea082966d7 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -1,16 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; use std::mem::MaybeUninit; +use std::ops::Range; use std::sync::Arc; use std::sync::LazyLock; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use futures::future::BoxFuture; use rstest::rstest; +use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -28,11 +33,13 @@ use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; use crate::arrays::FixedSizeListArray; +use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; use crate::assert_arrays_eq; use crate::buffer::BufferHandle; +use crate::buffer::DeviceBuffer; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; @@ -44,6 +51,8 @@ use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::operators::Operator; use crate::scalar_fn::unstable::row::FixedSizeListSink; use crate::scalar_fn::unstable::row::InitializedRow; use crate::scalar_fn::unstable::row::InputElement; @@ -61,6 +70,70 @@ use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); +#[derive(Clone)] +struct FilteredBoolean; + +impl RowFn for FilteredBoolean { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.filtered_boolean"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_bool::<(FilterOnlyI64, i64), false>(|(lhs, rhs)| lhs > rhs) + } +} + +#[derive(Debug, PartialEq, Eq, Hash)] +struct UnreadableDeviceBuffer { + len: usize, + alignment: Alignment, +} + +impl DeviceBuffer for UnreadableDeviceBuffer { + fn as_any(&self) -> &dyn Any { + self + } + + fn len(&self) -> usize { + self.len + } + + fn alignment(&self) -> Alignment { + self.alignment + } + + fn copy_to_host_sync(&self, _alignment: Alignment) -> VortexResult { + vortex_bail!("execution must not read the device buffer") + } + + fn copy_to_host( + &self, + _alignment: Alignment, + ) -> VortexResult>> { + vortex_bail!("execution must not read the device buffer") + } + + fn slice(&self, _range: Range) -> Arc { + panic!("execution must not slice the device buffer") + } + + fn aligned(self: Arc, alignment: Alignment) -> VortexResult> { + assert_eq!(alignment, self.alignment()); + Ok(self) + } +} + #[derive(Clone, Default)] struct DeferredAdd { /// The number of preparations across the dense attempt and any valid-row retry. @@ -807,6 +880,64 @@ fn test_bool_output_handles_partial_constants() -> VortexResult<()> { assert_bool_output_handles_partial_constants::() } +#[test] +fn test_bool_output_handles_null_literal_constant() -> VortexResult<()> { + let varying = PrimitiveArray::from_iter(0_i64..65).into_array(); + let null = Scalar::null(DType::Primitive(i64::PTYPE, Nullability::Nullable)); + let constant = ConstantArray::new(null, 65).into_array(); + let args = VecExecutionArgs::new(vec![varying, constant], 65); + let mut ctx = SESSION.create_execution_ctx(); + + let actual = execute_rows(&PackedGreaterThan::, &EmptyOptions, &args, &mut ctx)?; + let expected = BoolArray::from_iter(std::iter::repeat_n(None::, 65)).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_bool_output_handles_masked_constant() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let varying = PrimitiveArray::from_iter(0_i64..65).into_array(); + let valid: Vec = (0..65).map(|index| index % 3 != 0).collect(); + let constant = MaskedArray::try_new( + ConstantArray::new(32_i64, 65).into_array(), + Validity::from_iter(valid.iter().copied()), + &mut ctx, + )? + .into_array(); + let args = VecExecutionArgs::new(vec![varying, constant], 65); + + let actual = execute_rows(&PackedGreaterThan::, &EmptyOptions, &args, &mut ctx)?; + let expected = BoolArray::from_iter( + (0_i64..65) + .zip(valid) + .map(|(value, valid)| valid.then_some(value > 32)), + ) + .into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_bool_output_rejects_wrong_row_count() { + let short = PrimitiveArray::from_iter(0_i64..64).into_array(); + let long = PrimitiveArray::from_iter(0_i64..65).into_array(); + let args = VecExecutionArgs::new(vec![long, short], 65); + let mut ctx = SESSION.create_execution_ctx(); + + let error = execute_rows(&PackedGreaterThan::, &EmptyOptions, &args, &mut ctx) + .expect_err("a short input must fail"); + + assert!( + error + .to_string() + .contains("the test.packed_greater_than input 1 has the wrong row count"), + "{error}", + ); +} + #[test] fn test_bool_output_handles_all_constant_input() -> VortexResult<()> { let input = ConstantArray::new(7_i64, 65).into_array(); @@ -1537,3 +1668,99 @@ fn constant_output_retains_execution_allocator_payload() -> VortexResult<()> { tracker.assert_owns(value.inner().as_slice()); Ok(()) } + +#[test] +fn test_valid_only_bool_skips_invalid_decode_payloads() -> VortexResult<()> { + let lhs = PrimitiveArray::new(vec![i64::MIN, 7], Validity::from_iter([false, true])); + let rhs = ConstantArray::new(0_i64, 2); + let args = VecExecutionArgs::new(vec![lhs.into_array(), rhs.into_array()], 2); + let mut ctx = SESSION.create_execution_ctx(); + + let actual = execute_rows(&FilteredBoolean, &EmptyOptions, &args, &mut ctx)?; + assert_arrays_eq!(actual, BoolArray::from_iter([None, Some(true)]), &mut ctx); + Ok(()) +} + +#[test] +fn test_null_literal_skips_earlier_decode() -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter([i64::MIN, 7]); + let rhs = ConstantArray::new( + Scalar::null(DType::Primitive(i64::PTYPE, Nullability::Nullable)), + 2, + ); + let args = VecExecutionArgs::new(vec![lhs.into_array(), rhs.into_array()], 2); + let mut ctx = SESSION.create_execution_ctx(); + + let actual = execute_rows(&FilteredBoolean, &EmptyOptions, &args, &mut ctx)?; + assert_arrays_eq!(actual, BoolArray::from_iter([None::, None]), &mut ctx); + Ok(()) +} + +#[test] +fn test_null_literal_skips_lazy_input_execution() -> VortexResult<()> { + let lhs = Binary::try_new( + PrimitiveArray::from_iter([1_i64, 2]).into_array(), + ConstantArray::new(0_i64, 2).into_array(), + Operator::Div, + )? + .into_array(); + let rhs = ConstantArray::new( + Scalar::null(DType::Primitive(i64::PTYPE, Nullability::Nullable)), + 2, + ); + let args = VecExecutionArgs::new(vec![lhs, rhs.into_array()], 2); + let mut ctx = SESSION.create_execution_ctx(); + + let actual = execute_rows(&PackedGreaterThan::, &EmptyOptions, &args, &mut ctx)?; + assert_arrays_eq!(actual, BoolArray::from_iter([None::, None]), &mut ctx); + Ok(()) +} + +#[test] +fn test_device_validity_stays_lazy() -> VortexResult<()> { + let validity = BoolArray::try_new_from_handle( + BufferHandle::new_device(Arc::new(UnreadableDeviceBuffer { + len: 1, + alignment: Alignment::of::(), + })), + 0, + 2, + Validity::NonNullable, + )? + .into_array(); + let lhs = PrimitiveArray::new(vec![1_i64, 2], Validity::Array(validity.clone())); + let rhs = PrimitiveArray::new(vec![3_i64, 4], Validity::Array(validity)); + let args = VecExecutionArgs::new(vec![lhs.into_array(), rhs.into_array()], 2); + let mut ctx = SESSION.create_execution_ctx(); + + let actual = execute_rows(&DeferredAdd::default(), &EmptyOptions, &args, &mut ctx)?; + assert_eq!(actual.len(), 2); + assert_eq!( + actual.dtype(), + &DType::Primitive(i64::PTYPE, Nullability::Nullable) + ); + assert!(matches!(actual.validity()?, Validity::Array(_))); + Ok(()) +} + +#[test] +fn test_null_literal_skips_device_input_transfer() -> VortexResult<()> { + let input = PrimitiveArray::from_buffer_handle( + BufferHandle::new_device(Arc::new(UnreadableDeviceBuffer { + len: 2 * size_of::(), + alignment: Alignment::of::(), + })), + i64::PTYPE, + Validity::NonNullable, + ); + let null = ConstantArray::new( + Scalar::null(DType::Primitive(i64::PTYPE, Nullability::Nullable)), + 2, + ); + let args = VecExecutionArgs::new(vec![input.into_array(), null.into_array()], 2); + let mut ctx = SESSION.create_execution_ctx(); + + let actual = execute_rows(&PackedGreaterThan::, &EmptyOptions, &args, &mut ctx)?; + assert_arrays_eq!(actual, BoolArray::from_iter([None::, None]), &mut ctx); + Ok(()) +}