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+
139use vortex_error:: VortexResult ;
1410use vortex_error:: vortex_bail;
1511
1612use crate :: ArrayRef ;
1713use crate :: ExecutionCtx ;
18- use crate :: IntoArray ;
19- use crate :: arrays:: BoolArray ;
20- use crate :: arrays:: ConstantArray ;
2114use crate :: dtype:: DType ;
2215use crate :: dtype:: NativePType ;
23- use crate :: dtype:: Nullability ;
2416use crate :: dtype:: PType ;
2517use 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 ;
2921use 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}
0 commit comments