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

vortex_runend/compute/
compare.rs
