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/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(()) +} 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/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/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/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/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]; 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"