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 97e248a1a7d..1a4000b6a9e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native comparison of primitive arrays via bit-packing lane kernels. +//! Native comparison of primitive arrays with byte-oriented bitmap packing. 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; @@ -19,8 +21,7 @@ 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::bit_buffer_from_words; use crate::scalar_fn::fns::binary::compare::compare_validity; use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; use crate::scalar_fn::fns::operators::CompareOperator; @@ -140,3 +141,40 @@ fn compare_slice_constant( CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs), allocator), } } + +fn collect_bits( + 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])); + } + bit_buffer_from_words(words, values.len()) +} + +fn collect_zip_bits( + 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()) +} 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 85260afbd99..14a80ef5b07 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -883,3 +883,102 @@ fn struct_of_map_compare() -> VortexResult<()> { Ok(()) } + +#[rstest] +fn primitive_comparisons_across_bitmap_words( + #[values( + PType::I8, PType::I16, PType::I32, PType::I64, PType::U8, PType::U16, PType::U32, + PType::U64, PType::F16, PType::F32, PType::F64 + )] + ptype: PType, + #[values(0, 1, 63, 64, 65, 129)] len: usize, + #[values( + CompareOperator::Eq, + CompareOperator::NotEq, + CompareOperator::Lt, + CompareOperator::Lte, + CompareOperator::Gt, + CompareOperator::Gte + )] + op: CompareOperator, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let value_for_type = |value: Option| { + value.map(|value| { + if ptype.is_unsigned_int() { + value.abs() + } else { + value + } + }) + }; + let left: Vec<_> = [Some(-128i16), Some(-1), None, Some(127), Some(42), Some(2)] + .into_iter() + .cycle() + .take(len) + .map(value_for_type) + .collect(); + let right: Vec<_> = [ + Some(-1i16), + Some(-128), + Some(127), + None, + Some(42), + Some(3), + Some(100), + ] + .into_iter() + .cycle() + .take(len) + .map(value_for_type) + .collect(); + let dtype = DType::Primitive(ptype, Nullability::Nullable); + let lhs = PrimitiveArray::from_option_iter(left.iter().copied()) + .into_array() + .cast(dtype.clone())? + .execute::(&mut ctx)? + .into_array(); + let rhs = PrimitiveArray::from_option_iter(right.iter().copied()) + .into_array() + .cast(dtype.clone())? + .execute::(&mut ctx)? + .into_array(); + let predicate = |a: i16, b: i16| match op { + CompareOperator::Eq => a == b, + CompareOperator::NotEq => a != b, + CompareOperator::Lt => a < b, + CompareOperator::Lte => a <= b, + CompareOperator::Gt => a > b, + CompareOperator::Gte => a >= b, + }; + let expected = BoolArray::from_iter( + left.iter() + .zip(&right) + .map(|(a, b)| a.zip(*b).map(|(a, b)| predicate(a, b))), + ); + assert_arrays_eq!(lhs.binary(rhs, op.into())?, expected, &mut ctx); + + let constant = ConstantArray::new( + Scalar::primitive(42u8, Nullability::Nullable).cast(&dtype)?, + len, + ) + .into_array(); + for swapped in [false, true] { + let actual = if swapped { + constant.binary(lhs.clone(), op.into())? + } else { + lhs.binary(constant.clone(), op.into())? + }; + let expected = BoolArray::from_iter(left.iter().map(|value| { + value.map(|value| { + if swapped { + predicate(42, value) + } else { + predicate(value, 42) + } + }) + })); + assert_arrays_eq!(actual, expected, &mut ctx); + } + Ok(()) +}