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 num_traits::Bounded;
7use vortex_array::ArrayRef;
8use vortex_array::ArrayView;
9use vortex_array::IntoArray;
10use vortex_array::arrays::ConstantArray;
11use vortex_array::builtins::ArrayBuiltins;
12use vortex_array::dtype::NativeDType;
13use vortex_array::dtype::NativePType;
14use vortex_array::dtype::Nullability;
15use vortex_array::scalar::PValue;
16use vortex_array::scalar::Scalar;
17use vortex_array::scalar_fn::fns::between::BetweenOptions;
18use vortex_array::scalar_fn::fns::between::BetweenReduce;
19use vortex_array::scalar_fn::fns::between::StrictComparison;
20use vortex_error::VortexResult;
21
22use crate::ALP;
23use crate::ALPFloat;
24use crate::Exponents;
25use crate::alp::array::ALPArrayExt;
26use crate::alp::array::ALPArraySlotsExt;
27use crate::match_each_alp_float_ptype;
28
29impl BetweenReduce for ALP {
30    fn between(
31        array: ArrayView<'_, Self>,
32        lower: &ArrayRef,
33        upper: &ArrayRef,
34        options: &BetweenOptions,
35    ) -> VortexResult<Option<ArrayRef>> {
36        let (Some(lower), Some(upper)) = (lower.as_constant(), upper.as_constant()) else {
37            return Ok(None);
38        };
39
40        if array.patches().is_some() {
41            return Ok(None);
42        }
43
44        let nullability =
45            array.dtype().nullability() | lower.dtype().nullability() | upper.dtype().nullability();
46        match_each_alp_float_ptype!(array.dtype().as_ptype(), |F| {
47            between_impl::<F>(
48                array,
49                F::try_from(&lower)?,
50                F::try_from(&upper)?,
51                nullability,
52                options,
53            )
54        })
55        .map(Some)
56    }
57}
58
59fn between_impl<T: NativePType + ALPFloat>(
60    array: ArrayView<'_, ALP>,
61    lower: T,
62    upper: T,
63    nullability: Nullability,
64    options: &BetweenOptions,
65) -> VortexResult<ArrayRef>
66where
67    Scalar: From<T::ALPInt>,
68    <T as ALPFloat>::ALPInt: NativeDType + NativePType + Into<PValue> + Debug,
69{
70    let exponents = array.exponents();
71
72    // There are always compared
73    // the below bound is `value {< | <=} x`, either value encodes into the ALPInt domain
74    // in which case we can leave the comparison unchanged `enc(value) {< | <=} x` or it doesn't
75    // and we encode into value below enc_below(value) < value < x, in which case the comparison
76    // becomes enc(value) < x. See `alp_scalar_compare` for more details.
77    // note that if the value doesn't encode than value != x, so must use strict comparison.
78    let Some((lower_enc, lower_strict)) =
79        encode_lower_bound::<T>(lower, exponents, options.lower_strict)
80    else {
81        return Ok(ConstantArray::new(Scalar::bool(false, nullability), array.len()).into_array());
82    };
83
84    // the upper value `x { < | <= } value` similarly encodes or `x < value < enc_above(value())`
85    let Some((upper_enc, upper_strict)) =
86        encode_upper_bound::<T>(upper, exponents, options.upper_strict)
87    else {
88        return Ok(ConstantArray::new(Scalar::bool(false, nullability), array.len()).into_array());
89    };
90
91    let options = BetweenOptions {
92        lower_strict,
93        upper_strict,
94    };
95
96    array.encoded().clone().between(
97        ConstantArray::new(Scalar::primitive(lower_enc, nullability), array.len()).into_array(),
98        ConstantArray::new(Scalar::primitive(upper_enc, nullability), array.len()).into_array(),
99        options,
100    )
101}
102
103fn encode_lower_bound<T: ALPFloat + NativePType>(
104    lower: T,
105    exponents: Exponents,
106    strict: StrictComparison,
107) -> Option<(T::ALPInt, StrictComparison)> {
108    if NativePType::is_nan(lower) || NativePType::is_infinite(lower) {
109        return NativePType::is_lt(lower, T::zero())
110            .then_some((T::ALPInt::min_value(), StrictComparison::NonStrict));
111    }
112
113    Some(
114        T::encode_single(lower, exponents)
115            .map(|x| (x, strict))
116            .unwrap_or_else(|| (T::encode_below(lower, exponents), StrictComparison::Strict)),
117    )
118}
119
120fn encode_upper_bound<T: ALPFloat + NativePType>(
121    upper: T,
122    exponents: Exponents,
123    strict: StrictComparison,
124) -> Option<(T::ALPInt, StrictComparison)> {
125    if NativePType::is_nan(upper) || NativePType::is_infinite(upper) {
126        return NativePType::is_gt(upper, T::zero())
127            .then_some((T::ALPInt::max_value(), StrictComparison::NonStrict));
128    }
129
130    Some(
131        T::encode_single(upper, exponents)
132            .map(|x| (x, strict))
133            .unwrap_or_else(|| (T::encode_above(upper, exponents), StrictComparison::Strict)),
134    )
135}
136
137#[cfg(test)]
138mod tests {
139    use std::sync::LazyLock;
140
141    use vortex_array::IntoArray;
142    use vortex_array::VortexSessionExecute;
143    use vortex_array::arrays::ConstantArray;
144    use vortex_array::arrays::PrimitiveArray;
145    use vortex_array::assert_arrays_eq;
146    use vortex_array::dtype::Nullability;
147    use vortex_array::scalar::Scalar;
148    use vortex_array::scalar_fn::fns::between::BetweenOptions;
149    use vortex_array::scalar_fn::fns::between::StrictComparison;
150    use vortex_session::VortexSession;
151
152    use crate::ALPArray;
153    use crate::alp::array::ALPArrayExt;
154    use crate::alp::compute::between::between_impl;
155    use crate::alp_encode;
156
157    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
158        let session = vortex_array::array_session();
159        crate::initialize(&session);
160        session
161    });
162
163    fn assert_between(
164        arr: &ALPArray,
165        lower: f32,
166        upper: f32,
167        options: &BetweenOptions,
168        expected: bool,
169    ) {
170        let res =
171            between_impl(arr.as_view(), lower, upper, Nullability::Nullable, options).unwrap();
172        assert_arrays_eq!(
173            res,
174            ConstantArray::new(Scalar::bool(expected, res.dtype().nullability()), arr.len())
175                .into_array(),
176            &mut SESSION.create_execution_ctx()
177        );
178    }
179
180    #[test]
181    fn comparison_range() {
182        let value = 0.0605_f32;
183        let array = PrimitiveArray::from_iter([value; 1]);
184        let encoded =
185            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
186        assert!(encoded.patches().is_none());
187
188        assert_between(
189            &encoded,
190            0.0605_f32,
191            0.0605,
192            &BetweenOptions {
193                lower_strict: StrictComparison::NonStrict,
194                upper_strict: StrictComparison::NonStrict,
195            },
196            true,
197        );
198
199        assert_between(
200            &encoded,
201            0.0605_f32,
202            0.0605,
203            &BetweenOptions {
204                lower_strict: StrictComparison::Strict,
205                upper_strict: StrictComparison::NonStrict,
206            },
207            false,
208        );
209
210        assert_between(
211            &encoded,
212            0.0605_f32,
213            0.0605,
214            &BetweenOptions {
215                lower_strict: StrictComparison::NonStrict,
216                upper_strict: StrictComparison::Strict,
217            },
218            false,
219        );
220
221        assert_between(
222            &encoded,
223            0.060499_f32,
224            0.06051,
225            &BetweenOptions {
226                lower_strict: StrictComparison::NonStrict,
227                upper_strict: StrictComparison::NonStrict,
228            },
229            true,
230        );
231
232        assert_between(
233            &encoded,
234            0.06_f32,
235            0.06051,
236            &BetweenOptions {
237                lower_strict: StrictComparison::NonStrict,
238                upper_strict: StrictComparison::Strict,
239            },
240            true,
241        );
242    }
243
244    #[test]
245    fn non_finite_bounds_use_total_order() {
246        let mut ctx = SESSION.create_execution_ctx();
247        let array = PrimitiveArray::from_iter([1.234f32; 10]);
248        let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
249        assert!(encoded.patches().is_none());
250
251        let options = BetweenOptions {
252            lower_strict: StrictComparison::NonStrict,
253            upper_strict: StrictComparison::Strict,
254        };
255
256        assert_between(&encoded, f32::from_bits(0xffffff5e), 2.0, &options, true);
257        assert_between(&encoded, f32::NAN, 2.0, &options, false);
258        assert_between(&encoded, f32::NEG_INFINITY, 2.0, &options, true);
259        assert_between(&encoded, f32::INFINITY, 2.0, &options, false);
260
261        assert_between(&encoded, 0.0, f32::NAN, &options, true);
262        assert_between(&encoded, 0.0, f32::from_bits(0xffffff5e), &options, false);
263        assert_between(&encoded, 0.0, f32::INFINITY, &options, true);
264        assert_between(&encoded, 0.0, f32::NEG_INFINITY, &options, false);
265    }
266}
vortex_alp/alp/compute/between.rs

vortex_alp/alp/compute/
between.rs
