fix bug in date/datetime/time arithmetic dispatch

pola-rs · Oct 20, 2021 · 1a6db29 · 1a6db29
1 parent 98ba743
commit 1a6db29
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 21 deletions.
diff --git a/polars/polars-core/src/frame/groupby/mod.rs b/polars/polars-core/src/frame/groupby/mod.rs
@@ -359,7 +359,7 @@ impl DataFrame {
                     _ => {
                         // is date like
                         if !s.is_numeric() && s.is_numeric_physical() {
-                            s.to_physical_repr()
+                            s.to_physical_repr().into_owned()
                         } else {
                             s.clone()
                         }

diff --git a/polars/polars-core/src/frame/groupby/resample.rs b/polars/polars-core/src/frame/groupby/resample.rs
@@ -216,7 +216,7 @@ impl DataFrame {
         let temp_key = "__POLAR_TEMP_NAME";
 
         let key_dtype = key.dtype();
-        let mut key_phys = key.to_physical_repr();
+        let mut key_phys = key.to_physical_repr().into_owned();
         let mut key = key.clone();
         let key_name = key.name().to_string();
         let wrong_key_dtype = || {

diff --git a/polars/polars-core/src/series/arithmetic.rs b/polars/polars-core/src/series/arithmetic.rs
@@ -396,13 +396,14 @@ where
     type Output = Series;
 
     fn sub(self, rhs: T) -> Self::Output {
+        let s = self.to_physical_repr();
         macro_rules! sub {
             ($ca:expr) => {{
                 $ca.sub(rhs).into_series()
             }};
         }
 
-        match_arrow_data_type_apply_macro_ca_logical_num!(self, sub)
+        match_arrow_data_type_apply_macro_ca_logical_num!(s, sub)
     }
 }
 
@@ -424,12 +425,13 @@ where
     type Output = Series;
 
     fn add(self, rhs: T) -> Self::Output {
+        let s = self.to_physical_repr();
         macro_rules! add {
             ($ca:expr) => {{
                 $ca.add(rhs).into_series()
             }};
         }
-        match_arrow_data_type_apply_macro_ca_logical_num!(self, add)
+        match_arrow_data_type_apply_macro_ca_logical_num!(s, add)
     }
 }
 
@@ -451,13 +453,14 @@ where
     type Output = Series;
 
     fn div(self, rhs: T) -> Self::Output {
+        let s = self.to_physical_repr();
         macro_rules! div {
             ($ca:expr) => {{
                 $ca.div(rhs).into_series()
             }};
         }
 
-        match_arrow_data_type_apply_macro_ca_logical_num!(self, div)
+        match_arrow_data_type_apply_macro_ca_logical_num!(s, div)
     }
 }
 
@@ -479,12 +482,13 @@ where
     type Output = Series;
 
     fn mul(self, rhs: T) -> Self::Output {
+        let s = self.to_physical_repr();
         macro_rules! mul {
             ($ca:expr) => {{
                 $ca.mul(rhs).into_series()
             }};
         }
-        match_arrow_data_type_apply_macro_ca_logical_num!(self, mul)
+        match_arrow_data_type_apply_macro_ca_logical_num!(s, mul)
     }
 }
 
@@ -506,12 +510,13 @@ where
     type Output = Series;
 
     fn rem(self, rhs: T) -> Self::Output {
+        let s = self.to_physical_repr();
         macro_rules! rem {
             ($ca:expr) => {{
                 $ca.rem(rhs).into_series()
             }};
         }
-        match_arrow_data_type_apply_macro_ca_logical_num!(self, rem)
+        match_arrow_data_type_apply_macro_ca_logical_num!(s, rem)
     }
 }
 

