fix: Allow sub on date/datetime (#17207)

pola-rs · Jun 26, 2024 · fcd4dbb · fcd4dbb
1 parent 30e4d63
commit fcd4dbb
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/crates/polars-plan/src/plans/aexpr/schema.rs b/crates/polars-plan/src/plans/aexpr/schema.rs
@@ -327,6 +327,7 @@ fn get_arithmetic_field(
                 | (Date, Duration(_))
                 | (Duration(_), Time)
                 | (Time, Duration(_)) => try_get_supertype(left_field.data_type(), &right_type)?,
+                (Datetime(tu, _), Date) | (Date, Datetime(tu, _)) => Duration(*tu),
                 // T - T != T if T is a datetime / date
                 (Datetime(tul, _), Datetime(tur, _)) => Duration(get_time_units(tul, tur)),
                 (_, Datetime(_, _)) | (Datetime(_, _), _) => {

diff --git a/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py b/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py
@@ -748,3 +748,15 @@ def test_invalid_shapes_err() -> None:
         match=r"cannot do arithmetic operation on series of different lengths: got 2 and 3",
     ):
         pl.Series([1, 2]) + pl.Series([1, 2, 3])
+
+
+def test_date_datetime_sub() -> None:
+    df = pl.DataFrame({"foo": [date(2020, 1, 1)], "bar": [datetime(2020, 1, 5)]})
+
+    assert df.select(
+        pl.col("foo") - pl.col("bar"),
+        pl.col("bar") - pl.col("foo"),
+    ).to_dict(as_series=False) == {
+        "foo": [timedelta(days=-4)],
+        "bar": [timedelta(days=4)],
+    }