enh: support offset_by for cudf (#2823)

Author: MarcoGorelli

narwhals/_arrow/series_dt.py

@@ -213,13 +213,14 @@ def offset_by(self, by: str) -> ArrowSeries:
         if interval.unit in {"y", "q", "mo"}:
             msg = f"Offsetting by {interval.unit} is not yet supported for pyarrow."
             raise NotImplementedError(msg)
-        if interval.unit == "d":
+        dtype = self.compliant.dtype
+        datetime_dtype = self.version.dtypes.Datetime
+        if interval.unit == "d" and isinstance(dtype, datetime_dtype) and dtype.time_zone:
             offset: pa.DurationScalar[Any] = lit(interval.to_timedelta())
-            if time_zone := native.type.tz:
-                native_naive = pc.local_timestamp(native)
-                result = pc.assume_timezone(pc.add(native_naive, offset), time_zone)
-                return self.with_native(result)
-        elif interval.unit == "ns":  # pragma: no cover
+            native_naive = pc.local_timestamp(native)
+            result = pc.assume_timezone(pc.add(native_naive, offset), dtype.time_zone)
+            return self.with_native(result)
+        if interval.unit == "ns":  # pragma: no cover
             offset = lit(interval.multiple, pa.duration("ns"))  # type: ignore[assignment]
         else:
             offset = lit(interval.to_timedelta())

narwhals/_pandas_like/series_dt.py

@@ -2,8 +2,6 @@
 
 from typing import TYPE_CHECKING, Any
 
-import pandas as pd
-
 from narwhals._compliant.any_namespace import DateTimeNamespace
 from narwhals._constants import (
     EPOCH_YEAR,
@@ -25,6 +23,10 @@
 )
 
 if TYPE_CHECKING:
+    from datetime import timedelta
+
+    import pandas as pd
+
     from narwhals._pandas_like.series import PandasLikeSeries
     from narwhals.typing import TimeUnit
 
@@ -246,10 +248,8 @@ def truncate(self, every: str) -> PandasLikeSeries:
         )
 
     def offset_by(self, by: str) -> PandasLikeSeries:
-        if self.implementation.is_cudf():
-            msg = "Not implemented for cuDF."
-            raise NotImplementedError(msg)
         native = self.native
+        pdx = self.compliant.__native_namespace__()
         if self._is_pyarrow():
             import pyarrow as pa  # ignore-banned-import
 
@@ -270,19 +270,21 @@ def offset_by(self, by: str) -> PandasLikeSeries:
             if unit == "q":
                 multiple *= 3
                 unit = "mo"
-            offset: pd.DateOffset | pd.Timedelta
+            offset: pd.DateOffset | timedelta
             if unit == "y":
-                offset = pd.DateOffset(years=multiple)
+                offset = pdx.DateOffset(years=multiple)
             elif unit == "mo":
-                offset = pd.DateOffset(months=multiple)
+                offset = pdx.DateOffset(months=multiple)
+            elif unit == "ns":
+                offset = pdx.Timedelta(multiple, unit=UNITS_DICT[unit])
             else:
-                offset = pd.Timedelta(multiple, unit=UNITS_DICT[unit])  # type: ignore[arg-type]
-            if unit == "d":
-                original_timezone = native.dt.tz
+                offset = interval.to_timedelta()
+            dtype = self.compliant.dtype
+            datetime_dtype = self.version.dtypes.Datetime
+            if unit == "d" and isinstance(dtype, datetime_dtype) and dtype.time_zone:
                 native_without_timezone = native.dt.tz_localize(None)
                 result_pd = native_without_timezone + offset
-                if original_timezone is not None:
-                    result_pd = result_pd.dt.tz_localize(original_timezone)
+                result_pd = result_pd.dt.tz_localize(dtype.time_zone)
             else:
                 result_pd = native + offset
 

tests/expr_and_series/dt/offset_by_test.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from datetime import datetime, timezone
+from datetime import date, datetime, timezone
 
 import pytest
 
