remainder operation on series/ chunkedarrays

pola-rs · Oct 3, 2020 · fd3822a · fd3822a
1 parent b923894
commit fd3822a
Show file tree

Hide file tree

Showing 6 changed files with 69 additions and 21 deletions.
diff --git a/polars/src/chunked_array/arithmetic.rs b/polars/src/chunked_array/arithmetic.rs
@@ -3,7 +3,7 @@ use crate::prelude::*;
 use crate::utils::Xob;
 use arrow::{array::ArrayRef, compute};
 use num::{Num, NumCast, ToPrimitive};
-use std::ops::{Add, Div, Mul, Sub};
+use std::ops::{Add, Div, Mul, Rem, Sub};
 use std::sync::Arc;
 
 // TODO: Add Boolean arithmetic
@@ -134,6 +134,18 @@ where
     }
 }
 
+impl<T> Rem for &ChunkedArray<T>
+where
+    T: PolarsNumericType,
+    T::Native: Rem<Output = T::Native>,
+{
+    type Output = ChunkedArray<T>;
+
+    fn rem(self, rhs: Self) -> Self::Output {
+        apply_operand_on_chunkedarray_by_iter!(self, rhs, %)
+    }
+}
+
 impl<T> Sub for &ChunkedArray<T>
 where
     T: PolarsNumericType,
@@ -167,11 +179,7 @@ where
     type Output = Self;
 
     fn add(self, rhs: Self) -> Self::Output {
-        if self.chunk_id == rhs.chunk_id {
-            (&self).add(&rhs)
-        } else {
-            apply_operand_on_chunkedarray_by_iter!(self, rhs, +)
-        }
+        (&self).add(&rhs)
     }
 }
 
@@ -224,6 +232,18 @@ where
     }
 }
 
+impl<T> Rem for ChunkedArray<T>
+where
+    T: PolarsNumericType,
+    T::Native: Rem<Output = T::Native>,
+{
+    type Output = ChunkedArray<T>;
+
+    fn rem(self, rhs: Self) -> Self::Output {
+        (&self).rem(&rhs)
+    }
+}
+
 // Operands on ChunkedArray & Num
 
 impl<T, N> Add<N> for &ChunkedArray<T>
@@ -333,6 +353,28 @@ where
     }
 }
 
+impl<T, N> Rem<N> for &ChunkedArray<T>
+where
+    T: PolarsNumericType,
+    T::Native: NumCast,
+    N: Num + ToPrimitive,
+    T::Native: Rem<Output = T::Native>,
+{
+    type Output = ChunkedArray<T>;
+
+    fn rem(self, rhs: N) -> Self::Output {
+        let operand: T::Native = NumCast::from(rhs).unwrap();
+        if self.is_optimal_aligned() {
+            let intermed: Xob<_> = self.into_no_null_iter().map(|val| val % operand).collect();
+            intermed.into_inner()
+        } else {
+            self.into_iter()
+                .map(|opt_val| opt_val.map(|val| val % operand))
+                .collect()
+        }
+    }
+}
+
 pub trait Pow {
     fn pow_f32(&self, exp: f32) -> Float32Chunked;
     fn pow_f64(&self, exp: f64) -> Float64Chunked;

diff --git a/polars/src/chunked_array/comparison.rs b/polars/src/chunked_array/comparison.rs
@@ -865,14 +865,8 @@ mod test {
     fn test_bitwise_ops() {
         let a = BooleanChunked::new_from_slice("a", &[true, false, false]);
         let b = BooleanChunked::new_from_opt_slice("b", &[Some(true), Some(true), None]);
-        assert_eq!(
-            Vec::from((&a | &b).unwrap()),
-            &[Some(true), Some(true), None]
-        );
-        assert_eq!(
-            Vec::from((&a & &b).unwrap()),
-            &[Some(true), Some(false), None]
-        );
+        assert_eq!(Vec::from(&a | &b), &[Some(true), Some(true), None]);
+        assert_eq!(Vec::from(&a & &b), &[Some(true), Some(false), None]);
         assert_eq!(Vec::from(!b), &[Some(false), Some(false), None]);
     }
 

diff --git a/polars/src/frame/mod.rs b/polars/src/frame/mod.rs
@@ -828,10 +828,8 @@ impl DataFrame {
     /// let mut df = DataFrame::new(vec![s0, s1]).unwrap();
     ///
     /// // create a mask
-    /// let mask = || {
-    ///     (df.column("values")?.lt_eq(1) | df.column("values")?.gt_eq(5))
-    /// };
-    /// let mask = mask().unwrap();
+    /// let values = df.column("values").unwrap();
+    /// let mask = values.lt_eq(1) | values.gt_eq(5);
     ///
     /// df.may_apply("foo", |s| {
     ///     s.utf8()?

diff --git a/polars/src/lazy/physical_plan/expressions.rs b/polars/src/lazy/physical_plan/expressions.rs
@@ -92,7 +92,9 @@ impl PhysicalExpr for BinaryExpr {
             Operator::Not => Ok(apply_method_all_series!(left, eq_series, &right).into_series()),
             Operator::Like => todo!(),
             Operator::NotLike => todo!(),
-            Operator::Modulus => todo!(),
+            Operator::Modulus => {
+                apply_method_all_series!(left, remainder, &right).map(|ca| ca.into_series())
+            }
         }
     }
     fn to_field(&self, _input_schema: &Schema) -> Result<Field> {

diff --git a/polars/src/prelude.rs b/polars/src/prelude.rs
@@ -30,7 +30,10 @@ pub use crate::{
         },
         DataFrame, IntoSeries,
     },
-    series::{arithmetic::LhsNumOps, NamedFrom, Series},
+    series::{
+        arithmetic::{LhsNumOps, NumOpsDispatch},
+        NamedFrom, Series,
+    },
     testing::*,
 };
 pub use arrow::datatypes::{ArrowPrimitiveType, Field, Schema};

diff --git a/polars/src/series/arithmetic.rs b/polars/src/series/arithmetic.rs
@@ -2,7 +2,7 @@ use crate::prelude::*;
 use num::{Num, NumCast, ToPrimitive};
 use std::ops;
 
-pub(super) trait NumOpsDispatch {
+pub trait NumOpsDispatch {
     fn subtract(&self, _rhs: &Series) -> Result<Series> {
         Err(PolarsError::InvalidOperation)
     }
@@ -15,6 +15,9 @@ pub(super) trait NumOpsDispatch {
     fn divide(&self, _rhs: &Series) -> Result<Series> {
         Err(PolarsError::InvalidOperation)
     }
+    fn remainder(&self, _rhs: &Series) -> Result<Series> {
+        Err(PolarsError::InvalidOperation)
+    }
 }
 
 impl<T> NumOpsDispatch for ChunkedArray<T>
@@ -24,6 +27,7 @@ where
         + ops::Sub<Output = T::Native>
         + ops::Mul<Output = T::Native>
         + ops::Div<Output = T::Native>
+        + ops::Rem<Output = T::Native>
         + num::Zero
         + num::One,
 {
@@ -47,6 +51,11 @@ where
         let out = self / rhs;
         Ok(out.into_series())
     }
+    fn remainder(&self, rhs: &Series) -> Result<Series> {
+        let rhs = self.unpack_series_matching_type(rhs)?;
+        let out = self % rhs;
+        Ok(out.into_series())
+    }
 }
 
 impl NumOpsDispatch for Utf8Chunked {}