vortex_array/arrays/primitive/compute/
between.rs

1use arrow_buffer::BooleanBuffer;
2use vortex_dtype::{NativePType, Nullability, match_each_native_ptype};
3use vortex_error::VortexResult;
4
5use crate::arrays::{BoolArray, PrimitiveArray, PrimitiveVTable};
6use crate::compute::{BetweenKernel, BetweenKernelAdapter, BetweenOptions, StrictComparison};
7use crate::vtable::ValidityHelper;
8use crate::{Array, ArrayRef, IntoArray, register_kernel};
9
10impl BetweenKernel for PrimitiveVTable {
11    fn between(
12        &self,
13        arr: &PrimitiveArray,
14        lower: &dyn Array,
15        upper: &dyn Array,
16        options: &BetweenOptions,
17    ) -> VortexResult<Option<ArrayRef>> {
18        let (Some(lower), Some(upper)) = (lower.as_constant(), upper.as_constant()) else {
19            return Ok(None);
20        };
21
22        // Note, we know that have checked before that the lower and upper bounds are not constant
23        // null values
24
25        let nullability =
26            arr.dtype.nullability() | lower.dtype().nullability() | upper.dtype().nullability();
27
28        Ok(Some(match_each_native_ptype!(arr.ptype(), |$P| {
29            between_impl::<$P>(arr, $P::try_from(lower)?, $P::try_from(upper)?, nullability, options)
30        })))
31    }
32}
33
34register_kernel!(BetweenKernelAdapter(PrimitiveVTable).lift());
35
36fn between_impl<T: NativePType + Copy>(
37    arr: &PrimitiveArray,
38    lower: T,
39    upper: T,
40    nullability: Nullability,
41    options: &BetweenOptions,
42) -> ArrayRef {
43    match (options.lower_strict, options.upper_strict) {
44        // Note: these comparisons are explicitly passed in to allow function impl inlining
45        (StrictComparison::Strict, StrictComparison::Strict) => between_impl_(
46            arr,
47            lower,
48            NativePType::is_lt,
49            upper,
50            NativePType::is_lt,
51            nullability,
52        ),
53        (StrictComparison::Strict, StrictComparison::NonStrict) => between_impl_(
54            arr,
55            lower,
56            NativePType::is_lt,
57            upper,
58            NativePType::is_le,
59            nullability,
60        ),
61        (StrictComparison::NonStrict, StrictComparison::Strict) => between_impl_(
62            arr,
63            lower,
64            NativePType::is_le,
65            upper,
66            NativePType::is_lt,
67            nullability,
68        ),
69        (StrictComparison::NonStrict, StrictComparison::NonStrict) => between_impl_(
70            arr,
71            lower,
72            NativePType::is_le,
73            upper,
74            NativePType::is_le,
75            nullability,
76        ),
77    }
78}
79
80fn between_impl_<T>(
81    arr: &PrimitiveArray,
82    lower: T,
83    lower_fn: impl Fn(T, T) -> bool,
84    upper: T,
85    upper_fn: impl Fn(T, T) -> bool,
86    nullability: Nullability,
87) -> ArrayRef
88where
89    T: NativePType + Copy,
90{
91    let slice = arr.as_slice::<T>();
92    BoolArray::new(
93        BooleanBuffer::collect_bool(slice.len(), |idx| {
94            // We only iterate upto arr len and |arr| == |slice|.
95            let i = unsafe { *slice.get_unchecked(idx) };
96            lower_fn(lower, i) & upper_fn(i, upper)
97        }),
98        arr.validity().clone().union_nullability(nullability),
99    )
100    .into_array()
101}
vortex_array/arrays/primitive/compute/between.rs

vortex_array/arrays/primitive/compute/
between.rs
