Add missing arithmetic ops for ChunkedArray

pola-rs · Sep 4, 2020 · 0e7d82b · 0e7d82b
1 parent afb7a7b
commit 0e7d82b
Showing 1 changed file with 112 additions and 1 deletion.
diff --git a/polars/src/chunked_array/arithmetic.rs b/polars/src/chunked_array/arithmetic.rs
@@ -2,7 +2,7 @@
 use crate::prelude::*;
 use crate::utils::Xob;
 use arrow::{array::ArrayRef, compute};
-use num::ToPrimitive;
+use num::{Num, NumCast, ToPrimitive};
 use std::ops::{Add, Div, Mul, Sub};
 use std::sync::Arc;
 
@@ -56,6 +56,8 @@ macro_rules! apply_operand_on_chunkedarray_by_iter {
     }
 }
 
+// Operands on ChunkedArray & ChunkedArray
+
 impl<T> Add for &ChunkedArray<T>
 where
     T: PolarsNumericType,
@@ -210,6 +212,115 @@ where
     }
 }
 
+// Operands on ChunkedArray & Num
+
+impl<T, N> Add<N> for &ChunkedArray<T>
+where
+    T: PolarsNumericType,
+    T::Native: NumCast,
+    N: Num + ToPrimitive,
+    T::Native: Add<Output = T::Native>
+        + Sub<Output = T::Native>
+        + Mul<Output = T::Native>
+        + Div<Output = T::Native>
+        + num::Zero,
+{
+    type Output = ChunkedArray<T>;
+
+    fn add(self, rhs: N) -> Self::Output {
+        let adder: T::Native = NumCast::from(rhs).unwrap();
+        if self.is_optimal_aligned() {
+            let intermed: Xob<_> = self.into_no_null_iter().map(|val| val + adder).collect();
+            intermed.into_inner()
+        } else {
+            self.into_iter()
+                .map(|opt_val| opt_val.map(|val| val + adder))
+                .collect()
+        }
+    }
+}
+
+impl<T, N> Sub<N> for &ChunkedArray<T>
+where
+    T: PolarsNumericType,
+    T::Native: NumCast,
+    N: Num + ToPrimitive,
+    T::Native: Add<Output = T::Native>
+        + Sub<Output = T::Native>
+        + Mul<Output = T::Native>
+        + Div<Output = T::Native>
+        + num::Zero,
+{
+    type Output = ChunkedArray<T>;
+
+    fn sub(self, rhs: N) -> Self::Output {
+        let subber: T::Native = NumCast::from(rhs).unwrap();
+        if self.is_optimal_aligned() {
+            let intermed: Xob<_> = self.into_no_null_iter().map(|val| val - subber).collect();
+            intermed.into_inner()
+        } else {
+            self.into_iter()
+                .map(|opt_val| opt_val.map(|val| val - subber))
+                .collect()
+        }
+    }
+}
+
+impl<T, N> Div<N> for &ChunkedArray<T>
+where
+    T: PolarsNumericType,
+    T::Native: NumCast,
+    N: Num + ToPrimitive,
+    T::Native: Add<Output = T::Native>
+        + Sub<Output = T::Native>
+        + Mul<Output = T::Native>
+        + Div<Output = T::Native>
+        + num::Zero,
+{
+    type Output = ChunkedArray<T>;
+
+    fn div(self, rhs: N) -> Self::Output {
+        let divider: T::Native = NumCast::from(rhs).unwrap();
+        if self.is_optimal_aligned() {
+            let intermed: Xob<_> = self.into_no_null_iter().map(|val| val / divider).collect();
+            intermed.into_inner()
+        } else {
+            self.into_iter()
+                .map(|opt_val| opt_val.map(|val| val / divider))
+                .collect()
+        }
+    }
+}
+
+impl<T, N> Mul<N> for &ChunkedArray<T>
+where
+    T: PolarsNumericType,
+    T::Native: NumCast,
+    N: Num + ToPrimitive,
+    T::Native: Add<Output = T::Native>
+        + Sub<Output = T::Native>
+        + Mul<Output = T::Native>
+        + Div<Output = T::Native>
+        + num::Zero,
+{
+    type Output = ChunkedArray<T>;
+
+    fn mul(self, rhs: N) -> Self::Output {
+        let multiplier: T::Native = NumCast::from(rhs).unwrap();
+        if self.is_optimal_aligned() {
+            let intermed: Xob<_> = self
+                .into_no_null_iter()
+                .map(|val| val * multiplier)
+                .collect();
+            intermed.into_inner()
+        } else {
+            self.into_iter()
+                .map(|opt_val| opt_val.map(|val| val * multiplier))
+                .collect()
+        }
+    }
+}
+
 pub trait Pow {
     fn pow_f32(&self, exp: f32) -> Float32Chunked;
     fn pow_f64(&self, exp: f64) -> Float64Chunked;