vortex_alp/alp/compute/
compare.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::ExecutionCtx;
9use vortex_array::IntoArray;
10use vortex_array::arrays::ConstantArray;
11use vortex_array::builtins::ArrayBuiltins;
12use vortex_array::dtype::NativePType;
13use vortex_array::scalar::Scalar;
14use vortex_array::scalar_fn::fns::binary::CompareKernel;
15use vortex_array::scalar_fn::fns::operators::CompareOperator;
16use vortex_array::scalar_fn::fns::operators::Operator;
17use vortex_error::VortexResult;
18use vortex_error::vortex_bail;
19use vortex_error::vortex_err;
20
21use crate::ALP;
22use crate::ALPArrayExt;
23use crate::ALPArraySlotsExt;
24use crate::ALPFloat;
25use crate::match_each_alp_float_ptype;
26
27// TODO(joe): add fuzzing.
28
29impl CompareKernel for ALP {
30    fn compare(
31        lhs: ArrayView<'_, Self>,
32        rhs: &ArrayRef,
33        operator: CompareOperator,
34        _ctx: &mut ExecutionCtx,
35    ) -> VortexResult<Option<ArrayRef>> {
36        if lhs.patches().is_some() {
37            // TODO(joe): support patches
38            return Ok(None);
39        }
40        if lhs.dtype().is_nullable() || rhs.dtype().is_nullable() {
41            // TODO(joe): support nullability
42            return Ok(None);
43        }
44
45        if let Some(const_scalar) = rhs.as_constant() {
46            let pscalar = const_scalar.as_primitive_opt().ok_or_else(|| {
47                vortex_err!(
48                    "ALP Compare RHS had the wrong type {}, expected {}",
49                    const_scalar,
50                    const_scalar.dtype()
51                )
52            })?;
53
54            match_each_alp_float_ptype!(pscalar.ptype(), |T| {
55                match pscalar.typed_value::<T>() {
56                    Some(value) => return alp_scalar_compare(lhs, value, operator),
57                    None => vortex_bail!(
58                        "Failed to convert scalar {:?} to ALP type {:?}",
59                        pscalar,
60                        pscalar.ptype()
61                    ),
62                }
63            });
64        }
65
66        Ok(None)
67    }
68}
69
70/// We can compare a scalar to an ALPArray by encoding the scalar into the ALP domain and comparing
71/// the encoded value to the encoded values in the ALPArray. There are fixups when the value doesn't
72/// encode into the ALP domain.
73fn alp_scalar_compare<F: ALPFloat + Into<Scalar>>(
74    alp: ArrayView<ALP>,
75    value: F,
76    operator: CompareOperator,
77) -> VortexResult<Option<ArrayRef>>
78where
79    F::ALPInt: Into<Scalar>,
80    <F as ALPFloat>::ALPInt: Debug,
81{
82    // TODO(joe): support patches, this is checked above.
83    if alp.patches().is_some() {
84        return Ok(None);
85    }
86
87    let exponents = alp.exponents();
88    // If the scalar doesn't fit into the ALP domain,
89    // it cannot be equal to any values in the encoded array.
90    let encoded = F::encode_single(value, alp.exponents());
91    match encoded {
92        Some(encoded) => {
93            let s = ConstantArray::new(encoded, alp.len());
94            Ok(Some(
95                alp.encoded()
96                    .binary(s.into_array(), Operator::from(operator))?,
97            ))
98        }
99        None => match operator {
100            // Since this value is not encodable it cannot be equal to any value in the encoded
101            // array.
102            CompareOperator::Eq => Ok(Some(ConstantArray::new(false, alp.len()).into_array())),
103            // Since this value is not encodable it cannot be equal to any value in the encoded
104            // array, hence != to all values in the encoded array.
105            CompareOperator::NotEq => Ok(Some(ConstantArray::new(true, alp.len()).into_array())),
106            CompareOperator::Gt | CompareOperator::Gte => {
107                // Per IEEE 754 totalOrder semantics the ordering is -Nan < -Inf < Inf < Nan.
108                // All values in the encoded array are definitely finite
109                let is_not_finite = NativePType::is_infinite(value) || NativePType::is_nan(value);
110                if is_not_finite {
111                    Ok(Some(
112                        ConstantArray::new(value.is_sign_negative(), alp.len()).into_array(),
113                    ))
114                } else {
115                    Ok(Some(
116                        alp.encoded().binary(
117                            ConstantArray::new(F::encode_above(value, exponents), alp.len())
118                                .into_array(),
119                            // Since the encoded value is unencodable gte is equivalent to gt.
120                            // Consider a value v, between two encodable values v_l (just less) and
121                            // v_a (just above), then for all encodable values (u), v > u <=> v_g >= u
122                            Operator::Gte,
123                        )?,
124                    ))
125                }
126            }
127            CompareOperator::Lt | CompareOperator::Lte => {
128                // Per IEEE 754 totalOrder semantics the ordering is -Nan < -Inf < Inf < Nan.
129                // All values in the encoded array are definitely finite
130                let is_not_finite = NativePType::is_infinite(value) || NativePType::is_nan(value);
131                if is_not_finite {
132                    Ok(Some(
133                        ConstantArray::new(value.is_sign_positive(), alp.len()).into_array(),
134                    ))
135                } else {
136                    Ok(Some(
137                        alp.encoded().binary(
138                            ConstantArray::new(F::encode_below(value, exponents), alp.len())
139                                .into_array(),
140                            // Since the encoded values unencodable lt is equivalent to lte.
141                            // See Gt | Gte for further explanation.
142                            Operator::Lte,
143                        )?,
144                    ))
145                }
146            }
147        },
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use std::f32;
154    use std::sync::LazyLock;
155
156    use rstest::rstest;
157    use vortex_array::ArrayRef;
158    use vortex_array::VortexSessionExecute;
159    use vortex_array::arrays::BoolArray;
160    use vortex_array::arrays::ConstantArray;
161    use vortex_array::arrays::PrimitiveArray;
162    use vortex_array::assert_arrays_eq;
163    use vortex_array::builtins::ArrayBuiltins;
164    use vortex_array::dtype::DType;
165    use vortex_array::dtype::Nullability;
166    use vortex_array::dtype::PType;
167    use vortex_array::scalar::Scalar;
168    use vortex_array::scalar_fn::fns::operators::CompareOperator;
169    use vortex_array::scalar_fn::fns::operators::Operator;
170    use vortex_session::VortexSession;
171
172    use super::*;
173    use crate::alp_encode;
174
175    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
176        let session = vortex_array::array_session();
177        crate::initialize(&session);
178        session
179    });
180
181    fn test_alp_compare<F: ALPFloat + Into<Scalar>>(
182        alp: ArrayView<ALP>,
183        value: F,
184        operator: CompareOperator,
185    ) -> Option<ArrayRef>
186    where
187        F::ALPInt: Into<Scalar>,
188        <F as ALPFloat>::ALPInt: Debug,
189    {
190        alp_scalar_compare(alp, value, operator).unwrap()
191    }
192
193    #[test]
194    fn basic_comparison_test() {
195        let mut ctx = SESSION.create_execution_ctx();
196        let array = PrimitiveArray::from_iter([1.234f32; 1025]);
197        let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
198        assert!(encoded.patches().is_none());
199        let encoded_prim = encoded
200            .encoded()
201            .clone()
202            .execute::<PrimitiveArray>(&mut ctx)
203            .unwrap();
204        assert_eq!(encoded_prim.as_slice::<i32>(), vec![1234; 1025]);
205
206        let r = alp_scalar_compare(encoded.as_view(), 1.3_f32, CompareOperator::Eq)
207            .unwrap()
208            .unwrap();
209        let expected = BoolArray::from_iter([false; 1025]);
210        assert_arrays_eq!(r, expected, &mut ctx);
211
212        let r = alp_scalar_compare(encoded.as_view(), 1.234f32, CompareOperator::Eq)
213            .unwrap()
214            .unwrap();
215        let expected = BoolArray::from_iter([true; 1025]);
216        assert_arrays_eq!(r, expected, &mut ctx);
217    }
218
219    #[test]
220    fn comparison_with_unencodable_value() {
221        let mut ctx = SESSION.create_execution_ctx();
222        let array = PrimitiveArray::from_iter([1.234f32; 1025]);
223        let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
224        assert!(encoded.patches().is_none());
225        let encoded_prim = encoded
226            .encoded()
227            .clone()
228            .execute::<PrimitiveArray>(&mut ctx)
229            .unwrap();
230        assert_eq!(encoded_prim.as_slice::<i32>(), vec![1234; 1025]);
231
232        let r_eq = alp_scalar_compare(encoded.as_view(), 1.234444_f32, CompareOperator::Eq)
233            .unwrap()
234            .unwrap();
235        let expected = BoolArray::from_iter([false; 1025]);
236        assert_arrays_eq!(r_eq, expected, &mut ctx);
237
238        let r_neq = alp_scalar_compare(encoded.as_view(), 1.234444f32, CompareOperator::NotEq)
239            .unwrap()
240            .unwrap();
241        let expected = BoolArray::from_iter([true; 1025]);
242        assert_arrays_eq!(r_neq, expected, &mut ctx);
243    }
244
245    #[test]
246    fn comparison_range() {
247        let mut ctx = SESSION.create_execution_ctx();
248        let array = PrimitiveArray::from_iter([0.0605_f32; 10]);
249        let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
250        assert!(encoded.patches().is_none());
251        let encoded_prim = encoded
252            .encoded()
253            .clone()
254            .execute::<PrimitiveArray>(&mut ctx)
255            .unwrap();
256        assert_eq!(encoded_prim.as_slice::<i32>(), vec![605; 10]);
257
258        // !(0.0605_f32 >= 0.06051_f32);
259        let r_gte = alp_scalar_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Gte)
260            .unwrap()
261            .unwrap();
262        let expected = BoolArray::from_iter([false; 10]);
263        assert_arrays_eq!(r_gte, expected, &mut ctx);
264
265        // (0.0605_f32 > 0.06051_f32);
266        let r_gt = alp_scalar_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Gt)
267            .unwrap()
268            .unwrap();
269        let expected = BoolArray::from_iter([false; 10]);
270        assert_arrays_eq!(r_gt, expected, &mut ctx);
271
272        // 0.0605_f32 <= 0.06051_f32;
273        let r_lte = alp_scalar_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lte)
274            .unwrap()
275            .unwrap();
276        let expected = BoolArray::from_iter([true; 10]);
277        assert_arrays_eq!(r_lte, expected, &mut ctx);
278
279        // 0.0605_f32 < 0.06051_f32;
280        let r_lt = alp_scalar_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lt)
281            .unwrap()
282            .unwrap();
283        let expected = BoolArray::from_iter([true; 10]);
284        assert_arrays_eq!(r_lt, expected, &mut ctx);
285    }
286
287    #[test]
288    fn comparison_zeroes() {
289        let mut ctx = SESSION.create_execution_ctx();
290        let array = PrimitiveArray::from_iter([0.0_f32; 10]);
291        let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
292        assert!(encoded.patches().is_none());
293        let encoded_prim = encoded
294            .encoded()
295            .clone()
296            .execute::<PrimitiveArray>(&mut ctx)
297            .unwrap();
298        assert_eq!(encoded_prim.as_slice::<i32>(), vec![0; 10]);
299
300        let r_gte =
301            test_alp_compare(encoded.as_view(), -0.00000001_f32, CompareOperator::Gte).unwrap();
302        let expected = BoolArray::from_iter([true; 10]);
303        assert_arrays_eq!(r_gte, expected, &mut ctx);
304
305        let r_gte = test_alp_compare(encoded.as_view(), -0.0_f32, CompareOperator::Gte).unwrap();
306        let expected = BoolArray::from_iter([true; 10]);
307        assert_arrays_eq!(r_gte, expected, &mut ctx);
308
309        let r_gt =
310            test_alp_compare(encoded.as_view(), -0.0000000001f32, CompareOperator::Gt).unwrap();
311        let expected = BoolArray::from_iter([true; 10]);
312        assert_arrays_eq!(r_gt, expected, &mut ctx);
313
314        let r_gte = test_alp_compare(encoded.as_view(), -0.0_f32, CompareOperator::Gt).unwrap();
315        let expected = BoolArray::from_iter([true; 10]);
316        assert_arrays_eq!(r_gte, expected, &mut ctx);
317
318        let r_lte = test_alp_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lte).unwrap();
319        let expected = BoolArray::from_iter([true; 10]);
320        assert_arrays_eq!(r_lte, expected, &mut ctx);
321
322        let r_lt = test_alp_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lt).unwrap();
323        let expected = BoolArray::from_iter([true; 10]);
324        assert_arrays_eq!(r_lt, expected, &mut ctx);
325
326        let r_lt = test_alp_compare(encoded.as_view(), -0.00001_f32, CompareOperator::Lt).unwrap();
327        let expected = BoolArray::from_iter([false; 10]);
328        assert_arrays_eq!(r_lt, expected, &mut ctx);
329    }
330
331    #[test]
332    fn compare_with_patches() {
333        let array = PrimitiveArray::from_iter([1.234f32, 1.5, 19.0, f32::consts::E, 1_000_000.9]);
334        let encoded =
335            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
336        assert!(encoded.patches().is_some());
337
338        // Not supported!
339        assert!(
340            alp_scalar_compare(encoded.as_view(), 1_000_000.9_f32, CompareOperator::Eq)
341                .unwrap()
342                .is_none()
343        )
344    }
345
346    #[test]
347    fn compare_to_null() {
348        let array = PrimitiveArray::from_iter([1.234f32; 10]);
349        let encoded =
350            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
351
352        let other = ConstantArray::new(
353            Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable)),
354            array.len(),
355        );
356
357        let r = encoded
358            .into_array()
359            .binary(other.into_array(), Operator::Eq)
360            .unwrap();
361        // Comparing to null yields null results
362        let expected = BoolArray::from_iter([None::<bool>; 10]);
363        assert_arrays_eq!(r, expected, &mut SESSION.create_execution_ctx());
364    }
365
366    #[rstest]
367    #[case(f32::NAN, false)]
368    #[case(-1.0f32 / 0.0f32, true)]
369    #[case(f32::INFINITY, false)]
370    #[case(f32::NEG_INFINITY, true)]
371    fn compare_to_non_finite_gt(#[case] value: f32, #[case] result: bool) {
372        let array = PrimitiveArray::from_iter([1.234f32; 10]);
373        let encoded =
374            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
375
376        let r = test_alp_compare(encoded.as_view(), value, CompareOperator::Gt).unwrap();
377        let expected = BoolArray::from_iter([result; 10]);
378        assert_arrays_eq!(r, expected, &mut SESSION.create_execution_ctx());
379    }
380
381    #[rstest]
382    #[case(f32::NAN, true)]
383    #[case(-1.0f32 / 0.0f32, false)]
384    #[case(f32::INFINITY, true)]
385    #[case(f32::NEG_INFINITY, false)]
386    fn compare_to_non_finite_lt(#[case] value: f32, #[case] result: bool) {
387        let array = PrimitiveArray::from_iter([1.234f32; 10]);
388        let encoded =
389            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
390
391        let r = test_alp_compare(encoded.as_view(), value, CompareOperator::Lt).unwrap();
392        let expected = BoolArray::from_iter([result; 10]);
393        assert_arrays_eq!(r, expected, &mut SESSION.create_execution_ctx());
394    }
395}
vortex_alp/alp/compute/compare.rs

vortex_alp/alp/compute/
compare.rs
