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            .map(|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::{ConstantArray, PrimitiveArray};
48    use vortex_array::compute::{Operator, compare};
49    use vortex_array::{IntoArray, ToCanonical};
50    use vortex_buffer::BitBuffer;
51
52    use crate::RunEndArray;
53
54    fn ree_array() -> RunEndArray {
55        RunEndArray::encode(
56            PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(),
57        )
58        .unwrap()
59    }
60
61    #[test]
62    fn compare_run_end() {
63        let arr = ree_array();
64        let res = compare(
65            arr.as_ref(),
66            ConstantArray::new(5, 12).as_ref(),
67            Operator::Eq,
68        )
69        .unwrap();
70        let res_canon = res.to_bool();
71        assert_eq!(
72            res_canon.bit_buffer(),
73            &BitBuffer::from(vec![
74                false, false, false, false, false, false, false, false, true, true, true, true
75            ])
76        );
77    }
78}
vortex_runend/compute/compare.rs

vortex_runend/compute/
compare.rs
