diff --git a/crates/polars-core/src/utils/mod.rs b/crates/polars-core/src/utils/mod.rs index c85c807096f57..3be69c2cced80 100644 --- a/crates/polars-core/src/utils/mod.rs +++ b/crates/polars-core/src/utils/mod.rs @@ -554,8 +554,8 @@ macro_rules! df { pub fn get_time_units(tu_l: &TimeUnit, tu_r: &TimeUnit) -> TimeUnit { use TimeUnit::*; match (tu_l, tu_r) { - (Nanoseconds, Microseconds) => Microseconds, - (_, Milliseconds) => Milliseconds, + (_, Nanoseconds) => Nanoseconds, + (Milliseconds, Microseconds) => Microseconds, _ => *tu_l, } } diff --git a/crates/polars/tests/it/core/date_like.rs b/crates/polars/tests/it/core/date_like.rs index 8d97b5e87d1de..db4850726a22a 100644 --- a/crates/polars/tests/it/core/date_like.rs +++ b/crates/polars/tests/it/core/date_like.rs @@ -155,6 +155,57 @@ fn test_duration() -> PolarsResult<()> { .into_duration(TimeUnit::Nanoseconds) .into_series() ); + + Ok(()) +} + +#[test] +#[cfg(feature = "dtype-duration")] +fn test_duration_arithmetic_cross_timeunit() -> PolarsResult<()> { + let datetime_ns = Int64Chunked::new("", &[2, 3, 4]) + .into_datetime(TimeUnit::Nanoseconds, None) + .into_series(); + + let datetime_us = Int64Chunked::new("", &[2, 3, 4]) + .into_datetime(TimeUnit::Microseconds, None) + .into_series(); + + let datetime_ms = Int64Chunked::new("", &[2, 3, 4]) + .into_datetime(TimeUnit::Milliseconds, None) + .into_series(); + + let duration_ns = Int64Chunked::new("", &[1, 1, 1]) + .into_duration(TimeUnit::Nanoseconds) + .into_series(); + + let duration_us = Int64Chunked::new("", &[1, 1, 1]) + .into_duration(TimeUnit::Microseconds) + .into_series(); + + let duration_ms = Int64Chunked::new("", &[1, 1, 1]) + .into_duration(TimeUnit::Milliseconds) + .into_series(); + + assert_eq!( + *(datetime_ns.clone() + duration_us).dtype(), + DataType::Datetime(TimeUnit::Nanoseconds, None) + ); + + assert_eq!( + *(datetime_ns + duration_ms.clone()).dtype(), + DataType::Datetime(TimeUnit::Nanoseconds, None) + ); + + assert_eq!( + *(datetime_us + duration_ms).dtype(), + DataType::Datetime(TimeUnit::Microseconds, None) + ); + + assert_eq!( + *(datetime_ms + duration_ns).dtype(), + DataType::Datetime(TimeUnit::Nanoseconds, None) + ); + Ok(()) } diff --git a/py-polars/tests/unit/functions/as_datatype/test_duration.py b/py-polars/tests/unit/functions/as_datatype/test_duration.py index 07dc165a8a0fb..bb8f1afc443ed 100644 --- a/py-polars/tests/unit/functions/as_datatype/test_duration.py +++ b/py-polars/tests/unit/functions/as_datatype/test_duration.py @@ -143,6 +143,49 @@ def test_add_duration_3786() -> None: } +def test_datetime_ms_add_duration_us() -> None: + duration_us = pl.duration(microseconds=15) + df_datetimes = pl.DataFrame( + { + "datetime": [ + datetime(2000, 1, 1, 1, second=30), + datetime(2000, 1, 1, 1, second=35), + datetime(2000, 1, 1, 1, second=40), + ] + }, + schema={"datetime": pl.Datetime(time_unit="ms")}, + ) + + df_add_us = df_datetimes.select(pl.col("datetime") + duration_us) + + assert df_add_us["datetime"].dt.microsecond().to_list() == [15, 15, 15] + + assert df_add_us["datetime"].dtype.time_unit == "us" + + +def test_datetime_us_add_duration_ns() -> None: + duration_ns = pl.duration(nanoseconds=15, time_unit="ns") + df_datetimes = pl.DataFrame( + { + "datetime": [ + datetime(2000, 1, 1, 1, microsecond=30), + datetime(2000, 1, 1, 1, microsecond=35), + datetime(2000, 1, 1, 1, microsecond=40), + ] + } + ) + + df_add_ns = df_datetimes.select(pl.col("datetime") + duration_ns) + + assert df_add_ns["datetime"].dt.nanosecond().to_list() == [ + 15 + 30 * 1000, + 15 + 35 * 1000, + 15 + 40 * 1000, + ] + + assert df_add_ns["datetime"].dtype.time_unit == "ns" + + @pytest.mark.parametrize( ("time_unit", "ms", "us", "ns"), [