vortex_array/arrays/primitive/compute/
between.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::BitBuffer;
5use vortex_error::VortexExpect;
6use vortex_error::VortexResult;
7
8use crate::ArrayRef;
9use crate::ExecutionCtx;
10use crate::IntoArray;
11use crate::array::ArrayView;
12use crate::arrays::BoolArray;
13use crate::arrays::Primitive;
14use crate::dtype::NativePType;
15use crate::dtype::Nullability;
16use crate::match_each_native_ptype;
17use crate::scalar_fn::fns::between::BetweenKernel;
18use crate::scalar_fn::fns::between::BetweenOptions;
19use crate::scalar_fn::fns::between::StrictComparison;
20
21impl BetweenKernel for Primitive {
22    fn between(
23        arr: ArrayView<'_, Primitive>,
24        lower: &ArrayRef,
25        upper: &ArrayRef,
26        options: &BetweenOptions,
27        ctx: &mut ExecutionCtx,
28    ) -> VortexResult<Option<ArrayRef>> {
29        let (Some(lower), Some(upper)) = (lower.as_constant(), upper.as_constant()) else {
30            return Ok(None);
31        };
32
33        // Note, we know that have checked before that the lower and upper bounds are not constant
34        // null values
35
36        let nullability =
37            arr.dtype().nullability() | lower.dtype().nullability() | upper.dtype().nullability();
38
39        Ok(Some(match_each_native_ptype!(arr.ptype(), |P| {
40            between_impl::<P>(
41                arr,
42                P::try_from(&lower)?,
43                P::try_from(&upper)?,
44                nullability,
45                options,
46                ctx,
47            )
48        })))
49    }
50}
51
52fn between_impl<T: NativePType + Copy>(
53    arr: ArrayView<'_, Primitive>,
54    lower: T,
55    upper: T,
56    nullability: Nullability,
57    options: &BetweenOptions,
58    ctx: &mut ExecutionCtx,
59) -> ArrayRef {
60    match (options.lower_strict, options.upper_strict) {
61        // Note: these comparisons are explicitly passed in to allow function impl inlining
62        (StrictComparison::Strict, StrictComparison::Strict) => between_impl_(
63            arr,
64            lower,
65            NativePType::is_lt,
66            upper,
67            NativePType::is_lt,
68            nullability,
69            ctx,
70        ),
71        (StrictComparison::Strict, StrictComparison::NonStrict) => between_impl_(
72            arr,
73            lower,
74            NativePType::is_lt,
75            upper,
76            NativePType::is_le,
77            nullability,
78            ctx,
79        ),
80        (StrictComparison::NonStrict, StrictComparison::Strict) => between_impl_(
81            arr,
82            lower,
83            NativePType::is_le,
84            upper,
85            NativePType::is_lt,
86            nullability,
87            ctx,
88        ),
89        (StrictComparison::NonStrict, StrictComparison::NonStrict) => between_impl_(
90            arr,
91            lower,
92            NativePType::is_le,
93            upper,
94            NativePType::is_le,
95            nullability,
96            ctx,
97        ),
98    }
99}
100
101fn between_impl_<T>(
102    arr: ArrayView<'_, Primitive>,
103    lower: T,
104    lower_fn: impl Fn(T, T) -> bool,
105    upper: T,
106    upper_fn: impl Fn(T, T) -> bool,
107    nullability: Nullability,
108    ctx: &mut ExecutionCtx,
109) -> ArrayRef
110where
111    T: NativePType + Copy,
112{
113    let slice = arr.as_slice::<T>();
114    BoolArray::new(
115        BitBuffer::collect_bool_multiversioned_in(
116            slice.len(),
117            |idx| {
118                // We only iterate upto arr len and |arr| == |slice|.
119                let i = unsafe { *slice.get_unchecked(idx) };
120                lower_fn(lower, i) & upper_fn(i, upper)
121            },
122            ctx.allocator().clone(),
123        ),
124        arr.validity()
125            .vortex_expect("validity should be derivable")
126            .union_nullability(nullability),
127    )
128    .into_array()
129}
vortex_array/arrays/primitive/compute/between.rs

vortex_array/arrays/primitive/compute/
between.rs
