vortex_array/arrays/primitive/compute/
between.rs

1use arrow_buffer::BooleanBuffer;
2use vortex_dtype::{NativePType, 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        Ok(Some(match_each_native_ptype!(arr.ptype(), |$P| {
23            between_impl::<$P>(arr, $P::try_from(lower)?, $P::try_from(upper)?, options)
24        })))
25    }
26}
27
28fn between_impl<T: NativePType + Copy>(
29    arr: &PrimitiveArray,
30    lower: T,
31    upper: T,
32    options: &BetweenOptions,
33) -> ArrayRef {
34    match (options.lower_strict, options.upper_strict) {
35        // Note: these comparisons are explicitly passed in to allow function impl inlining
36        (StrictComparison::Strict, StrictComparison::Strict) => {
37            between_impl_(arr, lower, NativePType::is_lt, upper, NativePType::is_lt)
38        }
39        (StrictComparison::Strict, StrictComparison::NonStrict) => {
40            between_impl_(arr, lower, NativePType::is_lt, upper, NativePType::is_le)
41        }
42        (StrictComparison::NonStrict, StrictComparison::Strict) => {
43            between_impl_(arr, lower, NativePType::is_le, upper, NativePType::is_lt)
44        }
45        (StrictComparison::NonStrict, StrictComparison::NonStrict) => {
46            between_impl_(arr, lower, NativePType::is_le, upper, NativePType::is_le)
47        }
48    }
49}
50
51fn between_impl_<T>(
52    arr: &PrimitiveArray,
53    lower: T,
54    lower_fn: impl Fn(T, T) -> bool,
55    upper: T,
56    upper_fn: impl Fn(T, T) -> bool,
57) -> ArrayRef
58where
59    T: NativePType + Copy,
60{
61    let slice = arr.as_slice::<T>();
62    BoolArray::new(
63        BooleanBuffer::collect_bool(slice.len(), |idx| {
64            // We only iterate upto arr len and |arr| == |slice|.
65            let i = unsafe { *slice.get_unchecked(idx) };
66            lower_fn(lower, i) & upper_fn(i, upper)
67        }),
68        arr.validity().clone(),
69    )
70    .into_array()
71}
vortex_array/arrays/primitive/compute/between.rs

vortex_array/arrays/primitive/compute/
between.rs
