fix(rust, python): error on addition with datetime/time (#5931)

pola-rs · Dec 29, 2022 · 7163747 · 7163747
1 parent 13f5cad
commit 7163747
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 18 deletions.
diff --git a/polars/polars-core/src/series/arithmetic/borrowed.rs b/polars/polars-core/src/series/arithmetic/borrowed.rs
@@ -426,22 +426,27 @@ impl ops::Sub for &Series {
     }
 }
 
-impl ops::Add for &Series {
-    type Output = Series;
-
-    fn add(self, rhs: Self) -> Self::Output {
+impl Series {
+    pub fn try_add(&self, rhs: &Series) -> PolarsResult<Series> {
         match (self.dtype(), rhs.dtype()) {
             #[cfg(feature = "dtype-struct")]
             (DataType::Struct(_), DataType::Struct(_)) => {
-                struct_arithmetic(self, rhs, |a, b| a.add(b))
+                Ok(struct_arithmetic(self, rhs, |a, b| a.add(b)))
             }
             _ => {
-                let (lhs, rhs) = coerce_lhs_rhs(self, rhs).expect("cannot coerce datatypes");
-                lhs.add_to(rhs.as_ref()).expect("data types don't match")
+                let (lhs, rhs) = coerce_lhs_rhs(self, rhs)?;
+                lhs.add_to(rhs.as_ref())
             }
         }
     }
 }
+impl ops::Add for &Series {
+    type Output = Series;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        self.try_add(rhs).unwrap()
+    }
+}
 
 impl ops::Mul for &Series {
     type Output = Series;

diff --git a/polars/polars-core/src/utils/supertype.rs b/polars/polars-core/src/utils/supertype.rs
@@ -196,15 +196,6 @@ pub fn get_supertype(l: &DataType, r: &DataType) -> Option<DataType> {
             #[cfg(feature = "dtype-time")]
             (Time, Float64) => Some(Float64),
 
-            #[cfg(all(feature = "dtype-time", feature = "dtype-datetime"))]
-            (Time, Datetime(_, _)) => Some(Int64),
-            #[cfg(all(feature = "dtype-datetime", feature = "dtype-time"))]
-            (Datetime(_, _), Time) => Some(Int64),
-            #[cfg(all(feature = "dtype-time", feature = "dtype-date"))]
-            (Time, Date) => Some(Int64),
-            #[cfg(all(feature = "dtype-date", feature = "dtype-time"))]
-            (Date, Time) => Some(Int64),
-
             // every known type can be casted to a string except binary
             #[cfg(feature = "dtype-binary")]
             (dt, Utf8) if dt != &DataType::Unknown && dt != &DataType::Binary => Some(Utf8),

diff --git a/py-polars/src/series.rs b/py-polars/src/series.rs
@@ -536,8 +536,12 @@ impl PySeries {
         }
     }
 
-    pub fn add(&self, other: &PySeries) -> Self {
-        (&self.series + &other.series).into()
+    pub fn add(&self, other: &PySeries) -> PyResult<Self> {
+        let out = self
+            .series
+            .try_add(&other.series)
+            .map_err(PyPolarsErr::from)?;
+        Ok(out.into())
     }
 
     pub fn sub(&self, other: &PySeries) -> Self {

diff --git a/py-polars/tests/unit/test_errors.py b/py-polars/tests/unit/test_errors.py
@@ -311,3 +311,8 @@ def test_duplicate_columns_arg_csv() -> None:
         ValueError, match=r"'columns' arg should only have unique values"
     ):
         pl.read_csv(f, columns=["x", "x", "y"])
+
+
+def test_datetime_time_add_err() -> None:
+    with pytest.raises(pl.ComputeError):
+        pl.Series([datetime(1970, 1, 1, 0, 0, 1)]) + pl.Series([time(0, 0, 2)])