feat: Add Expr.dt.offset_by() (#2733)

---------

Co-authored-by: dangotbanned <[email protected]>
Co-authored-by: Marco Edward Gorelli <[email protected]>

Author: 3 people

docs/api-reference/expr_dt.md

@@ -13,6 +13,7 @@
         - minute
         - month
         - nanosecond
+        - offset_by
         - ordinal_day
         - replace_time_zone
         - second

docs/api-reference/series_dt.md

@@ -13,6 +13,7 @@
         - minute
         - month
         - nanosecond
+        - offset_by
         - ordinal_day
         - replace_time_zone
         - second

narwhals/_arrow/series_dt.py

@@ -18,7 +18,7 @@
     US_PER_MINUTE,
     US_PER_SECOND,
 )
-from narwhals._duration import parse_interval_string
+from narwhals._duration import Interval
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
@@ -202,7 +202,25 @@ def total_nanoseconds(self) -> ArrowSeries:
         return self.with_native(pc.multiply(self.native, factor).cast(pa.int64()))
 
     def truncate(self, every: str) -> ArrowSeries:
-        multiple, unit = parse_interval_string(every)
+        interval = Interval.parse(every)
         return self.with_native(
-            pc.floor_temporal(self.native, multiple=multiple, unit=UNITS_DICT[unit])
+            pc.floor_temporal(self.native, interval.multiple, UNITS_DICT[interval.unit])
         )
+
+    def offset_by(self, by: str) -> ArrowSeries:
+        interval = Interval.parse_no_constraints(by)
+        native = self.native
+        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":
+            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
+            offset = lit(interval.multiple, pa.duration("ns"))  # type: ignore[assignment]
+        else:
+            offset = lit(interval.to_timedelta())
+        return self.with_native(pc.add(native, offset))

narwhals/_compliant/any_namespace.py

@@ -48,6 +48,7 @@ def total_milliseconds(self) -> CompliantT_co: ...
     def total_microseconds(self) -> CompliantT_co: ...
     def total_nanoseconds(self) -> CompliantT_co: ...
     def truncate(self, every: str) -> CompliantT_co: ...
+    def offset_by(self, by: str) -> CompliantT_co: ...
 
 
 class ListNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):

narwhals/_compliant/expr.py

@@ -1112,6 +1112,9 @@ def total_nanoseconds(self) -> EagerExprT:
     def truncate(self, every: str) -> EagerExprT:
         return self.compliant._reuse_series_namespace("dt", "truncate", every=every)
 
+    def offset_by(self, by: str) -> EagerExprT:
+        return self.compliant._reuse_series_namespace("dt", "offset_by", by=by)
+
 
 class EagerExprListNamespace(
     EagerExprNamespace[EagerExprT], ListNamespace[EagerExprT], Generic[EagerExprT]

narwhals/_dask/expr_dt.py

@@ -5,9 +5,9 @@
 from narwhals._compliant import LazyExprNamespace
 from narwhals._compliant.any_namespace import DateTimeNamespace
 from narwhals._constants import MS_PER_SECOND, NS_PER_SECOND, US_PER_SECOND
-from narwhals._duration import parse_interval_string
+from narwhals._duration import Interval
 from narwhals._pandas_like.utils import (
-    UNIT_DICT,
+    ALIAS_DICT,
     calculate_timestamp_date,
     calculate_timestamp_datetime,
     native_to_narwhals_dtype,
@@ -154,9 +154,22 @@ def total_nanoseconds(self) -> DaskExpr:
         )
 
     def truncate(self, every: str) -> DaskExpr:
-        multiple, unit = parse_interval_string(every)
+        interval = Interval.parse(every)
+        unit = interval.unit
         if unit in {"mo", "q", "y"}:
-            msg = f"Truncating to {unit} is not supported yet for dask."
+            msg = f"Truncating to {unit} is not yet supported for dask."
             raise NotImplementedError(msg)
-        freq = f"{multiple}{UNIT_DICT.get(unit, unit)}"
+        freq = f"{interval.multiple}{ALIAS_DICT.get(unit, unit)}"
         return self.compliant._with_callable(lambda expr: expr.dt.floor(freq), "truncate")
+
+    def offset_by(self, by: str) -> DaskExpr:
+        def func(s: dx.Series, by: str) -> dx.Series:
+            interval = Interval.parse_no_constraints(by)
+            unit = interval.unit
+            if unit in {"y", "q", "mo", "d", "ns"}:
+                msg = f"Offsetting by {unit} is not yet supported for dask."
+                raise NotImplementedError(msg)
+            offset = interval.to_timedelta()
+            return s.add(offset)
+
+        return self.compliant._with_callable(func, "offset_by", by=by)

narwhals/_duckdb/expr_dt.py

@@ -13,7 +13,7 @@
     US_PER_SECOND,
 )
 from narwhals._duckdb.utils import UNITS_DICT, F, fetch_rel_time_zone, lit
