Skip to content

Commit 33ccf9b

Browse files
committed
Compare primitive values through the shared RowFn executor
Keep the primitive kernel focused on type and operator dispatch. Delegate borrowed inputs, constants, strict validity, and packed Boolean output to the lower stack layers while preserving Binary identity and floating-point total ordering. Replaces the comparison migration in #9979. Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent d3afc24 commit 33ccf9b

6 files changed

Lines changed: 87 additions & 250 deletions

File tree

‎vortex-array/benches/compare.rs‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,11 @@
33

44
//! Benchmarks for the binary comparison path, over every array kind it accepts.
55
//!
6-
//! The primitive cases carry `#[cpu_features]`, so they are measured on every walltime
7-
//! CPU-feature leg rather than in simulation. Each is written once and compiled differently
8-
//! per leg: today the primitive comparison path is a portable lane kernel, and how well it
9-
//! auto-vectorizes is decided by the build. That is the baseline a hand-written kernel
10-
//! selected through `cfg(target_feature)` has to beat, measured on the silicon it would run
11-
//! on.
6+
//! The primitive cases carry `#[cpu_features]`, so they are measured on every walltime CPU-feature
7+
//! leg rather than in simulation. They all exercise the same [`RowFn`] comparison path, including
8+
//! its runtime-selected packed Boolean collector.
9+
//!
10+
//! [`RowFn`]: vortex_array::scalar_fn::unstable::row::RowFn
1211
//!
1312
//! Every case here compares one array against another, which is the shape the path is tuned for.
1413
//! The three constant cases that remain — boolean, integer, and string — are a regression guard on

‎vortex-array/src/scalar_fn/fns/binary/compare/mod.rs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@
44
//! Native comparison kernels.
55
//!
66
//! [`execute_compare`] dispatches on the logical [`DType`] of its operands and evaluates every
7-
//! comparison directly over Vortex canonical arrays — bit buffers for booleans, lane kernels from
8-
//! `vortex-compute` for primitives and decimals, binary views for strings/bytes, and a row-wise
9-
//! comparator for nested types. There is no Arrow fallback.
7+
//! comparison directly over Vortex canonical arrays: bit buffers for booleans, [`RowFn`] for
8+
//! primitives, lane kernels from `vortex-compute` for decimals, binary views for strings/bytes, and
9+
//! a row-wise comparator for nested types. There is no Arrow fallback.
10+
//!
11+
//! [`RowFn`]: crate::scalar_fn::unstable::row::RowFn
1012
//!
1113
//! Floating point values compare with Vortex's total ordering (`NaN` is the largest value,
1214
//! `-0.0 < +0.0`, and equality is bitwise), matching [`Scalar`] comparison semantics.
@@ -209,7 +211,7 @@ fn compare_arrays(
209211
)
210212
.into_array()),
211213
DType::Bool(_) => boolean::compare_bool(lhs, rhs, op, nullability, ctx),
212-
DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx),
214+
DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, ctx),
213215
DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx),
214216
DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx),
215217
DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Map(..) => {
Lines changed: 74 additions & 170 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,28 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4-
//! Native comparison of primitive arrays with specialized bitmap packing for 8-bit inputs.
5-
6-
use vortex_buffer::BitBuffer;
7-
use vortex_buffer::BufferAllocatorRef;
8-
use vortex_buffer::BufferMut;
9-
use vortex_buffer::collect_bool_word;
10-
use vortex_compute::lane_kernels::bit_buffer_from_words;
11-
use vortex_compute::lane_kernels::collect_bits;
12-
use vortex_compute::lane_kernels::collect_zip_bits;
4+
//! Primitive comparison execution through [`RowFn`].
5+
//!
6+
//! [`PrimitiveCompare`] delegates decoding, constant handling, validity, and packed Boolean output
7+
//! to the row executor. Its row kernel contains only the comparison selected by [`CompareOperator`].
8+
139
use vortex_error::VortexResult;
1410
use vortex_error::vortex_bail;
1511

1612
use crate::ArrayRef;
1713
use crate::ExecutionCtx;
18-
use crate::IntoArray;
19-
use crate::arrays::BoolArray;
20-
use crate::arrays::ConstantArray;
2114
use crate::dtype::DType;
2215
use crate::dtype::NativePType;
23-
use crate::dtype::Nullability;
2416
use crate::dtype::PType;
2517
use crate::match_each_native_ptype;
26-
use crate::scalar::Scalar;
27-
use crate::scalar_fn::fns::binary::compare::compare_validity;
28-
use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand;
18+
use crate::scalar_fn::ScalarFnId;
19+
use crate::scalar_fn::ScalarFnVTable;
20+
use crate::scalar_fn::fns::binary::Binary;
2921
use crate::scalar_fn::fns::operators::CompareOperator;
22+
use crate::scalar_fn::unstable::row::RowFn;
23+
use crate::scalar_fn::unstable::row::RowInputs;
24+
use crate::scalar_fn::unstable::row::RowVisitor;
25+
use crate::scalar_fn::unstable::row::execute_row_inputs;
3026

3127
/// Compare two primitive arrays of the same [`PType`].
3228
///
@@ -36,178 +32,86 @@ pub(super) fn compare_primitive(
3632
lhs: &ArrayRef,
3733
rhs: &ArrayRef,
3834
op: CompareOperator,
39-
nullability: Nullability,
4035
ctx: &mut ExecutionCtx,
4136
) -> VortexResult<ArrayRef> {
42-
let ptype = PType::try_from(lhs.dtype())?;
43-
match_each_native_ptype!(ptype, |T| {
44-
compare_primitive_typed::<T>(lhs, rhs, op, nullability, ctx)
45-
})
37+
execute_row_inputs(&PrimitiveCompare, &op, &CompareArgs([lhs, rhs]), ctx)
4638
}
4739

48-
fn compare_primitive_typed<T: NativePType>(
49-
lhs: &ArrayRef,
50-
rhs: &ArrayRef,
51-
op: CompareOperator,
52-
nullability: Nullability,
53-
ctx: &mut ExecutionCtx,
54-
) -> VortexResult<ArrayRef> {
55-
let len = lhs.len();
56-
let lhs = PrimitiveOperand::<T>::try_new(lhs, ctx)?;
57-
let rhs = PrimitiveOperand::<T>::try_new(rhs, ctx)?;
58-
if lhs.len() != rhs.len() {
59-
vortex_bail!(
60-
"compare operator requires equal lengths, got {} and {}",
61-
lhs.len(),
62-
rhs.len()
63-
);
64-
}
40+
/// Borrowed comparison operands, so a call neither allocates nor clones them.
41+
struct CompareArgs<'a>([&'a ArrayRef; 2]);
6542