diff --git a/polars/polars-core/src/series/implementations/dates_time.rs b/polars/polars-core/src/series/implementations/dates_time.rs
@@ -85,7 +85,7 @@ macro_rules! impl_dyn_series {
 
             #[cfg(feature = "zip_with")]
             fn zip_with_same_type(&self, mask: &BooleanChunked, other: &Series) -> Result<Series> {
-                let other = other.to_physical_repr();
+                let other = other.to_physical_repr().into_owned();
                 self.0
                     .zip_with(mask, &other.as_ref().as_ref())
                     .map(|ca| ca.$into_logical().into_series())
@@ -188,23 +188,23 @@ macro_rules! impl_dyn_series {
                 self.0.pivot_count(pivot_series, keys, groups)
             }
             fn hash_join_inner(&self, other: &Series) -> Vec<(u32, u32)> {
-                let other = other.to_physical_repr();
+                let other = other.to_physical_repr().into_owned();
                 self.0.hash_join_inner(&other.as_ref().as_ref())
             }
             fn hash_join_left(&self, other: &Series) -> Vec<(u32, Option<u32>)> {
-                let other = other.to_physical_repr();
+                let other = other.to_physical_repr().into_owned();
                 self.0.hash_join_left(&other.as_ref().as_ref())
             }
             fn hash_join_outer(&self, other: &Series) -> Vec<(Option<u32>, Option<u32>)> {
-                let other = other.to_physical_repr();
+                let other = other.to_physical_repr().into_owned();
                 self.0.hash_join_outer(&other.as_ref().as_ref())
             }
             fn zip_outer_join_column(
                 &self,
                 right_column: &Series,
                 opt_join_tuples: &[(Option<u32>, Option<u32>)],
             ) -> Series {
-                let right_column = right_column.to_physical_repr();
+                let right_column = right_column.to_physical_repr().into_owned();
                 self.0
                     .zip_outer_join_column(&right_column, opt_join_tuples)
                     .$into_logical()
@@ -353,7 +353,7 @@ macro_rules! impl_dyn_series {
 
             fn append(&mut self, other: &Series) -> Result<()> {
                 if self.0.dtype() == other.dtype() {
-                    let other = other.to_physical_repr();
+                    let other = other.to_physical_repr().into_owned();
                     self.0.append(other.as_ref().as_ref());
                     Ok(())
                 } else {
@@ -754,4 +754,17 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn test_arithmetic_dispatch() {
+        let s = Int64Chunked::new_from_slice("", &[1, 2, 3])
+            .into_date()
+            .into_series();
+
+        // check if we don't panic.
+        let _ = &s * 100;
+        let _ = &s / 100;
+        let _ = &s + 100;
+        let _ = &s - 100;
+    }
 }
diff --git a/polars/polars-core/src/series/mod.rs b/polars/polars-core/src/series/mod.rs
@@ -24,6 +24,7 @@ pub use from::*;
 use num::NumCast;
 use rayon::prelude::*;
 pub use series_trait::*;
+use std::borrow::Cow;
 #[cfg(feature = "groupby_list")]
 use std::hash::{Hash, Hasher};
 use std::ops::Deref;
@@ -323,14 +324,13 @@ impl Series {
     /// * Date -> Int32
     /// * Datetime-> Int64
     ///
-    pub fn to_physical_repr(&self) -> Series {
+    pub fn to_physical_repr(&self) -> Cow<Series> {
         use DataType::*;
-        let out = match self.dtype() {
-            Date => self.cast(&DataType::Int32),
-            Datetime => self.cast(&DataType::Int64),
-            _ => return self.clone(),
-        };
-        out.unwrap()
+        match self.dtype() {
+            Date => Cow::Owned(self.cast(&DataType::Int32).unwrap()),
+            Datetime => Cow::Owned(self.cast(&DataType::Int64).unwrap()),
+            _ => Cow::Borrowed(self),
+        }
     }
 
     /// Take by index if ChunkedArray contains a single chunk.

diff --git a/polars/polars-core/src/utils.rs b/polars/polars-core/src/utils.rs
@@ -331,7 +331,7 @@ macro_rules! match_arrow_data_type_apply_macro_ca_logical_num {
             DataType::Int64 => $macro!($self.i64().unwrap() $(, $opt_args)*),
             DataType::Float32 => $macro!($self.f32().unwrap() $(, $opt_args)*),
             DataType::Float64 => $macro!($self.f64().unwrap() $(, $opt_args)*),
-            _ => unimplemented!(),
+            dt => panic!("not implemented for {:?}", dt),
         }
     }};
 }