vortex_dict/compute/
binary_numeric.rs

1use vortex_array::arrays::ConstantArray;
2use vortex_array::compute::{BinaryNumericFn, binary_numeric};
3use vortex_array::{Array, ArrayRef};
4use vortex_error::VortexResult;
5use vortex_scalar::BinaryNumericOperator;
6
7use crate::{DictArray, DictEncoding};
8
9impl BinaryNumericFn<&DictArray> for DictEncoding {
10    fn binary_numeric(
11        &self,
12        array: &DictArray,
13        rhs: &dyn Array,
14        op: BinaryNumericOperator,
15    ) -> VortexResult<Option<ArrayRef>> {
16        let Some(rhs_scalar) = rhs.as_constant() else {
17            return Ok(None);
18        };
19        let rhs_const_array = ConstantArray::new(rhs_scalar, array.values().len()).into_array();
20
21        Ok(Some(
22            DictArray::try_new(
23                array.codes().clone(),
24                binary_numeric(array.values(), &rhs_const_array, op)?,
25            )?
26            .into_array(),
27        ))
28    }
29}
30
31#[cfg(test)]
32mod tests {
33    use vortex_array::ArrayRef;
34    use vortex_array::arrays::PrimitiveArray;
35    use vortex_array::compute::slice;
36    use vortex_array::compute::test_harness::test_binary_numeric;
37
38    use crate::builders::dict_encode;
39
40    fn sliced_dict_array() -> ArrayRef {
41        let reference = PrimitiveArray::from_option_iter([
42            Some(42),
43            Some(-9),
44            None,
45            Some(42),
46            Some(1),
47            Some(5),
48        ]);
49        let dict = dict_encode(&reference).unwrap();
50        slice(&dict, 1, 4).unwrap()
51    }
52
53    #[test]
54    fn test_dict_binary_numeric() {
55        let array = sliced_dict_array();
56        test_binary_numeric::<i32>(array)
57    }
58}
vortex_dict/compute/binary_numeric.rs

vortex_dict/compute/
binary_numeric.rs