@@ -32,6 +32,13 @@
                 datetime(2020, 1, 2, 2, 4, 14, 715125),
             ],
         ),
+        (
+            "2000ns",
+            [
+                datetime(2021, 3, 1, 12, 34, 56, 49014),
+                datetime(2020, 1, 2, 2, 4, 14, 715125),
+            ],
+        ),
         (
             "2ms",
             [
@@ -107,29 +114,27 @@ def test_offset_by(
         x in str(constructor) for x in ("dask", "pyarrow", "ibis")
     ):
         request.applymarker(pytest.mark.xfail())
+    if "ns" in by and any(
+        x in str(constructor) for x in ("dask", "pyspark", "ibis", "cudf", "duckdb")
+    ):
+        request.applymarker(pytest.mark.xfail())
     if by.endswith("d") and any(x in str(constructor) for x in ("dask", "ibis")):
         request.applymarker(pytest.mark.xfail())
-    if "cudf" in str(constructor):
-        # https://github.com/rapidsai/cudf/issues/19292
-        request.applymarker(pytest.mark.xfail)
     result = df.select(nw.col("a").dt.offset_by(by))
     assert_equal_data(result, {"a": expected})
 
 
 @pytest.mark.parametrize(
     ("by", "expected"),
     [
-        ("2d", ["2024-01-03T05:45+0545"]),
-        ("5mo", ["2024-06-01T05:45+0545"]),
-        ("7q", ["2025-10-01T05:45+0545"]),
-        ("5y", ["2029-01-01T05:45+0545"]),
+        ("2d", "2024-01-03T05:45+0545"),
+        ("5mo", "2024-06-01T05:45+0545"),
+        ("7q", "2025-10-01T05:45+0545"),
+        ("5y", "2029-01-01T05:45+0545"),
     ],
 )
 def test_offset_by_tz(
-    request: pytest.FixtureRequest,
-    constructor: Constructor,
-    by: str,
-    expected: list[datetime],
+    request: pytest.FixtureRequest, constructor: Constructor, by: str, expected: str
 ) -> None:
     if (
         ("pyarrow" in str(constructor) and is_windows())
@@ -142,35 +147,31 @@ def test_offset_by_tz(
         # pyspark,duckdb don't support changing time zones.
         # convert_time_zone is not supported for ibis.
         request.applymarker(pytest.mark.xfail())
+    if any(x in str(constructor) for x in ("cudf",)) and "d" not in by:
+        # cudf: https://github.com/rapidsai/cudf/issues/19363
+        request.applymarker(pytest.mark.xfail())
     if any(x in by for x in ("y", "q", "mo")) and any(
         x in str(constructor) for x in ("dask", "pyarrow", "ibis")
     ):
         request.applymarker(pytest.mark.xfail())
     if by.endswith("d") and any(x in str(constructor) for x in ("dask",)):
         request.applymarker(pytest.mark.xfail())
-    if "cudf" in str(constructor):
-        # https://github.com/rapidsai/cudf/issues/19292
-        request.applymarker(pytest.mark.xfail)
     df = nw.from_native(constructor(data_tz))
     df = df.select(nw.col("a").dt.convert_time_zone("Asia/Kathmandu"))
     result = df.select(nw.col("a").dt.offset_by(by))
-    result_str = result.select(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M%z"))
-    assert_equal_data(result_str, {"a": expected})
+    assert_equal_data(result, {"a": [datetime.strptime(expected, "%Y-%m-%dT%H:%M%z")]})
 
 
 @pytest.mark.parametrize(
     ("by", "expected"),
     [
-        ("2d", ["2020-10-27T02:00+0100"]),
-        ("5mo", ["2021-03-25T02:00+0100"]),
-        ("1q", ["2021-01-25T02:00+0100"]),
+        ("2d", "2020-10-27T02:00+0100"),
+        ("5mo", "2021-03-25T02:00+0100"),
+        ("1q", "2021-01-25T02:00+0100"),
     ],
 )
 def test_offset_by_dst(
-    request: pytest.FixtureRequest,
-    constructor: Constructor,
-    by: str,
-    expected: list[datetime],
+    request: pytest.FixtureRequest, constructor: Constructor, by: str, expected: str
 ) -> None:
     if (
         ("pyarrow" in str(constructor) and is_windows())
@@ -179,12 +180,12 @@ def test_offset_by_dst(
         or ("modin_pyarrow" in str(constructor) and PANDAS_VERSION < (2, 1))
     ):
         pytest.skip()
-    if any(
-        x in str(constructor) for x in ("duckdb", "sqlframe", "pyspark", "ibis", "cudf")
-    ):
+    if any(x in str(constructor) for x in ("duckdb", "sqlframe", "pyspark", "ibis")):
         # pyspark,duckdb don't support changing time zones.
         # convert_time_zone is not supported for ibis.
-        # cudf https://github.com/rapidsai/cudf/issues/19292
+        request.applymarker(pytest.mark.xfail())
+    if any(x in str(constructor) for x in ("cudf",)) and "d" not in by:
+        # cudf: https://github.com/rapidsai/cudf/issues/19363
         request.applymarker(pytest.mark.xfail())
     if any(x in by for x in ("y", "q", "mo")) and any(
         x in str(constructor) for x in ("dask", "pyarrow")
@@ -195,16 +196,10 @@ def test_offset_by_dst(
     df = nw.from_native(constructor(data_dst))
     df = df.with_columns(a=nw.col("a").dt.convert_time_zone("Europe/Amsterdam"))
     result = df.select(nw.col("a").dt.offset_by(by))
-    result_str = result.select(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M%z"))
-    assert_equal_data(result_str, {"a": expected})
+    assert_equal_data(result, {"a": [datetime.strptime(expected, "%Y-%m-%dT%H:%M%z")]})
 
 
-def test_offset_by_series(
-    constructor_eager: ConstructorEager, request: pytest.FixtureRequest
-) -> None:
-    if "cudf" in str(constructor_eager):
-        # https://github.com/rapidsai/cudf/issues/19292
-        request.applymarker(pytest.mark.xfail)
+def test_offset_by_series(constructor_eager: ConstructorEager) -> None:
     df = nw.from_native(constructor_eager(data), eager_only=True)
     result = df.select(df["a"].dt.offset_by("1h"))
     expected = {
@@ -216,13 +211,23 @@ def test_offset_by_series(
     assert_equal_data(result, expected)
 
 
-def test_offset_by_invalid_interval(
-    constructor: Constructor, request: pytest.FixtureRequest
-) -> None:
-    if "cudf" in str(constructor):
-        # https://github.com/rapidsai/cudf/issues/19292
-        request.applymarker(pytest.mark.xfail)
+def test_offset_by_invalid_interval(constructor: Constructor) -> None:
     df = nw.from_native(constructor(data))
     msg = "Invalid `every` string"
     with pytest.raises(ValueError, match=msg):
         df.select(nw.col("a").dt.offset_by("1r"))
+
+
+@pytest.mark.skipif(PANDAS_VERSION < (2, 2), reason="too old for pyarrow date type")
+def test_offset_by_date_pandas() -> None:
+    pytest.importorskip("pandas")
+    import pandas as pd
+
+    df = nw.from_native(pd.DataFrame({"a": [date(2020, 1, 1)]}, dtype="date32[pyarrow]"))
+    result = df.select(nw.col("a").dt.offset_by("1d"))
+    expected = {"a": [date(2020, 1, 2)]}
+    assert_equal_data(result, expected)
+    df = nw.from_native(pd.DataFrame({"a": [date(2020, 1, 1)]}))
+    result = df.select(nw.col("a").dt.offset_by("1d"))
+    expected = {"a": [date(2020, 1, 2)]}
+    assert_equal_data(result, expected)