vortex_runend/compute/
compare.rs

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

vortex_runend/compute/
compare.rs
