vortex_runend/compute/
compare.rs

1use vortex_array::arrays::ConstantArray;
2use vortex_array::compute::{CompareFn, Operator, compare};
3use vortex_array::{Array, ArrayRef, ToCanonical};
4use vortex_error::VortexResult;
5
6use crate::compress::runend_decode_bools;
7use crate::{RunEndArray, RunEndEncoding};
8
9impl CompareFn<&RunEndArray> 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#[cfg(test)]
40mod test {
41    use vortex_array::arrays::{BooleanBuffer, ConstantArray, PrimitiveArray};
42    use vortex_array::compute::{Operator, compare};
43    use vortex_array::{Array, ToCanonical};
44
45    use crate::RunEndArray;
46
47    fn ree_array() -> RunEndArray {
48        RunEndArray::encode(
49            PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(),
50        )
51        .unwrap()
52    }
53
54    #[test]
55    fn compare_run_end() {
56        let arr = ree_array();
57        let res = compare(&arr, &ConstantArray::new(5, 12), Operator::Eq).unwrap();
58        let res_canon = res.to_bool().unwrap();
59        assert_eq!(
60            res_canon.boolean_buffer(),
61            &BooleanBuffer::from(vec![
62                false, false, false, false, false, false, false, false, true, true, true, true
63            ])
64        );
65    }
66}
vortex_runend/compute/compare.rs

vortex_runend/compute/
compare.rs
