Merge remote-tracking branch 'camriddell/main' into simp-pandas-group-by

Author: dangotbanned

docs/api-reference/expr_dt.md

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

docs/api-reference/narwhals.md

@@ -9,6 +9,7 @@ Here are the top-level functions available in Narwhals.
         - all
         - all_horizontal
         - any_horizontal
+        - coalesce
         - col
         - concat
         - concat_str

docs/api-reference/series_dt.md

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

narwhals/__init__.py

@@ -49,6 +49,7 @@
     all_ as all,
     all_horizontal,
     any_horizontal,
+    coalesce,
     col,
     concat,
     concat_str,
@@ -127,6 +128,7 @@
     "all",
     "all_horizontal",
     "any_horizontal",
+    "coalesce",
     "col",
     "concat",
     "concat_str",

narwhals/_arrow/namespace.py

@@ -262,6 +262,27 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             context=self,
         )
 
+    def coalesce(self, *exprs: ArrowExpr) -> ArrowExpr:
+        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+            align = self._series._align_full_broadcast
+            init_series, *series = align(*chain.from_iterable(expr(df) for expr in exprs))
+            return [
+                ArrowSeries(
+                    pc.coalesce(init_series.native, *(s.native for s in series)),
+                    name=init_series.name,
+                    version=self._version,
+                )
+            ]
+
+        return self._expr._from_callable(
+            func=func,
+            depth=max(x._depth for x in exprs) + 1,
+            function_name="coalesce",
+            evaluate_output_names=combine_evaluate_output_names(*exprs),
+            alias_output_names=combine_alias_output_names(*exprs),
+            context=self,
+        )
+
 
 class ArrowWhen(EagerWhen[ArrowDataFrame, ArrowSeries, ArrowExpr, "ChunkedArrayAny"]):
     @property

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/_compliant/namespace.py

@@ -96,6 +96,7 @@ def concat_str(
     def selectors(self) -> CompliantSelectorNamespace[Any, Any]: ...
     @property
     def _expr(self) -> type[CompliantExprT]: ...
+    def coalesce(self, *exprs: CompliantExprT) -> CompliantExprT: ...
 
 
 class DepthTrackingNamespace(

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)