From 0c2eb01bb3815ee37219a06c03812fe9584c011b Mon Sep 17 00:00:00 2001
From: Simon Braun <48355812+ptiza@users.noreply.github.com>
Date: Sun, 15 Dec 2024 09:04:12 +0000
Subject: [PATCH] feat: Allow substraction of time dtype columns (#20300)

Co-authored-by: ritchie <ritchie46@gmail.com>
---
 .../src/series/implementations/time.rs        |  9 ++++++-
 crates/polars-plan/src/plans/aexpr/schema.rs  |  1 +
 .../operations/arithmetic/test_arithmetic.py  | 25 +++++++++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/crates/polars-core/src/series/implementations/time.rs b/crates/polars-core/src/series/implementations/time.rs
index 74b57a48f3c2..5fe2c64bd542 100644
--- a/crates/polars-core/src/series/implementations/time.rs
+++ b/crates/polars-core/src/series/implementations/time.rs
@@ -88,7 +88,14 @@ impl private::PrivateSeries for SeriesWrap<TimeChunked> {
     }
 
     fn subtract(&self, rhs: &Series) -> PolarsResult<Series> {
-        polars_bail!(opq = sub, DataType::Time, rhs.dtype());
+        let rhs = rhs.time().map_err(|_| polars_err!(InvalidOperation: "cannot subtract a {} dtype with a series of type: {}", self.dtype(), rhs.dtype()))?;
+
+        let phys = self
+            .0
+            .physical()
+            .subtract(&rhs.physical().clone().into_series())?;
+
+        Ok(phys.into_duration(TimeUnit::Nanoseconds))
     }
 
     fn add_to(&self, rhs: &Series) -> PolarsResult<Series> {
diff --git a/crates/polars-plan/src/plans/aexpr/schema.rs b/crates/polars-plan/src/plans/aexpr/schema.rs
index fad154c216ee..024cf95981f6 100644
--- a/crates/polars-plan/src/plans/aexpr/schema.rs
+++ b/crates/polars-plan/src/plans/aexpr/schema.rs
@@ -467,6 +467,7 @@ fn get_arithmetic_field(
                 (_, Duration(_)) | (Duration(_), _) => {
                     polars_bail!(InvalidOperation: "{} not allowed on {} and {}", op, left_field.dtype, right_type)
                 },
+                (Time, Time) => Duration(TimeUnit::Nanoseconds),
                 (_, Time) | (Time, _) => {
                     polars_bail!(InvalidOperation: "{} not allowed on {} and {}", op, left_field.dtype, right_type)
                 },
diff --git a/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py b/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py
index b003219bf6d0..a1f5f5c249a6 100644
--- a/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py
+++ b/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py
@@ -791,6 +791,31 @@ def test_date_datetime_sub() -> None:
     }
 
 
+def test_time_time_sub() -> None:
+    df = pl.DataFrame(
+        {
+            "foo": pl.Series([-1, 0, 10]).cast(pl.Datetime("us")),
+            "bar": pl.Series([1, 0, 1]).cast(pl.Datetime("us")),
+        }
+    )
+
+    assert df.select(
+        pl.col("foo").dt.time() - pl.col("bar").dt.time(),
+        pl.col("bar").dt.time() - pl.col("foo").dt.time(),
+    ).to_dict(as_series=False) == {
+        "foo": [
+            timedelta(days=1, microseconds=-2),
+            timedelta(0),
+            timedelta(microseconds=9),
+        ],
+        "bar": [
+            timedelta(days=-1, microseconds=2),
+            timedelta(0),
+            timedelta(microseconds=-9),
+        ],
+    }
+
+
 def test_raise_invalid_shape() -> None:
     with pytest.raises(InvalidOperationError):
         pl.DataFrame([[1, 2], [3, 4]]) * pl.DataFrame([1, 2, 3])