vortex_compute/arithmetic/
pscalar.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::NativePType;
5use vortex_vector::primitive::PScalar;
6
7use crate::arithmetic::Arithmetic;
8use crate::arithmetic::CheckedArithmetic;
9use crate::arithmetic::CheckedOperator;
10use crate::arithmetic::Operator;
11
12impl<Op, T> Arithmetic<Op> for &PScalar<T>
13where
14    T: NativePType,
15    Op: Operator<T>,
16{
17    type Output = PScalar<T>;
18
19    fn eval(self, rhs: &PScalar<T>) -> Self::Output {
20        match (self.value(), rhs.value()) {
21            (Some(a), Some(b)) => {
22                let value = Op::apply(&a, &b);
23                PScalar::new(Some(value))
24            }
25            (..) => {
26                // At least one side is null, so result is null
27                PScalar::new(None)
28            }
29        }
30    }
31}
32
33impl<Op, T> CheckedArithmetic<Op> for &PScalar<T>
34where
35    T: NativePType,
36    Op: CheckedOperator<T>,
37{
38    type Output = PScalar<T>;
39
40    fn checked_eval(self, rhs: Self) -> Option<Self::Output> {
41        match (self.value(), rhs.value()) {
42            (Some(a), Some(b)) => {
43                let value = Op::apply(&a, &b)?;
44                Some(PScalar::new(Some(value)))
45            }
46            (..) => {
47                // At least one side is null, so result is null
48                Some(PScalar::new(None))
49            }
50        }
51    }
52}
53
54#[cfg(test)]
55mod test {
56    use vortex_vector::primitive::PScalar;
57
58    use super::*;
59    use crate::arithmetic::Add;
60    use crate::arithmetic::CheckedArithmetic;
61    use crate::arithmetic::WrappingSub;
62
63    #[test]
64    fn test_add() {
65        let left = PScalar::new(Some(5u32));
66        let right = PScalar::new(Some(3u32));
67
68        let result = CheckedArithmetic::<Add>::checked_eval(&left, &right).unwrap();
69        assert_eq!(result.value(), Some(8u32));
70
71        let left_null = PScalar::new(None);
72        let result_null = CheckedArithmetic::<Add>::checked_eval(&left_null, &right).unwrap();
73        assert_eq!(result_null.value(), None);
74    }
75
76    #[test]
77    fn test_subtract() {
78        let left = PScalar::new(Some(10u32));
79        let right = PScalar::new(Some(4u32));
80
81        let result = Arithmetic::<WrappingSub>::eval(&left, &right);
82        assert_eq!(result.value(), Some(6u32));
83
84        let right_null = PScalar::new(None);
85        let result_null = Arithmetic::<WrappingSub>::eval(&left, &right_null);
86        assert_eq!(result_null.value(), None);
87    }
88}