nd_arrow_array/compute/
ops.rs

1//! Operator overloads: `&a + &b` calls [`add`](crate::compute::add) and returns a `Result`.
2
3use std::ops;
4
5use crate::NdArrowArray;
6use crate::compute::elementwise;
7use crate::error::Result;
8
9macro_rules! binary_operator {
10    ($($trait:ident :: $method:ident => $kernel:path),* $(,)?) => {$(
11        impl ops::$trait<&NdArrowArray> for &NdArrowArray {
12            type Output = Result<NdArrowArray>;
13
14            fn $method(self, rhs: &NdArrowArray) -> Self::Output {
15                $kernel(self, rhs)
16            }
17        }
18    )*};
19}
20
21binary_operator!(
22    Add::add => elementwise::add,
23    Sub::sub => elementwise::sub,
24    Mul::mul => elementwise::mul,
25    Div::div => elementwise::div,
26    Rem::rem => elementwise::rem,
27    BitAnd::bitand => elementwise::logical_and,
28    BitOr::bitor => elementwise::logical_or,
29    BitXor::bitxor => elementwise::logical_xor,
30);
31
32impl ops::Neg for &NdArrowArray {
33    type Output = Result<NdArrowArray>;
34
35    fn neg(self) -> Self::Output {
36        elementwise::neg(self)
37    }
38}
39
40impl ops::Not for &NdArrowArray {
41    type Output = Result<NdArrowArray>;
42
43    fn not(self) -> Self::Output {
44        elementwise::logical_not(self)
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use std::sync::Arc;
51
52    use arrow::array::{AsArray, Int32Array};
53    use arrow::datatypes::Int32Type;
54
55    use crate::{Dimensions, IntoScalar, NdArrowArray};
56
57    #[test]
58    fn operators_call_kernels() {
59        let a = NdArrowArray::try_new(
60            Arc::new(Int32Array::from(vec![1, 2, 3])),
61            Dimensions::try_new([("x", 3)]).unwrap(),
62        )
63        .unwrap();
64        let two = NdArrowArray::try_new_scalar(2i32.into_scalar()).unwrap();
65        let b = (&(&a * &two).unwrap() - &a).unwrap();
66        assert_eq!(b.values().as_primitive::<Int32Type>().values(), &[1, 2, 3]);
67        let n = (-&a).unwrap();
68        assert_eq!(
69            n.values().as_primitive::<Int32Type>().values(),
70            &[-1, -2, -3]
71        );
72    }
73}
nd_arrow_array/compute/ops.rs

nd_arrow_array/compute/
ops.rs
