feat: dt.offset_by supports broadcasting lhs (#11095)

pola-rs · Sep 14, 2023 · 6886949 · 6886949
1 parent 7ba1e82
commit 6886949
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 2 deletions.
diff --git a/crates/polars-plan/src/dsl/function_expr/temporal.rs b/crates/polars-plan/src/dsl/function_expr/temporal.rs
@@ -126,8 +126,14 @@ fn apply_offsets_to_datetime(
     offset_fn: fn(&Duration, i64, Option<&Tz>) -> PolarsResult<i64>,
     time_zone: Option<&Tz>,
 ) -> PolarsResult<Int64Chunked> {
-    match offsets.len() {
-        1 => match offsets.get(0) {
+    match (datetime.len(), offsets.len()) {
+        (1, _) => match datetime.0.get(0) {
+            Some(dt) => offsets.try_apply_values_generic(|offset| {
+                offset_fn(&Duration::parse(offset), dt, time_zone)
+            }),
+            _ => Ok(Int64Chunked::full_null(datetime.0.name(), offsets.len())),
+        },
+        (_, 1) => match offsets.get(0) {
             Some(offset) => datetime
                 .0
                 .try_apply(|v| offset_fn(&Duration::parse(offset), v, time_zone)),

diff --git a/py-polars/tests/unit/namespaces/test_datetime.py b/py-polars/tests/unit/namespaces/test_datetime.py
@@ -690,6 +690,80 @@ def test_offset_by_truncate_sorted_flag() -> None:
     assert s2.flags["SORTED_ASC"]
 
 
+def test_offset_by_broadcasting() -> None:
+    # test broadcast lhs
+    df = pl.DataFrame(
+        {
+            "offset": ["1d", "10d", "3d", None],
+        }
+    )
+    result = df.select(
+        d1=pl.lit(datetime(2020, 10, 25)).dt.offset_by(pl.col("offset")),
+        d2=pl.lit(datetime(2020, 10, 25))
+        .dt.cast_time_unit("ms")
+        .dt.offset_by(pl.col("offset")),
+        d3=pl.lit(datetime(2020, 10, 25))
+        .dt.replace_time_zone("Europe/London")
+        .dt.offset_by(pl.col("offset")),
+        d4=pl.lit(datetime(2020, 10, 25)).dt.date().dt.offset_by(pl.col("offset")),
+        d5=pl.lit(None, dtype=pl.Datetime).dt.offset_by(pl.col("offset")),
+    )
+    expected_dict = {
+        "d1": [
+            datetime(2020, 10, 26),
+            datetime(2020, 11, 4),
+            datetime(2020, 10, 28),
+            None,
+        ],
+        "d2": [
+            datetime(2020, 10, 26),
+            datetime(2020, 11, 4),
+            datetime(2020, 10, 28),
+            None,
+        ],
+        "d3": [
+            datetime(2020, 10, 26, tzinfo=ZoneInfo(key="Europe/London")),
+            datetime(2020, 11, 4, tzinfo=ZoneInfo(key="Europe/London")),
+            datetime(2020, 10, 28, tzinfo=ZoneInfo(key="Europe/London")),
+            None,
+        ],
+        "d4": [
+            datetime(2020, 10, 26).date(),
+            datetime(2020, 11, 4).date(),
+            datetime(2020, 10, 28).date(),
+            None,
+        ],
+        "d5": [None, None, None, None],
+    }
+    assert result.to_dict(False) == expected_dict
+
+    # test broadcast rhs
+    df = pl.DataFrame({"dt": [datetime(2020, 10, 25), datetime(2021, 1, 2), None]})
+    result = df.select(
+        d1=pl.col("dt").dt.offset_by(pl.lit("1mo3d")),
+        d2=pl.col("dt").dt.cast_time_unit("ms").dt.offset_by(pl.lit("1y1mo")),
+        d3=pl.col("dt")
+        .dt.replace_time_zone("Europe/London")
+        .dt.offset_by(pl.lit("3d")),
+        d4=pl.col("dt").dt.date().dt.offset_by(pl.lit("1y1mo1d")),
+    )
+    expected_dict = {
+        "d1": [datetime(2020, 11, 28), datetime(2021, 2, 5), None],
+        "d2": [datetime(2021, 11, 25), datetime(2022, 2, 2), None],
+        "d3": [
+            datetime(2020, 10, 28, tzinfo=ZoneInfo(key="Europe/London")),
+            datetime(2021, 1, 5, tzinfo=ZoneInfo(key="Europe/London")),
+            None,
+        ],
+        "d4": [datetime(2021, 11, 26).date(), datetime(2022, 2, 3).date(), None],
+    }
+    assert result.to_dict(False) == expected_dict
+
+    # test all literal
+    result = df.select(d=pl.lit(datetime(2021, 11, 26)).dt.offset_by("1mo1d"))
+    assert result.to_dict(False) == {"d": [datetime(2021, 12, 27)]}
+
+
 def test_offset_by_expressions() -> None:
     df = pl.DataFrame(
         {