vortex_dict/compute/
binary_numeric.rs

1use vortex_array::arrays::ConstantArray;
2use vortex_array::compute::{NumericKernel, NumericKernelAdapter, numeric};
3use vortex_array::{Array, ArrayRef, IntoArray, register_kernel};
4use vortex_error::VortexResult;
5use vortex_scalar::NumericOperator;
6
7use crate::{DictArray, DictVTable};
8
9impl NumericKernel for DictVTable {
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(DictVTable).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
44    use crate::builders::dict_encode;
45
46    fn sliced_dict_array() -> ArrayRef {
47        let reference = PrimitiveArray::from_option_iter([
48            Some(42),
49            Some(-9),
50            None,
51            Some(42),
52            Some(1),
53            Some(5),
54        ]);
55        let dict = dict_encode(reference.as_ref()).unwrap();
56        dict.slice(1, 4).unwrap()
57    }
58
59    #[test]
60    fn test_dict_binary_numeric() {
61        let array = sliced_dict_array();
62        test_numeric::<i32>(array)
63    }
64}
vortex_dict/compute/binary_numeric.rs

vortex_dict/compute/
binary_numeric.rs
