vortex_fastlanes/bitpacking/compute/
compare.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Block-streaming compare kernel for [`BitPackedArray`] against a constant.
5//!
6//! Avoids materialising the full primitive: the array is walked one 1024-element FastLanes
7//! block at a time through a reusable scratch buffer, and a per-element bool is folded into
8//! a [`BitBuffer`]. Patches are re-applied at the end by overwriting bits at the patched
9//! indices with `predicate(patch_value)`.
10//!
11//! [`BitPackedArray`]: crate::BitPackedArray
12//! [`BitBuffer`]: vortex_buffer::BitBuffer
13
14use fastlanes::BitPacking;
15use fastlanes::BitPackingCompare;
16use fastlanes::FastLanesComparable;
17use vortex_array::ArrayRef;
18use vortex_array::ArrayView;
19use vortex_array::ExecutionCtx;
20use vortex_array::dtype::NativePType;
21use vortex_array::dtype::Nullability;
22use vortex_array::dtype::PhysicalPType;
23use vortex_array::match_each_integer_ptype;
24use vortex_array::scalar_fn::fns::binary::CompareKernel;
25use vortex_array::scalar_fn::fns::operators::CompareOperator;
26use vortex_error::VortexExpect;
27use vortex_error::VortexResult;
28
29use crate::BitPacked;
30use crate::bitpacking::compute::compare_fused::stream_compare_fused;
31use crate::unpack_iter::BitPacked as BitPackedIter;
32
33impl CompareKernel for BitPacked {
34    fn compare(
35        lhs: ArrayView<'_, Self>,
36        rhs: &ArrayRef,
37        operator: CompareOperator,
38        ctx: &mut ExecutionCtx,
39    ) -> VortexResult<Option<ArrayRef>> {
40        // Only accelerate compare-against-constant.
41        let Some(constant) = rhs.as_constant() else {
42            return Ok(None);
43        };
44        let Some(constant_prim) = constant.as_primitive_opt() else {
45            return Ok(None);
46        };
47
48        // Adaptor strips null-constant RHS, and the binary scalar-fn coerce_args step has
49        // already promoted both sides to a common ptype.
50        let nullability = lhs.dtype().nullability() | rhs.dtype().nullability();
51        let lhs_ptype = lhs.dtype().as_ptype();
52        if constant_prim.ptype() != lhs_ptype {
53            return Ok(None);
54        }
55
56        let result = match_each_integer_ptype!(lhs_ptype, |T| {
57            let rhs: T = constant_prim
58                .typed_value::<T>()
59                .vortex_expect("compare adaptor strips null constants");
60            compare_constant_typed::<T>(lhs, rhs, operator, nullability, ctx)?
61        });
62        Ok(Some(result))
63    }
64}
65
66/// Compare every value against the constant via the fused FastLanes `unpack_cmp` kernel.
67///
68/// `NativePType::is_eq` / `is_lt` etc. provide total comparison (matching the primitive between
69/// kernel's dispatch shape). `NotEq` has no direct method, so use `!is_eq`.
70fn compare_constant_typed<T>(
71    lhs: ArrayView<'_, BitPacked>,
72    rhs: T,
73    operator: CompareOperator,
74    nullability: Nullability,
75    ctx: &mut ExecutionCtx,
76) -> VortexResult<ArrayRef>
77where
78    T: NativePType
79        + BitPackedIter
80        + FastLanesComparable<Bitpacked = <T as PhysicalPType>::Physical>,
81    <T as PhysicalPType>::Physical: BitPacking + NativePType + BitPackingCompare,
82{
83    match operator {
84        CompareOperator::Eq => {
85            stream_compare_fused::<T, _>(lhs, rhs, nullability, |a, b| a.is_eq(b), ctx)
86        }
87        CompareOperator::NotEq => {
88            stream_compare_fused::<T, _>(lhs, rhs, nullability, |a, b| !a.is_eq(b), ctx)
89        }
90        CompareOperator::Lt => {
91            stream_compare_fused::<T, _>(lhs, rhs, nullability, |a, b| a.is_lt(b), ctx)
92        }
93        CompareOperator::Lte => {
94            stream_compare_fused::<T, _>(lhs, rhs, nullability, |a, b| a.is_le(b), ctx)
95        }
96        CompareOperator::Gt => {
97            stream_compare_fused::<T, _>(lhs, rhs, nullability, |a, b| a.is_gt(b), ctx)
98        }
99        CompareOperator::Gte => {
100            stream_compare_fused::<T, _>(lhs, rhs, nullability, |a, b| a.is_ge(b), ctx)
101        }
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use std::sync::LazyLock;
108
109    use rstest::rstest;
110    use vortex_array::IntoArray;
111    use vortex_array::VortexSessionExecute;
112    use vortex_array::arrays::BoolArray;
113    use vortex_array::arrays::ConstantArray;
114    use vortex_array::arrays::PrimitiveArray;
115    use vortex_array::arrays::slice::SliceKernel;
116    use vortex_array::assert_arrays_eq;
117    use vortex_array::builtins::ArrayBuiltins;
118    use vortex_array::scalar_fn::fns::binary::CompareKernel;
119    use vortex_array::scalar_fn::fns::operators::CompareOperator;
120    use vortex_array::scalar_fn::fns::operators::Operator;
121    use vortex_error::VortexResult;
122    use vortex_session::VortexSession;
123
124    use crate::BitPacked;
125    use crate::BitPackedArrayExt;
126    use crate::BitPackedData;
127
128    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
129        let session = vortex_array::array_session();
130        crate::initialize(&session);
131        session
132    });
133
134    /// All six operators on a small in-range input.
135    #[rstest]
136    #[case(Operator::Eq, vec![false, false, false, true, false, false, true])]
137    #[case(Operator::NotEq, vec![true, true, true, false, true, true, false])]
138    #[case(Operator::Lt, vec![true, true, true, false, false, false, false])]
139    #[case(Operator::Lte, vec![true, true, true, true, false, false, true])]
140    #[case(Operator::Gt, vec![false, false, false, false, true, true, false])]
141    #[case(Operator::Gte, vec![false, false, false, true, true, true, true])]
142    fn small(#[case] op: Operator, #[case] expected: Vec<bool>) {
143        let mut ctx = SESSION.create_execution_ctx();
144        let values = PrimitiveArray::from_iter([0u32, 1, 2, 3, 4, 5, 3]);
145        let packed = BitPackedData::encode(&values.into_array(), 3, &mut ctx).unwrap();
146        let rhs = ConstantArray::new(3u32, packed.len()).into_array();
147        let result = packed
148            .into_array()
149            .binary(rhs, op)
150            .unwrap()
151            .execute::<BoolArray>(&mut ctx)
152            .unwrap();
153        assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx);
154    }
155
156    /// Sweep every native int type across several bit-widths. 2048 elements spans two
157    /// FastLanes blocks, exercising the per-type monomorphised inner loop. The kernel is
158    /// invoked *directly* and asserted `Some`, proving the streaming path engages (rather
159    /// than silently falling back to the Arrow compare), and its output is checked against
160    /// the Primitive fallback.
161    macro_rules! sweep {
162        ($name:ident, $T:ty, $($bw:expr),+) => {
163            #[test]
164            fn $name() -> VortexResult<()> {
165                let mut ctx = SESSION.create_execution_ctx();
166                for bw in [$($bw),+] {
167                    let cap: u128 = 1u128 << bw;
168                    let values: Vec<$T> = (0..2048u128).map(|i| (i % cap) as $T).collect();
169                    let prim = PrimitiveArray::from_iter(values);
170                    let packed = BitPackedData::encode(&prim.clone().into_array(), bw, &mut ctx)?;
171                    let rhs_val = (cap.min(2048) / 2) as $T;
172                    let rhs = ConstantArray::new(rhs_val, prim.len()).into_array();
173                    for op in [CompareOperator::Eq, CompareOperator::Lt, CompareOperator::Gte] {
174                        let got = <BitPacked as CompareKernel>::compare(
175                            packed.as_view(), &rhs, op, &mut ctx,
176                        )?
177                        .expect("streaming compare kernel must engage")
178                        .execute::<BoolArray>(&mut ctx)?;
179                        let want = prim
180                            .clone()
181                            .into_array()
182                            .binary(rhs.clone(), Operator::from(op))?
183                            .execute::<BoolArray>(&mut ctx)?;
184                        assert_arrays_eq!(got, want, &mut SESSION.create_execution_ctx());
185                    }
186                }
187                Ok(())
188            }
189        };
190    }
191
192    sweep!(sweep_u8, u8, 1, 4, 7);
193    sweep!(sweep_u16, u16, 1, 8, 15);
194    sweep!(sweep_u32, u32, 1, 16, 31);
195    sweep!(sweep_u64, u64, 1, 32, 63);
196    sweep!(sweep_i8, i8, 1, 4, 7);
197    sweep!(sweep_i16, i16, 1, 8, 15);
198    sweep!(sweep_i32, i32, 1, 16, 31);
199    sweep!(sweep_i64, i64, 1, 32, 63);
200
201    /// Inline-patch path: encode signed i32 values that exceed the bit-width range so they
202    /// end up in `Patches`. The streaming kernel must splice the patches in before the
203    /// predicate runs.
204    #[test]
205    fn signed_with_patches_matches_primitive() -> VortexResult<()> {
206        let mut ctx = SESSION.create_execution_ctx();
207        let values: Vec<i32> = (0..1500)
208            .map(|i| if i % 73 == 0 { 100_000 + i } else { i % 100 })
209            .collect();
210        let prim = PrimitiveArray::from_iter(values);
211        let packed = BitPackedData::encode(&prim.clone().into_array(), 7, &mut ctx)?;
212        assert!(packed.patches().is_some(), "test setup expects patches");
213        let rhs = ConstantArray::new(50i32, prim.len()).into_array();
214        let expected = prim
215            .into_array()
216            .binary(rhs.clone(), Operator::Eq)?
217            .execute::<BoolArray>(&mut ctx)?;
218        let actual = packed
219            .into_array()
220            .binary(rhs, Operator::Eq)?
221            .execute::<BoolArray>(&mut ctx)?;
222        assert_arrays_eq!(actual, expected, &mut ctx);
223        Ok(())
224    }
225
226    /// Sliced inputs: a non-zero block offset (and a length spanning several blocks) must still go
227    /// through the fused kernel and agree with the primitive fallback. Sweeps slice starts that
228    /// land both inside the first block and past it, with lengths that end mid-block and on a block
229    /// boundary.
230    #[rstest]
231    #[case(1, 4000)] // start mid-first-block, multi-block length
232    #[case(1023, 2)] // start at the last row of the first block
233    #[case(1024, 1024)] // start exactly on a block boundary, exactly one block long
234    #[case(1500, 1000)] // start mid-second-block
235    #[case(3, 1021)] // ends exactly on the first block boundary
236    fn sliced_matches_primitive(
237        #[case] start: usize,
238        #[case] slice_len: usize,
239    ) -> VortexResult<()> {
240        let mut ctx = SESSION.create_execution_ctx();
241        let values: Vec<u32> = (0..5000u32).map(|i| i % 128).collect();
242        let prim = PrimitiveArray::from_iter(values);
243        let packed = BitPackedData::encode(&prim.clone().into_array(), 7, &mut ctx)?;
244
245        let sliced = packed.into_array().slice(start..start + slice_len)?;
246        let rhs = ConstantArray::new(50u32, slice_len).into_array();
247        for op in [
248            CompareOperator::Eq,
249            CompareOperator::Lt,
250            CompareOperator::Gte,
251        ] {
252            let got = <BitPacked as CompareKernel>::compare(
253                sliced.as_::<BitPacked>(),
254                &rhs,
255                op,
256                &mut ctx,
257            )?
258            .expect("fused compare kernel must engage for sliced arrays")
259            .execute::<BoolArray>(&mut ctx)?;
260            let want = prim
261                .clone()
262                .into_array()
263                .slice(start..start + slice_len)?
264                .binary(rhs.clone(), Operator::from(op))?
265                .execute::<BoolArray>(&mut ctx)?;
266            assert_arrays_eq!(got, want, &mut ctx);
267        }
268        Ok(())
269    }
270
271    /// Sliced *and* patched: combine a non-zero offset with out-of-range values that land in
272    /// `Patches`, exercising the `offset + (global - p_off)` patch-position math.
273    #[test]
274    fn sliced_with_patches_matches_primitive() -> VortexResult<()> {
275        let mut ctx = SESSION.create_execution_ctx();
276        let values: Vec<i32> = (0..4096)
277            .map(|i| if i % 91 == 0 { 100_000 + i } else { i % 100 })
278            .collect();
279        let prim = PrimitiveArray::from_iter(values);
280        let packed = BitPackedData::encode(&prim.clone().into_array(), 7, &mut ctx)?;
281        assert!(packed.patches().is_some(), "test setup expects patches");
282
283        let (start, end) = (700usize, 3500usize);
284        // `ArrayRef::slice` leaves a lazy `SliceArray` over a patched `BitPacked` (the
285        // `SliceReduce` path bails when patches are present), so go through the `SliceKernel`,
286        // which reads the buffers and produces a sliced `BitPacked` with sliced patches.
287        let sliced = <BitPacked as SliceKernel>::slice(packed.as_view(), start..end, &mut ctx)?
288            .expect("slice kernel produces a sliced bitpacked array");
289        let rhs = ConstantArray::new(50i32, end - start).into_array();
290        let got = <BitPacked as CompareKernel>::compare(
291            sliced.as_::<BitPacked>(),
292            &rhs,
293            CompareOperator::Eq,
294            &mut ctx,
295        )?
296        .expect("fused compare kernel must engage for sliced arrays with patches")
297        .execute::<BoolArray>(&mut ctx)?;
298        let want = prim
299            .into_array()
300            .slice(start..end)?
301            .binary(rhs, Operator::Eq)?
302            .execute::<BoolArray>(&mut ctx)?;
303        assert_arrays_eq!(got, want, &mut ctx);
304        Ok(())
305    }
306
307    /// Nullable input — the result must carry the array's validity.
308    #[test]
309    fn nullable_propagates_validity() -> VortexResult<()> {
310        let mut ctx = SESSION.create_execution_ctx();
311        let prim = PrimitiveArray::from_option_iter([Some(1u32), None, Some(3), Some(4), None]);
312        let packed = BitPackedData::encode(&prim.clone().into_array(), 3, &mut ctx)?;
313        let rhs = ConstantArray::new(3u32, packed.len()).into_array();
314        let actual = packed
315            .into_array()
316            .binary(rhs.clone(), Operator::Eq)?
317            .execute::<BoolArray>(&mut ctx)?;
318        let expected = prim
319            .into_array()
320            .binary(rhs, Operator::Eq)?
321            .execute::<BoolArray>(&mut ctx)?;
322        assert_arrays_eq!(actual, expected, &mut ctx);
323        Ok(())
324    }
325}
vortex_fastlanes/bitpacking/compute/compare.rs

vortex_fastlanes/bitpacking/compute/
compare.rs
