vortex_runend/compute/
compare.rs

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

vortex_runend/compute/
compare.rs
