cleanup old arithmetic code

pola-rs · Apr 1, 2021 · 5261fd8 · 5261fd8
1 parent 47f7598
commit 5261fd8
Show file tree

Hide file tree

Showing 5 changed files with 228 additions and 154 deletions.
diff --git a/polars/polars-core/src/chunked_array/arithmetic.rs b/polars/polars-core/src/chunked_array/arithmetic.rs
@@ -2,8 +2,9 @@
 use crate::prelude::*;
 use crate::utils::{align_chunks_binary, NoNull};
 use arrow::array::PrimitiveArray;
+use arrow::compute::divide_scalar;
 use arrow::{array::ArrayRef, compute};
-use num::{Num, NumCast, ToPrimitive};
+use num::{Num, NumCast, One, ToPrimitive, Zero};
 use std::ops::{Add, Div, Mul, Rem, Sub};
 use std::sync::Arc;
 
@@ -301,15 +302,14 @@ where
 impl<T, N> Div<N> for &ChunkedArray<T>
 where
     T: PolarsNumericType,
-    T::Native: NumCast,
+    T::Native: NumCast + Div<Output = T::Native> + One + Zero + Sub<Output = T::Native>,
     N: Num + ToPrimitive,
-    T::Native: Div<Output = T::Native>,
 {
     type Output = ChunkedArray<T>;
 
     fn div(self, rhs: N) -> Self::Output {
-        let divider: T::Native = NumCast::from(rhs).unwrap();
-        self.apply(|val| val / divider)
+        let rhs: T::Native = NumCast::from(rhs).expect("could not cast");
+        self.apply_kernel(|arr| Arc::new(divide_scalar(arr, rhs).unwrap()))
     }
 }
 

diff --git a/polars/polars-core/src/chunked_array/ops/apply.rs b/polars/polars-core/src/chunked_array/ops/apply.rs
@@ -75,22 +75,17 @@ where
     where
         F: Fn(T::Native) -> T::Native + Copy,
     {
-        if let Ok(slice) = self.cont_slice() {
-            let new: NoNull<ChunkedArray<T>> = slice.iter().copied().map(f).collect();
-            new.into_inner()
-        } else {
-            let mut ca: ChunkedArray<T> = self
-                .data_views()
-                .into_iter()
-                .zip(self.null_bits())
-                .map(|(slice, (_null_count, opt_buffer))| {
-                    let vec: AlignedVec<_> = slice.iter().copied().map(f).collect();
-                    (vec, opt_buffer)
-                })
-                .collect();
-            ca.rename(self.name());
-            ca
-        }
+        let mut ca: ChunkedArray<T> = self
+            .data_views()
+            .into_iter()
+            .zip(self.null_bits())
+            .map(|(slice, (_null_count, opt_buffer))| {
+                let vec: AlignedVec<_> = slice.iter().copied().map(f).collect();
+                (vec, opt_buffer)
+            })
+            .collect();
+        ca.rename(self.name());
+        ca
     }
 
     fn apply_with_idx<F>(&'a self, f: F) -> Self