From 1db27f22053a5aa421d1cfd9b2bb571b21eb83d9 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 18 Sep 2026 16:55:32 -0400 Subject: [PATCH 1/4] perf: pack primitive comparison results in fixed-size chunks Reuse collect_bool_word for primitive comparison bitmaps while preserving execution-context allocation and comparison semantics. Cover bitmap boundaries, all primitive types, nulls, and constant operand orientation. Signed-off-by: "Matt Katz" Signed-off-by: Matt Katz --- .../scalar_fn/fns/binary/compare/primitive.rs | 44 ++++++++- .../src/scalar_fn/fns/binary/compare/tests.rs | 99 +++++++++++++++++++ 2 files changed, 140 insertions(+), 3 deletions(-) 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(()) +} From 95b15506fd31d10503b84825289be95b8e6871d9 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 22 Sep 2026 10:34:36 -0400 Subject: [PATCH 2/4] perf: specialize comparison bitmap packing for 8-bit inputs Use byte packing only for i8/u8 and preserve the existing lane collectors for wider primitive types. The type selection folds away during monomorphization and avoids the wider AVX2 comparison regressions. Signed-off-by: "Matt Katz" Signed-off-by: Matt Katz --- .../src/scalar_fn/fns/binary/compare/primitive.rs | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) 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 1a4000b6a9e..c0d3c85a68c 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native comparison of primitive arrays with byte-oriented bitmap packing. +//! Native comparison of primitive arrays with specialized bitmap packing for 8-bit inputs. use vortex_buffer::BitBuffer; use vortex_buffer::BufferAllocatorRef; @@ -142,11 +142,16 @@ fn compare_slice_constant( } } -fn collect_bits( +fn collect_bits( 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) { + return super::collect_bits(values, f, allocator); + } 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. @@ -159,12 +164,15 @@ fn collect_bits( bit_buffer_from_words(words, values.len()) } -fn collect_zip_bits( +fn collect_zip_bits( lhs: &[T], rhs: &[T], f: impl Fn(T, T) -> bool, allocator: &BufferAllocatorRef, ) -> BitBuffer { + if !matches!(T::PTYPE, PType::I8 | PType::U8) { + return super::collect_zip_bits(lhs, rhs, f, allocator); + } assert_eq!(lhs.len(), rhs.len()); let (left_chunks, left_tail) = lhs.as_chunks::<64>(); let (right_chunks, right_tail) = rhs.as_chunks::<64>(); From 4483ef94e7f6968a527e544883dcb33ac72217f2 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 22 Sep 2026 10:55:03 -0400 Subject: [PATCH 3/4] refactor: separate primitive comparison bitmap dispatch Keep type selection in collect_bits_dispatch and collect_zip_bits_dispatch, with independent narrow collectors and the existing wider collectors. Signed-off-by: "Matt Katz" Signed-off-by: Matt Katz --- .../scalar_fn/fns/binary/compare/primitive.rs | 63 +++++++++++++------ 1 file changed, 44 insertions(+), 19 deletions(-) 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 c0d3c85a68c..d7607953651 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -22,6 +22,8 @@ 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::fns::operators::CompareOperator; @@ -117,12 +119,16 @@ fn compare_slices( // 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), allocator), - CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b), allocator), - CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt, allocator), - CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge, allocator), - CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt, allocator), - CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le, allocator), + 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), } } @@ -133,25 +139,47 @@ fn compare_slice_constant( allocator: &BufferAllocatorRef, ) -> BitBuffer { match op { - CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs), allocator), - CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs), allocator), - CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs), allocator), - CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs), allocator), - CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs), allocator), - CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs), allocator), + 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 collect_bits( +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) { - return super::collect_bits(values, f, allocator); + if matches!(T::PTYPE, PType::I8 | PType::U8) { + collect_bits_narrow(values, f, allocator) + } else { + collect_bits(values, f, allocator) } +} + +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) + } +} + +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. @@ -164,15 +192,12 @@ fn collect_bits( bit_buffer_from_words(words, values.len()) } -fn collect_zip_bits( +fn collect_zip_bits_narrow( lhs: &[T], rhs: &[T], f: impl Fn(T, T) -> bool, allocator: &BufferAllocatorRef, ) -> BitBuffer { - if !matches!(T::PTYPE, PType::I8 | PType::U8) { - return super::collect_zip_bits(lhs, rhs, f, allocator); - } assert_eq!(lhs.len(), rhs.len()); let (left_chunks, left_tail) = lhs.as_chunks::<64>(); let (right_chunks, right_tail) = rhs.as_chunks::<64>(); From c215b0e423f558ef83ac0f99841160d70237a9b7 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 22 Sep 2026 10:59:48 -0400 Subject: [PATCH 4/4] test: focus bitmap comparison coverage on 8-bit inputs Limit the new boundary matrix to the i8/u8 specialization. Existing tests cover the unchanged wider primitive comparison paths. Signed-off-by: "Matt Katz" Signed-off-by: Matt Katz --- vortex-array/src/scalar_fn/fns/binary/compare/tests.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) 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 14a80ef5b07..4688d323c90 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -885,12 +885,8 @@ fn struct_of_map_compare() -> VortexResult<()> { } #[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, +fn byte_comparisons_across_bitmap_words( + #[values(PType::I8, PType::U8)] ptype: PType, #[values(0, 1, 63, 64, 65, 129)] len: usize, #[values( CompareOperator::Eq,