diff --git a/crates/polars-ops/src/series/ops/duration.rs b/crates/polars-ops/src/series/ops/duration.rs index bdfb17114459..b839fda8d375 100644 --- a/crates/polars-ops/src/series/ops/duration.rs +++ b/crates/polars-ops/src/series/ops/duration.rs @@ -72,19 +72,24 @@ pub fn impl_duration(s: &[Column], time_unit: TimeUnit) -> PolarsResult TimeUnit::Milliseconds => MILLISECONDS, }; if !is_zero_scalar(&seconds) { - duration = ((duration + seconds)? * multiplier)?; + let units = seconds * multiplier; + duration = (duration + units?)?; } if !is_zero_scalar(&minutes) { - duration = ((duration + minutes)? * (multiplier * 60))?; + let units = minutes * (multiplier * 60); + duration = (duration + units?)?; } if !is_zero_scalar(&hours) { - duration = ((duration + hours)? * (multiplier * 60 * 60))?; + let units = hours * (multiplier * 60 * 60); + duration = (duration + units?)?; } if !is_zero_scalar(&days) { - duration = ((duration + days)? * (multiplier * SECONDS_IN_DAY))?; + let units = days * (multiplier * SECONDS_IN_DAY); + duration = (duration + units?)?; } if !is_zero_scalar(&weeks) { - duration = ((duration + weeks)? * (multiplier * SECONDS_IN_DAY * 7))?; + let units = weeks * (multiplier * SECONDS_IN_DAY * 7); + duration = (duration + units?)?; } duration diff --git a/py-polars/tests/unit/datatypes/test_duration.py b/py-polars/tests/unit/datatypes/test_duration.py index 899d896def12..597ac1c4a624 100644 --- a/py-polars/tests/unit/datatypes/test_duration.py +++ b/py-polars/tests/unit/datatypes/test_duration.py @@ -60,3 +60,36 @@ def test_series_duration_var_overflow() -> None: s = pl.Series([timedelta(days=10), timedelta(days=20), timedelta(days=40)]) with pytest.raises(PanicException, match="OverflowError"): s.var() + + +def test_series_duration_units() -> None: + td = timedelta + + assert_frame_equal( + pl.DataFrame({"x": [0, 1, 2, 3]}).select(x=pl.duration(weeks=pl.col("x"))), + pl.DataFrame({"x": [td(weeks=i) for i in range(4)]}), + ) + assert_frame_equal( + pl.DataFrame({"x": [0, 1, 2, 3]}).select(x=pl.duration(days=pl.col("x"))), + pl.DataFrame({"x": [td(days=i) for i in range(4)]}), + ) + assert_frame_equal( + pl.DataFrame({"x": [0, 1, 2, 3]}).select(x=pl.duration(hours=pl.col("x"))), + pl.DataFrame({"x": [td(hours=i) for i in range(4)]}), + ) + assert_frame_equal( + pl.DataFrame({"x": [0, 1, 2, 3]}).select(x=pl.duration(minutes=pl.col("x"))), + pl.DataFrame({"x": [td(minutes=i) for i in range(4)]}), + ) + assert_frame_equal( + pl.DataFrame({"x": [0, 1, 2, 3]}).select( + x=pl.duration(milliseconds=pl.col("x")) + ), + pl.DataFrame({"x": [td(milliseconds=i) for i in range(4)]}), + ) + assert_frame_equal( + pl.DataFrame({"x": [0, 1, 2, 3]}).select( + x=pl.duration(microseconds=pl.col("x")) + ), + pl.DataFrame({"x": [td(microseconds=i) for i in range(4)]}), + )