diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 8e040ce1b91..82971cf5d02 100644 --- a/vortex-array/benches/compare.rs +++ b/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 //! //! Every case here compares one array against another, which is the shape the path is tuned for. //! The three constant cases that remain — boolean, integer, and string — are a regression guard on diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index a1635c718ff..5a3c37b4704 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/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. @@ -212,7 +214,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(..) => { diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index d7607953651..7e224a4df8a 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -1,32 +1,28 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native comparison of primitive arrays with specialized bitmap packing for 8-bit inputs. +//! 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_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; -use vortex_buffer::collect_bool_word; 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::bit_buffer_from_words; -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`]. /// @@ -36,178 +32,71 @@ pub(super) fn compare_primitive( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, - nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - match_each_native_ptype!(ptype, |T| { - compare_primitive_typed::(lhs, rhs, op, nullability, ctx) - }) -} - -fn compare_primitive_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: CompareOperator, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let len = lhs.len(); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - if lhs.len() != rhs.len() { - vortex_bail!( - "compare operator requires equal lengths, got {} and {}", - lhs.len(), - rhs.len() - ); - } + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); - let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + execute_rows(&PrimitiveCompare, &op, &args, ctx) +} - let bits = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slices(lhs, rhs, op, ctx.allocator()), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => compare_slice_constant(lhs, *rhs, op, ctx.allocator()), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slice_constant(rhs, *lhs, op.swap(), ctx.allocator()), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Unreachable through `execute_compare` (constant-constant is folded there), but - // cheap to answer anyway. - BitBuffer::full_in(apply_op(*lhs, *rhs, op), len, ctx.allocator().clone()) - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { - return Ok( - ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) - .into_array(), - ); - } - }; +/// Internal row execution for primitive comparison operators. +#[derive(Clone)] +struct PrimitiveCompare; - Ok(BoolArray::try_new(bits, validity)?.into_array()) -} +impl RowFn for PrimitiveCompare { + type Options = CompareOperator; -#[allow(clippy::inline_always)] -#[inline(always)] -fn apply_op(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), - } -} + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; -fn compare_slices( - lhs: &[T], - rhs: &[T], - op: CompareOperator, - allocator: &BufferAllocatorRef, -) -> BitBuffer { - // Dispatch the operator outside the lane loop so each instantiation vectorizes a single - // branch-free predicate. - match op { - CompareOperator::Eq => { - collect_zip_bits_dispatch(lhs, rhs, |a: T, b: T| a.is_eq(b), allocator) - } - CompareOperator::NotEq => { - collect_zip_bits_dispatch(lhs, rhs, |a: T, b: T| !a.is_eq(b), allocator) - } - CompareOperator::Gt => collect_zip_bits_dispatch(lhs, rhs, T::is_gt, allocator), - CompareOperator::Gte => collect_zip_bits_dispatch(lhs, rhs, T::is_ge, allocator), - CompareOperator::Lt => collect_zip_bits_dispatch(lhs, rhs, T::is_lt, allocator), - CompareOperator::Lte => collect_zip_bits_dispatch(lhs, rhs, T::is_le, allocator), - } -} + const INFALLIBLE: bool = true; -fn compare_slice_constant( - lhs: &[T], - rhs: T, - op: CompareOperator, - allocator: &BufferAllocatorRef, -) -> BitBuffer { - match op { - CompareOperator::Eq => collect_bits_dispatch(lhs, |a: T| a.is_eq(rhs), allocator), - CompareOperator::NotEq => collect_bits_dispatch(lhs, |a: T| !a.is_eq(rhs), allocator), - CompareOperator::Gt => collect_bits_dispatch(lhs, |a: T| a.is_gt(rhs), allocator), - CompareOperator::Gte => collect_bits_dispatch(lhs, |a: T| a.is_ge(rhs), allocator), - CompareOperator::Lt => collect_bits_dispatch(lhs, |a: T| a.is_lt(rhs), allocator), - CompareOperator::Lte => collect_bits_dispatch(lhs, |a: T| a.is_le(rhs), allocator), + 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 collect_bits_dispatch( - values: &[T], - f: impl Fn(T) -> bool, - allocator: &BufferAllocatorRef, -) -> BitBuffer { - // This type check folds away during monomorphization. Wider masks keep the lane kernel: - // byte packing regresses 64-bit comparisons on AVX2. - if matches!(T::PTYPE, PType::I8 | PType::U8) { - collect_bits_narrow(values, f, allocator) - } else { - collect_bits(values, f, allocator) - } -} + fn dispatch( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let [lhs_dtype, _] = args else { + vortex_bail!( + "a primitive comparison requires two operands, got {}", + args.len(), + ); + }; + let ptype = PType::try_from(lhs_dtype)?; -fn collect_zip_bits_dispatch( - lhs: &[T], - rhs: &[T], - f: impl Fn(T, T) -> bool, - allocator: &BufferAllocatorRef, -) -> BitBuffer { - if matches!(T::PTYPE, PType::I8 | PType::U8) { - collect_zip_bits_narrow(lhs, rhs, f, allocator) - } else { - collect_zip_bits(lhs, rhs, f, allocator) + match_each_native_ptype!(ptype, |T| { visit_compare::(*op, visitor) }) } } -fn collect_bits_narrow( - values: &[T], - f: impl Fn(T) -> bool, - allocator: &BufferAllocatorRef, -) -> BitBuffer { - let (chunks, tail) = values.as_chunks::<64>(); - let mut words = BufferMut::::zeroed_in(values.len().div_ceil(64), allocator.clone()); - // Fixed-size chunks let the compiler prove the predicate's indexing stays in bounds. - for (word, chunk) in words.iter_mut().zip(chunks) { - *word = collect_bool_word(64, |i| f(chunk[i])); - } - if !tail.is_empty() { - words[chunks.len()] = collect_bool_word(tail.len(), |i| f(tail[i])); +fn visit_compare(op: CompareOperator, visitor: V) -> VortexResult +where + T: NativePType, + V: RowVisitor, +{ + match op { + CompareOperator::Eq => visit_compare_with::(visitor, T::is_eq), + CompareOperator::NotEq => visit_compare_with::(visitor, |lhs, rhs| !lhs.is_eq(rhs)), + CompareOperator::Gt => visit_compare_with::(visitor, T::is_gt), + CompareOperator::Gte => visit_compare_with::(visitor, T::is_ge), + CompareOperator::Lt => visit_compare_with::(visitor, T::is_lt), + CompareOperator::Lte => visit_compare_with::(visitor, T::is_le), } - bit_buffer_from_words(words, values.len()) } -fn collect_zip_bits_narrow( - lhs: &[T], - rhs: &[T], - f: impl Fn(T, T) -> bool, - allocator: &BufferAllocatorRef, -) -> BitBuffer { - assert_eq!(lhs.len(), rhs.len()); - let (left_chunks, left_tail) = lhs.as_chunks::<64>(); - let (right_chunks, right_tail) = rhs.as_chunks::<64>(); - let mut words = BufferMut::::zeroed_in(lhs.len().div_ceil(64), allocator.clone()); - for ((word, left), right) in words.iter_mut().zip(left_chunks).zip(right_chunks) { - *word = collect_bool_word(64, |i| f(left[i], right[i])); - } - if !left_tail.is_empty() { - words[left_chunks.len()] = - collect_bool_word(left_tail.len(), |i| f(left_tail[i], right_tail[i])); - } - bit_buffer_from_words(words, lhs.len()) +fn visit_compare_with( + visitor: V, + compare: impl Fn(T, T) -> bool, +) -> VortexResult +where + T: NativePType, + V: RowVisitor, +{ + visitor.visit_bool::<(T, T), true>(move |(lhs, rhs)| compare(lhs, rhs)) } diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index e790bc4d18f..6ca84fa2e9c 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/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; diff --git a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs b/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs deleted file mode 100644 index 71d1122fc79..00000000000 --- a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs +++ /dev/null @@ -1,69 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Decoding shared by primitive binary operators. - -use vortex_buffer::Buffer; -use vortex_error::VortexResult; - -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrays::Constant; -use crate::arrays::PrimitiveArray; -use crate::dtype::NativePType; -use crate::validity::Validity; - -/// A materialized primitive column, a non-null constant, or an all-null constant. -pub(super) enum PrimitiveOperand { - Array { - values: Buffer, - validity: Validity, - }, - Constant { - value: T, - len: usize, - validity: Validity, - }, - Null(usize), -} - -impl PrimitiveOperand { - pub(super) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - if let Some(constant) = array.as_opt::() { - return Ok( - match constant.scalar().as_primitive().try_typed_value::()? { - Some(value) => Self::Constant { - value, - len: array.len(), - validity: if constant.scalar().dtype().is_nullable() { - Validity::AllValid - } else { - Validity::NonNullable - }, - }, - None => Self::Null(array.len()), - }, - ); - } - - let array = array.clone().execute::(ctx)?; - let validity = array.validity()?; - let values = array.into_buffer::(); - Ok(Self::Array { values, validity }) - } - - pub(super) fn len(&self) -> usize { - match self { - Self::Array { values, .. } => values.len(), - Self::Constant { len, .. } | Self::Null(len) => *len, - } - } - - pub(super) fn validity(&self) -> Validity { - match self { - Self::Array { validity, .. } => validity.clone(), - Self::Constant { validity, .. } => validity.clone(), - Self::Null(_) => Validity::AllInvalid, - } - } -}