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

vortex_runend/compute/
compare.rs
