vortex_runend/compute/
compare.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::arrays::ConstantArray;
5use vortex_array::compute::{CompareKernel, CompareKernelAdapter, Operator, compare};
6use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical, register_kernel};
7use vortex_error::VortexResult;
8
9use crate::compress::runend_decode_bools;
10use crate::{RunEndArray, RunEndVTable};
11
12impl CompareKernel for RunEndVTable {
13    fn compare(
14        &self,
15        lhs: &RunEndArray,
16        rhs: &dyn Array,
17        operator: Operator,
18    ) -> VortexResult<Option<ArrayRef>> {
19        // If the RHS is constant, then we just need to compare against our encoded values.
20        if let Some(const_scalar) = rhs.as_constant() {
21            return compare(
22                lhs.values(),
23                ConstantArray::new(const_scalar, lhs.values().len()).as_ref(),
24                operator,
25            )
26            .and_then(|values| {
27                runend_decode_bools(
28                    lhs.ends().to_primitive()?,
29                    values.to_bool()?,
30                    lhs.offset(),
31                    lhs.len(),
32                )
33            })
34            .map(|a| a.into_array())
35            .map(Some);
36        }
37
38        // Otherwise, fall back
39        Ok(None)
40    }
41}
42
43register_kernel!(CompareKernelAdapter(RunEndVTable).lift());
44
45#[cfg(test)]
46mod test {
47    use vortex_array::arrays::{BooleanBuffer, ConstantArray, PrimitiveArray};
48    use vortex_array::compute::{Operator, compare};
49    use vortex_array::{IntoArray, ToCanonical};
50
51    use crate::RunEndArray;
52
53    fn ree_array() -> RunEndArray {
54        RunEndArray::encode(
55            PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(),
56        )
57        .unwrap()
58    }
59
60    #[test]
61    fn compare_run_end() {
62        let arr = ree_array();
63        let res = compare(
64            arr.as_ref(),
65            ConstantArray::new(5, 12).as_ref(),
66            Operator::Eq,
67        )
68        .unwrap();
69        let res_canon = res.to_bool().unwrap();
70        assert_eq!(
71            res_canon.boolean_buffer(),
72            &BooleanBuffer::from(vec![
73                false, false, false, false, false, false, false, false, true, true, true, true
74            ])
75        );
76    }
77}
vortex_runend/compute/compare.rs

vortex_runend/compute/
compare.rs
