Compare primitive values with RowFn

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-array/benches/compare.rs

@@ -3,12 +3,11 @@
 
 //! Benchmarks for the binary comparison path, over every array kind it accepts.
 //!
-//! The primitive cases carry `#[cpu_features]`, so they are measured on every walltime
-//! CPU-feature leg rather than in simulation. Each is written once and compiled differently
-//! per leg: today the primitive comparison path is a portable lane kernel, and how well it
-//! auto-vectorizes is decided by the build. That is the baseline a hand-written kernel
-//! selected through `cfg(target_feature)` has to beat, measured on the silicon it would run
-//! on.
+//! The primitive cases carry `#[cpu_features]`, so they are measured on every walltime CPU-feature
+//! leg rather than in simulation. They all exercise the same [`RowFn`] comparison path, including
+//! its runtime-selected packed Boolean collector.
+//!
+//! [`RowFn`]: vortex_array::scalar_fn::unstable::row::RowFn
 //!
 //! The boolean, decimal, string, and struct cases are not tagged. A wider vector register is
 //! not what decides them: booleans are already word-at-a-time over a bitmap, decimals are

vortex-array/src/scalar_fn/fns/binary/compare/mod.rs

@@ -4,9 +4,11 @@
 //! Native comparison kernels.
 //!
 //! [`execute_compare`] dispatches on the logical [`DType`] of its operands and evaluates every
-//! comparison directly over Vortex canonical arrays — bit buffers for booleans, lane kernels from
-//! `vortex-compute` for primitives and decimals, binary views for strings/bytes, and a row-wise
-//! comparator for nested types. There is no Arrow fallback.
+//! comparison directly over Vortex canonical arrays: bit buffers for booleans, [`RowFn`] for
+//! primitives, lane kernels from `vortex-compute` for decimals, binary views for strings/bytes, and
+//! a row-wise comparator for nested types. There is no Arrow fallback.
+//!
+//! [`RowFn`]: crate::scalar_fn::unstable::row::RowFn
 //!
 //! Floating point values compare with Vortex's total ordering (`NaN` is the largest value,
 //! `-0.0 < +0.0`, and equality is bitwise), matching [`Scalar`] comparison semantics.
@@ -211,7 +213,7 @@ fn compare_arrays(
         )
         .into_array()),
         DType::Bool(_) => boolean::compare_bool(lhs, rhs, op, nullability, ctx),
-        DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx),
+        DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, ctx),
         DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx),
         DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx),
         DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Map(..) => {

vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs

@@ -1,28 +1,28 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! Native comparison of primitive arrays via bit-packing lane kernels.
+//! Primitive comparison execution through [`RowFn`].
+//!
+//! [`PrimitiveCompare`] delegates decoding, constant handling, validity, and packed Boolean output
+//! to the row executor. Its row kernel contains only the comparison selected by [`CompareOperator`].
 
-use vortex_buffer::BitBuffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
 use crate::ArrayRef;
 use crate::ExecutionCtx;
-use crate::IntoArray;
-use crate::arrays::BoolArray;
-use crate::arrays::ConstantArray;
 use crate::dtype::DType;
 use crate::dtype::NativePType;
-use crate::dtype::Nullability;
 use crate::dtype::PType;
 use crate::match_each_native_ptype;
-use crate::scalar::Scalar;
-use crate::scalar_fn::fns::binary::compare::collect_bits;
-use crate::scalar_fn::fns::binary::compare::collect_zip_bits;
-use crate::scalar_fn::fns::binary::compare::compare_validity;
-use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand;
+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::fns::operators::CompareOperator;
+use crate::scalar_fn::unstable::row::RowFn;
+use crate::scalar_fn::unstable::row::RowVisitor;
+use crate::scalar_fn::unstable::row::execute_rows;
 
 /// Compare two primitive arrays of the same [`PType`].
 ///
@@ -32,99 +32,71 @@ pub(super) fn compare_primitive(
     lhs: &ArrayRef,
     rhs: &ArrayRef,
     op: CompareOperator,
-    nullability: Nullability,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
-    let ptype = PType::try_from(lhs.dtype())?;
-    match_each_native_ptype!(ptype, |T| {
-        compare_primitive_typed::<T>(lhs, rhs, op, nullability, ctx)
-    })
+    let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len());
+
+    execute_rows(&PrimitiveCompare, &op, &args, ctx)
 }
 
