diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index ddf6343e9ab..99ad1c02092 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 6593af28c6e..91f4eb77603 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. @@ -209,7 +211,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 329af3c61d6..a491b301a5c 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -1,239 +1,160 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native comparison of primitive arrays with packed bitmap output. - -use vortex_buffer::BitBuffer; -use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; -use vortex_buffer::collect_bool_word; -use vortex_compute::lane_kernels::bit_buffer_from_words; -use vortex_compute::lane_kernels::collect_bits; -use vortex_compute::lane_kernels::collect_zip_bits; +//! 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_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::array::VTable; +use crate::arrays::Constant; 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::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::fns::binary::Binary; use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowInputs; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_row_inputs; /// Compare two primitive arrays of the same [`PType`]. /// -/// Floats use a total ordering over their bit patterns, including NaN signs and payloads. -/// Equality is bitwise, so signed zeros and distinct NaN payloads remain distinct. +/// Floats compare with Vortex's total ordering over their bit patterns: a `NaN` with the sign bit +/// set sorts below `-inf`, any other `NaN` sorts above `+inf`, and `-0.0 < +0.0`. Equality is +/// bitwise, so distinct `NaN` payloads are not equal. 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() - ); + // The row loop vectorizes worse with a constant first operand, so keep the constant on the + // right. Comparing encodings avoids two dynamic downcasts on every call. + let constant = VTable::id(&Constant); + if lhs.encoding_id() == constant && rhs.encoding_id() != constant { + let args = CompareArgs::new(rhs, lhs, ptype); + return execute_row_inputs(&PrimitiveCompare, &op.swap(), &args, ctx); } - let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; - - 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(), - ); - } - }; - - Ok(BoolArray::try_new(bits, validity)?.into_array()) + execute_row_inputs( + &PrimitiveCompare, + &op, + &CompareArgs::new(lhs, rhs, ptype), + ctx, + ) } -#[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), - } +/// Borrowed comparison operands with their dtypes, so a call does not allocate. +struct CompareArgs<'a> { + arrays: [&'a ArrayRef; 2], + dtypes: [DType; 2], } -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) +impl<'a> CompareArgs<'a> { + /// Pair two primitive operands of the same [`PType`] with their dtypes. + /// + /// The dtypes are rebuilt as primitive values rather than cloned, so the compiler knows their + /// variant and drops them without a runtime match. + #[expect(clippy::inline_always)] + #[inline(always)] + fn new(lhs: &'a ArrayRef, rhs: &'a ArrayRef, ptype: PType) -> Self { + Self { + arrays: [lhs, rhs], + dtypes: [ + DType::Primitive(ptype, lhs.dtype().nullability()), + DType::Primitive(ptype, rhs.dtype().nullability()), + ], } - 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), } } -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), +impl RowInputs for CompareArgs<'_> { + fn num_inputs(&self) -> usize { + 2 } -} -fn collect_bits_dispatch( - values: &[T], - f: impl Fn(T) -> bool, - allocator: &BufferAllocatorRef, -) -> BitBuffer { - // This type check folds away during monomorphization. Word packing improves float32 - // comparisons; wider types retain the lane kernel to avoid AVX2 constant-case regressions. - if matches!(T::PTYPE, PType::I8 | PType::U8 | PType::F32) { - collect_bits_narrow(values, f, allocator) - } else { - collect_bits(values, f, allocator) + fn row_count(&self) -> usize { + self.arrays[0].len() } -} -fn collect_zip_bits_dispatch( - lhs: &[T], - rhs: &[T], - f: impl Fn(T, T) -> bool, - allocator: &BufferAllocatorRef, -) -> BitBuffer { - #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] - if matches!(T::PTYPE, PType::I64 | PType::U64 | PType::F64) { - // Fill-and-pack improves wide array pairs on NEON. Constants keep their existing loop. - return collect_zip_bits_wide(lhs, rhs, f, allocator); + fn input(&self, index: usize) -> &ArrayRef { + self.arrays[index] } - if matches!(T::PTYPE, PType::I8 | PType::U8 | PType::F32) { - collect_zip_bits_narrow(lhs, rhs, f, allocator) - } else { - collect_zip_bits(lhs, rhs, f, allocator) + fn dtypes(&self) -> Option<&[DType]> { + Some(&self.dtypes) } } -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])); +/// Internal row execution for primitive comparison operators. +#[derive(Clone)] +struct PrimitiveCompare; + +impl RowFn for PrimitiveCompare { + type Options = CompareOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + const INFALLIBLE: bool = true; + + 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) } - if !tail.is_empty() { - words[chunks.len()] = collect_bool_word(tail.len(), |i| f(tail[i])); + + 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)?; + + match_each_native_ptype!(ptype, |T| { visit_compare::(*op, visitor) }) } - 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])); +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, lhs.len()) } -/// Keep the predicate free of bounds checks so NEON can fuse comparison and bitmap packing. -#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] -fn collect_zip_bits_wide( - lhs: &[T], - rhs: &[T], - f: impl Fn(T, T) -> bool, - allocator: &BufferAllocatorRef, -) -> BitBuffer { - assert_eq!(lhs.len(), rhs.len()); - BitBuffer::collect_bool_multiversioned_in( - lhs.len(), - |index| { - // SAFETY: the collector calls with `index < lhs.len()`, and the slice lengths match. - let (lhs, rhs) = unsafe { (*lhs.get_unchecked(index), *rhs.get_unchecked(index)) }; - f(lhs, rhs) - }, - allocator.clone(), - ) +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/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs index a555466519e..06959e6903a 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -447,6 +447,36 @@ fn int_constant_lhs_and_rhs(#[case] op: Operator, #[case] expected: [bool; 4]) { assert_arrays_eq!(swapped, BoolArray::from_iter(expected_swapped), &mut ctx); } +/// A constant left operand is moved to the right with the swapped operator. The 65 rows span one +/// full bitmap word and a partial word. +#[rstest] +#[case(Operator::Eq)] +#[case(Operator::NotEq)] +#[case(Operator::Lt)] +#[case(Operator::Lte)] +#[case(Operator::Gt)] +#[case(Operator::Gte)] +fn int_constant_lhs_all_operators(#[case] op: Operator) { + let mut ctx = SESSION.create_execution_ctx(); + let values: Vec = (0..65).collect(); + let array = PrimitiveArray::from_iter(values.iter().copied()).into_array(); + let constant = ConstantArray::new(32i64, values.len()).into_array(); + + let compare = |lhs: i64, rhs: i64| match op { + Operator::Eq => lhs == rhs, + Operator::NotEq => lhs != rhs, + Operator::Lt => lhs < rhs, + Operator::Lte => lhs <= rhs, + Operator::Gt => lhs > rhs, + Operator::Gte => lhs >= rhs, + _ => unreachable!(), + }; + let expected = BoolArray::from_iter(values.iter().map(|&value| compare(32, value))); + + let result = execute_compare_test(constant, array, op); + assert_arrays_eq!(result, expected, &mut ctx); +} + /// Floats compare with Vortex's total ordering: NaN is the largest value, equality is bitwise, /// and -0.0 < +0.0. This matches `Scalar` comparison semantics. #[rstest] diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index ffea35a5bff..65d6f584b4d 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -44,7 +44,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, - } - } -} diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index f6f5479578c..3345a04c2bd 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -24,6 +24,7 @@ mod execute; mod batch; +pub(crate) use batch::RowInputs; mod row_fn; pub use row_fn::RowFn; @@ -52,5 +53,6 @@ mod visitor; pub use visitor::RowVisitor; mod vtable; +pub(crate) use vtable::execute_row_inputs; pub use vtable::execute_rows; pub use vtable::row_fn_return_dtype;