66-
let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?;
67-
68-
let bits = match (&lhs, &rhs) {
69-
(
70-
PrimitiveOperand::Array { values: lhs, .. },
71-
PrimitiveOperand::Array { values: rhs, .. },
72-
) => compare_slices(lhs, rhs, op, ctx.allocator()),
73-
(
74-
PrimitiveOperand::Array { values: lhs, .. },
75-
PrimitiveOperand::Constant { value: rhs, .. },
76-
) => compare_slice_constant(lhs, *rhs, op, ctx.allocator()),
77-
(
78-
PrimitiveOperand::Constant { value: lhs, .. },
79-
PrimitiveOperand::Array { values: rhs, .. },
80-
) => compare_slice_constant(rhs, *lhs, op.swap(), ctx.allocator()),
81-
(
82-
PrimitiveOperand::Constant { value: lhs, .. },
83-
PrimitiveOperand::Constant { value: rhs, .. },
84-
) => {
85-
// Unreachable through `execute_compare` (constant-constant is folded there), but
86-
// cheap to answer anyway.
87-
BitBuffer::full_in(apply_op(*lhs, *rhs, op), len, ctx.allocator().clone())
88-
}
89-
(PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => {
90-
return Ok(
91-
ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len)
92-
.into_array(),
93-
);
94-
}
95-
};
96-
97-
Ok(BoolArray::try_new(bits, validity)?.into_array())
98-
}
99-
100-
#[allow(clippy::inline_always)]
101-
#[inline(always)]
102-
fn apply_op<T: NativePType>(lhs: T, rhs: T, op: CompareOperator) -> bool {
103-
match op {
104-
CompareOperator::Eq => lhs.is_eq(rhs),
105-
CompareOperator::NotEq => !lhs.is_eq(rhs),
106-
CompareOperator::Gt => lhs.is_gt(rhs),
107-
CompareOperator::Gte => lhs.is_ge(rhs),
108-
CompareOperator::Lt => lhs.is_lt(rhs),
109-
CompareOperator::Lte => lhs.is_le(rhs),
43+
impl RowInputs for CompareArgs<'_> {
44+
fn num_inputs(&self) -> usize {
45+
2
11046
}
111-
}
11247

113-
fn compare_slices<T: NativePType>(
114-
lhs: &[T],
115-
rhs: &[T],
116-
op: CompareOperator,
117-
allocator: &BufferAllocatorRef,
118-
) -> BitBuffer {
119-
// Dispatch the operator outside the lane loop so each instantiation vectorizes a single
120-
// branch-free predicate.
121-
match op {
122-
CompareOperator::Eq => {
123-
collect_zip_bits_dispatch(lhs, rhs, |a: T, b: T| a.is_eq(b), allocator)
124-
}
125-
CompareOperator::NotEq => {
126-
collect_zip_bits_dispatch(lhs, rhs, |a: T, b: T| !a.is_eq(b), allocator)
127-
}
128-
CompareOperator::Gt => collect_zip_bits_dispatch(lhs, rhs, T::is_gt, allocator),
129-
CompareOperator::Gte => collect_zip_bits_dispatch(lhs, rhs, T::is_ge, allocator),
130-
CompareOperator::Lt => collect_zip_bits_dispatch(lhs, rhs, T::is_lt, allocator),
131-
CompareOperator::Lte => collect_zip_bits_dispatch(lhs, rhs, T::is_le, allocator),
48+
fn row_count(&self) -> usize {
49+
self.0[0].len()
13250
}
133-
}
13451

135-
fn compare_slice_constant<T: NativePType>(
136-
lhs: &[T],
137-
rhs: T,
138-
op: CompareOperator,
139-
allocator: &BufferAllocatorRef,
140-
) -> BitBuffer {
141-
match op {
142-
CompareOperator::Eq => collect_bits_dispatch(lhs, |a: T| a.is_eq(rhs), allocator),
143-
CompareOperator::NotEq => collect_bits_dispatch(lhs, |a: T| !a.is_eq(rhs), allocator),
144-
CompareOperator::Gt => collect_bits_dispatch(lhs, |a: T| a.is_gt(rhs), allocator),
145-
CompareOperator::Gte => collect_bits_dispatch(lhs, |a: T| a.is_ge(rhs), allocator),
146-
CompareOperator::Lt => collect_bits_dispatch(lhs, |a: T| a.is_lt(rhs), allocator),
147-
CompareOperator::Lte => collect_bits_dispatch(lhs, |a: T| a.is_le(rhs), allocator),
52+
fn input(&self, index: usize) -> &ArrayRef {
53+
self.0[index]
14854
}
14955
}
15056

151-
fn collect_bits_dispatch<T: NativePType>(
152-
values: &[T],
153-
f: impl Fn(T) -> bool,
154-
allocator: &BufferAllocatorRef,
155-
) -> BitBuffer {
156-
// This type check folds away during monomorphization. Wider masks keep the lane kernel:
157-
// byte packing regresses 64-bit comparisons on AVX2.
158-
if matches!(T::PTYPE, PType::I8 | PType::U8) {
159-
collect_bits_narrow(values, f, allocator)
160-
} else {
161-
collect_bits(values, f, allocator)
57+
/// Internal row execution for primitive comparison operators.
58+
#[derive(Clone)]
59+
struct PrimitiveCompare;
60+
61+
impl RowFn for PrimitiveCompare {
62+
type Options = CompareOperator;
63+
64+
const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"];
65+
66+
const INFALLIBLE: bool = true;
67+
68+
fn id(&self) -> ScalarFnId {
69+
// `PrimitiveCompare` is a private implementation detail of `Binary`: it is never registered
70+
// or serialized independently. Reusing the public ID keeps execution errors attributed to
71+
// `Binary`. If this type becomes registrable, it needs its own ID and persistence contract.
72+
ScalarFnVTable::id(&Binary)
16273
}
163-
}
16474

165-
fn collect_zip_bits_dispatch<T: NativePType>(
166-
lhs: &[T],
167-
rhs: &[T],
168-
f: impl Fn(T, T) -> bool,
169-
allocator: &BufferAllocatorRef,
170-
) -> BitBuffer {
171-
if matches!(T::PTYPE, PType::I8 | PType::U8) {
172-
collect_zip_bits_narrow(lhs, rhs, f, allocator)
173-
} else {
174-
collect_zip_bits(lhs, rhs, f, allocator)
75+
fn dispatch<V: RowVisitor>(
76+
&self,
77+
op: &Self::Options,
78+
args: &[DType],
79+
visitor: V,
80+
) -> VortexResult<V::VisitResult> {
81+
let [lhs_dtype, _] = args else {
82+
vortex_bail!(
83+
"a primitive comparison requires two operands, got {}",
84+
args.len(),
85+
);
86+
};
87+
let ptype = PType::try_from(lhs_dtype)?;
88+
89+
match_each_native_ptype!(ptype, |T| { visit_compare::<T, V>(*op, visitor) })
17590
}
17691
}
17792

178-
fn collect_bits_narrow<T: Copy>(
179-
values: &[T],
180-
f: impl Fn(T) -> bool,
181-
allocator: &BufferAllocatorRef,
182-
) -> BitBuffer {
183-
let (chunks, tail) = values.as_chunks::<64>();
184-
let mut words = BufferMut::<u64>::zeroed_in(values.len().div_ceil(64), allocator.clone());
185-
// Fixed-size chunks let the compiler prove the predicate's indexing stays in bounds.
186-
for (word, chunk) in words.iter_mut().zip(chunks) {
187-
*word = collect_bool_word(64, |i| f(chunk[i]));
188-
}
189-
if !tail.is_empty() {
190-
words[chunks.len()] = collect_bool_word(tail.len(), |i| f(tail[i]));
93+
fn visit_compare<T, V>(op: CompareOperator, visitor: V) -> VortexResult<V::VisitResult>
94+
where
95+
T: NativePType,
96+
V: RowVisitor,
97+
{
98+
match op {
99+
CompareOperator::Eq => visit_compare_with::<T, V>(visitor, T::is_eq),
100+
CompareOperator::NotEq => visit_compare_with::<T, V>(visitor, |lhs, rhs| !lhs.is_eq(rhs)),
101+
CompareOperator::Gt => visit_compare_with::<T, V>(visitor, T::is_gt),
102+
CompareOperator::Gte => visit_compare_with::<T, V>(visitor, T::is_ge),
103+
CompareOperator::Lt => visit_compare_with::<T, V>(visitor, T::is_lt),
104+
CompareOperator::Lte => visit_compare_with::<T, V>(visitor, T::is_le),
191105
}
192-
bit_buffer_from_words(words, values.len())
193106
}
194107

195-
fn collect_zip_bits_narrow<T: Copy>(
196-
lhs: &[T],
197-
rhs: &[T],
198-
f: impl Fn(T, T) -> bool,
199-
allocator: &BufferAllocatorRef,
200-
) -> BitBuffer {
201-
assert_eq!(lhs.len(), rhs.len());
202-
let (left_chunks, left_tail) = lhs.as_chunks::<64>();
203-
let (right_chunks, right_tail) = rhs.as_chunks::<64>();
204-
let mut words = BufferMut::<u64>::zeroed_in(lhs.len().div_ceil(64), allocator.clone());
205-
for ((word, left), right) in words.iter_mut().zip(left_chunks).zip(right_chunks) {
206-
*word = collect_bool_word(64, |i| f(left[i], right[i]));
207-
}
208-
if !left_tail.is_empty() {
209-
words[left_chunks.len()] =
210-
collect_bool_word(left_tail.len(), |i| f(left_tail[i], right_tail[i]));
211-
}
212-
bit_buffer_from_words(words, lhs.len())
108+
fn visit_compare_with<T, V>(
109+
visitor: V,
110+
compare: impl Fn(T, T) -> bool,
111+
) -> VortexResult<V::VisitResult>
112+
where
113+
T: NativePType,
114+
V: RowVisitor,
115+
{
116+
visitor.visit_bool::<(T, T), true>(move |(lhs, rhs)| compare(lhs, rhs))
213117
}

‎vortex-array/src/scalar_fn/fns/binary/mod.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@ mod compare;
4444
pub use compare::*;
4545
mod numeric;
4646
pub(crate) use numeric::*;
47-
mod primitive_operand;
4847

4948
use crate::scalar::NumericOperator;
5049
use crate::scalar::Scalar;

0 commit comments

Comments
 (0)