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        Ok(Some(
26            RunEndArray::with_offset_and_length(
27                array.ends().clone(),
28                numeric(array.values(), &rhs_const_array, op)?,
29                array.offset(),
30                array.len(),
31            )?
32            .into_array(),
33        ))
34    }
35}
36
37register_kernel!(NumericKernelAdapter(RunEndVTable).lift());
38
39#[cfg(test)]
40mod tests {
41    use rstest::rstest;
42    use vortex_array::IntoArray;
43    use vortex_array::arrays::PrimitiveArray;
44    use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array;
45
46    use crate::RunEndArray;
47
48    #[rstest]
49    #[case::runend_i32_basic(RunEndArray::encode(
50        PrimitiveArray::from_iter([10i32, 10, 10, 20, 20, 30, 30, 30, 30]).into_array()
51    ).unwrap())]
52    #[case::runend_u32_basic(RunEndArray::encode(
53        PrimitiveArray::from_iter([100u32, 100, 200, 200, 200]).into_array()
54    ).unwrap())]
55    #[case::runend_i64_basic(RunEndArray::encode(
56        PrimitiveArray::from_iter([1000i64, 1000, 2000, 2000, 3000, 3000]).into_array()
57    ).unwrap())]
58    #[case::runend_u64_basic(RunEndArray::encode(
59        PrimitiveArray::from_iter([5000u64, 5000, 5000, 6000, 6000]).into_array()
60    ).unwrap())]
61    #[case::runend_f32_basic(RunEndArray::encode(
62        PrimitiveArray::from_iter([1.5f32, 1.5, 2.5, 2.5, 3.5]).into_array()
63    ).unwrap())]
64    #[case::runend_f64_basic(RunEndArray::encode(
65        PrimitiveArray::from_iter([10.1f64, 10.1, 20.2, 20.2, 20.2]).into_array()
66    ).unwrap())]
67    #[case::runend_i32_large(RunEndArray::encode(
68        PrimitiveArray::from_iter((0..100).map(|i| i / 5)).into_array()
69    ).unwrap())]
70    fn test_runend_binary_numeric(#[case] array: RunEndArray) {
71        test_binary_numeric_array(array.into_array());
72    }
73}