feat: Allow mode(..., keep="any") (#3019)

Author: FBruzzesi

docs/api-reference/typing.md

@@ -30,6 +30,7 @@ Narwhals comes fully statically typed. In addition to `nw.DataFrame`, `nw.Expr`,
         - ConcatMethod
         - FillNullStrategy
         - JoinStrategy
+        - ModeKeepStrategy
         - PivotAgg
         - RankMethod
         - RollingInterpolationMethod

narwhals/_arrow/series.py

@@ -66,6 +66,7 @@
         FillNullStrategy,
         Into1DArray,
         IntoDType,
+        ModeKeepStrategy,
         NonNestedLiteral,
         NumericLiteral,
         PythonLiteral,
@@ -856,16 +857,17 @@ def clip(
     def to_arrow(self) -> ArrayAny:
         return self.native.combine_chunks()
 
-    def mode(self) -> ArrowSeries:
+    def mode(self, *, keep: ModeKeepStrategy) -> ArrowSeries:
         plx = self.__narwhals_namespace__()
         col_token = generate_temporary_column_name(n_bytes=8, columns=[self.name])
         counts = self.value_counts(
             name=col_token, normalize=False, sort=False, parallel=False
         )
-        return counts.filter(
+        result = counts.filter(
             plx.col(col_token)
             == plx.col(col_token).max().broadcast(kind=ExprKind.AGGREGATION)
         ).get_column(self.name)
+        return result.head(1) if keep == "any" else result
 
     def is_finite(self) -> Self:
         return self._with_native(pc.is_finite(self.native))

narwhals/_compliant/column.py

@@ -20,6 +20,7 @@
         ClosedInterval,
         FillNullStrategy,
         IntoDType,
+        ModeKeepStrategy,
         NonNestedLiteral,
         NumericLiteral,
         RankMethod,
@@ -174,7 +175,7 @@ def is_nan(self) -> Self: ...
     def is_null(self) -> Self: ...
     def is_unique(self) -> Self: ...
     def log(self, base: float) -> Self: ...
-    def mode(self) -> Self: ...
+    def mode(self, *, keep: ModeKeepStrategy) -> Self: ...
     def rank(self, method: RankMethod, *, descending: bool) -> Self: ...
     def replace_strict(
         self,

narwhals/_compliant/expr.py

@@ -45,6 +45,7 @@
         ClosedInterval,
         FillNullStrategy,
         IntoDType,
+        ModeKeepStrategy,
         NonNestedLiteral,
         NumericLiteral,
         RankMethod,
@@ -702,8 +703,8 @@ def len(self) -> Self:
     def gather_every(self, n: int, offset: int) -> Self:
         return self._reuse_series("gather_every", n=n, offset=offset)
 
-    def mode(self) -> Self:
-        return self._reuse_series("mode")
+    def mode(self, *, keep: ModeKeepStrategy) -> Self:
+        return self._reuse_series("mode", scalar_kwargs={"keep": keep})
 
     def is_finite(self) -> Self:
         return self._reuse_series("is_finite")

narwhals/_compliant/typing.py

@@ -24,6 +24,7 @@
     from narwhals.typing import (
         FillNullStrategy,
         IntoLazyFrame,
+        ModeKeepStrategy,
         NativeDataFrame,
         NativeFrame,
         NativeSeries,
@@ -43,6 +44,7 @@ class ScalarKwargs(TypedDict, total=False):
         half_life: float | None
         ignore_nulls: bool
         interpolation: RollingInterpolationMethod
+        keep: ModeKeepStrategy
         limit: int | None
         method: RankMethod
         min_samples: int
@@ -180,6 +182,7 @@ class ScalarKwargs(TypedDict, total=False):
     "median",
     "max",
     "min",
+    "mode",
     "std",
     "var",
     "len",

narwhals/_dask/expr.py

@@ -34,6 +34,7 @@
     from narwhals.typing import (
         FillNullStrategy,
         IntoDType,
+        ModeKeepStrategy,
         NonNestedLiteral,
         NumericLiteral,
         RollingInterpolationMethod,
@@ -648,6 +649,14 @@ def sqrt(self) -> Self:
 
         return self._with_callable(da.sqrt, "sqrt")
 
+    def mode(self, *, keep: ModeKeepStrategy) -> Self:
+        def func(expr: dx.Series) -> dx.Series:
+            _name = expr.name
+            result = expr.to_frame().mode()[_name]
+            return result.head(1) if keep == "any" else result
+
+        return self._with_callable(func, "mode", scalar_kwargs={"keep": keep})
+
     @property
     def str(self) -> DaskExprStringNamespace:
         return DaskExprStringNamespace(self)
@@ -663,7 +672,6 @@ def dt(self) -> DaskExprDateTimeNamespace:
     gather_every: not_implemented = not_implemented()
     head: not_implemented = not_implemented()
     map_batches: not_implemented = not_implemented()
-    mode: not_implemented = not_implemented()
     sample: not_implemented = not_implemented()
     rank: not_implemented = not_implemented()
     replace_strict: not_implemented = not_implemented()

narwhals/_pandas_like/group_by.py

@@ -39,6 +39,7 @@
     "mean",
     "median",
     "min",
+    "mode",
     "nunique",
     "prod",
     "quantile",
@@ -115,6 +116,35 @@ def _getitem_aggs(
             result = ns._concat_horizontal(
                 [ns.from_native(result_single).alias(name).native for name in names]
             )
+        elif self.is_mode():
+            compliant = group_by.compliant
+            if (keep := self.kwargs.get("keep")) != "any":  # pragma: no cover
+                msg = (
+                    f"`Expr.mode(keep='{keep}')` is not implemented in group by context for "
+                    f"backend {compliant._implementation}\n\n"
+                    "Hint: Use `nw.col(...).mode(keep='any')` instead."
+                )
+                raise NotImplementedError(msg)
+
+            cols = list(names)
+            native = compliant.native
+            keys, kwargs = group_by._keys, group_by._kwargs
+
+            # Implementation based on the following suggestion:
+            # https://github.com/pandas-dev/pandas/issues/19254#issuecomment-778661578
+            ns = compliant.__narwhals_namespace__()
+            result = ns._concat_horizontal(
+                [
+                    native.groupby([*keys, col], **kwargs)
+                    .size()
+                    .sort_values(ascending=False)
+                    .reset_index(col)
+                    .groupby(keys, **kwargs)[col]
+                    .head(1)
+                    .sort_index()
+                    for col in cols
+                ]
+            )
         else:
             select = names[0] if len(names) == 1 else list(names)
             result = self.native_agg()(group_by._grouped[select])
@@ -127,6 +157,9 @@ def _getitem_aggs(
     def is_len(self) -> bool:
         return self.leaf_name == "len"
 
+    def is_mode(self) -> bool:
+        return self.leaf_name == "mode"
+
     def is_top_level_function(self) -> bool:
         # e.g. `nw.len()`.
         return self.expr._depth == 0
@@ -158,6 +191,7 @@ class PandasLikeGroupBy(
         "median": "median",
         "max": "max",
         "min": "min",
+        "mode": "mode",
         "std": "std",
         "var": "var",
         "len": "size",
@@ -176,6 +210,9 @@ class PandasLikeGroupBy(
     _output_key_names: list[str]
     """Stores the **original** version of group keys."""
 
+    _kwargs: Mapping[str, bool]
+    """Stores keyword arguments for `DataFrame.groupby` other than `by`."""
+
     @property
     def exclude(self) -> tuple[str, ...]:
         """Group keys to ignore when expanding multi-output aggregations."""
@@ -200,13 +237,14 @@ def __init__(
         native = self.compliant.native
         if set(native.index.names).intersection(self.compliant.columns):
             native = native.reset_index(drop=True)
-        self._grouped: NativeGroupBy = native.groupby(
-            self._keys.copy(),
-            sort=False,
-            as_index=True,
-            dropna=drop_null_keys,
-            observed=True,
-        )
+
+        self._kwargs = {
+            "sort": False,
+            "as_index": True,
+            "dropna": drop_null_keys,
+            "observed": True,
+        }
+        self._grouped: NativeGroupBy = native.groupby(self._keys.copy(), **self._kwargs)
 
     def agg(self, *exprs: PandasLikeExpr) -> PandasLikeDataFrame:
         all_aggs_are_simple = True

narwhals/_pandas_like/series.py

@@ -47,6 +47,7 @@
         FillNullStrategy,
         Into1DArray,
         IntoDType,
+        ModeKeepStrategy,
         NonNestedLiteral,
         NumericLiteral,
         RankMethod,
@@ -849,10 +850,10 @@ def to_arrow(self) -> pa.Array[Any]:
 
         return pa.Array.from_pandas(self.native)
 
-    def mode(self) -> Self:
+    def mode(self, *, keep: ModeKeepStrategy) -> Self:
         result = self.native.mode()
         result.name = self.name
-        return self._with_native(result)
+        return self._with_native(result.head(1) if keep == "any" else result)
 
     def cum_count(self, *, reverse: bool) -> Self:
         not_na_series = ~self.native.isna()

narwhals/_polars/expr.py

@@ -27,7 +27,7 @@
     from narwhals._polars.dataframe import Method, PolarsDataFrame
     from narwhals._polars.namespace import PolarsNamespace
     from narwhals._utils import Version, _LimitedContext
-    from narwhals.typing import IntoDType, NumericLiteral
+    from narwhals.typing import IntoDType, ModeKeepStrategy, NumericLiteral
 
 
 class PolarsExpr:
@@ -280,6 +280,10 @@ def is_close(
             )
         return self._with_native(result)
 
+    def mode(self, *, keep: ModeKeepStrategy) -> Self:
+        result = self.native.mode()
+        return self._with_native(result.first() if keep == "any" else result)
+
     @property
     def dt(self) -> PolarsExprDateTimeNamespace:
         return PolarsExprDateTimeNamespace(self)
@@ -366,7 +370,6 @@ def _eval_names_indices(indices: Sequence[int], /) -> EvalNames[PolarsDataFrame]
     mean: Method[Self]
     median: Method[Self]
     min: Method[Self]
-    mode: Method[Self]
     n_unique: Method[Self]
     null_count: Method[Self]
     quantile: Method[Self]

narwhals/_polars/series.py

@@ -41,6 +41,7 @@
     from narwhals.typing import (
         Into1DArray,
         IntoDType,
+        ModeKeepStrategy,
         MultiIndexSelector,
         NonNestedLiteral,
         NumericLiteral,
@@ -518,6 +519,10 @@ def is_close(
         )
         return self._with_native(result)
 
+    def mode(self, *, keep: ModeKeepStrategy) -> Self:
+        result = self.native.mode()
+        return self._with_native(result.head(1) if keep == "any" else result)
+
     def hist_from_bins(
         self, bins: list[float], *, include_breakpoint: bool
     ) -> PolarsDataFrame:
@@ -702,7 +707,6 @@ def struct(self) -> PolarsSeriesStructNamespace:
     max: Method[Any]
     mean: Method[float]
     min: Method[Any]
-    mode: Method[Self]
     n_unique: Method[int]
     null_count: Method[int]
     quantile: Method[float]