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, PrimitiveEncoding};
6use crate::compute::{BetweenFn, BetweenOptions, StrictComparison};
7use crate::variants::PrimitiveArrayTrait;
8use crate::{Array, ArrayRef};
9
10impl BetweenFn<&PrimitiveArray> for PrimitiveEncoding {
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
34fn between_impl<T: NativePType + Copy>(
35    arr: &PrimitiveArray,
36    lower: T,
37    upper: T,
38    nullability: Nullability,
39    options: &BetweenOptions,
40) -> ArrayRef {
41    match (options.lower_strict, options.upper_strict) {
42        // Note: these comparisons are explicitly passed in to allow function impl inlining
43        (StrictComparison::Strict, StrictComparison::Strict) => between_impl_(
44            arr,
45            lower,
46            NativePType::is_lt,
47            upper,
48            NativePType::is_lt,
49            nullability,
50        ),
51        (StrictComparison::Strict, StrictComparison::NonStrict) => between_impl_(
52            arr,
53            lower,
54            NativePType::is_lt,
55            upper,
56            NativePType::is_le,
57            nullability,
58        ),
59        (StrictComparison::NonStrict, StrictComparison::Strict) => between_impl_(
60            arr,
61            lower,
62            NativePType::is_le,
63            upper,
64            NativePType::is_lt,
65            nullability,
66        ),
67        (StrictComparison::NonStrict, StrictComparison::NonStrict) => between_impl_(
68            arr,
69            lower,
70            NativePType::is_le,
71            upper,
72            NativePType::is_le,
73            nullability,
74        ),
75    }
76}
77
78fn between_impl_<T>(
79    arr: &PrimitiveArray,
80    lower: T,
81    lower_fn: impl Fn(T, T) -> bool,
82    upper: T,
83    upper_fn: impl Fn(T, T) -> bool,
84    nullability: Nullability,
85) -> ArrayRef
86where
87    T: NativePType + Copy,
88{
89    let slice = arr.as_slice::<T>();
90    BoolArray::new(
91        BooleanBuffer::collect_bool(slice.len(), |idx| {
92            // We only iterate upto arr len and |arr| == |slice|.
93            let i = unsafe { *slice.get_unchecked(idx) };
94            lower_fn(lower, i) & upper_fn(i, upper)
95        }),
96        arr.validity().clone().union_nullability(nullability),
97    )
98    .into_array()
99}
vortex_array/arrays/primitive/compute/between.rs

vortex_array/arrays/primitive/compute/
between.rs
