vortex_runend/compute/
compare.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayRef;
5use vortex_array::ArrayView;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::arrays::BoolArray;
9use vortex_array::arrays::ConstantArray;
10use vortex_array::arrays::PrimitiveArray;
11use vortex_array::builtins::ArrayBuiltins;
12use vortex_array::scalar_fn::fns::binary::CompareKernel;
13use vortex_array::scalar_fn::fns::operators::CompareOperator;
14use vortex_array::scalar_fn::fns::operators::Operator;
15use vortex_error::VortexResult;
16
17use crate::RunEnd;
18use crate::array::RunEndArrayExt;
19use crate::decompress_bool::runend_decode_bools;
20
21impl CompareKernel for RunEnd {
22    fn compare(
23        lhs: ArrayView<'_, Self>,
24        rhs: &ArrayRef,
25        operator: CompareOperator,
26        ctx: &mut ExecutionCtx,
27    ) -> VortexResult<Option<ArrayRef>> {
28        // If the RHS is constant, then we just need to compare against our encoded values.
29        if let Some(const_scalar) = rhs.as_constant() {
30            let values = lhs.values().binary(
31                ConstantArray::new(const_scalar, lhs.values().len()).into_array(),
32                Operator::from(operator),
33            )?;
34            return runend_decode_bools(
35                lhs.ends().clone().execute::<PrimitiveArray>(ctx)?,
36                values.execute::<BoolArray>(ctx)?,
37                lhs.offset(),
38                lhs.len(),
39                ctx,
40            )
41            .map(Some);
42        }
43
44        // Otherwise, fall back
45        Ok(None)
46    }
47}
48
49#[cfg(test)]
50mod test {
51    use std::sync::LazyLock;
52
53    use vortex_array::ExecutionCtx;
54    use vortex_array::IntoArray;
55    use vortex_array::VortexSessionExecute;
56    use vortex_array::arrays::BoolArray;
57    use vortex_array::arrays::ConstantArray;
58    use vortex_array::arrays::PrimitiveArray;
59    use vortex_array::assert_arrays_eq;
60    use vortex_array::builtins::ArrayBuiltins;
61    use vortex_array::scalar_fn::fns::operators::Operator;
62    use vortex_session::VortexSession;
63
64    use crate::RunEnd;
65    use crate::RunEndArray;
66
67    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
68        let session = vortex_array::array_session();
69        crate::initialize(&session);
70        session
71    });
72
73    fn ree_array(ctx: &mut ExecutionCtx) -> RunEndArray {
74        RunEnd::encode(
75            PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(),
76            ctx,
77        )
78        .unwrap()
79    }
80
81    #[test]
82    fn compare_run_end() {
83        let mut ctx = SESSION.create_execution_ctx();
84        let arr = ree_array(&mut ctx);
85        let res = arr
86            .into_array()
87            .binary(ConstantArray::new(5, 12).into_array(), Operator::Eq)
88            .unwrap();
89        let expected = BoolArray::from_iter([
90            false, false, false, false, false, false, false, false, true, true, true, true,
91        ]);
92        assert_arrays_eq!(res, expected, &mut ctx);
93    }
94}
vortex_runend/compute/compare.rs

vortex_runend/compute/
compare.rs
