vortex_alp/alp/compute/
between.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5
6use vortex_array::ArrayRef;
7use vortex_array::ArrayView;
8use vortex_array::IntoArray;
9use vortex_array::arrays::ConstantArray;
10use vortex_array::builtins::ArrayBuiltins;
11use vortex_array::dtype::NativeDType;
12use vortex_array::dtype::NativePType;
13use vortex_array::dtype::Nullability;
14use vortex_array::scalar::Scalar;
15use vortex_array::scalar_fn::fns::between::BetweenOptions;
16use vortex_array::scalar_fn::fns::between::BetweenReduce;
17use vortex_array::scalar_fn::fns::between::StrictComparison;
18use vortex_error::VortexResult;
19
20use crate::ALP;
21use crate::ALPFloat;
22use crate::alp::array::ALPArrayExt;
23use crate::alp::array::ALPArraySlotsExt;
24use crate::match_each_alp_float_ptype;
25
26impl BetweenReduce for ALP {
27    fn between(
28        array: ArrayView<'_, Self>,
29        lower: &ArrayRef,
30        upper: &ArrayRef,
31        options: &BetweenOptions,
32    ) -> VortexResult<Option<ArrayRef>> {
33        let (Some(lower), Some(upper)) = (lower.as_constant(), upper.as_constant()) else {
34            return Ok(None);
35        };
36
37        if array.patches().is_some() {
38            return Ok(None);
39        }
40
41        let nullability =
42            array.dtype().nullability() | lower.dtype().nullability() | upper.dtype().nullability();
43        match_each_alp_float_ptype!(array.dtype().as_ptype(), |F| {
44            between_impl::<F>(
45                array,
46                F::try_from(&lower)?,
47                F::try_from(&upper)?,
48                nullability,
49                options,
50            )
51        })
52        .map(Some)
53    }
54}
55
56fn between_impl<T: NativePType + ALPFloat>(
57    array: ArrayView<'_, ALP>,
58    lower: T,
59    upper: T,
60    nullability: Nullability,
61    options: &BetweenOptions,
62) -> VortexResult<ArrayRef>
63where
64    Scalar: From<T::ALPInt>,
65    <T as ALPFloat>::ALPInt: NativeDType + Debug,
66{
67    let exponents = array.exponents();
68
69    // There are always compared
70    // the below bound is `value {< | <=} x`, either value encodes into the ALPInt domain
71    // in which case we can leave the comparison unchanged `enc(value) {< | <=} x` or it doesn't
72    // and we encode into value below enc_below(value) < value < x, in which case the comparison
73    // becomes enc(value) < x. See `alp_scalar_compare` for more details.
74    // note that if the value doesn't encode than value != x, so must use strict comparison.
75    let (lower_enc, lower_strict) = T::encode_single(lower, exponents)
76        .map(|x| (x, options.lower_strict))
77        .unwrap_or_else(|| (T::encode_below(lower, exponents), StrictComparison::Strict));
78
79    // the upper value `x { < | <= } value` similarly encodes or `x < value < enc_above(value())`
80    let (upper_enc, upper_strict) = T::encode_single(upper, exponents)
81        .map(|x| (x, options.upper_strict))
82        .unwrap_or_else(|| (T::encode_above(upper, exponents), StrictComparison::Strict));
83
84    let options = BetweenOptions {
85        lower_strict,
86        upper_strict,
87    };
88
89    array.encoded().clone().between(
90        ConstantArray::new(Scalar::primitive(lower_enc, nullability), array.len()).into_array(),
91        ConstantArray::new(Scalar::primitive(upper_enc, nullability), array.len()).into_array(),
92        options,
93    )
94}
95
96#[cfg(test)]
97mod tests {
98    use vortex_array::LEGACY_SESSION;
99    use vortex_array::VortexSessionExecute;
100    use vortex_array::arrays::BoolArray;
101    use vortex_array::arrays::PrimitiveArray;
102    use vortex_array::assert_arrays_eq;
103    use vortex_array::dtype::Nullability;
104    use vortex_array::scalar_fn::fns::between::BetweenOptions;
105    use vortex_array::scalar_fn::fns::between::StrictComparison;
106
107    use crate::ALPArray;
108    use crate::alp::array::ALPArrayExt;
109    use crate::alp::compute::between::between_impl;
110    use crate::alp_encode;
111
112    fn assert_between(
113        arr: &ALPArray,
114        lower: f32,
115        upper: f32,
116        options: &BetweenOptions,
117        expected: bool,
118    ) {
119        let res =
120            between_impl(arr.as_view(), lower, upper, Nullability::Nullable, options).unwrap();
121        assert_arrays_eq!(res, BoolArray::from_iter([Some(expected)]));
122    }
123
124    #[test]
125    fn comparison_range() {
126        let value = 0.0605_f32;
127        let array = PrimitiveArray::from_iter([value; 1]);
128        let encoded = alp_encode(
129            array.as_view(),
130            None,
131            &mut LEGACY_SESSION.create_execution_ctx(),
132        )
133        .unwrap();
134        assert!(encoded.patches().is_none());
135
136        assert_between(
137            &encoded,
138            0.0605_f32,
139            0.0605,
140            &BetweenOptions {
141                lower_strict: StrictComparison::NonStrict,
142                upper_strict: StrictComparison::NonStrict,
143            },
144            true,
145        );
146
147        assert_between(
148            &encoded,
149            0.0605_f32,
150            0.0605,
151            &BetweenOptions {
152                lower_strict: StrictComparison::Strict,
153                upper_strict: StrictComparison::NonStrict,
154            },
155            false,
156        );
157
158        assert_between(
159            &encoded,
160            0.0605_f32,
161            0.0605,
162            &BetweenOptions {
163                lower_strict: StrictComparison::NonStrict,
164                upper_strict: StrictComparison::Strict,
165            },
166            false,
167        );
168
169        assert_between(
170            &encoded,
171            0.060499_f32,
172            0.06051,
173            &BetweenOptions {
174                lower_strict: StrictComparison::NonStrict,
175                upper_strict: StrictComparison::NonStrict,
176            },
177            true,
178        );
179
180        assert_between(
181            &encoded,
182            0.06_f32,
183            0.06051,
184            &BetweenOptions {
185                lower_strict: StrictComparison::NonStrict,
186                upper_strict: StrictComparison::Strict,
187            },
188            true,
189        );
190    }
191}
vortex_alp/alp/compute/between.rs

vortex_alp/alp/compute/
between.rs
