Execute primitive numeric operators with RowFn

vortex-array/benches/binary_ops.rs

@@ -38,10 +38,60 @@ static SESSION: LazyLock<VortexSession> = LazyLock::new(array_session);
 
 const LEN: usize = 32_768;
 
+const ROWFN_MATRIX_CASES: &[(usize, RowFnShape)] = &[
+    (128, RowFnShape::PerRowPerRow),
+    (128, RowFnShape::PerRowConstant),
+    (128, RowFnShape::ConstantPerRow),
+    (128, RowFnShape::PerRowNullableConstant),
+    (LEN, RowFnShape::PerRowPerRow),
+    (LEN, RowFnShape::PerRowConstant),
+    (LEN, RowFnShape::ConstantPerRow),
+    (LEN, RowFnShape::PerRowNullableConstant),
+];
+
+#[derive(Clone, Copy, Debug)]
+enum RowFnShape {
+    PerRowPerRow,
+    PerRowConstant,
+    ConstantPerRow,
+    PerRowNullableConstant,
+}
+
 /// Decimal Mul and Div cost far more per lane than Add, so they run over a shorter array to keep
 /// the instrumented CodSpeed runs quick.
 const DECIMAL_MUL_DIV_LEN: usize = 8_192;
 
+#[divan::bench(args = ROWFN_MATRIX_CASES)]
+fn rowfn_add(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) {
+    bench_rowfn_shape(bencher, len, shape, Operator::Add);
+}
+
+#[divan::bench(args = ROWFN_MATRIX_CASES)]
+fn rowfn_subtract(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) {
+    bench_rowfn_shape(bencher, len, shape, Operator::Sub);
+}
+
+#[divan::bench(args = ROWFN_MATRIX_CASES)]
+fn rowfn_multiply(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) {
+    bench_rowfn_shape(bencher, len, shape, Operator::Mul);
+}
+
+fn bench_rowfn_shape(bencher: Bencher, len: usize, shape: RowFnShape, operator: Operator) {
+    let per_row =
+        || PrimitiveArray::from_iter((0..len).map(|index| (index % 1_024) as i64 + 1)).into_array();
+    let constant = || ConstantArray::new(17_i64, len).into_array();
+    let nullable_constant = || ConstantArray::new(Some(17_i64), len).into_array();
+
+    let (lhs, rhs) = match shape {
+        RowFnShape::PerRowPerRow => (per_row(), per_row()),
+        RowFnShape::PerRowConstant => (per_row(), constant()),
+        RowFnShape::ConstantPerRow => (constant(), per_row()),
+        RowFnShape::PerRowNullableConstant => (per_row(), nullable_constant()),
+    };
+
+    bench_primitive(bencher, lhs, rhs, operator);
+}
+
 #[divan::bench]
 fn add_i64_nonnull(bencher: Bencher) {
     let lhs = primitive_nonnull(0).into_array();
@@ -170,6 +220,14 @@ fn div_i64_nonnull(bencher: Bencher) {
     bench_primitive(bencher, lhs, rhs, Operator::Div);
 }
 
+#[divan::bench]
+fn div_i64_nullable(bencher: Bencher) {
+    let lhs = primitive_nullable(1_000_000, 7).into_array();
+    let rhs = primitive_nullable(17, 5).into_array();
+
+    bench_primitive(bencher, lhs, rhs, Operator::Div);
+}
+
 #[divan::bench]
 fn sub_i64_constant(bencher: Bencher) {
     let lhs = primitive_nonnull(0).into_array();

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.

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<S, T, Apply>(
     source: S,
@@ -61,7 +48,6 @@ where
     };
 
     let mut values = BufferMut::<T>::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<S, T, Fail, Apply>(
-    source: S,
-    valid_rows: &Mask,
-    mut apply: Apply,
-) -> Result<Buffer<T>, usize>
-where
-    S: IndexedSource + Copy,
-    T: Copy + Default,
-    Fail: Failure,
-    Apply: FnMut(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::<T>::with_capacity(len);
-
-    let out = &mut values.spare_capacity_mut()[..len];
-    if source.map_checked_into(out, &mut apply) != Fail::default() {
-        let mut checked = |item: S::Item| {
-            let (value, failure) = apply(item);
-            (failure == Fail::default()).then_some(value)
-        };
-        match valid_bits {
-            None => source.try_map_into(out, &mut checked)?,
-            Some(valid_bits) => source.try_map_masked_into(valid_bits, out, &mut checked)?,
-        }
-    }
-
-    // SAFETY: the kernels initialize every lane in `out`.
-    unsafe { values.set_len(len) };
-
-    Ok(values.freeze())
-}

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::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;