diff --git a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs index 6b389a0554b..d7bcc8d0b4c 100644 --- a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs +++ b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_err; use crate::scalar_fn::fns::operators::Operator; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] /// Binary element-wise operations. pub enum NumericOperator { /// Binary element-wise addition of two arrays or of two scalars. diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index 46ff0f36e15..1901f0260c0 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Checked-lane execution for numeric kernels, driven by the shared +//! Checked-lane execution for the decimal kernels, driven by the shared //! `vortex-compute` lane kernels. - -use std::ops::BitOrAssign; +//! +//! The primitive widths do not come through here: they are computed one row at a time by +//! [`row`](super::row), which writes a value for every row and reduces failure evidence without +//! scanning the finished output. use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -13,33 +15,18 @@ use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_mask::AllOr; use vortex_mask::Mask; -/// Evidence that a lane failed, anything other than [`Default`] meaning failure. -/// -/// `bool` is the ordinary choice; [`map_checked_into`] explains why the wider members exist and -/// asserts the width bound that membership here does **not** imply. -/// -/// [`map_checked_into`]: IndexedSourceExt::map_checked_into -pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - -impl Failure for bool {} -impl Failure for u8 {} -impl Failure for u16 {} -impl Failure for u32 {} -impl Failure for u64 {} - /// Apply the fallible `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it returns `None` on a _valid_ lane. +/// `Err(first_failing_valid_lane)` only when it returns `None` on a valid lane. /// /// `apply` also runs on invalid lanes, whose failures are masked out and whose values are /// unspecified, so it must be total: no panics or side effects on any stored lane value. /// -/// This drives the one-pass early-exit kernels, which abort at the end of the enclosing 64-lane -/// chunk. Use it when per-lane failure handling is cheap relative to the operation, as in integer -/// division. Prefer [`checked_apply_lanes`] when the failure check itself vectorizes. +/// This drives the one-pass early-exit kernels: failures abort at the end of the enclosing +/// 64-lane chunk. It suits an operation whose per-lane failure handling is cheap relative to the +/// operation itself, which is what the decimal kernels and their per-lane casts are. /// -/// `#[inline]`: this must inline into the caller that builds the closure, so that a captured -/// constant operand flattens into a register rather than living behind a pointer the loop reloads -/// on every lane, which blocks vectorization. +/// Keep this wrapper inlineable so captured constants can become loop invariants in the caller. +/// The lane kernels retain their own inlining decisions. #[inline] pub(super) fn checked_lanes( source: S, @@ -61,7 +48,6 @@ where }; let mut values = BufferMut::::with_capacity(len); - let out = &mut values.spare_capacity_mut()[..len]; match valid_bits { None => source.try_map_into(out, apply)?, @@ -73,55 +59,3 @@ where Ok(values.freeze()) } - -/// Apply the split value/failure `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it flags a _valid_ lane. -/// -/// The hot pass writes every value unconditionally and OR-reduces one piece of evidence, leaving -/// the loop free of per-lane selects and per-chunk exit branches. Only if a lane flagged does a -/// cold second pass re-run `apply` through the early-exit kernels, which drop null-lane failures -/// and attribute the first valid one. -/// -/// Like [`checked_lanes`], `apply` runs on invalid lanes and must be total. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -pub(super) fn checked_apply_lanes( - source: S, - valid_rows: &Mask, - apply: Apply, -) -> Result, usize> -where - S: IndexedSource + Copy, - T: Copy + Default, - Fail: Failure, - Apply: Fn(S::Item) -> (T, Fail), -{ - let len = source.len(); - debug_assert_eq!(len, valid_rows.len()); - - let valid_bits = match valid_rows.bit_buffer() { - AllOr::All => None, - AllOr::None => return Ok(Buffer::zeroed(len)), - AllOr::Some(valid_bits) => Some(valid_bits), - }; - - let mut values = BufferMut::::with_capacity(len); - - let out = &mut values.spare_capacity_mut()[..len]; - if source.map_checked_into(out, &apply) != Fail::default() { - let checked = |item: S::Item| { - let (value, failure) = apply(item); - (failure == Fail::default()).then_some(value) - }; - match valid_bits { - None => source.try_map_into(out, checked)?, - Some(valid_bits) => source.try_map_masked_into(valid_bits, out, checked)?, - } - } - - // SAFETY: the kernels initialize every lane in `out`. - unsafe { values.set_len(len) }; - - Ok(values.freeze()) -} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index 6dc0de0fbea..a0c427b142e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -4,16 +4,19 @@ //! Native execution of the arithmetic operators (Add/Sub/Mul/Div) of the [`Binary`] scalar //! function. There is no Arrow fallback. //! +//! The primitive widths are computed by a +//! [`RowFn`](crate::scalar_fn::unstable::row::RowFn), which owns null handling, constants, and +//! validity for them; see [`row`]. Decimal keeps its own columnar implementation in [`decimal`]. +//! //! [`Binary`]: super::Binary mod checked; mod decimal; mod primitive; -#[cfg(test)] -mod tests; +mod row; use decimal::execute_numeric_decimal; -use primitive::execute_numeric_primitive; +use row::execute_numeric_primitive; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -81,3 +84,6 @@ fn build_empty_result( Ok(Canonical::empty(&result_dtype).into_array()) } + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 3a3c3f1b585..7cbe88c7c8e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -1,73 +1,52 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; -use vortex_compute::lane_kernels::IndexedSource; -use vortex_compute::lane_kernels::LaneZip; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::Mask; - -use super::checked::Failure; -use super::checked::checked_apply_lanes; -use super::checked::checked_lanes; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::PrimitiveArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; +//! Checked arithmetic for one primitive row. + +use std::ops::BitOrAssign; + use crate::dtype::NativePType; -use crate::dtype::PType; use crate::dtype::half::f16; -use crate::match_each_native_ptype; -use crate::scalar::NumericOperator; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; -use crate::validity::Validity; - -struct CheckedAdd; -struct CheckedSub; +/// Checked addition, failing on integer overflow. +pub(super) struct CheckedAdd; -struct CheckedMul; +/// Checked subtraction, failing on integer overflow. +pub(super) struct CheckedSub; -struct CheckedDiv; +/// Checked multiplication, failing on integer overflow. +pub(super) struct CheckedMul; -trait CheckedPrimitiveOp: Sized { - /// Message for the error raised when this operation fails on a valid lane. - const ERROR: &'static str; +/// Checked division, failing on integer division by zero and on `MIN / -1`. +pub(super) struct CheckedDiv; - /// Whether to check inside the value loop and exit early, rather than reducing evidence over - /// the whole batch and re-running only once some lane has flagged. - const CHECKED_VALUE_LOOP: bool = false; +/// OR-reducible evidence that a row failed, with [`Default`] meaning success. +pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - /// How this operation reports a failing lane. See [`Failure`]. - type Failure: Failure; +impl Failure for bool {} +impl Failure for u8 {} +impl Failure for u16 {} +impl Failure for u32 {} +impl Failure for u64 {} - /// Compute a lane's value and its failure evidence together. - /// - /// Overflowing-style rather than `Option`, so the vectorizable kernels can write every - /// lane's value unconditionally and reduce the evidence separately. [`Self::checked`] is the - /// shape for the scalar and early-exit paths. - fn apply(lhs: T, rhs: T) -> (T, Self::Failure); +/// One arithmetic operator at one width, split into its value and failure evidence. +pub(super) trait CheckedPrimitiveOp: 'static + Sized { + /// The error reported for a batch in which some valid row failed. + const ERROR: &'static str; - /// [`Self::apply`] folded into the `Option` that the early-exit kernels take. - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - let (value, failed) = Self::apply(lhs, rhs); + /// How this operation reports a failing row. See [`Failure`]. + type Fail: Failure; - (failed == Self::Failure::default()).then_some(value) - } + /// The result of this operation, paired with evidence of whether the row failed. + fn apply(lhs: T, rhs: T) -> (T, Self::Fail); } impl CheckedPrimitiveOp for CheckedAdd { const ERROR: &'static str = "integer overflow in checked add"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.add_value(rhs), lhs.add_error(rhs)) } @@ -76,9 +55,9 @@ impl CheckedPrimitiveOp for CheckedAdd { impl CheckedPrimitiveOp for CheckedSub { const ERROR: &'static str = "integer overflow in checked sub"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.sub_value(rhs), lhs.sub_error(rhs)) } @@ -87,9 +66,9 @@ impl CheckedPrimitiveOp for CheckedSub { impl CheckedPrimitiveOp for CheckedMul { const ERROR: &'static str = "integer overflow in checked mul"; - type Failure = T::MulFailure; + type Fail = T::MulFailure; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, T::MulFailure) { (lhs.mul_value(rhs), lhs.mul_failure(rhs)) } @@ -97,16 +76,10 @@ impl CheckedPrimitiveOp for CheckedMul { impl CheckedPrimitiveOp for CheckedDiv { const ERROR: &'static str = "integer division by zero or overflow in checked div"; - // Integer division still lowers to scalar divides, so the split - // value/error-scan loop used to auto-vectorize add/sub/mul only adds a - // second full scan. Use the one-pass early-exit checked kernel for integer - // division, matching Arrow/Velox. Float division has no checked errors and - // stays on the split/vectorizable default path. - const CHECKED_VALUE_LOOP: bool = T::DIV_CHECKS_IN_VALUE_LOOP; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { let failed = lhs.div_error(rhs); let value = if failed { @@ -116,151 +89,16 @@ impl CheckedPrimitiveOp for CheckedDiv { }; (value, failed) } - - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - lhs.div_checked(rhs) - } } -/// Execute a numeric operation between two primitive-typed arrays. -pub(super) fn execute_numeric_primitive( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: NumericOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - - match_each_native_ptype!(ptype, |T| { - match op { - NumericOperator::Add => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Sub => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Mul => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Div => execute_checked_typed::(lhs, rhs, ctx), - } - }) -} - -fn execute_checked_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: NativePType, - Op: CheckedPrimitiveOp, - Scalar: From, - Scalar: From>, -{ - let result_dtype = lhs - .dtype() - .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - let len = lhs.len(); - debug_assert_eq!(len, rhs.len()); - - let validity = lhs.validity().and(rhs.validity())?; - let valid_rows = validity.execute_mask(len, ctx)?; - - let values = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_op_lanes::<_, T, Op>( - LaneZip::new(lhs.as_slice(), rhs.as_slice()), - &valid_rows, - |(lhs, rhs)| (lhs, rhs), - ), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Capture the constant by value so it stays hoisted out of the lane loop. - let rhs = *rhs; - checked_op_lanes::<_, T, Op>(lhs.as_slice(), &valid_rows, move |lhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => { - let lhs = *lhs; - checked_op_lanes::<_, T, Op>(rhs.as_slice(), &valid_rows, move |rhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - let value = Op::checked(*lhs, *rhs) - .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - return Ok(constant_result_array(value, len, &result_dtype)); - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => Ok(Buffer::zeroed(len)), - } - .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - - primitive_result_array::(values, validity, &result_dtype) -} - -/// Run `Op` over the lanes of `source` in the loop shape it declares: the one-pass early-exit -/// kernel when `Op::CHECKED_VALUE_LOOP` is set, and the split value/failure kernel otherwise. +/// Per-width arithmetic used to compute values and failure evidence. /// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -fn checked_op_lanes( - source: S, - valid_rows: &Mask, - to_operands: impl Fn(S::Item) -> (T, T), -) -> Result, usize> -where - S: IndexedSource + Copy, - T: NativePType, - Op: CheckedPrimitiveOp, -{ - if Op::CHECKED_VALUE_LOOP { - checked_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::checked(lhs, rhs) - }) - } else { - checked_apply_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::apply(lhs, rhs) - }) - } -} - -fn primitive_result_array( - values: Buffer, - validity: Validity, - dtype: &DType, -) -> VortexResult { - let array = PrimitiveArray::new(values, validity).into_array(); - if array.dtype() == dtype { - return Ok(array); - } - array.cast(dtype.clone()) -} - -fn constant_result_array(value: T, len: usize, dtype: &DType) -> ArrayRef -where - T: NativePType, - Scalar: From + From>, -{ - if dtype.is_nullable() { - ConstantArray::new(Some(value), len).into_array() - } else { - ConstantArray::new(value, len).into_array() - } -} - -trait CheckedArithmetic: NativePType { - const DIV_CHECKS_IN_VALUE_LOOP: bool; - - /// How multiplication reports a failing lane, which is width-dependent in a way the other - /// operations are not. `impl_checked_unsigned!` and `impl_checked_signed!` say why each family - /// reports what it reports. +/// The add, subtract, and multiply value methods **must** be total over every stored lane value. +/// [`Self::div_value`] may assume that [`Self::div_error`] returned `false` for the same operands. +pub(super) trait CheckedArithmetic: NativePType { + /// How multiplication reports a failing row. + /// + /// This may be a word rather than `bool` when narrowing evidence would block vectorization. type MulFailure: Failure; fn add_value(self, rhs: Self) -> Self; @@ -269,18 +107,15 @@ trait CheckedArithmetic: NativePType { fn sub_error(self, rhs: Self) -> bool; fn mul_value(self, rhs: Self) -> Self; fn mul_failure(self, rhs: Self) -> Self::MulFailure; + + /// Divide operands that [`Self::div_error`] accepted. fn div_value(self, rhs: Self) -> Self; + + /// Return whether [`Self::div_value`] would trap for these operands. fn div_error(self, rhs: Self) -> bool; - fn div_checked(self, rhs: Self) -> Option; } -/// The integer arithmetic every width shares, given the two things that actually differ between -/// them: how multiplication reports a failing lane, and how add/sub/div detect one. -/// -/// Only `mul_failure` genuinely varies per width, so it is the macro's parameter and everything -/// else is written once. The `$mul_failure_ty` and body are supplied by the caller because the -/// choice between a widened high half and a `bool` is exactly the vectorization decision described -/// on [`Failure`]. +/// Generate the shared integer operations from their failure predicates. macro_rules! impl_checked_integer { ( $ty:ty, @@ -291,67 +126,57 @@ macro_rules! impl_checked_integer { = |$mf_lhs:ident, $mf_rhs:ident| $mul_failure:expr, ) => { impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = true; - type MulFailure = $mul_failure_ty; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self.wrapping_add(rhs) } - #[inline(always)] + #[inline] fn add_error(self, rhs: Self) -> bool { let ($add_lhs, $add_rhs) = (self, rhs); $add_error } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self.wrapping_sub(rhs) } - #[inline(always)] + #[inline] fn sub_error(self, rhs: Self) -> bool { let ($sub_lhs, $sub_rhs) = (self, rhs); $sub_error } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self.wrapping_mul(rhs) } - #[inline(always)] + #[inline] $(#[$mul_failure_attr])* fn mul_failure(self, rhs: Self) -> $mul_failure_ty { let ($mf_lhs, $mf_rhs) = (self, rhs); $mul_failure } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, rhs: Self) -> bool { let ($div_lhs, $div_rhs) = (self, rhs); $div_error } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - self.checked_div(rhs) - } } }; } -/// The unsigned widths. `add`, `sub` and `div` are written once here, and `widening_mul` covers -/// every width including the 64-bit one, which widens into `u128`: the discarded high half of the -/// widened product is the failure evidence, and costs none of the comparison LLVM folds into -/// `umul.with.overflow`. +/// Unsigned multiplication reports its discarded high half as failure evidence. macro_rules! impl_checked_unsigned { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_integer!( @@ -364,12 +189,7 @@ macro_rules! impl_checked_unsigned { }; } -/// The signed widths, on the same principle. `widening_mul` is the shorthand for the narrow widths, -/// whose two-sided range check over a wider product is not the shape LLVM folds into an overflow -/// intrinsic, so they vectorize while reporting a plain `bool`. The 64-bit width cannot: deriving a -/// `bool` there costs the comparison that scalarizes the loop, so `high_half_mul` hands back the -/// discarded high half as a word. Both arms derive every shift and bound from `$ty`, so -/// instantiating one at a new width cannot silently keep another width's constants. +/// Signed widths use a range check or discarded high-half evidence. macro_rules! impl_checked_signed { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_signed!($ty, mul_failure: bool = |lhs, rhs| { @@ -377,9 +197,6 @@ macro_rules! impl_checked_signed { product < <$ty>::MIN as $wide || product > <$ty>::MAX as $wide }); }; - // Zero exactly when the product fits: a signed multiply overflows iff the high half of the true - // product differs from the sign extension of the half that was kept, so the two XOR to zero on - // the lanes that fit. `tests::test_multiply_overflow_boundaries` pins the boundaries. ($ty:ty, high_half_mul: $wide:ty => $failure:ty) => { impl_checked_signed!($ty, mul_failure: #[expect( clippy::cast_possible_truncation, @@ -389,13 +206,17 @@ macro_rules! impl_checked_signed { let kept = wide as $ty; let discarded = (wide >> <$ty>::BITS) as $ty; + // A product fits exactly when its discarded half is the sign extension of the kept + // half. XOR reduces that comparison to zero evidence for success and nonzero evidence + // for overflow without converting the wide product to a branch. + (discarded ^ (kept >> (<$ty>::BITS - 1))) as $failure }); }; ( $ty:ty, mul_failure: $(#[$mul_failure_attr:meta])* $mul_failure_ty:ty - = |$l:ident, $r:ident| $mul_failure:expr + = |$lhs:ident, $rhs:ident| $mul_failure:expr ) => { impl_checked_integer!( $ty, @@ -408,7 +229,7 @@ macro_rules! impl_checked_signed { ((lhs ^ rhs) & (lhs ^ value)) < 0 }, div_error: |lhs, rhs| rhs == 0 || (lhs == <$ty>::MIN && rhs == -1), - mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$l, $r| $mul_failure, + mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$lhs, $rhs| $mul_failure, ); }; } @@ -417,54 +238,47 @@ macro_rules! impl_checked_float { ($($ty:ty),+ $(,)?) => { $( impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = false; - type MulFailure = bool; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self + rhs } - #[inline(always)] + #[inline] fn add_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self - rhs } - #[inline(always)] + #[inline] fn sub_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self * rhs } - #[inline(always)] + #[inline] fn mul_failure(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, _rhs: Self) -> bool { false } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - Some(self / rhs) - } } )+ }; @@ -484,30 +298,27 @@ impl_checked_float!(f16, f32, f64); mod tests { use super::CheckedArithmetic; - /// Values whose pairwise products are worth probing: the saturating boundaries, the sign-change - /// pivots, and a spread of magnitudes that straddles the 64-bit split. + /// Values around zero, signed extrema, and 32- and 64-bit boundaries where the discarded + /// multiplication half or its sign extension changes. const PROBES: &[i64] = &[ - 0, // - 1, // - -1, // - 2, // - -2, // - 3, // - i64::MIN, // - i64::MIN + 1, // - i64::MAX, // - i64::MAX - 1, // - 1 << 31, // - 1 << 32, // - 1 << 62, // - -(1 << 62), // - 0x7FFF_FFFF, // - -0x8000_0000, // + 0, // Additive identity. + 1, // Smallest positive value. + -1, // All sign bits set. + 2, // Small positive power of two. + -2, // Small negative power of two. + 3, // Small non-power of two. + i64::MIN, // Minimum signed value. + i64::MIN + 1, // Minimum signed value's neighbor. + i64::MAX, // Maximum signed value. + i64::MAX - 1, // Maximum signed value's neighbor. + 1 << 31, // First positive value outside i32. + 1 << 32, // First value with bit 32 set. + 1 << 62, // Largest positive power of two in i64. + -(1 << 62), // Negative counterpart of the largest power of two. + 0x7FFF_FFFF, // Maximum i32 represented as i64. + -0x8000_0000, // Minimum i32 represented as i64. ]; - /// Every `mul_failure` impl is either a bit trick or a two-sided range check, so hold each - /// against the obvious reference: `reference` is the width's own `checked_mul`, whose `None` - /// _is_ the definition of overflow. #[track_caller] fn assert_agrees_with_checked_mul(lhs: T, rhs: T, reference: Option) { let failed = lhs.mul_failure(rhs) != ::default(); @@ -522,14 +333,11 @@ mod tests { assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); let (lhs, rhs) = (lhs as u64, rhs as u64); - assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); } } } - /// The 8-bit widths are cheap enough to check exhaustively, which pins the shift of the - /// unsigned formula and the range check of the signed one against every product that exists. #[test] fn mul_failure_agrees_with_checked_mul_exhaustively_at_8_bits() { for lhs in u8::MIN..=u8::MAX { diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs new file mode 100644 index 00000000000..7f7a3f0e441 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive arithmetic execution through [`RowFn`]. +//! +//! `Binary` keeps its registered contract; [`NumericBinary`] is only an execution helper. Decimal +//! arithmetic remains on its existing columnar path. + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::primitive::CheckedAdd; +use super::primitive::CheckedArithmetic; +use super::primitive::CheckedDiv; +use super::primitive::CheckedMul; +use super::primitive::CheckedPrimitiveOp; +use super::primitive::CheckedSub; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::NumericOperator; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::unstable::row::InitializedElement; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::UninitElementSink; +use crate::scalar_fn::unstable::row::execute_rows; + +pub(super) fn execute_numeric_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + execute_rows(&NumericBinary, &op, &args, ctx) +} + +/// Internal row execution for the primitive arithmetic operators. +#[derive(Clone)] +struct NumericBinary; + +impl RowFn for NumericBinary { + type Options = NumericOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + // Fallibility is queried without input dtypes, so this conservatively covers integer widths. + const INFALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + // `NumericBinary` is a private implementation detail of `Binary`: it is never registered or + // serialized independently. Reusing the public ID keeps execution errors attributed to + // `Binary`. If this type becomes registrable, it needs its own ID and persistence contract. + ScalarFnVTable::id(&Binary) + } + + fn dispatch>( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = PType::try_from( + args.first() + .ok_or_else(|| vortex_err!("a numeric operator takes two operands, got none"))?, + )?; + + match_each_native_ptype!(ptype, |T| { + match op { + NumericOperator::Add => visit_checked::(visitor), + NumericOperator::Sub => visit_checked::(visitor), + NumericOperator::Mul => visit_checked::(visitor), + NumericOperator::Div => visit_div::(visitor), + } + }) + } +} + +fn visit_checked(visitor: V) -> VortexResult +where + T: NativePType, + Op: CheckedPrimitiveOp, + V: RowVisitor, +{ + visitor.visit_deferred::<(T, T), T, Op::Fail>( + |(lhs, rhs)| Op::apply(lhs, rhs), + |failure| { + if failure != ::default() { + return Err(numeric_error(Op::ERROR)); + } + + Ok(()) + }, + ) +} + +fn visit_div(visitor: V) -> VortexResult +where + T: CheckedArithmetic, + V: RowVisitor, +{ + if T::PTYPE.is_float() { + return visit_checked::(visitor); + } + + // Integer division is scalar and expensive, so deferring its cheap failure check preserves no + // vectorization. Check each divide immediately and stop at the first failure. + // Dense execution leaves output uninitialized. Nullable branches fill placeholders only when + // they need to skip invalid rows. + visitor.visit_into::<(T, T), UninitElementSink, _>(|(lhs, rhs), output| { + let (value, failed) = CheckedDiv::apply(lhs, rhs); + if failed { + return Err(numeric_error(>::ERROR)); + } + + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + Ok(unsafe { InitializedElement::write(output, value) }) + }) +} + +/// Keep rich error construction out of row closures so the closures remain inlineable. +#[cold] +#[inline(never)] +fn numeric_error(message: &'static str) -> VortexError { + vortex_err!(InvalidArgument: "{message}") +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs index 7ee98ae717e..3813c8612b3 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -201,8 +201,7 @@ fn test_integer_array_array_errors_on_valid_lanes() { assert!(result.is_err()); } -/// Multiply two non-nullable lanes of `lhs` by two of `rhs`, expecting `Some(product)` where the -/// product fits and `None` where the checked kernel must report overflow. +/// Assert one checked multiplication through the complete array execution path. #[track_caller] fn assert_multiply(lhs: T, rhs: T, expected: Option) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -297,13 +296,11 @@ fn test_multiply_overflow_on_null_lane_ignored( Ok(()) } -/// The hot pass OR-reduces evidence across whole 64-lane chunks before anything looks at it, so an -/// overflow in a late chunk must still be caught, and must still be suppressed when its lane is -/// null. Every other test here fits in a single chunk and cannot show either. +/// An overflow late in the batch must still be reported, unless its row is null. #[rstest] #[case::reported(true)] #[case::suppressed_by_null(false)] -fn test_multiply_overflow_survives_chunk_reduction( +fn test_multiply_overflow_survives_batch_reduction( #[case] lane_is_valid: bool, ) -> VortexResult<()> { const LEN: u32 = 1000; 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 9855e74b0bc..31179b029cb 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 @@ -1,13 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexError; use vortex_error::VortexResult; +use vortex_error::vortex_panic; +use vortex_mask::Mask; use super::super::RowFnExecutionArgs; use super::super::args::BorrowedRowFnArgs; use crate::ArrayRef; use crate::ExecutionCtx; use crate::builtins::ArrayBuiltins; +use crate::scalar_fn::unstable::row::execute::DenseAttempt; use crate::validity::Validity; impl RowFnExecutionArgs { @@ -18,6 +22,82 @@ impl RowFnExecutionArgs { ctx: &mut ExecutionCtx, ) -> VortexResult { let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?; + + self.finalize_dense_output(values, ctx) + } + + /// Run all stored payloads, retrying valid rows only when reduced failure evidence is rejected. + /// + /// Validity stays lazy on success because masking makes null-row values unobservable. Rejected + /// evidence loses which row failed, so validity decides whether to return, suppress, or + /// recompute the error. + pub(super) fn execute_dense_with_retry( + &self, + execute_dense_attempt: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut ExecutionCtx, + ) -> VortexResult, + try_valid_rows: impl FnOnce( + BorrowedRowFnArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let attempt = + execute_dense_attempt(self.execution_args(&self.inputs, self.row_count), ctx)?; + + match attempt { + DenseAttempt::Values(values) => self.finalize_dense_output(values, ctx), + DenseAttempt::DeferredError(error) => { + self.resolve_deferred_error(error, try_valid_rows, ctx) + } + } + } + + fn resolve_deferred_error( + &self, + deferred_error: VortexError, + try_valid_rows: impl FnOnce( + BorrowedRowFnArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid_rows = self.validity.execute_mask(self.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. Check all-true before all-false because an + // empty mask is both. + if valid_rows.all_true() { + return Err(deferred_error); + } + + if valid_rows.all_false() { + return Ok(self.all_null()); + } + + // Reduced evidence does not identify which rows failed. Discard the dense error before + // retrying only observable rows. + drop(deferred_error); + + if let Some(result) = self.try_execute_valid_rows(try_valid_rows, &valid_rows, ctx)? { + return Ok(result); + } + + vortex_panic!( + "dense retry requires direct valid-row support after a deferred error, but {} declined it", + self.id, + ) + } + + fn finalize_dense_output( + &self, + values: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { let values = self.validate_kernel_output(values, self.row_count, ctx)?; match self.validity.clone() { 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 d80091844c7..ca0d27b599f 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 @@ -15,6 +15,7 @@ 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; @@ -33,6 +34,10 @@ impl RowFnExecutionArgs { pub(crate) fn execute( &self, kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + execute_dense_attempt: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut ExecutionCtx, + ) -> VortexResult, try_valid_rows: impl FnOnce( BorrowedRowFnArgs<'_>, &Mask, @@ -70,6 +75,9 @@ impl RowFnExecutionArgs { match self.policy { RowPolicy::Dense => self.execute_dense(kernel, ctx), + RowPolicy::DenseWithRetry => { + self.execute_dense_with_retry(execute_dense_attempt, try_valid_rows, ctx) + } RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_valid_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 c614e04dc3e..26fc62fb178 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 @@ -82,7 +82,7 @@ impl RowFnExecutionArgs { } /// Try execution against the original inputs, then mask a returned full-length result. - fn try_execute_valid_rows( + pub(super) fn try_execute_valid_rows( &self, try_valid_rows: impl FnOnce( BorrowedRowFnArgs<'_>, 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 dfa83cc7017..78dcb554a07 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use vortex_buffer::BufferMut; use vortex_error::VortexResult; @@ -30,8 +32,17 @@ use crate::scalar_fn::unstable::row::RowVisitor; use crate::scalar_fn::unstable::row::execute_rows; use crate::validity::Validity; -#[derive(Clone)] -struct DeferredAdd; +#[derive(Clone, Default)] +struct DeferredAdd { + /// The number of preparations across the dense attempt and any valid-row retry. + prepare_count: Arc, +} + +impl DeferredAdd { + fn prepare_count(&self) -> usize { + self.prepare_count.load(Ordering::Relaxed) + } +} #[derive(Clone)] struct ValidOnlyIdentity; @@ -104,8 +115,13 @@ impl RowFn for DeferredAdd { _args: &[DType], visitor: V, ) -> VortexResult { - visitor.visit_deferred::<(i64, i64), i64, bool>( - |(lhs, rhs)| lhs.overflowing_add(rhs), + let prepare_count = Arc::clone(&self.prepare_count); + + visitor.visit_prepared_deferred::<(i64, i64), i64, (), bool>( + move |_| { + prepare_count.fetch_add(1, Ordering::Relaxed); + }, + |&(), (lhs, rhs)| lhs.overflowing_add(rhs), |overflowed| { if overflowed { vortex_bail!(InvalidArgument: "deferred addition overflowed"); @@ -209,17 +225,78 @@ fn test_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { } #[test] -fn test_deferred_owned_execution_skips_invalid_rows() -> VortexResult<()> { +fn test_deferred_owned_execution_retries_null_row_failure() -> VortexResult<()> { + let function = DeferredAdd::default(); let validity = Validity::from_iter([true, false]); let lhs = PrimitiveArray::new(vec![1_i64, i64::MAX], validity.clone()).into_array(); let rhs = ConstantArray::new(1_i64, 2).into_array(); let args = VecExecutionArgs::new(vec![lhs, rhs], 2); let mut ctx = array_session().create_execution_ctx(); - let actual = execute_rows(&DeferredAdd, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::new(vec![2_i64, 0], validity).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(function.prepare_count(), 2); + Ok(()) +} + +#[test] +fn test_deferred_owned_execution_does_not_retry_partially_valid_success() -> VortexResult<()> { + let function = DeferredAdd::default(); + let validity = Validity::from_iter([true, false]); + let lhs = PrimitiveArray::new(vec![1_i64, 2], validity.clone()).into_array(); + let rhs = ConstantArray::new(1_i64, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; let expected = PrimitiveArray::new(vec![2_i64, 0], validity).into_array(); assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(function.prepare_count(), 1); + Ok(()) +} + +#[test] +fn test_deferred_owned_execution_retries_and_reports_valid_row_failure() -> VortexResult<()> { + let function = DeferredAdd::default(); + let validity = Validity::from_iter([true, false]); + let lhs = PrimitiveArray::new(vec![i64::MAX, 1], validity).into_array(); + let rhs = ConstantArray::new(1_i64, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = execute_rows(&function, &EmptyOptions, &args, &mut ctx) + .expect_err("a valid-row overflow must remain observable"); + let error = error.to_string(); + + assert!( + error.contains("deferred addition overflowed"), + "the valid-row retry must report its deferred error, got {error}", + ); + assert_eq!(function.prepare_count(), 2); + Ok(()) +} + +#[test] +fn test_deferred_owned_array_backed_all_valid_error_does_not_retry() -> VortexResult<()> { + let function = DeferredAdd::default(); + let validity = Validity::Array(ConstantArray::new(true, 2).into_array()); + let lhs = PrimitiveArray::new(vec![i64::MAX, 1], validity).into_array(); + let rhs = ConstantArray::new(1_i64, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = execute_rows(&function, &EmptyOptions, &args, &mut ctx) + .expect_err("an all-valid deferred error must remain observable"); + let error = error.to_string(); + + assert!( + error.contains("deferred addition overflowed"), + "the dense attempt must return its original deferred error, got {error}", + ); + assert_eq!(function.prepare_count(), 1); Ok(()) } 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 8adecd4206d..3ed0b4801f4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -12,6 +12,10 @@ pub(super) use owned::execute_owned_infallible; pub(super) use owned::execute_owned_infallible_valid_rows; pub(super) use owned::execute_owned_valid_rows; +mod retry; +pub(super) use retry::DenseAttempt; +pub(super) use retry::execute_owned_dense_attempt; + mod sink; pub(super) use sink::execute_sink; pub(super) use sink::execute_sink_valid_rows; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs b/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs new file mode 100644 index 00000000000..13cc8734f56 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs @@ -0,0 +1,112 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Dense owned execution that can request a valid-row retry. +//! +//! [`execute_owned_dense_attempt`] decodes and evaluates every row while reducing compact failure +//! evidence. Accepted evidence returns dense values. Rejected evidence discards those values and +//! preserves the error while batch execution resolves input validity. Decode, validation, and +//! allocation errors remain terminal. + +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ExecutionCtx; +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::OutputElement; +use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; + +/// The result of attempting dense execution for a deferred row kernel. +pub(in crate::scalar_fn::unstable::row) enum DenseAttempt { + /// Dense values whose reduced failure evidence was accepted. + /// + /// Batch execution must still validate these values and attach input validity. + Values(ArrayRef), + + /// An error produced from the reduced failure evidence. + /// + /// Batch execution must resolve input validity before deciding whether this error is + /// observable. + DeferredError(VortexError), +} + +/// Decode every input column and report rejected failure evidence to the batch executor. +/// +/// Returns [`DenseAttempt::DeferredError`] only when `finish_failure` rejects the reduced +/// evidence. Decode, validation, and allocation errors remain ordinary [`VortexResult`] errors. +pub(in crate::scalar_fn::unstable::row) fn execute_owned_dense_attempt( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + 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, +{ + // The output vector stays at length zero until every slot is initialized so that an unwind + // abandons partially initialized spare capacity. This no-drop assertion proves that no + // initialized value requires a destructor to run. + const { assert_owned_output_needs_no_drop::() }; + + // Keep this dense row loop separate from `execute_owned`. Factoring their shared state into a + // helper changes LLVM's optimized dense kernel even when the helper is inlined. + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::const_values(&columns)); + + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let output = &mut values.spare_capacity_mut()[..row_count]; + + let failure_evidence = if let Some(views) = Args::views_if_no_consts(&columns) { + // Keep this validation beside the views so LLVM sees their common length here. + vortex_ensure!( + Args::view_lens_match(&views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // SAFETY: `view_lens_match` checked that these exact retained views address `row_count` + // rows. + let source = unsafe { Args::indexed_source(views, row_count) }; + + source.map_checked_into(output, |elements| apply(&prepared, elements)) + } else { + // Keep this proof branch-local. Shared validation prevents LLVM from specializing this + // loop for each batch-constant arrangement, leaving it scalar under multiple CGUs without + // LTO. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + let mut accumulated_failure = Fail::default(); + + // Iterate over `output` directly. A `0..row_count` range reuses the address-taken value + // from the validation error formatter and retains an output bounds check. + for (index, slot) in output.iter_mut().enumerate() { + // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the loop. + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + + slot.write(value); + accumulated_failure |= row_failure; + } + + accumulated_failure + }; + + // SAFETY: normal completion of either execution path initializes `0..row_count` exactly + // once, and `values` was allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + match finish_failure(failure_evidence) { + Ok(()) => Ok(DenseAttempt::Values(Out::build(values))), + Err(error) => Ok(DenseAttempt::DeferredError(error)), + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index 5551684dd4b..bd73916c85c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -17,7 +17,9 @@ //! only for valid rows without changing row positions. //! //! Prepared visits move work derived from constant operands outside the hot loop. Deferred visits -//! reduce compact failure evidence without constructing errors in that loop. +//! reduce compact failure evidence without constructing errors in that loop. Eligible kernels +//! first evaluate all payloads. If the reduced evidence reports an error for a partially valid +//! batch, execution retries only valid rows. mod execute; 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 0b7a8078cfe..148f784bcba 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 @@ -54,8 +54,9 @@ pub unsafe trait InputElement: 'static { /// Decode `array` into its column representation. /// - /// Called once per row-kernel invocation. Hoist dtype checks, downcasts, and other - /// invocation-invariant work into this method. + /// Called once per row-kernel invocation. Retrying a partially valid batch after a dense + /// deferred error starts a second invocation over valid rows. Hoist dtype checks, downcasts, + /// and other invocation-invariant work into this method. fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; /// Whether [`decode_null_tolerant`](Self::decode_null_tolerant) can decode this array. 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 7016aedaf23..493394ae982 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -8,7 +8,6 @@ //! signature cannot execute over the original inputs. use vortex_error::VortexResult; -use vortex_error::vortex_ensure_eq; use vortex_mask::Mask; use super::RowPolicy; @@ -18,6 +17,7 @@ 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::ensure_plan; use super::row_visitor::private; use crate::ArrayRef; use crate::ExecutionCtx; @@ -291,23 +291,3 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { ) } } - -fn ensure_plan( - planned_output: &DType, - planned_policy: RowPolicy, - actual_output: DType, - actual_policy: RowPolicy, -) -> VortexResult<()> { - vortex_ensure_eq!( - actual_policy, - planned_policy, - "row dispatch must select the planned nullable execution policy: planned {planned_policy:?}, got {actual_policy:?}", - ); - vortex_ensure_eq!( - actual_output, - *planned_output, - "row dispatch must select the planned output dtype: planned {planned_output}, got {actual_output}", - ); - - Ok(()) -} 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 c7f9baf6a62..8c0bfea0027 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -5,6 +5,11 @@ //! //! [`RowFn::dispatch`]: crate::scalar_fn::unstable::row::RowFn::dispatch +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use crate::dtype::DType; + mod check; pub(super) use check::assert_owned_output_needs_no_drop; @@ -12,6 +17,9 @@ mod execute; pub(super) use execute::ExecuteRows; pub(super) use execute::ExecuteValidRows; +mod retry; +pub(super) use retry::ExecuteDenseWithRetry; + mod plan; pub(super) use plan::BatchPlan; pub(super) use plan::BatchPlanner; @@ -19,3 +27,23 @@ pub(super) use plan::RowPolicy; mod row_visitor; pub use row_visitor::RowVisitor; + +fn ensure_plan( + planned_output: &DType, + planned_policy: RowPolicy, + actual_output: DType, + actual_policy: RowPolicy, +) -> VortexResult<()> { + vortex_ensure_eq!( + actual_policy, + planned_policy, + "row dispatch must select the planned nullable execution policy: planned {planned_policy:?}, got {actual_policy:?}", + ); + vortex_ensure_eq!( + actual_output, + *planned_output, + "row dispatch must select the planned output dtype: planned {planned_output}, got {actual_output}", + ); + + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs index 75d5852d83f..9299c1d59fc 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -133,6 +133,10 @@ pub(crate) enum RowPolicy { /// Evaluate all rows and mask the result. Dense, + /// Evaluate all rows, then retry a partially valid batch if reduced failure evidence reports an + /// error. + DenseWithRetry, + /// Execute only valid rows over the original inputs. ValidOnly, } @@ -149,8 +153,11 @@ impl RowPolicy { /// The policy for an owned output carrying batch-deferred failure evidence. pub(crate) const fn for_deferred_output() -> Self { - let _ = PhantomData::; - Self::ValidOnly + if Args::DENSE_SAFE && Args::DECODE_INFALLIBLE { + Self::DenseWithRetry + } else { + Self::ValidOnly + } } /// The policy for a sink-writing output. diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs new file mode 100644 index 00000000000..e7925dcbe62 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Executes the dense attempt for deferred kernels that can retry valid rows. +//! +//! [`ExecuteDenseWithRetry`] replays the dispatch selected by the planning visitor. It returns a +//! dense attempt containing either unvalidated values or an error that batch execution resolves +//! against input validity. + +use vortex_error::VortexResult; + +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::ensure_plan; +use super::row_visitor::private; +use crate::ExecutionCtx; +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::BorrowedRowFnArgs; +use crate::scalar_fn::unstable::row::execute::DenseAttempt; +use crate::scalar_fn::unstable::row::execute::execute_owned_dense_attempt; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible; +use crate::scalar_fn::unstable::row::execute::execute_sink; + +/// The runtime visit for a planned [`RowPolicy::DenseWithRetry`] dispatch. +pub(in crate::scalar_fn::unstable::row) struct ExecuteDenseWithRetry< + 'visit, + 'inputs, + 'ctx, + F: RowFn, +> { + /// The inputs and planning metadata for this dense attempt. + args: &'visit BorrowedRowFnArgs<'inputs>, + + /// The function options used to derive a sink's runtime dtype. + options: &'visit F::Options, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, +} + +impl<'visit, 'inputs, 'ctx, F: RowFn> ExecuteDenseWithRetry<'visit, 'inputs, 'ctx, F> { + pub(in crate::scalar_fn::unstable::row) fn new( + args: &'visit BorrowedRowFnArgs<'inputs>, + options: &'visit F::Options, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { args, options, ctx } + } +} + +impl private::Sealed for ExecuteDenseWithRetry<'_, '_, '_, F> {} + +impl RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> { + type VisitResult = DenseAttempt; + + 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::() }; + ensure_plan( + self.args.output_dtype(), + self.args.policy(), + validate_owned_visit::(self.args.dtypes())?, + RowPolicy::for_owned_output::(), + )?; + + execute_owned_infallible::(self.args, self.ctx, prepare, apply) + .map(DenseAttempt::Values) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_plan( + self.args.output_dtype(), + self.args.policy(), + validate_sink_visit::(self.options, self.args.dtypes())?, + RowPolicy::for_sink::(), + )?; + + execute_sink::( + self.args, self.ctx, prepare, apply, + ) + .map(DenseAttempt::Values) + } + + 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::() }; + ensure_plan( + self.args.output_dtype(), + self.args.policy(), + validate_owned_visit::(self.args.dtypes())?, + RowPolicy::for_deferred_output::(), + )?; + + execute_owned_dense_attempt::( + self.args, + self.ctx, + prepare, + apply, + finish_failure, + ) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index 3add5e918d1..23a0d602d79 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -73,8 +73,8 @@ pub trait RowVisitor: private::Sealed + Sized { /// # Examples /// /// Test whether each string occurs in its allowed-values list. The prepare closure builds one - /// lookup table for a batch-constant list. The row closure scans the current list from the input - /// column directly. + /// lookup table for a batch-constant list. The row closure scans the current list from the + /// input column directly. /// /// ```ignore /// visitor.visit_prepared::< @@ -199,9 +199,14 @@ pub trait RowVisitor: private::Sealed + Sized { /// The executor OR-reduces [`FailureEvidence`] across rows and passes the result to /// `finish_failure`. /// - /// [`RowFn::INFALLIBLE`](crate::scalar_fn::unstable::row::RowFn::INFALLIBLE) **must** be `false`. - /// `Out` must not require drop glue. `Fail` must be no wider than `Out`, or failure tracking - /// reduces the vector width. The framework checks these requirements. + /// If `finish_failure` rejects dense evidence, reduction has lost which row failed. Batch + /// execution therefore resolves input validity: it preserves the error for all-valid input, + /// suppresses it for all-null input, and retries `apply` over valid rows for partially valid + /// input. A prepared retry also runs `prepare` again. + /// + /// [`RowFn::INFALLIBLE`](crate::scalar_fn::unstable::row::RowFn::INFALLIBLE) **must** be + /// `false`. `Out` must not require drop glue. `Fail` must be no wider than `Out`, or failure + /// tracking reduces the vector width. The framework checks these requirements. /// /// # Examples /// diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index bb0d7815e22..c342780653a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -18,6 +18,7 @@ use super::batch::RowFnExecutionArgs; use super::batch::finalize_kernel_output; use super::row_fn::RowFn; use super::visitor::BatchPlanner; +use super::visitor::ExecuteDenseWithRetry; use super::visitor::ExecuteRows; use super::visitor::ExecuteValidRows; use crate::ArrayRef; @@ -30,6 +31,7 @@ use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::unstable::row::execute::DenseAttempt; impl ScalarFnVTable for F { type Options = F::Options; @@ -122,6 +124,7 @@ pub fn execute_rows( let batch = prepare_batch(function, options, args)?; batch.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), ctx, ) @@ -175,6 +178,19 @@ fn execute_row_kernel( ) } +fn execute_dense_attempt( + function: &F, + options: &F::Options, + args: BorrowedRowFnArgs<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + function.dispatch( + options, + args.dtypes(), + ExecuteDenseWithRetry::::new(&args, options, ctx), + ) +} + fn try_execute_valid_rows( function: &F, options: &F::Options, @@ -382,7 +398,7 @@ mod tests { "unexpected error: {error}", ); assert!( - message.contains("planned Dense, got ValidOnly"), + message.contains("planned Dense, got DenseWithRetry"), "unexpected error: {error}", ); Ok(())