-from narwhals._duration import parse_interval_string
+from narwhals._duration import Interval
 from narwhals._utils import not_implemented
 
 if TYPE_CHECKING:
@@ -99,7 +99,8 @@ def total_microseconds(self) -> DuckDBExpr:
         )
 
     def truncate(self, every: str) -> DuckDBExpr:
-        multiple, unit = parse_interval_string(every)
+        interval = Interval.parse(every)
+        multiple, unit = interval.multiple, interval.unit
         if multiple != 1:
             # https://github.com/duckdb/duckdb/issues/17554
             msg = f"Only multiple 1 is currently supported for DuckDB.\nGot {multiple!s}."
@@ -114,6 +115,15 @@ def _truncate(expr: Expression) -> Expression:
 
         return self.compliant._with_elementwise(_truncate)
 
+    def offset_by(self, by: str) -> DuckDBExpr:
+        interval = Interval.parse_no_constraints(by)
+        format = lit(f"{interval.multiple!s} {UNITS_DICT[interval.unit]}")
+
+        def _offset_by(expr: Expression) -> Expression:
+            return F("date_add", format, expr)
+
+        return self.compliant._with_callable(_offset_by)
+
     def _no_op_time_zone(self, time_zone: str) -> DuckDBExpr:
         def func(df: DuckDBLazyFrame) -> Sequence[Expression]:
             native_series_list = self.compliant(df)

narwhals/_duration.py

@@ -2,13 +2,16 @@
 
 from __future__ import annotations
 
+import datetime as dt
 import re
 from typing import TYPE_CHECKING, Literal, cast, get_args
 
 if TYPE_CHECKING:
+    from collections.abc import Container, Mapping
+
     from typing_extensions import TypeAlias
 
-__all__ = ["IntervalUnit", "parse_interval_string"]
+__all__ = ["IntervalUnit"]
 
 IntervalUnit: TypeAlias = Literal["ns", "us", "ms", "s", "m", "h", "d", "mo", "q", "y"]
 """A Polars duration string interval unit.
@@ -24,23 +27,43 @@
 - 'q': quarter.
 - 'y': year.
 """
+TimedeltaKwd: TypeAlias = Literal[
+    "days", "hours", "minutes", "seconds", "milliseconds", "microseconds"
+]
 
 PATTERN_INTERVAL: re.Pattern[str] = re.compile(
     r"^(?P<multiple>\d+)(?P<unit>ns|us|ms|mo|m|s|h|d|q|y)\Z"
 )
 MONTH_MULTIPLES = frozenset([1, 2, 3, 4, 6, 12])
 QUARTER_MULTIPLES = frozenset([1, 2, 4])
+UNIT_TO_TIMEDELTA: Mapping[IntervalUnit, TimedeltaKwd] = {
+    "d": "days",
+    "h": "hours",
+    "m": "minutes",
+    "s": "seconds",
+    "ms": "milliseconds",
+    "us": "microseconds",
+}
+
 
+class Interval:
+    def __init__(self, multiple: int, unit: IntervalUnit, /) -> None:
+        self.multiple: int = multiple
+        self.unit: IntervalUnit = unit
 
