vortex_dict/compute/
binary_numeric.rs

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

vortex_dict/compute/
binary_numeric.rs