-fn compare_primitive_typed<T: NativePType>(
-    lhs: &ArrayRef,
-    rhs: &ArrayRef,
-    op: CompareOperator,
-    nullability: Nullability,
-    ctx: &mut ExecutionCtx,
-) -> VortexResult<ArrayRef> {
-    let len = lhs.len();
-    let lhs = PrimitiveOperand::<T>::try_new(lhs, ctx)?;
-    let rhs = PrimitiveOperand::<T>::try_new(rhs, ctx)?;
-    if lhs.len() != rhs.len() {
-        vortex_bail!(
-            "compare operator requires equal lengths, got {} and {}",
-            lhs.len(),
-            rhs.len()
-        );
-    }
+/// Internal row execution for primitive comparison operators.
+#[derive(Clone)]
+struct PrimitiveCompare;
 
-    let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?;
+impl RowFn for PrimitiveCompare {
+    type Options = CompareOperator;
 
-    let bits = match (&lhs, &rhs) {
-        (
-            PrimitiveOperand::Array { values: lhs, .. },
-            PrimitiveOperand::Array { values: rhs, .. },
-        ) => compare_slices(lhs, rhs, op),
-        (
-            PrimitiveOperand::Array { values: lhs, .. },
-            PrimitiveOperand::Constant { value: rhs, .. },
-        ) => compare_slice_constant(lhs, *rhs, op),
-        (
-            PrimitiveOperand::Constant { value: lhs, .. },
-            PrimitiveOperand::Array { values: rhs, .. },
-        ) => compare_slice_constant(rhs, *lhs, op.swap()),
-        (
-            PrimitiveOperand::Constant { value: lhs, .. },
-            PrimitiveOperand::Constant { value: rhs, .. },
-        ) => {
-            // Unreachable through `execute_compare` (constant-constant is folded there), but
-            // cheap to answer anyway.
-            BitBuffer::full(apply_op(*lhs, *rhs, op), len)
-        }
-        (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => {
-            return Ok(
-                ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len)
-                    .into_array(),
-            );
-        }
-    };
+    const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"];
 
-    Ok(BoolArray::try_new(bits, validity)?.into_array())
-}
+    const INFALLIBLE: bool = true;
 
-#[inline(always)]
-fn apply_op<T: NativePType>(lhs: T, rhs: T, op: CompareOperator) -> bool {
-    match op {
-        CompareOperator::Eq => lhs.is_eq(rhs),
-        CompareOperator::NotEq => !lhs.is_eq(rhs),
-        CompareOperator::Gt => lhs.is_gt(rhs),
-        CompareOperator::Gte => lhs.is_ge(rhs),
-        CompareOperator::Lt => lhs.is_lt(rhs),
-        CompareOperator::Lte => lhs.is_le(rhs),
+    fn id(&self) -> ScalarFnId {
+        // `PrimitiveCompare` 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 compare_slices<T: NativePType>(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer {
-    // Dispatch the operator outside the lane loop so each instantiation vectorizes a single
-    // branch-free predicate.
-    match op {
-        CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a.is_eq(b)),
-        CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b)),
-        CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt),
-        CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge),
-        CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt),
-        CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le),
+    fn dispatch<V: RowVisitor>(
+        &self,
+        op: &Self::Options,
+        args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        let [lhs_dtype, _] = args else {
+            vortex_bail!(
+                "a primitive comparison requires two operands, got {}",
+                args.len(),
+            );
+        };
+        let ptype = PType::try_from(lhs_dtype)?;
+
+        match_each_native_ptype!(ptype, |T| { visit_compare::<T, V>(*op, visitor) })
     }
 }
 
-fn compare_slice_constant<T: NativePType>(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer {
+fn visit_compare<T, V>(op: CompareOperator, visitor: V) -> VortexResult<V::VisitResult>
+where
+    T: NativePType,
+    V: RowVisitor,
+{
     match op {
-        CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs)),
-        CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs)),
-        CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs)),
-        CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs)),
-        CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs)),
-        CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)),
+        CompareOperator::Eq => visit_compare_with::<T, V>(visitor, T::is_eq),
+        CompareOperator::NotEq => visit_compare_with::<T, V>(visitor, |lhs, rhs| !lhs.is_eq(rhs)),
+        CompareOperator::Gt => visit_compare_with::<T, V>(visitor, T::is_gt),
+        CompareOperator::Gte => visit_compare_with::<T, V>(visitor, T::is_ge),
+        CompareOperator::Lt => visit_compare_with::<T, V>(visitor, T::is_lt),
+        CompareOperator::Lte => visit_compare_with::<T, V>(visitor, T::is_le),
     }
 }
+
+fn visit_compare_with<T, V>(
+    visitor: V,
+    compare: impl Fn(T, T) -> bool,
+) -> VortexResult<V::VisitResult>
+where
+    T: NativePType,
+    V: RowVisitor,
+{
+    visitor.visit_bool::<(T, T), true>(move |(lhs, rhs)| compare(lhs, rhs))
+}

vortex-array/src/scalar_fn/fns/binary/mod.rs

@@ -45,7 +45,6 @@ mod compare;
 pub use compare::*;
 mod numeric;
 pub(crate) use numeric::*;
-mod primitive_operand;
 
 use crate::scalar::NumericOperator;
 use crate::scalar::Scalar;