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        // 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                binary_numeric(array.values(), &rhs_const_array, op)?,
30            )?
31            .into_array(),
32        ))
33    }
34}
35
36#[cfg(test)]
37mod tests {
38    use vortex_array::ArrayRef;
39    use vortex_array::arrays::PrimitiveArray;
40    use vortex_array::compute::conformance::binary_numeric::test_binary_numeric;
41    use vortex_array::compute::slice;
42
43    use crate::builders::dict_encode;
44
45    fn sliced_dict_array() -> ArrayRef {
46        let reference = PrimitiveArray::from_option_iter([
47            Some(42),
48            Some(-9),
49            None,
50            Some(42),
51            Some(1),
52            Some(5),
53        ]);
54        let dict = dict_encode(&reference).unwrap();
55        slice(&dict, 1, 4).unwrap()
56    }
57
58    #[test]
59    fn test_dict_binary_numeric() {
60        let array = sliced_dict_array();
61        test_binary_numeric::<i32>(array)
62    }
63}
vortex_dict/compute/binary_numeric.rs

vortex_dict/compute/
binary_numeric.rs
