vortex_runend/compute/
binary_numeric.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::arrays::ConstantArray;
5use vortex_array::compute::{NumericKernel, NumericKernelAdapter, numeric};
6use vortex_array::{Array, ArrayRef, IntoArray, register_kernel};
7use vortex_error::VortexResult;
8use vortex_scalar::NumericOperator;
9
10use crate::{RunEndArray, RunEndVTable};
11
12impl NumericKernel for RunEndVTable {
13    fn numeric(
14        &self,
15        array: &RunEndArray,
16        rhs: &dyn Array,
17        op: NumericOperator,
18    ) -> VortexResult<Option<ArrayRef>> {
19        let Some(rhs_scalar) = rhs.as_constant() else {
20            return Ok(None);
21        };
22
23        let rhs_const_array = ConstantArray::new(rhs_scalar, array.values().len()).into_array();
24
25        // SAFETY: ends are preserved.
26        unsafe {
27            Ok(Some(
28                RunEndArray::new_unchecked(
29                    array.ends().clone(),
30                    numeric(array.values(), &rhs_const_array, op)?,
31                    array.offset(),
32                    array.len(),
33                )
34                .into_array(),
35            ))
36        }
37    }
38}
39
40register_kernel!(NumericKernelAdapter(RunEndVTable).lift());
41
42#[cfg(test)]
43mod tests {
44    use rstest::rstest;
45    use vortex_array::IntoArray;
46    use vortex_array::arrays::PrimitiveArray;
47    use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array;
48
49    use crate::RunEndArray;
50
51    #[rstest]
52    #[case::runend_i32_basic(RunEndArray::encode(
53        PrimitiveArray::from_iter([10i32, 10, 10, 20, 20, 30, 30, 30, 30]).into_array()
54    ).unwrap())]
55    #[case::runend_u32_basic(RunEndArray::encode(
56        PrimitiveArray::from_iter([100u32, 100, 200, 200, 200]).into_array()
57    ).unwrap())]
58    #[case::runend_i64_basic(RunEndArray::encode(
59        PrimitiveArray::from_iter([1000i64, 1000, 2000, 2000, 3000, 3000]).into_array()
60    ).unwrap())]
61    #[case::runend_u64_basic(RunEndArray::encode(
62        PrimitiveArray::from_iter([5000u64, 5000, 5000, 6000, 6000]).into_array()
63    ).unwrap())]
64    #[case::runend_f32_basic(RunEndArray::encode(
65        PrimitiveArray::from_iter([1.5f32, 1.5, 2.5, 2.5, 3.5]).into_array()
66    ).unwrap())]
67    #[case::runend_f64_basic(RunEndArray::encode(
68        PrimitiveArray::from_iter([10.1f64, 10.1, 20.2, 20.2, 20.2]).into_array()
69    ).unwrap())]
70    #[case::runend_i32_large(RunEndArray::encode(
71        PrimitiveArray::from_iter((0..100).map(|i| i / 5)).into_array()
72    ).unwrap())]
73    fn test_runend_binary_numeric(#[case] array: RunEndArray) {
74        test_binary_numeric_array(array.into_array());
75    }
76}
vortex_runend/compute/binary_numeric.rs

vortex_runend/compute/
binary_numeric.rs
