perf: pack primitive comparison results in fixed-size chunks

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<T: NativePType>(
         CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs), allocator),
     }
 }
+
+fn collect_bits<T: Copy>(
+    values: &[T],
+    f: impl Fn(T) -> bool,
+    allocator: &BufferAllocatorRef,
+) -> BitBuffer {
+    let (chunks, tail) = values.as_chunks::<64>();
+    let mut words = BufferMut::<u64>::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<T: Copy>(
+    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::<u64>::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())
+}

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<i16>| {
+        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::<PrimitiveArray>(&mut ctx)?
+        .into_array();
+    let rhs = PrimitiveArray::from_option_iter(right.iter().copied())
+        .into_array()
+        .cast(dtype.clone())?
+        .execute::<PrimitiveArray>(&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(())
+}