From 35335b1aaf26cb7adf1dca040f67e60e2e4e221b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 27 Oct 2026 15:35:24 -0400 Subject: [PATCH 1/2] Plan and execute common RowFn batches in one dispatch Finish all-null, all-constant, null-free, and dense batches in the planning visit, and hand every other batch back as a PartiallyValidBatch for dense retry or valid-row execution. Decode a direct Boolean batch straight from the borrowed inputs, construct Boolean output with the batch validity attached, and mask dense output that does not carry it. Signed-off-by: Connor Tsui --- .../unstable/row/batch/execute/constant.rs | 9 +- .../unstable/row/batch/execute/dense.rs | 25 +- .../unstable/row/batch/execute/mod.rs | 101 ++++--- .../unstable/row/batch/execute/valid_only.rs | 2 +- .../scalar_fn/unstable/row/batch/inputs.rs | 7 +- .../src/scalar_fn/unstable/row/batch/mod.rs | 29 +- .../scalar_fn/unstable/row/batch/planning.rs | 44 ++- .../src/scalar_fn/unstable/row/batch/tests.rs | 8 +- .../src/scalar_fn/unstable/row/execute/mod.rs | 1 + .../unstable/row/execute/packed_bool.rs | 86 +++++- .../scalar_fn/unstable/row/execute/tests.rs | 77 ++++- .../unstable/row/types/element/mod.rs | 1 + .../row/types/element/tuple/element_tuple.rs | 22 ++ .../row/types/element/tuple/indexed.rs | 94 ++++++ .../unstable/row/types/element/tuple/mod.rs | 1 + .../src/scalar_fn/unstable/row/types/mod.rs | 1 + .../scalar_fn/unstable/row/visitor/execute.rs | 8 +- .../scalar_fn/unstable/row/visitor/fused.rs | 282 ++++++++++++++++++ .../src/scalar_fn/unstable/row/visitor/mod.rs | 4 + .../src/scalar_fn/unstable/row/vtable.rs | 98 +++++- 20 files changed, 819 insertions(+), 81 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/visitor/fused.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs index 62411f48b1a..abcfbc670bf 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs @@ -9,16 +9,21 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; +use crate::validity::Validity; impl RowFnExecutionArgs<'_> { /// Execute all-constant inputs by evaluating one row and broadcasting the validated result. pub(super) fn execute_all_constant( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { let result = self.validate_kernel_output( - kernel(self.execution_args(&self.batch.inputs, 1), ctx)?, + kernel(self.execution_args(&self.batch.inputs, 1), &mut None, ctx)?, 1, ctx, )?; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs index ac125b81709..4e151be81ee 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 @@ -22,14 +22,31 @@ impl RowFnExecutionArgs<'_> { /// Run every stored payload, then attach the input validity without materializing its mask. pub(super) fn execute_dense( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { + // Offer a partial validity to the kernel. A collector that takes it builds the output with + // that validity attached, so the output needs neither a second construction nor an + // all-valid check. + let offered = matches!(self.batch.probe.validity, Validity::Array(_)); + let mut validity = offered.then(|| self.batch.probe.validity.clone()); let values = kernel( self.execution_args(&self.batch.inputs, self.batch.probe.row_count), + &mut validity, ctx, )?; + // A taken validity proves nothing alone. Row kernels build non-nullable storage, so only a + // nullable output carries the validity. Mask every other output, including one from a + // collector that took the validity without attaching it. + if offered && validity.is_none() && values.dtype().is_nullable() { + return self.finalize_output(values, self.batch.probe.row_count); + } + self.finalize_dense_output(values, ctx) } @@ -91,9 +108,9 @@ impl RowFnExecutionArgs<'_> { .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 - // and all-null, so preserve the all-valid behavior. + // in `BatchInputs::try_execute_without_retry` could not prove that. The deferred error + // therefore came from an observable row and remains terminal. An empty mask is both + // all-valid and all-null, so preserve the all-valid behavior. Mask::AllTrue(_) | Mask::AllFalse(0) => return Err(deferred_error), Mask::AllFalse(_) => return Ok(self.all_null()), Mask::Values(valid_rows) => valid_rows, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs index c1e3c58664e..eebe73288df 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -3,17 +3,21 @@ //! Selects a batch execution strategy. //! -//! [`RowFnExecutionArgs::execute`] handles universal fast paths, then delegates to dense or -//! valid-only execution. +//! [`BatchInputs::try_execute_without_retry`] handles universal fast paths and dense execution, so +//! the visit that planned a batch can finish it without dispatching again. It hands every other +//! batch back as a [`PartiallyValidBatch`], which delegates to dense retry or valid-only execution. use vortex_error::VortexResult; use vortex_mask::MaskValuesRef; -use super::RowFnExecutionArgs; +use super::BatchInputs; +use super::BatchPlan; +use super::PartiallyValidBatch; use super::RowPolicy; use super::args::BorrowedRowFnArgs; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::dtype::DType; use crate::scalar_fn::unstable::row::execute::DenseAttempt; use crate::validity::Validity; @@ -25,16 +29,64 @@ mod valid_only; mod output; pub(crate) use output::finalize_kernel_output; -impl RowFnExecutionArgs<'_> { - /// Apply constant folding and null handling around `kernel`. +impl<'a> BatchInputs<'a> { + /// Execute this batch under `plan` when `kernel` alone can finish it. Otherwise return it as a + /// [`PartiallyValidBatch`] that needs a dense retry or valid-row execution. + /// + /// `kernel` reports a rejected deferred failure as a terminal error. That is correct for each + /// batch this method executes: no null row can be the cause of the failure. + pub(crate) fn try_execute_without_retry( + self, + arg_dtypes: &[DType], + plan: BatchPlan, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult>> { + let args = self.with_plan(arg_dtypes, plan); + + // Strictness: an all-null batch has no observable row work. Keep the literal-constant + // check explicit alongside the conjoined validity invariant. + if matches!(self.probe.validity, Validity::AllInvalid) || self.probe.has_null_const { + return Ok(Ok(args.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.probe.row_count > 0 + && self.probe.validity.definitely_no_nulls() + && self.all_batch_const() + { + return args.execute_all_constant(kernel, ctx).map(Ok); + } + + // 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.probe.validity.definitely_no_nulls() || args.plan.policy() == RowPolicy::Dense { + return args.execute_dense(kernel, ctx).map(Ok); + } + + let plan = args.into_plan(); + + Ok(Err(PartiallyValidBatch { batch: self, plan })) + } +} + +impl PartiallyValidBatch<'_> { + /// Execute the partially valid batch with the dense retry or valid-row strategy of its plan. /// /// For a partially valid batch, `try_valid_rows` can avoid filtering. `Ok(None)` filters the /// inputs to the valid rows and `execute_filtered_rows` runs the kernel over them while /// writing into the original row domain. Every result is checked against the planned shape /// and dtype. pub(crate) fn execute( - &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + self, + arg_dtypes: &[DType], + kernel: impl FnOnce(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, execute_dense_attempt: impl FnOnce( BorrowedRowFnArgs<'_>, &mut ExecutionCtx, @@ -51,40 +103,21 @@ impl RowFnExecutionArgs<'_> { ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - // Strictness: an all-null batch has no observable row work. Keep the literal-constant - // check explicit alongside the conjoined validity invariant. - if matches!(self.batch.probe.validity, Validity::AllInvalid) - || self.batch.probe.has_null_const - { - return Ok(self.all_null()); - } + let Self { batch, plan } = self; + let args = batch.with_plan(arg_dtypes, plan); - // All inputs are constant, and their conjoined validity proves that every row is non-null. - // The constant check sees through extension and masked wrappers, just like argument - // decoding. - if self.batch.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.batch.probe.validity.definitely_no_nulls() { - return self.execute_dense(kernel, ctx); - } - - match self.plan.policy() { - RowPolicy::Dense => self.execute_dense(kernel, ctx), - RowPolicy::DenseWithRetry => self.execute_dense_with_retry( + match args.plan.policy() { + RowPolicy::Dense => { + unreachable!("BatchInputs::try_execute_without_retry executes every dense batch") + } + RowPolicy::DenseWithRetry => args.execute_dense_with_retry( execute_dense_attempt, try_valid_rows, execute_filtered_rows, ctx, ), RowPolicy::ValidOnly => { - self.execute_valid_only(kernel, try_valid_rows, execute_filtered_rows, ctx) + args.execute_valid_only(kernel, try_valid_rows, execute_filtered_rows, ctx) } } } 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 b593249c4e3..0ca4f61ff2d 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 @@ -16,7 +16,7 @@ impl RowFnExecutionArgs<'_> { /// Resolve validity and try direct valid-row execution before filtered execution. pub(super) fn execute_valid_only( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, try_valid_rows: impl FnOnce( BorrowedRowFnArgs<'_>, MaskValuesRef, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs b/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs index 1ea379e8ff9..befee73b30e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs @@ -11,7 +11,10 @@ use crate::ArrayRef; use crate::dtype::DType; /// The same-length input columns of one row-function execution, borrowed by index. -pub(crate) trait RowInputs { +/// +/// The trait is `pub` only so that the sealed tuple decoding can bound on it. The private module +/// keeps it out of the public API. +pub trait RowInputs { /// Return the number of input columns. fn num_inputs(&self) -> usize; @@ -31,7 +34,7 @@ pub(crate) trait RowInputs { /// 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. + /// otherwise clones the columns when a visit cannot finish the batch directly. fn as_slice(&self) -> Option<&[ArrayRef]> { None } 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 f37cf0b32f4..ba1423e2cbf 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -10,14 +10,17 @@ //! output. //! //! [`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. +//! [`BatchInputs`] keeps them when a visit cannot finish the batch directly. [`BatchPlan`] carries +//! the nullable execution strategy selected by a concrete dispatch. [`RowFnExecutionArgs`] applies +//! that strategy, and [`BorrowedRowFnArgs`] pairs each kernel invocation with its planning +//! metadata. A [`PartiallyValidBatch`] carries a batch that needs a dense retry or valid-row +//! execution after the planning visit. use std::borrow::Cow; use crate::ArrayRef; use crate::dtype::DType; +use crate::dtype::Nullability; use crate::scalar_fn::ScalarFnId; use crate::validity::Validity; @@ -38,7 +41,8 @@ pub(super) use super::visitor::RowPolicy; /// What batch execution derives from its borrowed inputs before a dispatch plans them. /// -/// A probe holds no input column. [`BatchInputs`] pairs it with the columns. +/// A probe holds no input column. A visit that can finish the batch directly decodes the borrowed +/// inputs. Any other visit keeps them in [`BatchInputs`]. pub(crate) struct BatchProbe { /// The number of rows in the original execution scope. row_count: usize, @@ -53,6 +57,9 @@ pub(crate) struct BatchProbe { /// Whether any input is a null literal constant. has_null_const: bool, + + /// The strict result nullability: nullable when any input dtype is nullable. + nullability: Nullability, } /// The same-length input columns of one row-function execution, with their [`BatchProbe`]. @@ -67,6 +74,20 @@ pub(crate) struct BatchInputs<'a> { probe: BatchProbe, } +/// A partially valid batch that [`BatchInputs::try_execute_without_retry`] could not finish, with +/// the plan that a dispatch selected for it. +/// +/// Only that method creates this type. [`PartiallyValidBatch::execute`] therefore never receives +/// an all-null, all-constant, or provably null-free batch, or a batch planned with +/// [`RowPolicy::Dense`]. +pub(crate) struct PartiallyValidBatch<'a> { + /// The input columns and their probe. + batch: BatchInputs<'a>, + + /// The plan that a dispatch selected for `batch`. + plan: BatchPlan, +} + /// The same-length input columns and metadata for one row-function execution. pub(crate) struct RowFnExecutionArgs<'a> { /// The input columns and their conjoined validity. 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 9f23ca1fe60..ce00816423c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs @@ -21,6 +21,7 @@ use super::args::BorrowedRowFnArgs; use crate::ArrayRef; use crate::arrays::Constant; use crate::dtype::DType; +use crate::dtype::Nullability; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::unstable::row::types::is_batch_const; use crate::validity::Validity; @@ -47,6 +48,7 @@ impl BatchProbe { let mut const_mask = 0u32; let mut has_null_const = false; + let mut nullability = Nullability::NonNullable; for index in 0..num_inputs { let input = inputs.input(index); vortex_ensure_eq!( @@ -62,6 +64,10 @@ impl BatchProbe { .as_opt::() .is_some_and(|constant| constant.scalar().is_null()); } + + if input.dtype().is_nullable() { + nullability = Nullability::Nullable; + } } // A null literal makes every row invalid, so skip conjoining the other validities: each @@ -83,8 +89,39 @@ impl BatchProbe { validity, const_mask, has_null_const, + nullability, }) } + + /// Return the output validity of a _direct_ batch, or return the probe unchanged for any other + /// batch. + /// + /// A visit finishes a direct batch from its kernel output alone, once that output carries + /// this validity. A direct batch has no null literal, at least one input that is not + /// batch-constant, and at least one possibly valid row. It therefore needs no all-null or + /// all-constant handling. Set `dense` when the visit plans + /// [`RowPolicy::Dense`](super::RowPolicy::Dense). Without it, only a batch that provably has + /// no null row is direct, because the kernel visits every row. + pub(crate) fn into_direct_validity( + self, + num_inputs: usize, + dense: bool, + ) -> Result { + if self.has_null_const + || self.const_mask.count_ones() as usize == num_inputs + || matches!(self.validity, Validity::AllInvalid) + || !(dense || self.validity.definitely_no_nulls()) + { + return Err(self); + } + + Ok(self.validity.union_nullability(self.nullability)) + } + + /// Return the batch-constant classification of the inputs: bit `i` marks input `i`. + pub(crate) fn const_mask(&self) -> u32 { + self.const_mask + } } impl<'a> BatchInputs<'a> { @@ -119,7 +156,7 @@ impl<'a> BatchInputs<'a> { } /// Attach the plan a dispatch selected for these inputs. - pub(crate) fn with_plan<'b>( + pub(super) fn with_plan<'b>( &'b self, arg_dtypes: &'b [DType], plan: BatchPlan, @@ -136,6 +173,11 @@ impl<'a> BatchInputs<'a> { } impl RowFnExecutionArgs<'_> { + /// Return the plan, discarding the batch view. + pub(super) fn into_plan(self) -> BatchPlan { + self.plan + } + /// Pair an input view with this batch's planning metadata. pub(super) fn execution_args<'b>( &'b self, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 0ea082966d7..45c89a21433 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -1393,7 +1393,7 @@ impl RowFn for ChangingOutputDType { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["value"]; - const INFALLIBLE: bool = true; + const INFALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.changing_output_dtype"); @@ -1412,7 +1412,8 @@ impl RowFn for ChangingOutputDType { visitor }; - visitor.visit::<(i64,), i64>(|(value,)| value) + // A deferred visit plans a dense retry, so a partially valid batch dispatches again. + visitor.visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())) } } @@ -1594,7 +1595,8 @@ fn test_execution_rejects_a_changed_output_dtype() -> VortexResult<()> { let function = ChangingOutputDType { dispatches: Arc::new(AtomicUsize::new(0)), }; - let input = PrimitiveArray::from_iter(vec![1_i64, 2]).into_array(); + let validity = Validity::from_iter([true, false]); + let input = PrimitiveArray::new(vec![1_i64, 2], validity).into_array(); let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs index c9e17da1d09..3ca76d6360c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -18,6 +18,7 @@ mod packed_bool; pub(super) use packed_bool::execute_bool_dense_attempt; pub(super) use packed_bool::execute_owned_bool; pub(super) use packed_bool::execute_owned_infallible_bool; +pub(super) use packed_bool::execute_rows_infallible_bool; mod retry; pub(super) use retry::DenseAttempt; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs index 5fb3d0cadc1..a77d903004e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs @@ -25,37 +25,103 @@ use crate::arrays::BoolArray; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::unstable::row::FailureEvidence; use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::batch::RowInputs; +use crate::scalar_fn::unstable::row::types::decode_rows; use crate::scalar_fn::unstable::row::types::decoded_source; use crate::validity::Validity; -/// Decode every input column, then pack infallible Boolean outputs. +/// Decode every input column, then pack infallible Boolean outputs with `validity` attached. pub(crate) fn execute_owned_infallible_bool( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, + validity: Validity, apply: impl Fn(Args::Elems<'_>) -> bool, ) -> VortexResult where Args: IndexedElementTuple, { let columns = Args::decode(args, ctx)?; - let row_count = args.row_count(); - let Some(source) = decoded_source::(&columns, row_count) else { - vortex_bail!("a decoded row input does not address exactly {row_count} rows"); + pack_infallible_bool::(&columns, args.row_count(), ctx, validity, apply) +} + +/// Decode one borrowed array per argument, then pack infallible Boolean outputs with `validity` +/// attached. +/// +/// Bit `i` of `const_mask` is set exactly when input `i` is batch-constant. +/// +/// Decoding and packing inline into this one frame, and this frame stays out of the dispatch +/// that calls it. A short batch then pays for one call boundary, not one per step, and each arm +/// of the dispatch stays small. +#[inline(never)] +pub(crate) fn execute_rows_infallible_bool( + inputs: &I, + const_mask: u32, + ctx: &mut ExecutionCtx, + validity: Validity, + apply: impl Fn(Args::Elems<'_>) -> bool, +) -> VortexResult +where + Args: IndexedElementTuple, + I: RowInputs + ?Sized, +{ + let columns = decode_rows::(inputs, const_mask, ctx)?; + + pack_infallible_bool::(&columns, inputs.row_count(), ctx, validity, apply) +} + +#[expect(clippy::inline_always)] +#[inline(always)] +fn pack_infallible_bool( + columns: &Args::Columns, + row_count: usize, + ctx: &mut ExecutionCtx, + validity: Validity, + apply: impl Fn(Args::Elems<'_>) -> bool, +) -> VortexResult +where + Args: IndexedElementTuple, +{ + // Without batch constants, read plain views. Selecting each argument's constant or row + // source per row stays in short scalar loops, which LLVM does not unswitch. + let values = if let Some(views) = Args::views_if_no_consts(columns) { + if !Args::view_lens_match(&views, row_count) { + vortex_bail!("a decoded row input does not address exactly {row_count} rows"); + } + + // SAFETY: `view_lens_match` checked that these exact views address `row_count` rows. + let source = unsafe { Args::indexed_source(views, row_count) }; + collect_source_bits::(&source, row_count, ctx, &apply) + } else { + let Some(source) = decoded_source::(columns, row_count) else { + vortex_bail!("a decoded row input does not address exactly {row_count} rows"); + }; + collect_source_bits::(&source, row_count, ctx, &apply) }; + Ok(BoolArray::try_new(values, validity)?.into_array()) +} + +/// Pack `apply` over the `row_count` rows of `source`. +#[expect(clippy::inline_always)] +#[inline(always)] +fn collect_source_bits( + source: &S, + row_count: usize, + ctx: &ExecutionCtx, + apply: &impl Fn(S::Item) -> bool, +) -> BitBuffer { let collect = |index| { - // SAFETY: Both collectors only invoke this closure with `index < row_count`, and the decoded - // source was constructed with exactly `row_count` rows. + // SAFETY: Both collectors only invoke this closure with `index < row_count`, and the + // source addresses exactly `row_count` rows. apply(unsafe { source.get_unchecked(index) }) }; - let values = if MULTIVERSIONED { + + if MULTIVERSIONED { BitBuffer::collect_bool_multiversioned_in(row_count, collect, ctx.allocator().clone()) } else { BitBuffer::collect_bool_in(row_count, collect, ctx.allocator().clone()) - }; - - Ok(BoolArray::new(values, Validity::NonNullable).into_array()) + } } /// Decode every input column, then pack Boolean outputs while combining failure evidence. diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs index 3fd72339f4d..2b5c862415c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs @@ -32,6 +32,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::Bool; +use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; use crate::arrays::Primitive; @@ -190,12 +191,18 @@ fn packed_boolean_payload_uses_allocator( .with_allocator(allocator); let output = match (multiversioned, traversal) { - (false, Traversal::Infallible) => { - execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, |(value,)| value > 2)? - } - (true, Traversal::Infallible) => { - execute_owned_infallible_bool::<(i64,), true>(&args, &mut ctx, |(value,)| value > 2)? - } + (false, Traversal::Infallible) => execute_owned_infallible_bool::<(i64,), false>( + &args, + &mut ctx, + Validity::NonNullable, + |(value,)| value > 2, + )?, + (true, Traversal::Infallible) => execute_owned_infallible_bool::<(i64,), true>( + &args, + &mut ctx, + Validity::NonNullable, + |(value,)| value > 2, + )?, (false, Traversal::Fallible) => execute_owned_bool::<(i64,), (), bool, false>( &args, &mut ctx, @@ -347,8 +354,12 @@ fn empty_outputs_and_zero_width_rows_do_not_allocate_payloads() -> VortexResult< let primitive = execute_owned_infallible::<(i64,), i64, ()>(&args, &mut ctx, |_| (), |_, (value,)| value)?; - let boolean = - execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, |(value,)| value > 0)?; + let boolean = execute_owned_infallible_bool::<(i64,), false>( + &args, + &mut ctx, + Validity::NonNullable, + |(value,)| value > 0, + )?; let scalar = execute_sink::<(i64,), (), UninitElementSink, _>( &args, &(), @@ -495,3 +506,53 @@ fn zero_sized_output_uses_its_own_storage(#[case] traversal: Traversal) -> Vorte Ok(()) } + +#[rstest] +fn packed_boolean_attaches_validity( + #[values(false, true)] multiversioned: bool, +) -> VortexResult<()> { + let values = PrimitiveArray::from_iter([1_i64, 2, 3, 4]).into_array(); + let args = VecExecutionArgs::new(vec![values], 4); + let validity = Validity::from_iter([true, false, true, true]); + let mut ctx = VortexSession::empty().create_execution_ctx(); + + let output = if multiversioned { + execute_owned_infallible_bool::<(i64,), true>(&args, &mut ctx, validity, |(value,)| { + value > 2 + })? + } else { + execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, validity, |(value,)| { + value > 2 + })? + }; + + assert!(output.is::()); + let expected = BoolArray::new( + [false, false, true, true].into_iter().collect(), + Validity::from_iter([true, false, true, true]), + ) + .into_array(); + assert_arrays_eq!(&output, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn packed_boolean_rejects_wrong_validity_length() { + let values = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let args = VecExecutionArgs::new(vec![values], 2); + let validity = Validity::from_iter([true, false, true]); + let mut ctx = VortexSession::empty().create_execution_ctx(); + + let error = + execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, validity, |(value,)| { + value > 0 + }) + .expect_err("a validity longer than the batch must fail"); + + assert!( + error + .to_string() + .contains("BoolArray validity length does not match array length"), + "{error}", + ); +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs index 26f114a559f..ef4f516c328 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -24,6 +24,7 @@ mod primitive; mod tuple; pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; +pub(in crate::scalar_fn::unstable::row) use tuple::decode_rows; pub(in crate::scalar_fn::unstable::row) use tuple::decoded_source; pub(in crate::scalar_fn::unstable::row) use tuple::is_batch_const; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index 4f7ea96c856..35493de055d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -51,6 +51,28 @@ impl ArgColumn { Ok(Self(ArgColumnKind::Column(T::decode(array, ctx)?))) } + /// Decode an input already classified as batch-constant or not by [`is_batch_const`]. + #[inline] + pub(super) fn decode_rows( + array: &ArrayRef, + is_const: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if is_const { + // A literal constant is its own batch constant. Wrapped constants unwrap in `decode`. + if array.is::() && !array.is_empty() { + return Ok(Self(ArgColumnKind::Const(T::decode_constant( + array.clone(), + ctx, + )?))); + } + + return Self::decode(array.clone(), ctx); + } + + Ok(Self(ArgColumnKind::Column(T::decode_ref(array, ctx)?))) + } + fn decode_null_tolerant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { // Batch execution short-circuits null constants before selecting a strategy. if let Some(const_array) = batch_const(&array) diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs index 73605aa7d33..808ead04d1a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs @@ -8,12 +8,16 @@ use vortex_compute::lane_kernels::IndexedSource; use vortex_compute::lane_kernels::LaneZip; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; use super::ElementTuple; use super::element_tuple::ArgColumn; use super::element_tuple::ArgColumnKind; +use crate::ExecutionCtx; use crate::scalar_fn::unstable::row::InputElement; use crate::scalar_fn::unstable::row::ViewLen; +use crate::scalar_fn::unstable::row::batch::RowInputs; /// An argument tuple that supports a validated dense indexed traversal. /// @@ -43,6 +47,44 @@ pub(in crate::scalar_fn::unstable::row) fn decoded_source<'a, Args: IndexedEleme ::decoded_source(columns, row_count) } +/// Decode one borrowed array per argument, like [`ElementTuple::decode`]. +/// +/// Bit `i` of `const_mask` is set exactly when input `i` is batch-constant, as classified by +/// `is_batch_const`. Only those inputs repeat the batch-constant checks. +#[expect(clippy::inline_always)] +#[inline(always)] +pub(in crate::scalar_fn::unstable::row) fn decode_rows( + inputs: &Inputs, + const_mask: u32, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + Args: IndexedElementTuple, + Inputs: RowInputs + ?Sized, +{ + vortex_ensure_eq!( + inputs.num_inputs(), + Args::ARITY, + "row input count must match the argument tuple arity", + ); + + ::decode_rows(inputs, const_mask, ctx) +} + +/// Decode input `index` as an argument of element type `T`. +#[expect(clippy::inline_always)] +#[inline(always)] +fn decode_input( + inputs: &Inputs, + index: usize, + const_mask: u32, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let is_const = const_mask & (1 << index) != 0; + + ArgColumn::decode_rows(inputs.input(index), is_const, ctx) +} + /// Indexed access to one decoded [`ArgColumn`]. /// /// [`ArgColumn`] already records whether an input is row-wise or batch-constant. Keeping that choice @@ -182,6 +224,14 @@ impl IndexedElementTuple for () { } impl private::DecodedSource for () { + fn decode_rows( + _inputs: &Inputs, + _const_mask: u32, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(()) + } + fn decoded_source<'a>( _columns: &'a Self::Columns, row_count: usize, @@ -202,6 +252,16 @@ impl IndexedElementTuple for (A,) { } impl private::DecodedSource for (A,) { + #[expect(clippy::inline_always)] + #[inline(always)] + fn decode_rows( + inputs: &Inputs, + const_mask: u32, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok((decode_input(inputs, 0, const_mask, ctx)?,)) + } + fn decoded_source<'a>( columns: &'a Self::Columns, row_count: usize, @@ -222,6 +282,19 @@ impl IndexedElementTuple for (A, B) { } impl private::DecodedSource for (A, B) { + #[expect(clippy::inline_always)] + #[inline(always)] + fn decode_rows( + inputs: &Inputs, + const_mask: u32, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(( + decode_input(inputs, 0, const_mask, ctx)?, + decode_input(inputs, 1, const_mask, ctx)?, + )) + } + fn decoded_source<'a>( columns: &'a Self::Columns, row_count: usize, @@ -283,6 +356,16 @@ macro_rules! indexed_element_tuple { } impl<$($t: InputElement),+> private::DecodedSource for ($($t,)+) { + #[allow(clippy::inline_always)] + #[inline(always)] + fn decode_rows( + inputs: &Inputs, + const_mask: u32, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(($(decode_input::<$t, Inputs>(inputs, $idx, const_mask, ctx)?,)+)) + } + fn decoded_source<'a>( columns: &'a Self::Columns, row_count: usize, @@ -312,14 +395,25 @@ indexed_element_tuple!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: mod private { use vortex_compute::lane_kernels::IndexedSource; + use vortex_error::VortexResult; use super::ElementTuple; + use crate::ExecutionCtx; + use crate::scalar_fn::unstable::row::batch::RowInputs; /// The crate-private half of [`super::IndexedElementTuple`]. /// /// Rust trait methods have the visibility of their trait. This companion keeps source /// construction from decoded columns out of the public unstable API. pub trait DecodedSource: ElementTuple { + /// Decode one array per argument. Bit `i` of `const_mask` marks input `i` as + /// batch-constant. + fn decode_rows( + inputs: &Inputs, + const_mask: u32, + ctx: &mut ExecutionCtx, + ) -> VortexResult; + fn decoded_source<'a>( columns: &'a Self::Columns, row_count: usize, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs index 29b5a1c5221..0054c36838e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -12,6 +12,7 @@ pub(in crate::scalar_fn::unstable::row) use element_tuple::is_batch_const; mod indexed; pub use indexed::IndexedElementTuple; +pub(in crate::scalar_fn::unstable::row) use indexed::decode_rows; pub(in crate::scalar_fn::unstable::row) use indexed::decoded_source; #[cfg(test)] diff --git a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs index 2ce2e9e3e69..ee7d87954bd 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -19,6 +19,7 @@ pub use element::InputElement; pub use element::OutputBuffer; pub use element::OutputElement; pub use element::Utf8Column; +pub(in crate::scalar_fn::unstable::row) use element::decode_rows; pub(in crate::scalar_fn::unstable::row) use element::decoded_source; pub(in crate::scalar_fn::unstable::row) use element::is_batch_const; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index b5539d5b29e..e4fb8d6d0b3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -44,6 +44,7 @@ use crate::scalar_fn::unstable::row::execute::execute_owned_valid_rows; use crate::scalar_fn::unstable::row::execute::execute_sink; use crate::scalar_fn::unstable::row::execute::execute_sink_filtered; use crate::scalar_fn::unstable::row::execute::execute_sink_valid_rows; +use crate::validity::Validity; /// The runtime visit that decodes every column once and runs the selected row loop. pub(crate) struct ExecuteRows<'args, 'ctx, F: RowFn> { @@ -129,7 +130,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { )?; self.plan.ensure_reproduced_by(&visited)?; - execute_owned_infallible_bool::(self.args, self.ctx, apply) + execute_owned_infallible_bool::( + self.args, + self.ctx, + Validity::NonNullable, + apply, + ) } fn visit_prepared_into( diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/fused.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/fused.rs new file mode 100644 index 00000000000..c44198e5e80 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/fused.rs @@ -0,0 +1,282 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Plans and executes a batch in a single dispatch. +//! +//! [`ExecuteFused`] borrows the inputs and owns their [`BatchProbe`]. Each visit plans its concrete +//! signature, then runs every batch that its own kernel can finish: an all-null, all-constant, or +//! null-free batch, or a partially valid batch under [`RowPolicy::Dense`]. Such a batch therefore +//! needs no second dispatch and no check that a second dispatch reproduced the plan. A visit hands +//! any other batch back as [`Fused::Partial`], and batch execution dispatches again for dense +//! retry or valid-row execution. +//! +//! A Boolean visit decodes a direct batch straight from the borrowed inputs. Every other visit +//! keeps the inputs in [`BatchInputs`]. + +use vortex_error::VortexResult; + +use super::BatchPlan; +use super::RowPolicy; +use super::RowVisitor; +use super::check::assert_deferred_visit_contract; +use super::check::assert_owned_visit_contract; +use super::check::assert_sink_visit_contract; +use super::check::validate_owned_visit; +use super::check::validate_sink_visit; +use super::row_visitor::private; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::FailureEvidence; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::SinkResult; +use crate::scalar_fn::unstable::row::batch::BatchInputs; +use crate::scalar_fn::unstable::row::batch::BatchProbe; +use crate::scalar_fn::unstable::row::batch::BorrowedRowFnArgs; +use crate::scalar_fn::unstable::row::batch::PartiallyValidBatch; +use crate::scalar_fn::unstable::row::batch::RowInputs; +use crate::scalar_fn::unstable::row::execute::execute_owned; +use crate::scalar_fn::unstable::row::execute::execute_owned_bool; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible_bool; +use crate::scalar_fn::unstable::row::execute::execute_rows_infallible_bool; +use crate::scalar_fn::unstable::row::execute::execute_sink; +use crate::validity::Validity; + +/// The result of a fused planning and execution visit. +pub(crate) enum Fused<'a> { + /// The finished output column. + Done(ArrayRef), + + /// A partially valid batch that needs dense retry or valid-row execution, with its plan. + Partial(PartiallyValidBatch<'a>), +} + +/// The visit that plans a batch and executes it when no retry can be necessary. +pub(crate) struct ExecuteFused<'a, 'ctx, F: RowFn, I: ?Sized> { + /// The borrowed input columns. + inputs: &'a I, + + /// What batch execution derived from the inputs. + probe: BatchProbe, + + /// The input dtypes passed to the dispatch. + dtypes: &'a [DType], + + /// The output dtype declared by [`RowVisitor::with_output_dtype`], if any. + output_dtype: Option, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, + + /// The function being executed, named in the errors of a collected batch. + function: &'a F, +} + +impl<'a, 'ctx, F: RowFn, I: RowInputs + ?Sized> ExecuteFused<'a, 'ctx, F, I> { + pub(crate) fn new( + function: &'a F, + inputs: &'a I, + probe: BatchProbe, + dtypes: &'a [DType], + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + inputs, + probe, + dtypes, + output_dtype: None, + ctx, + function, + } + } + + /// Plan the inputs, then finish the batch with `kernel` when no retry can be necessary. + fn execute( + self, + storage_dtype: DType, + policy: RowPolicy, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, + ) -> VortexResult> { + let plan = BatchPlan::new(storage_dtype, self.output_dtype, policy)?; + let batch = BatchInputs::new(RowFn::id(self.function), self.probe, self.inputs); + + let fused = match batch.try_execute_without_retry(self.dtypes, plan, kernel, self.ctx)? { + Ok(values) => Fused::Done(values), + Err(partial) => Fused::Partial(partial), + }; + + Ok(fused) + } +} + +impl private::Sealed for ExecuteFused<'_, '_, F, I> {} + +impl<'a, F: RowFn, I: RowInputs + ?Sized> RowVisitor for ExecuteFused<'a, '_, F, I> { + type VisitResult = Fused<'a>; + + fn with_output_dtype(mut self, dtype: DType) -> Self { + self.output_dtype = Some(dtype); + self + } + + fn visit_prepared( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + let storage_dtype = validate_owned_visit::(self.dtypes)?; + + self.execute( + storage_dtype, + RowPolicy::for_owned_output::(), + |args, _validity, ctx| { + execute_owned_infallible::(&args, ctx, prepare, apply) + }, + ) + } + + fn visit_bool( + mut self, + apply: impl Fn(Args::Elems<'_>) -> bool, + ) -> VortexResult + where + Args: IndexedElementTuple, + { + const { assert_owned_visit_contract::() }; + let storage_dtype = validate_owned_visit::(self.dtypes)?; + + // A direct batch packs straight into the finished column. + if self.output_dtype.is_none() { + let dense = RowPolicy::for_owned_output::() == RowPolicy::Dense; + let const_mask = self.probe.const_mask(); + match self + .probe + .into_direct_validity(self.inputs.num_inputs(), dense) + { + Ok(validity) => { + // The collector builds a canonical Boolean array of exactly the batch length + // with this validity attached, so the output needs no plan, label, or + // validation. + return execute_rows_infallible_bool::( + self.inputs, + const_mask, + self.ctx, + validity, + apply, + ) + .map(Fused::Done); + } + Err(probe) => self.probe = probe, + } + } + + self.execute( + storage_dtype, + RowPolicy::for_owned_output::(), + |args, validity, ctx| { + execute_owned_infallible_bool::( + &args, + ctx, + validity.take().unwrap_or(Validity::NonNullable), + apply, + ) + }, + ) + } + + fn visit_prepared_into( + self, + params: Sink::Params, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + const { assert_sink_visit_contract::() }; + let storage_dtype = validate_sink_visit::(self.dtypes, ¶ms)?; + + self.execute( + storage_dtype, + RowPolicy::for_sink::(), + |args, _validity, ctx| { + execute_sink::( + &args, ¶ms, ctx, prepare, apply, + ) + }, + ) + } + + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: FailureEvidence, + { + const { assert_deferred_visit_contract::() }; + let storage_dtype = validate_owned_visit::(self.dtypes)?; + + self.execute( + storage_dtype, + RowPolicy::for_deferred_output::(), + |args, _validity, ctx| { + execute_owned::( + &args, + ctx, + prepare, + apply, + finish_failure, + ) + }, + ) + } + + fn visit_prepared_deferred_bool( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (bool, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Fail: FailureEvidence, + { + const { assert_deferred_visit_contract::() }; + let storage_dtype = validate_owned_visit::(self.dtypes)?; + + self.execute( + storage_dtype, + RowPolicy::for_deferred_output::(), + |args, _validity, ctx| { + execute_owned_bool::( + &args, + ctx, + prepare, + apply, + finish_failure, + ) + }, + ) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs index dc8d0209374..f7a3188b45c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -13,6 +13,10 @@ pub(super) use execute::ExecuteFilteredRows; pub(super) use execute::ExecuteRows; pub(super) use execute::ExecuteValidRows; +mod fused; +pub(super) use fused::ExecuteFused; +pub(super) use fused::Fused; + mod retry; pub(super) use retry::ExecuteDenseWithRetry; diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 55e686d1b09..5eee9e8bde2 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -15,7 +15,6 @@ use vortex_mask::MaskValuesRef; use vortex_session::VortexSession; use super::batch::ArrayInputs; -use super::batch::BatchInputs; use super::batch::BatchProbe; use super::batch::BorrowedRowFnArgs; use super::batch::RowInputs; @@ -24,8 +23,10 @@ use super::row_fn::RowFn; use super::visitor::BatchPlanner; use super::visitor::ExecuteDenseWithRetry; use super::visitor::ExecuteFilteredRows; +use super::visitor::ExecuteFused; use super::visitor::ExecuteRows; use super::visitor::ExecuteValidRows; +use super::visitor::Fused; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; @@ -149,6 +150,8 @@ pub(crate) fn execute_row_inputs( return execute_nullary_rows(function, options, inputs.row_count(), ctx); } + // Classify every input before decoding. A later null literal can skip an earlier fallible + // expression or device transfer. let probe = BatchProbe::new(|| RowFn::id(function), inputs)?; let collected: SmallVec<[DType; 4]>; let arg_dtypes = match inputs.dtypes() { @@ -161,10 +164,19 @@ pub(crate) fn execute_row_inputs( } }; - let plan = function.dispatch(options, arg_dtypes, BatchPlanner::::new(arg_dtypes))?; - let batch = BatchInputs::new(RowFn::id(function), probe, inputs); + // Plan and execute in one dispatch. Only a partially valid batch that needs dense retry or + // valid-row execution dispatches again. + let batch = match function.dispatch( + options, + arg_dtypes, + ExecuteFused::::new(function, inputs, probe, arg_dtypes, ctx), + )? { + Fused::Done(values) => return Ok(values), + Fused::Partial(batch) => batch, + }; - batch.with_plan(arg_dtypes, plan).execute( + batch.execute( + arg_dtypes, |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), @@ -203,6 +215,7 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { Ok(()) } +/// Execute the row kernel over `args`. fn execute_row_kernel( function: &F, options: &F::Options, @@ -275,6 +288,7 @@ mod tests { use super::row_fn_return_dtype; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -367,13 +381,15 @@ mod tests { _args: &[DType], visitor: V, ) -> VortexResult { + // The first dispatch plans dense execution with a retry, so a partially valid batch + // dispatches again for its dense attempt. if self.dispatches.fetch_add(1, Ordering::Relaxed) == 0 { - visitor.visit::<(i64,), i64>(|(value,)| value) + visitor.visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())) } else { match self.change { - DispatchChange::Policy => visitor - .visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())), - DispatchChange::Element => visitor.visit::<(u64,), u64>(|(value,)| value), + DispatchChange::Policy => visitor.visit::<(i64,), i64>(|(value,)| value), + DispatchChange::Element => visitor + .visit_deferred::<(u64,), u64, bool>(|(value,)| (value, false), |_| Ok(())), } } } @@ -411,13 +427,32 @@ mod tests { Ok(()) } + #[test] + fn test_execute_dispatches_once_without_retry() -> VortexResult<()> { + let dispatches = Arc::new(AtomicUsize::new(0)); + let function = ChangingDispatchRowFn { + dispatches: Arc::clone(&dispatches), + change: DispatchChange::Policy, + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = SESSION.create_execution_ctx(); + + let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; + + assert_eq!(dispatches.load(Ordering::Relaxed), 1); + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) + } + #[test] fn test_execute_rejects_dispatch_that_changes_after_planning() -> VortexResult<()> { let function = ChangingDispatchRowFn { dispatches: Arc::new(AtomicUsize::new(0)), change: DispatchChange::Policy, }; - let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let validity = Validity::from_iter([true, false]); + let input = PrimitiveArray::new(vec![1_i64, 2], validity).into_array(); let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = SESSION.create_execution_ctx(); @@ -434,7 +469,7 @@ mod tests { "unexpected error: {error}", ); assert!( - message.contains("planned Dense, got DenseWithRetry"), + message.contains("planned DenseWithRetry, got Dense"), "unexpected error: {error}", ); Ok(()) @@ -446,7 +481,8 @@ mod tests { dispatches: Arc::new(AtomicUsize::new(0)), change: DispatchChange::Element, }; - let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let validity = Validity::from_iter([true, false]); + let input = PrimitiveArray::new(vec![1_i64, 2], validity).into_array(); let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = SESSION.create_execution_ctx(); @@ -466,6 +502,46 @@ mod tests { Ok(()) } + #[derive(Clone)] + struct GreaterThanTwo; + + impl RowFn for GreaterThanTwo { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.greater_than_two"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_bool::<(i64,), false>(|(value,)| value > 2) + } + } + + #[test] + fn test_partially_valid_packed_bool_output_is_canonical() -> VortexResult<()> { + let validity = Validity::from_iter([true, false, true, true]); + let input = PrimitiveArray::new(vec![1_i64, 3, 3, 0], validity.clone()).into_array(); + let args = VecExecutionArgs::new(vec![input], 4); + let mut ctx = SESSION.create_execution_ctx(); + + let actual = execute_rows(&GreaterThanTwo, &EmptyOptions, &args, &mut ctx)?; + let expected = + BoolArray::new([false, true, true, false].into_iter().collect(), validity).into_array(); + + assert!(actual.is_canonical(), "{actual}"); + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } + #[track_caller] fn assert_arity_error(error: VortexError) { assert!( From 712ed0948cb6779ccfe4276c3afbf737a255ce46 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 27 Oct 2026 13:37:36 -0400 Subject: [PATCH 2/2] Inline the multiversioned bit packing loop in static AVX-512 builds A static AVX-512 build inlines collect_bool_words_avx512 only when rustc emits that copy in the caller's codegen unit. This branch moved it out of the RowFn deferred Boolean kernel's unit, adding 15% more instructions per row. Such a build has nothing wider to detect, so use collect_bool_words_inline, which already packs with AVX-512. Signed-off-by: Connor Tsui --- vortex-buffer/src/bit/pack.rs | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/vortex-buffer/src/bit/pack.rs b/vortex-buffer/src/bit/pack.rs index cf180084667..67e03f92ecb 100644 --- a/vortex-buffer/src/bit/pack.rs +++ b/vortex-buffer/src/bit/pack.rs @@ -111,6 +111,10 @@ where /// conversions, or unchecked slice comparisons like the `between` kernels — where the fused /// AVX-512 loop is worth another ~2x over the baseline kernel. /// +/// A build that enables AVX-512F and AVX-512BW at compile time skips detection and uses +/// `collect_bool_words_inline`, which already packs with AVX-512 without a `#[target_feature]` +/// boundary. +/// /// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`, /// exactly once per index in ascending order. /// @@ -133,7 +137,15 @@ where return collect_bool_words_inline(words, len, f); } - #[cfg(all(target_arch = "x86_64", not(miri)))] + // NB: Static AVX-512 builds must not call the `#[target_feature]` copies. A copy is a generic + // without `#[inline]`, so rustc emits it in one codegen unit, and the caller inlines it only + // when both share a unit. Unrelated code decides that: in #10675 it left the RowFn deferred + // Boolean loop out of line, with 15% more instructions per row. + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx512f", target_feature = "avx512bw")), + not(miri) + ))] { if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") { // SAFETY: runtime detection guarantees the required target features.