vortex_runend/compute/
binary_numeric.rs

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