-def parse_interval_string(every: str) -> tuple[int, IntervalUnit]:
-    """Parse a string like "1d", "2h", "3m" into a tuple of (number, unit).
+    def to_timedelta(
+        self, *, unsupported: Container[IntervalUnit] = frozenset(("ns", "mo", "q", "y"))
+    ) -> dt.timedelta:
+        if self.unit in unsupported:  # pragma: no cover
+            msg = f"Creating timedelta with {self.unit} unit is not supported."
+            raise NotImplementedError(msg)
+        kwd = UNIT_TO_TIMEDELTA[self.unit]
+        # error: Keywords must be strings (bad mypy)
+        return dt.timedelta(**{kwd: self.multiple})  # type: ignore[misc]
 
-    Returns:
-        A tuple of multiple and unit parsed from the interval string.
-    """
-    if match := PATTERN_INTERVAL.match(every):
-        multiple = int(match["multiple"])
-        unit = cast("IntervalUnit", match["unit"])
+    @classmethod
+    def parse(cls, every: str) -> Interval:
+        multiple, unit = cls._parse(every)
         if unit == "mo" and multiple not in MONTH_MULTIPLES:
             msg = f"Only the following multiples are supported for 'mo' unit: {MONTH_MULTIPLES}.\nGot: {multiple}."
             raise ValueError(msg)
@@ -52,9 +75,20 @@ def parse_interval_string(every: str) -> tuple[int, IntervalUnit]:
                 f"Only multiple 1 is currently supported for 'y' unit.\nGot: {multiple}."
             )
             raise ValueError(msg)
-        return multiple, unit
-    msg = (
-        f"Invalid `every` string: {every}. Expected string of kind <number><unit>, "
-        f"where 'unit' is one of: {get_args(IntervalUnit)}."
-    )
-    raise ValueError(msg)
+        return cls(multiple, unit)
+
+    @classmethod
+    def parse_no_constraints(cls, every: str) -> Interval:
+        return cls(*cls._parse(every))
+
+    @staticmethod
+    def _parse(every: str) -> tuple[int, IntervalUnit]:
+        if match := PATTERN_INTERVAL.match(every):
+            multiple = int(match["multiple"])
+            unit = cast("IntervalUnit", match["unit"])
+            return multiple, unit
+        msg = (
+            f"Invalid `every` string: {every}. Expected string of kind <number><unit>, "
+            f"where 'unit' is one of: {get_args(IntervalUnit)}."
+        )
+        raise ValueError(msg)

narwhals/_ibis/expr_dt.py

@@ -4,7 +4,7 @@
 
 from narwhals._compliant import LazyExprNamespace
 from narwhals._compliant.any_namespace import DateTimeNamespace
-from narwhals._duration import parse_interval_string
+from narwhals._duration import Interval
 from narwhals._ibis.utils import UNITS_DICT_BUCKET, UNITS_DICT_TRUNCATE
 from narwhals._utils import not_implemented
 
@@ -68,7 +68,8 @@ def fn(expr: ir.TimestampValue) -> ir.TimestampValue:
         return fn
 
     def truncate(self, every: str) -> IbisExpr:
-        multiple, unit = parse_interval_string(every)
+        interval = Interval.parse(every)
+        multiple, unit = interval.multiple, interval.unit
         if unit == "q":
             multiple, unit = 3 * multiple, "mo"
         if multiple != 1:
@@ -80,6 +81,15 @@ def truncate(self, every: str) -> IbisExpr:
             fn = self._truncate(UNITS_DICT_TRUNCATE[unit])
         return self.compliant._with_callable(fn)
 
+    def offset_by(self, every: str) -> IbisExpr:
+        interval = Interval.parse_no_constraints(every)
+        unit = interval.unit
+        if unit in {"y", "q", "mo", "d", "ns"}:
+            msg = f"Offsetting by {unit} is not yet supported for ibis."
+            raise NotImplementedError(msg)
+        offset = interval.to_timedelta()
+        return self.compliant._with_callable(lambda expr: expr.add(offset))
+
     def replace_time_zone(self, time_zone: str | None) -> IbisExpr:
         if time_zone is None:
             return self.compliant._with_callable(lambda expr: expr.cast("timestamp"))

narwhals/_pandas_like/series_dt.py

@@ -2,6 +2,8 @@
 
 from typing import TYPE_CHECKING, Any
 
