From 2e52af243250576e1e0bd8ff6de1c630e666b810 Mon Sep 17 00:00:00 2001 From: ptiza Date: Fri, 13 Dec 2024 15:48:16 +0100 Subject: [PATCH 1/5] allow sub operation of time dtype --- crates/polars-core/src/series/implementations/time.rs | 10 +++++++++- crates/polars-plan/src/plans/aexpr/schema.rs | 4 +--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/crates/polars-core/src/series/implementations/time.rs b/crates/polars-core/src/series/implementations/time.rs index 74b57a48f3c2..5591f44a3db6 100644 --- a/crates/polars-core/src/series/implementations/time.rs +++ b/crates/polars-core/src/series/implementations/time.rs @@ -88,7 +88,15 @@ impl private::PrivateSeries for SeriesWrap { } fn subtract(&self, rhs: &Series) -> PolarsResult { - polars_bail!(opq = sub, DataType::Time, rhs.dtype()); + match rhs.dtype() { + DataType::Time => { + let dt = DataType::Duration(TimeUnit::Nanoseconds); + let lhs = self.cast(&dt, CastOptions::NonStrict)?; + let rhs = rhs.cast(&dt)?; + lhs.subtract(&rhs) + }, + dtr => polars_bail!(opq = sub, DataType::Time, dtr), + } } fn add_to(&self, rhs: &Series) -> PolarsResult { diff --git a/crates/polars-plan/src/plans/aexpr/schema.rs b/crates/polars-plan/src/plans/aexpr/schema.rs index fad154c216ee..76941374aaab 100644 --- a/crates/polars-plan/src/plans/aexpr/schema.rs +++ b/crates/polars-plan/src/plans/aexpr/schema.rs @@ -467,9 +467,7 @@ fn get_arithmetic_field( (_, Duration(_)) | (Duration(_), _) => { polars_bail!(InvalidOperation: "{} not allowed on {} and {}", op, left_field.dtype, right_type) }, - (_, Time) | (Time, _) => { - polars_bail!(InvalidOperation: "{} not allowed on {} and {}", op, left_field.dtype, right_type) - }, + (_, Time) | (Time, _) => Duration(TimeUnit::Nanoseconds), (l @ List(a), r @ List(b)) if ![a, b] .into_iter() From b1237c844611b8f5c5b19babc4aaee9673add0a8 Mon Sep 17 00:00:00 2001 From: ptiza Date: Sat, 14 Dec 2024 20:06:05 +0100 Subject: [PATCH 2/5] fix super_type --- crates/polars-plan/src/plans/aexpr/schema.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/polars-plan/src/plans/aexpr/schema.rs b/crates/polars-plan/src/plans/aexpr/schema.rs index 76941374aaab..024cf95981f6 100644 --- a/crates/polars-plan/src/plans/aexpr/schema.rs +++ b/crates/polars-plan/src/plans/aexpr/schema.rs @@ -467,7 +467,10 @@ 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) => Duration(TimeUnit::Nanoseconds), + (_, Time) | (Time, _) => { + polars_bail!(InvalidOperation: "{} not allowed on {} and {}", op, left_field.dtype, right_type) + }, (l @ List(a), r @ List(b)) if ![a, b] .into_iter() From c8ca35c369f3b0829241d5b01c60518645888d3e Mon Sep 17 00:00:00 2001 From: ptiza Date: Sat, 14 Dec 2024 20:48:19 +0100 Subject: [PATCH 3/5] add test --- .../operations/arithmetic/test_arithmetic.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py b/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py index b003219bf6d0..9d4557aa09a3 100644 --- a/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py +++ b/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py @@ -790,6 +790,29 @@ def test_date_datetime_sub() -> None: "bar": [timedelta(days=4)], } +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): From 8057dc8dbbe9244c15dc3242ee1d6684d5fbf274 Mon Sep 17 00:00:00 2001 From: ptiza Date: Sat, 14 Dec 2024 20:50:30 +0100 Subject: [PATCH 4/5] reformat --- py-polars/tests/unit/operations/arithmetic/test_arithmetic.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py b/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py index 9d4557aa09a3..a1f5f5c249a6 100644 --- a/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py +++ b/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py @@ -790,6 +790,7 @@ def test_date_datetime_sub() -> None: "bar": [timedelta(days=4)], } + def test_time_time_sub() -> None: df = pl.DataFrame( { @@ -814,6 +815,7 @@ def test_time_time_sub() -> None: ], } + def test_raise_invalid_shape() -> None: with pytest.raises(InvalidOperationError): pl.DataFrame([[1, 2], [3, 4]]) * pl.DataFrame([1, 2, 3]) From 32a2475f1a726f8ad72dd29875ba476d4a9d9bd4 Mon Sep 17 00:00:00 2001 From: ritchie Date: Sun, 15 Dec 2024 09:37:15 +0100 Subject: [PATCH 5/5] use physicals --- .../src/series/implementations/time.rs | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/crates/polars-core/src/series/implementations/time.rs b/crates/polars-core/src/series/implementations/time.rs index 5591f44a3db6..5fe2c64bd542 100644 --- a/crates/polars-core/src/series/implementations/time.rs +++ b/crates/polars-core/src/series/implementations/time.rs @@ -88,15 +88,14 @@ impl private::PrivateSeries for SeriesWrap { } fn subtract(&self, rhs: &Series) -> PolarsResult { - match rhs.dtype() { - DataType::Time => { - let dt = DataType::Duration(TimeUnit::Nanoseconds); - let lhs = self.cast(&dt, CastOptions::NonStrict)?; - let rhs = rhs.cast(&dt)?; - lhs.subtract(&rhs) - }, - dtr => polars_bail!(opq = sub, DataType::Time, dtr), - } + 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 {