vortex_runend/compute/
binary_numeric.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
use vortex_array::array::ConstantArray;
use vortex_array::compute::{binary_numeric, BinaryNumericFn};
use vortex_array::{ArrayData, ArrayLen, IntoArrayData};
use vortex_error::VortexResult;
use vortex_scalar::BinaryNumericOperator;

use crate::{RunEndArray, RunEndEncoding};

impl BinaryNumericFn<RunEndArray> for RunEndEncoding {
    fn binary_numeric(
        &self,
        array: &RunEndArray,
        rhs: &ArrayData,
        op: BinaryNumericOperator,
    ) -> VortexResult<Option<ArrayData>> {
        let Some(rhs_scalar) = rhs.as_constant() else {
            return Ok(None);
        };

        let rhs_const_array = ConstantArray::new(rhs_scalar, array.values().len()).into_array();

        RunEndArray::with_offset_and_length(
            array.ends(),
            binary_numeric(&array.values(), &rhs_const_array, op)?,
            array.offset(),
            array.len(),
        )
        .map(IntoArrayData::into_array)
        .map(Some)
    }
}

#[cfg(test)]
mod tests {
    use vortex_array::array::PrimitiveArray;
    use vortex_array::compute::test_harness::test_binary_numeric;
    use vortex_array::{IntoArrayData, ToArrayData};

    use crate::RunEndArray;

    fn ree_array() -> RunEndArray {
        RunEndArray::encode(
            PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).to_array(),
        )
        .unwrap()
    }

    #[test]
    fn test_runend_binary_numeric() {
        let array = ree_array().into_array();
        test_binary_numeric::<i32>(array)
    }
}