+import pandas as pd
+
 from narwhals._compliant.any_namespace import DateTimeNamespace
 from narwhals._constants import (
     EPOCH_YEAR,
@@ -10,9 +12,10 @@
     SECONDS_PER_DAY,
     US_PER_SECOND,
 )
-from narwhals._duration import parse_interval_string
+from narwhals._duration import Interval
 from narwhals._pandas_like.utils import (
-    UNIT_DICT,
+    ALIAS_DICT,
+    UNITS_DICT,
     PandasLikeSeriesNamespace,
     calculate_timestamp_date,
     calculate_timestamp_datetime,
@@ -203,13 +206,14 @@ def timestamp(self, time_unit: TimeUnit) -> PandasLikeSeries:
         return self.with_native(result)
 
     def truncate(self, every: str) -> PandasLikeSeries:
-        multiple, unit = parse_interval_string(every)
+        interval = Interval.parse(every)
+        multiple, unit = interval.multiple, interval.unit
         native = self.native
         if self.implementation.is_cudf():
             if multiple != 1:
                 msg = f"Only multiple `1` is supported for cuDF, got: {multiple}."
                 raise NotImplementedError(msg)
-            return self.with_native(self.native.dt.floor(UNIT_DICT.get(unit, unit)))
+            return self.with_native(self.native.dt.floor(ALIAS_DICT.get(unit, unit)))
         dtype_backend = get_dtype_backend(native.dtype, self.compliant._implementation)
         if unit in {"mo", "q", "y"}:
             if self.implementation.is_cudf():
@@ -218,8 +222,6 @@ def truncate(self, every: str) -> PandasLikeSeries:
             if dtype_backend == "pyarrow":
                 import pyarrow.compute as pc  # ignore-banned-import
 
-                from narwhals._arrow.utils import UNITS_DICT
-
                 ca = native.array._pa_array
                 result_arr = pc.floor_temporal(ca, multiple, UNITS_DICT[unit])
             else:
@@ -230,7 +232,7 @@ def truncate(self, every: str) -> PandasLikeSeries:
                     np_unit = "M"
                 else:
                     np_unit = "Y"
-                arr = native.values
+                arr = native.values  # noqa: PD011
                 arr_dtype = arr.dtype
                 result_arr = arr.astype(f"datetime64[{multiple}{np_unit}]").astype(
                     arr_dtype
@@ -240,5 +242,48 @@ def truncate(self, every: str) -> PandasLikeSeries:
             )
             return self.with_native(result_native)
         return self.with_native(
-            self.native.dt.floor(f"{multiple}{UNIT_DICT.get(unit, unit)}")
+            self.native.dt.floor(f"{multiple}{ALIAS_DICT.get(unit, unit)}")
         )
+
+    def offset_by(self, by: str) -> PandasLikeSeries:
+        if self.implementation.is_cudf():
+            msg = "Not implemented for cuDF."
+            raise NotImplementedError(msg)
+        native = self.native
+        if self._is_pyarrow():
+            import pyarrow as pa  # ignore-banned-import
+
+            compliant = self.compliant
+            ca = pa.chunked_array([compliant.to_arrow()])  # type: ignore[arg-type]
+            result = (
+                compliant._version.namespace.from_backend("pyarrow")
+                .compliant.from_native(ca)
+                .dt.offset_by(by)
+                .native
+            )
+            result_pd = native.__class__(
+                result, dtype=native.dtype, index=native.index, name=native.name
+            )
+        else:
+            interval = Interval.parse_no_constraints(by)
+            multiple, unit = interval.multiple, interval.unit
+            if unit == "q":
+                multiple *= 3
+                unit = "mo"
+            offset: pd.DateOffset | pd.Timedelta
+            if unit == "y":
+                offset = pd.DateOffset(years=multiple)
+            elif unit == "mo":
+                offset = pd.DateOffset(months=multiple)
+            else:
+                offset = pd.Timedelta(multiple, unit=UNITS_DICT[unit])  # type: ignore[arg-type]
+            if unit == "d":
+                original_timezone = native.dt.tz
+                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)
+            else:
+                result_pd = native + offset
+
+        return self.with_native(result_pd)