diff --git a/encodings/runend/src/trace_tests.rs b/encodings/runend/src/trace_tests.rs index 96f2afff50a..be8fd0ba762 100644 --- a/encodings/runend/src/trace_tests.rs +++ b/encodings/runend/src/trace_tests.rs @@ -73,6 +73,14 @@ fn trace_compare_on_runend() -> VortexResult<()> { iter 0 current=vortex.runend(bool, len=9) builder_active=false execute_until target=AnyCanonical root=vortex.binary(bool, len=3) iter 0 current=vortex.binary(bool, len=3) builder_active=false + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=3) iter 1 current=vortex.bool(bool, len=3) builder_active=false return output=vortex.bool(bool, len=3) diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 4a399760dc2..9e6dd3e4e5b 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -4,6 +4,7 @@ #![expect(clippy::unwrap_used)] use divan::Bencher; +use divan::counter::ItemsCount; use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; @@ -38,6 +39,7 @@ const ARRAY_SIZE: usize = 65_536; fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { let session = vortex_array::array_session(); bencher + .counter(ItemsCount::new(ARRAY_SIZE)) .with_inputs(|| (&lhs, &rhs, session.create_execution_ctx())) .bench_refs(|input| { input @@ -49,6 +51,31 @@ fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { }); } +fn u8_array(offset: u8) -> ArrayRef { + (0u8..=u8::MAX) + .cycle() + .take(ARRAY_SIZE) + .map(|value| value.wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn i32_array(offset: i32) -> ArrayRef { + (0i32..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn u64_array(offset: u64) -> ArrayRef { + (0u64..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + fn bool_array(rng: &mut StdRng) -> ArrayRef { BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.5))).into_array() } @@ -87,6 +114,13 @@ fn float_array(rng: &mut StdRng) -> ArrayRef { .into_array() } +fn f32_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random_range(0.0f32..1.0)) + .collect::>() + .into_array() +} + fn string_array(rng: &mut StdRng) -> ArrayRef { VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| { let len = rng.random_range(1usize..24); @@ -153,6 +187,14 @@ fn compare_int_constant(bencher: Bencher) { bench_compare(bencher, arr, constant, Operator::Gte); } +#[divan::bench] +fn compare_int_constant_lhs(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let constant = ConstantArray::new(50_000_000i64, ARRAY_SIZE).into_array(); + let arr = int_array(&mut rng); + bench_compare(bencher, constant, arr, Operator::Gte); +} + #[divan::bench] fn compare_int_eq(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -161,6 +203,55 @@ fn compare_int_eq(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Eq); } +#[divan::bench] +fn compare_i32(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = i32_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_i32_constant(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = ConstantArray::new(1_000_000i32, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = u8_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8_constant(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = ConstantArray::new(127u8, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_constant(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = ConstantArray::new(1_000_000u64, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_eq(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Eq); +} + #[divan::bench] fn compare_float(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -169,6 +260,30 @@ fn compare_float(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Gte); } +#[divan::bench] +fn compare_float_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = float_array(&mut rng); + let arr2 = float_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + +#[divan::bench] +fn compare_f32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_f32_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + #[divan::bench] fn compare_decimal(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); 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 d25a652ee57..a36d0a22bde 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,9 @@ //! 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, row or fused lane +//! kernels for primitives, lane kernels for decimals, binary views for strings and bytes, and a +//! row-wise comparator for nested types. There is no Arrow fallback. //! //! 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 +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 1247358dce1..93afcf538ed 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -1,28 +1,30 @@ // 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`]. + +mod columnar; +mod operand; -use vortex_buffer::BitBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_bail; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::ConstantArray; +#[cfg(target_arch = "x86_64")] +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::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 +34,125 @@ 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) - }) + compare_primitive_with_path(lhs, rhs, op, PrimitiveComparisonPath::Auto, ctx) +} + +/// Selects automatic production dispatch or a forced implementation in tests. +#[derive(Clone, Copy)] +pub(super) enum PrimitiveComparisonPath { + /// Use the architecture and operand-specific production policy. + Auto, + + /// Force row execution. + #[cfg(test)] + Row, + + /// Force fused columnar execution. + #[cfg(test)] + Columnar, } -fn compare_primitive_typed( +/// Compare primitives through the selected implementation. +pub(super) fn compare_primitive_with_path( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, - nullability: Nullability, + path: PrimitiveComparisonPath, 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 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), - ( - 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(), - ); + let use_columnar = match path { + PrimitiveComparisonPath::Auto => { + #[cfg(target_arch = "x86_64")] + { + use_columnar_comparison(lhs, rhs, op)? + } + #[cfg(not(target_arch = "x86_64"))] + { + false + } } + #[cfg(test)] + PrimitiveComparisonPath::Row => false, + #[cfg(test)] + PrimitiveComparisonPath::Columnar => true, }; + if use_columnar { + return columnar::compare_primitive(lhs, rhs, op, ctx); + } - Ok(BoolArray::try_new(bits, validity)?.into_array()) + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + execute_rows(&PrimitiveCompare, &op, &args, ctx) } -#[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), +/// 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 FALLIBLE: bool = false; + + 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(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>( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = + PType::try_from(args.first().ok_or_else(|| { + vortex_err!("a comparison operator takes two operands, got none") + })?)?; + + match_each_native_ptype!(ptype, |T| { visit_compare::(*op, visitor) }) } } -fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { +#[cfg(target_arch = "x86_64")] +fn use_columnar_comparison( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, +) -> VortexResult { + let ptype = PType::try_from(lhs.dtype())?; + Ok(match (ptype, op) { + // Equality bit-packs efficiently for every type supported by the columnar path. + (PType::I64 | PType::U64 | PType::F64, CompareOperator::Eq | CompareOperator::NotEq) => { + true + } + // The fused comparison and bit-packing loop produces better x86 code for signed 64-bit + // integers and f64. The RowFn byte-output loop remains faster for narrower lanes. + (PType::I64 | PType::F64, _) => true, + // LLVM 22 vectorizes the mixed-constant RowFn loop at 16 CGUs without LTO. However, the + // fused comparison and bit-packing path is still about 38% faster in + // `compare_u64_constant`. Recheck that benchmark before changing this dispatch. + (PType::U64, _) => lhs.is::() || rhs.is::(), + _ => false, + }) +} + +fn visit_compare(op: CompareOperator, visitor: V) -> VortexResult +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 => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_eq(rhs)), + CompareOperator::NotEq => visitor.visit::<(T, T), bool>(|(lhs, rhs)| !lhs.is_eq(rhs)), + CompareOperator::Gt => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_gt(rhs)), + CompareOperator::Gte => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_ge(rhs)), + CompareOperator::Lt => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_lt(rhs)), + CompareOperator::Lte => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_le(rhs)), } } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs new file mode 100644 index 00000000000..6728437e6a8 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fused comparison and bit-packing for wide primitive lanes. +//! +//! Production uses this implementation only for measured x86 paths. Keeping it portable lets the +//! semantic tests exercise the RowFn and fused paths on every target. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use super::operand::PrimitiveOperand; +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::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::operators::CompareOperator; + +/// Compare primitive operands with one fused comparison and bit-packing loop. +pub(super) fn compare_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match PType::try_from(lhs.dtype())? { + PType::I64 => compare_primitive_typed::(lhs, rhs, op, ctx), + PType::U64 => compare_primitive_typed::(lhs, rhs, op, ctx), + PType::F64 => compare_primitive_typed::(lhs, rhs, op, ctx), + ptype => vortex_bail!("columnar comparison is not selected for {ptype}"), + } +} + +fn compare_primitive_typed( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let nullability = Nullability::from(lhs.dtype().is_nullable() || rhs.dtype().is_nullable()); + 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 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), + ( + 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, .. }, + ) => 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(), + ); + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +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), + } +} + +fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| !lhs.is_eq(rhs)), + 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 compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_bits(lhs, |lhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_bits(lhs, |lhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_bits(lhs, |lhs: T| lhs.is_gt(rhs)), + CompareOperator::Gte => collect_bits(lhs, |lhs: T| lhs.is_ge(rhs)), + CompareOperator::Lt => collect_bits(lhs, |lhs: T| lhs.is_lt(rhs)), + CompareOperator::Lte => collect_bits(lhs, |lhs: T| lhs.is_le(rhs)), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs similarity index 78% rename from vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs rename to vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs index 71d1122fc79..1563b8e68e6 100644 --- a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Decoding shared by primitive binary operators. +//! Operand decoding for the fused primitive comparison path. use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -15,19 +15,33 @@ use crate::validity::Validity; /// A materialized primitive column, a non-null constant, or an all-null constant. pub(super) enum PrimitiveOperand { + /// A per-row primitive column and its validity. Array { + /// The materialized values. values: Buffer, + + /// The validity of the values. validity: Validity, }, + + /// A non-null value repeated for every row. Constant { + /// The repeated value. value: T, + + /// The number of repeated rows. len: usize, + + /// The validity implied by the constant's dtype. validity: Validity, }, + + /// An all-null constant with this row count. Null(usize), } impl PrimitiveOperand { + /// Decode an operand once for the fused comparison loop. pub(super) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { if let Some(constant) = array.as_opt::() { return Ok( @@ -49,9 +63,11 @@ impl PrimitiveOperand { let array = array.clone().execute::(ctx)?; let validity = array.validity()?; let values = array.into_buffer::(); + Ok(Self::Array { values, validity }) } + /// Return the logical row count. pub(super) fn len(&self) -> usize { match self { Self::Array { values, .. } => values.len(), @@ -59,6 +75,7 @@ impl PrimitiveOperand { } } + /// Return the operand validity. pub(super) fn validity(&self) -> Validity { match self { Self::Array { validity, .. } => validity.clone(), 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 9831a963354..d0f7a9b5e57 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -12,7 +12,9 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::VTable; use crate::array_session; +use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; @@ -21,6 +23,7 @@ use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::ScalarFn; use crate::arrays::StructArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; @@ -39,6 +42,8 @@ use crate::extension::datetime::Timestamp; use crate::extension::datetime::TimestampOptions; use crate::scalar::DecimalValue; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::primitive::PrimitiveComparisonPath; +use crate::scalar_fn::fns::binary::compare::primitive::compare_primitive_with_path; use crate::scalar_fn::fns::binary::scalar_cmp; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; @@ -429,6 +434,142 @@ fn float_total_order() { ); } +#[rstest] +#[case::row_eq(PrimitiveComparisonPath::Row, CompareOperator::Eq)] +#[case::row_not_eq(PrimitiveComparisonPath::Row, CompareOperator::NotEq)] +#[case::row_lt(PrimitiveComparisonPath::Row, CompareOperator::Lt)] +#[case::columnar_eq(PrimitiveComparisonPath::Columnar, CompareOperator::Eq)] +#[case::columnar_not_eq(PrimitiveComparisonPath::Columnar, CompareOperator::NotEq)] +#[case::columnar_lt(PrimitiveComparisonPath::Columnar, CompareOperator::Lt)] +fn test_primitive_comparison_paths_preserve_semantics_and_encoding( + #[case] path: PrimitiveComparisonPath, + #[case] op: CompareOperator, +) -> VortexResult<()> { + let lhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::NAN, // Null on the left. + -0.0, // Signed zero ordering. + 1.0, // A finite value below NaN. + f64::NAN, // Null on the right. + ], + Validity::from_iter([ + true, // + false, // + true, // + true, // + true, // + ]), + ) + .into_array(); + let rhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::INFINITY, // Null on the left. + 0.0, // Signed zero ordering. + f64::NAN, // A finite value below NaN. + 1.0, // Null on the right. + ], + Validity::from_iter([ + true, // + true, // + true, // + true, // + false, // + ]), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let actual = compare_primitive_with_path(&lhs, &rhs, op, path, &mut ctx)?; + let expected = match op { + CompareOperator::Eq => [ + Some(true), // Equal NaNs. + None, // Null on the left. + Some(false), // Distinct signed zeroes. + Some(false), // A finite value and NaN. + None, // Null on the right. + ], + CompareOperator::NotEq | CompareOperator::Lt => [ + Some(false), // Equal NaNs. + None, // Null on the left. + Some(true), // Distinct signed zeroes. + Some(true), // A finite value and NaN. + None, // Null on the right. + ], + _ => unreachable!(), + }; + let expected = BoolArray::from_iter(expected); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(actual.dtype(), &DType::Bool(Nullability::Nullable)); + + // This encoding difference is intentional: the fused path materializes bits and validity + // together, while the RowFn path keeps masking lazy. + match path { + PrimitiveComparisonPath::Columnar => assert_eq!(actual.encoding_id(), Bool.id()), + PrimitiveComparisonPath::Row => assert!(actual.as_opt::().is_some()), + PrimitiveComparisonPath::Auto => unreachable!(), + } + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[rstest] +#[case::i64_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::i64_not_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::u64_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::u64_not_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::f64_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::f64_not_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +fn test_primitive_equality_auto_uses_columnar_for_supported_ptype( + #[case] lhs: ArrayRef, + #[case] rhs: ArrayRef, + #[case] op: CompareOperator, + #[case] expected: [bool; 3], +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let actual = + compare_primitive_with_path(&lhs, &rhs, op, PrimitiveComparisonPath::Auto, &mut ctx)?; + + assert_eq!(actual.encoding_id(), Bool.id()); + assert_arrays_eq!(actual, BoolArray::from_iter(expected), &mut ctx); + + Ok(()) +} + #[rstest] #[case(Operator::Eq, [true, false, true, true])] #[case(Operator::Lt, [false, true, false, false])] diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 80faff20e0a..a5b9fe70539 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -43,7 +43,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/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs index 98043caec13..fd9685dcb35 100644 --- a/vortex-array/src/test_harness/trace/tests.rs +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -685,6 +685,14 @@ fn trace_compare_on_dict() -> VortexResult<()> { iter 0 current=vortex.dict(bool, len=5) builder_active=false ExecuteSlot slot=1 parent=vortex.dict(bool, len=5) child=vortex.binary(bool, len=3) iter 1 current=vortex.binary(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=3) iter 2 current=vortex.bool(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false pop_frame slot=1 output=vortex.dict(bool, len=5) diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index d779757bc77..bf2b00f563e 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -209,6 +209,14 @@ fn trace_scan_compare_on_compressed_shipdate() -> VortexResult<()> { Done array=vortex.primitive(i32, len=4096) iter 1 current=vortex.primitive(i32, len=4096) builder_active=false return output=vortex.primitive(i32, len=4096) + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=4096) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=4096) iter 2 current=vortex.bool(bool, len=4096) builder_active=false return output=vortex.bool(bool, len=4096) @@ -267,6 +275,14 @@ fn trace_scan_compare_on_compressed_quantity() -> VortexResult<()> { Done array=vortex.primitive(i16, len=50) iter 1 current=vortex.primitive(i16, len=50) builder_active=false return output=vortex.primitive(i16, len=50) + optimize root=vortex.slice(i16, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i16, len=1) child=vortex.constant(i16, len=50) -> vortex.constant(i16, len=1) + done output=vortex.constant(i16, len=1) + execute_until target=AnyCanonical root=vortex.constant(i16, len=1) + iter 0 current=vortex.constant(i16, len=1) builder_active=false + Done array=vortex.primitive(i16, len=1) + iter 1 current=vortex.primitive(i16, len=1) builder_active=false + return output=vortex.primitive(i16, len=1) Done array=vortex.bool(bool, len=50) iter 6 current=vortex.bool(bool, len=50) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false pop_frame slot=1 output=vortex.dict(bool, len=4096)