refactor(rust): make binary chunkedarray functions DRY (#10607)

pola-rs · Aug 19, 2023 · 7dff4c3 · 7dff4c3
1 parent d7bc251
commit 7dff4c3
Show file tree

Hide file tree

Showing 17 changed files with 378 additions and 413 deletions.
diff --git a/crates/polars-core/src/chunked_array/arithmetic/decimal.rs b/crates/polars-core/src/chunked_array/arithmetic/decimal.rs
@@ -2,6 +2,7 @@ use polars_arrow::compute::arithmetics::decimal;
 
 use super::*;
 use crate::prelude::DecimalChunked;
+use crate::utils::align_chunks_binary;
 
 // TODO: remove
 impl ArrayArithmetics for i128 {

diff --git a/crates/polars-core/src/chunked_array/arithmetic/mod.rs b/crates/polars-core/src/chunked_array/arithmetic/mod.rs
@@ -15,7 +15,7 @@ use polars_arrow::utils::combine_validities_and;
 
 use crate::prelude::*;
 use crate::series::IsSorted;
-use crate::utils::{align_chunks_binary, align_chunks_binary_owned};
+use crate::utils::align_chunks_binary_owned;
 
 pub trait ArrayArithmetics
 where
@@ -148,12 +148,7 @@ impl Add for &BinaryChunked {
             };
         }
 
-        let (lhs, rhs) = align_chunks_binary(self, rhs);
-        let chunks = lhs
-            .downcast_iter()
-            .zip(rhs.downcast_iter())
-            .map(|(a, b)| concat_binary(a, b));
-        ChunkedArray::from_chunk_iter(self.name(), chunks)
+        arity::binary_mut(self, rhs, concat_binary)
     }
 }
 
@@ -202,12 +197,7 @@ impl Add for &BooleanChunked {
         if self.len() == 1 {
             return rhs.add(self);
         }
-        let (lhs, rhs) = align_chunks_binary(self, rhs);
-        let chunks = lhs
-            .downcast_iter()
-            .zip(rhs.downcast_iter())
-            .map(|(a, b)| add_boolean(a, b));
-        ChunkedArray::from_chunk_iter(self.name(), chunks)
+        arity::binary_mut(self, rhs, add_boolean)
     }
 }
 

diff --git a/crates/polars-core/src/chunked_array/arithmetic/numeric.rs b/crates/polars-core/src/chunked_array/arithmetic/numeric.rs
@@ -12,15 +12,7 @@ where
     F: Fn(T::Native, T::Native) -> T::Native,
 {
     let mut ca = match (lhs.len(), rhs.len()) {
-        (a, b) if a == b => {
-            let (lhs, rhs) = align_chunks_binary(lhs, rhs);
-            let chunks = lhs
-                .downcast_iter()
-                .zip(rhs.downcast_iter())
-                .map(|(lhs, rhs)| Box::new(kernel(lhs, rhs)) as ArrayRef)
-                .collect();
-            unsafe { lhs.copy_with_chunks(chunks, false, false) }
-        },
+        (a, b) if a == b => arity::binary_mut(lhs, rhs, |lhs, rhs| kernel(lhs, rhs)),
         // broadcast right path
         (_, 1) => {
             let opt_rhs = rhs.get(0);

diff --git a/crates/polars-core/src/chunked_array/bitwise.rs b/crates/polars-core/src/chunked_array/bitwise.rs
@@ -6,7 +6,6 @@ use polars_arrow::utils::combine_validities_and;
 
 use super::arithmetic::arithmetic_helper;
 use super::*;
-use crate::utils::align_chunks_binary;
 
 impl<T> BitAnd for &ChunkedArray<T>
 where
@@ -73,12 +72,7 @@ impl BitOr for &BooleanChunked {
             _ => {},
         }
 
-        let (lhs, rhs) = align_chunks_binary(self, rhs);
-        let chunks = lhs
-            .downcast_iter()
-            .zip(rhs.downcast_iter())
-            .map(|(lhs, rhs)| compute::boolean_kleene::or(lhs, rhs));
-        BooleanChunked::from_chunk_iter(self.name(), chunks)
+        arity::binary_mut(self, rhs, compute::boolean_kleene::or)
     }
 }
 
@@ -123,16 +117,11 @@ impl BitXor for &BooleanChunked {
             _ => {},
         }
 
-        let (l, r) = align_chunks_binary(self, rhs);
-        let chunks = l
-            .downcast_iter()
-            .zip(r.downcast_iter())
-            .map(|(l_arr, r_arr)| {
-                let validity = combine_validities_and(l_arr.validity(), r_arr.validity());
-                let values = l_arr.values() ^ r_arr.values();
-                BooleanArray::from_data_default(values, validity)
-            });
-        ChunkedArray::from_chunk_iter(self.name(), chunks)
+        arity::binary_mut(self, rhs, |l_arr, r_arr| {
+            let validity = combine_validities_and(l_arr.validity(), r_arr.validity());
+            let values = l_arr.values() ^ r_arr.values();
+            BooleanArray::from_data_default(values, validity)
+        })
     }
 }
 
@@ -169,12 +158,7 @@ impl BitAnd for &BooleanChunked {
             _ => {},
         }
 
-        let (lhs, rhs) = align_chunks_binary(self, rhs);
-        let chunks = lhs
-            .downcast_iter()
-            .zip(rhs.downcast_iter())
-            .map(|(lhs, rhs)| compute::boolean_kleene::and(lhs, rhs));
-        BooleanChunked::from_chunk_iter(self.name(), chunks)
+        arity::binary_mut(self, rhs, compute::boolean_kleene::and)
     }
 }