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    use vortex_buffer::buffer;
49
50    use crate::RunEndArray;
51
52    #[rstest]
53    #[case::runend_i32_basic(RunEndArray::encode(
54        buffer![10i32, 10, 10, 20, 20, 30, 30, 30, 30].into_array()
55    ).unwrap())]
56    #[case::runend_u32_basic(RunEndArray::encode(
57        buffer![100u32, 100, 200, 200, 200].into_array()
58    ).unwrap())]
59    #[case::runend_i64_basic(RunEndArray::encode(
60        buffer![1000i64, 1000, 2000, 2000, 3000, 3000].into_array()
61    ).unwrap())]
62    #[case::runend_u64_basic(RunEndArray::encode(
63        buffer![5000u64, 5000, 5000, 6000, 6000].into_array()
64    ).unwrap())]
65    #[case::runend_f32_basic(RunEndArray::encode(
66        buffer![1.5f32, 1.5, 2.5, 2.5, 3.5].into_array()
67    ).unwrap())]
68    #[case::runend_f64_basic(RunEndArray::encode(
69        buffer![10.1f64, 10.1, 20.2, 20.2, 20.2].into_array()
70    ).unwrap())]
71    #[case::runend_i32_large(RunEndArray::encode(
72        PrimitiveArray::from_iter((0..100).map(|i| i / 5)).into_array()
73    ).unwrap())]
74    fn test_runend_binary_numeric(#[case] array: RunEndArray) {
75        test_binary_numeric_array(array.into_array());
76    }
77}