vortex_runend/compute/
compare.rs

1use vortex_array::arrays::ConstantArray;
2use vortex_array::compute::{CompareKernel, CompareKernelAdapter, Operator, compare};
3use vortex_array::{Array, ArrayRef, ToCanonical, register_kernel};
4use vortex_error::VortexResult;
5
6use crate::compress::runend_decode_bools;
7use crate::{RunEndArray, RunEndEncoding};
8
9impl CompareKernel for RunEndEncoding {
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()),
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(RunEndEncoding).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::{Array, 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(&arr, &ConstantArray::new(5, 12), Operator::Eq).unwrap();
61        let res_canon = res.to_bool().unwrap();
62        assert_eq!(
63            res_canon.boolean_buffer(),
64            &BooleanBuffer::from(vec![
65                false, false, false, false, false, false, false, false, true, true, true, true
66            ])
67        );
68    }
69}
vortex_runend/compute/compare.rs

vortex_runend/compute/
compare.rs
