Repository: narwhals-dev/narwhals

Issue:
[Bug]: n_unique in `groupby-agg` raises
### Describe the bug

reported here https://github.com/narwhals-dev/narwhals/pull/910/files#r1746497849

### Steps or code to reproduce the bug

https://github.com/narwhals-dev/narwhals/pull/910/files#r1746497849

### Expected results

no error

### Actual results

error
```
AttributeError                            Traceback (most recent call last)
File ~/scratch/.venv/lib/python3.12/site-packages/narwhals/_pandas_like/group_by.py:165, in agg_pandas(grouped, exprs, keys, output_names, from_dataframe, implementation, backend_version, dataframe_is_empty)
    164 try:
--> 165     result_simple = grouped.agg(aggs)
    166 except AttributeError as exc:

File ~/scratch/.venv/lib/python3.12/site-packages/pandas/core/groupby/generic.py:1432, in DataFrameGroupBy.aggregate(self, func, engine, engine_kwargs, *args, **kwargs)
   1431 op = GroupByApply(self, func, args=args, kwargs=kwargs)
-> 1432 result = op.agg()
   1433 if not is_dict_like(func) and result is not None:
   1434     # GH #52849

File ~/scratch/.venv/lib/python3.12/site-packages/pandas/core/apply.py:190, in Apply.agg(self)
    189 if is_dict_like(func):
--> 190     return self.agg_dict_like()
    191 elif is_list_like(func):
    192     # we require a list, but not a 'str'

File ~/scratch/.venv/lib/python3.12/site-packages/pandas/core/apply.py:423, in Apply.agg_dict_like(self)
    416 """
    417 Compute aggregation in the case of a dict-like argument.
    418
   (...)
    421 Result of aggregation.
    422 """
--> 423 return self.agg_or_apply_dict_like(op_name="agg")

File ~/scratch/.venv/lib/python3.12/site-packages/pandas/core/apply.py:1608, in GroupByApply.agg_or_apply_dict_like(self, op_name)
   1605 with com.temp_setattr(
   1606     obj, "as_index", True, condition=hasattr(obj, "as_index")
   1607 ):
-> 1608     result_index, result_data = self.compute_dict_like(
   1609         op_name, selected_obj, selection, kwargs
   1610     )
   1611 result = self.wrap_results_dict_like(selected_obj, result_index, result_data)

File ~/scratch/.venv/lib/python3.12/site-packages/pandas/core/apply.py:497, in Apply.compute_dict_like(self, op_name, selected_obj, selection, kwargs)
    494 else:
    495     # key used for column selection and output
    496     results = [
--> 497         getattr(obj._gotitem(key, ndim=1), op_name)(how, **kwargs)
    498         for key, how in func.items()
    499     ]
    500     keys = list(func.keys())

File ~/scratch/.venv/lib/python3.12/site-packages/pandas/core/groupby/generic.py:257, in SeriesGroupBy.aggregate(self, func, engine, engine_kwargs, *args, **kwargs)
    256 kwargs["engine_kwargs"] = engine_kwargs
--> 257 ret = self._aggregate_multiple_funcs(func, *args, **kwargs)
    258 if relabeling:
    259     # columns is not narrowed by mypy from relabeling flag

File ~/scratch/.venv/lib/python3.12/site-packages/pandas/core/groupby/generic.py:362, in SeriesGroupBy._aggregate_multiple_funcs(self, arg, *args, **kwargs)
    361         key = base.OutputKey(label=name, position=idx)
--> 362         results[key] = self.aggregate(func, *args, **kwargs)
    364 if any(isinstance(x, DataFrame) for x in results.values()):

File ~/scratch/.venv/lib/python3.12/site-packages/pandas/core/groupby/generic.py:249, in SeriesGroupBy.aggregate(self, func, engine, engine_kwargs, *args, **kwargs)
    248         kwargs["engine_kwargs"] = engine_kwargs
--> 249     return getattr(self, func)(*args, **kwargs)
    251 elif isinstance(func, abc.Iterable):
    252     # Catch instances of lists / tuples
    253     # but not the class list / tuple itself.

File ~/scratch/.venv/lib/python3.12/site-packages/pandas/core/groupby/groupby.py:1363, in GroupBy.__getattr__(self, attr)
   1361     return self[attr]
-> 1363 raise AttributeError(
   1364     f"'{type(self).__name__}' object has no attribute '{attr}'"
   1365 )

AttributeError: 'SeriesGroupBy' object has no attribute 'n_unique'

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
Cell In[20], line 1
----> 1 func(pd.DataFrame(data)); func(pl.DataFrame(data))

File ~/scratch/.venv/lib/python3.12/site-packages/narwhals/translate.py:717, in narwhalify.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
    714     msg = "Found multiple backends. Make sure that all dataframe/series inputs come from the same backend."
    715     raise ValueError(msg)
--> 717 result = func(*args, **kwargs)
    719 return to_native(result, strict=strict)

Cell In[15], line 3, in func(df)
      1 @nw.narwhalify
      2 def func(df):
----> 3     df.group_by('a').agg(nw.col('b').n_unique())

File ~/scratch/.venv/lib/python3.12/site-packages/narwhals/group_by.py:111, in GroupBy.agg(self, *aggs, **named_aggs)
     31 """
     32 Compute aggregations for each group of a group by operation.
     33
   (...)
    107     └─────┴─────┴─────┘
    108 """
    109 aggs, named_aggs = self._df._flatten_and_extract(*aggs, **named_aggs)
    110 return self._df._from_compliant_dataframe(  # type: ignore[return-value]
--> 111     self._grouped.agg(*aggs, **named_aggs),
    112 )

File ~/scratch/.venv/lib/python3.12/site-packages/narwhals/_pandas_like/group_by.py:73, in PandasLikeGroupBy.agg(self, *aggs, **named_aggs)
     70         raise ValueError(msg)
     71     output_names.extend(expr._output_names)
---> 73 return agg_pandas(
     74     self._grouped,
     75     exprs,
     76     self._keys,
     77     output_names,
     78     self._from_native_frame,
     79     dataframe_is_empty=self._df._native_frame.empty,
     80     implementation=implementation,
     81     backend_version=self._df._backend_version,
     82 )

File ~/scratch/.venv/lib/python3.12/site-packages/narwhals/_pandas_like/group_by.py:168, in agg_pandas(grouped, exprs, keys, output_names, from_dataframe, implementation, backend_version, dataframe_is_empty)
    166 except AttributeError as exc:
    167     msg = "Failed to aggregated - does your aggregation function return a scalar?"
--> 168     raise RuntimeError(msg) from exc
    169 result_simple.columns = [f"{a}_{b}" for a, b in result_simple.columns]
    170 result_simple = result_simple.rename(columns=name_mapping).reset_index()

RuntimeError: Failed to aggregated - does your aggregation function return a scalar?
```

### Please run narwhals.show_version() and enter the output below.

```shell
System:
    python: 3.12.4 (main, Jul 25 2024, 22:42:01) [Clang 18.1.8 ]
executable: /home/marcogorelli/scratch/.venv/bin/python3
   machine: Linux-5.15.153.1-microsoft-standard-WSL2-x86_64-with-glibc2.35

Python dependencies:
     narwhals: 1.6.2
       pandas: 2.2.2
       polars: 1.6.0
         cudf:
        modin:
      pyarrow: 17.0.0
        numpy: 2.1.1
```


### Relevant log output

_No response_

Discussion:

Patch:
diff --git a/narwhals/_arrow/group_by.py b/narwhals/_arrow/group_by.py
index 27c7ff36..78b241c9 100644
--- a/narwhals/_arrow/group_by.py
+++ b/narwhals/_arrow/group_by.py
@@ -15,6 +15,12 @@ if TYPE_CHECKING:
     from narwhals._arrow.expr import ArrowExpr
     from narwhals._arrow.typing import IntoArrowExpr
 
+POLARS_TO_ARROW_AGGREGATIONS = {
+    "n_unique": "count_distinct",
+    "std": "stddev",
+    "var": "variance",  # currently unused, we don't have `var` yet
+}
+
 
 class ArrowGroupBy:
     def __init__(self, df: ArrowDataFrame, keys: list[str]) -> None:
@@ -112,16 +118,27 @@ def agg_arrow(
                 raise AssertionError(msg)
 
             function_name = remove_prefix(expr._function_name, "col->")
+            function_name = POLARS_TO_ARROW_AGGREGATIONS.get(function_name, function_name)
             for root_name, output_name in zip(expr._root_names, expr._output_names):
-                if function_name != "len":
+                if function_name == "len":
                     simple_aggregations[output_name] = (
-                        (root_name, function_name),
-                        f"{root_name}_{function_name}",
+                        (root_name, "count", pc.CountOptions(mode="all")),
+                        f"{root_name}_count",
+                    )
+                elif function_name == "count_distinct":
+                    simple_aggregations[output_name] = (
+                        (root_name, "count_distinct", pc.CountOptions(mode="all")),
+                        f"{root_name}_count_distinct",
+                    )
+                elif function_name == "stddev":
+                    simple_aggregations[output_name] = (
+                        (root_name, "stddev", pc.VarianceOptions(ddof=1)),
+                        f"{root_name}_stddev",
                     )
                 else:
                     simple_aggregations[output_name] = (
-                        (root_name, "count", pc.CountOptions(mode="all")),
-                        f"{root_name}_count",
+                        (root_name, function_name),
+                        f"{root_name}_{function_name}",
                     )
 
         aggs: list[Any] = []
diff --git a/narwhals/_dask/group_by.py b/narwhals/_dask/group_by.py
index 8538c62d..463d6fc5 100644
--- a/narwhals/_dask/group_by.py
+++ b/narwhals/_dask/group_by.py
@@ -10,12 +10,33 @@ from narwhals._expression_parsing import parse_into_exprs
 from narwhals.utils import remove_prefix
 
 if TYPE_CHECKING:
+    import dask.dataframe as dd
+    import pandas as pd
+
     from narwhals._dask.dataframe import DaskLazyFrame
     from narwhals._dask.expr import DaskExpr
     from narwhals._dask.typing import IntoDaskExpr
 
-POLARS_TO_PANDAS_AGGREGATIONS = {
+
+def n_unique() -> dd.Aggregation:
+    import dask.dataframe as dd  # ignore-banned-import
+
+    def chunk(s: pd.core.groupby.generic.SeriesGroupBy) -> int:
+        return s.nunique(dropna=False)  # type: ignore[no-any-return]
+
+    def agg(s0: pd.core.groupby.generic.SeriesGroupBy) -> int:
+        return s0.sum()  # type: ignore[no-any-return]
+
+    return dd.Aggregation(
+        name="nunique",
+        chunk=chunk,
+        agg=agg,
+    )
+
+
+POLARS_TO_DASK_AGGREGATIONS = {
     "len": "size",
+    "n_unique": n_unique,
 }
 
 
@@ -85,7 +106,7 @@ def agg_dask(
             break
 
     if all_simple_aggs:
-        simple_aggregations: dict[str, tuple[str, str]] = {}
+        simple_aggregations: dict[str, tuple[str, str | dd.Aggregation]] = {}
         for expr in exprs:
             if expr._depth == 0:
                 # e.g. agg(nw.len()) # noqa: ERA001
@@ -93,7 +114,7 @@ def agg_dask(
                     msg = "Safety assertion failed, please report a bug to https://github.com/narwhals-dev/narwhals/issues"
                     raise AssertionError(msg)
 
-                function_name = POLARS_TO_PANDAS_AGGREGATIONS.get(
+                function_name = POLARS_TO_DASK_AGGREGATIONS.get(
                     expr._function_name, expr._function_name
                 )
                 for output_name in expr._output_names:
@@ -108,9 +129,11 @@ def agg_dask(
                 raise AssertionError(msg)
 
             function_name = remove_prefix(expr._function_name, "col->")
-            function_name = POLARS_TO_PANDAS_AGGREGATIONS.get(
-                function_name, function_name
-            )
+            function_name = POLARS_TO_DASK_AGGREGATIONS.get(function_name, function_name)
+
+            # deal with n_unique case in a "lazy" mode to not depend on dask globally
+            function_name = function_name() if callable(function_name) else function_name
+
             for root_name, output_name in zip(expr._root_names, expr._output_names):
                 simple_aggregations[output_name] = (root_name, function_name)
         try:
diff --git a/narwhals/_pandas_like/group_by.py b/narwhals/_pandas_like/group_by.py
index 11abc85c..97a477dc 100644
--- a/narwhals/_pandas_like/group_by.py
+++ b/narwhals/_pandas_like/group_by.py
@@ -21,6 +21,7 @@ if TYPE_CHECKING:
 
 POLARS_TO_PANDAS_AGGREGATIONS = {
     "len": "size",
+    "n_unique": "nunique",
 }
 
 
@@ -103,7 +104,7 @@ class PandasLikeGroupBy:
         yield from ((key, self._from_native_frame(sub_df)) for (key, sub_df) in iterator)
 
 
-def agg_pandas(
+def agg_pandas(  # noqa: PLR0915
     grouped: Any,
     exprs: list[PandasLikeExpr],
     keys: list[str],
@@ -120,13 +121,18 @@ def agg_pandas(
     - https://github.com/rapidsai/cudf/issues/15118
     - https://github.com/rapidsai/cudf/issues/15084
     """
-    all_simple_aggs = True
+    all_aggs_are_simple = True
     for expr in exprs:
         if not is_simple_aggregation(expr):
-            all_simple_aggs = False
+            all_aggs_are_simple = False
             break
 
-    if all_simple_aggs:
+    # dict of {output_name: root_name} that we count n_unique on
+    # We need to do this separately from the rest so that we
+    # can pass the `dropna` kwargs.
+    nunique_aggs: dict[str, str] = {}
+
+    if all_aggs_are_simple:
         simple_aggregations: dict[str, tuple[str, str]] = {}
         for expr in exprs:
             if expr._depth == 0:
@@ -154,21 +160,54 @@ def agg_pandas(
                 function_name, function_name
             )
             for root_name, output_name in zip(expr._root_names, expr._output_names):
-                simple_aggregations[output_name] = (root_name, function_name)
+                if function_name == "nunique":
+                    nunique_aggs[output_name] = root_name
+                else:
+                    simple_aggregations[output_name] = (root_name, function_name)
 
-        aggs = collections.defaultdict(list)
+        simple_aggs = collections.defaultdict(list)
         name_mapping = {}
         for output_name, named_agg in simple_aggregations.items():
-            aggs[named_agg[0]].append(named_agg[1])
+            simple_aggs[named_agg[0]].append(named_agg[1])
             name_mapping[f"{named_agg[0]}_{named_agg[1]}"] = output_name
-        try:
-            result_simple = grouped.agg(aggs)
-        except AttributeError as exc:
-            msg = "Failed to aggregated - does your aggregation function return a scalar?"
-            raise RuntimeError(msg) from exc
-        result_simple.columns = [f"{a}_{b}" for a, b in result_simple.columns]
-        result_simple = result_simple.rename(columns=name_mapping).reset_index()
-        return from_dataframe(result_simple.loc[:, output_names])
+        if simple_aggs:
+            try:
+                result_simple_aggs = grouped.agg(simple_aggs)
+            except AttributeError as exc:
+                msg = "Failed to aggregated - does your aggregation function return a scalar?"
+                raise RuntimeError(msg) from exc
+            result_simple_aggs.columns = [
+                f"{a}_{b}" for a, b in result_simple_aggs.columns
+            ]
+            result_simple_aggs = result_simple_aggs.rename(
+                columns=name_mapping
+            ).reset_index()
+        if nunique_aggs:
+            result_nunique_aggs = grouped[list(nunique_aggs.values())].nunique(
+                dropna=False
+            )
+            result_nunique_aggs.columns = list(nunique_aggs.keys())
+            result_nunique_aggs = result_nunique_aggs.reset_index()
+        if simple_aggs and nunique_aggs:
+            if (
+                set(result_simple_aggs.columns)
+                .difference(keys)
+                .intersection(result_nunique_aggs.columns)
+            ):
+                msg = (
+                    "Got two aggregations with the same output name. Please make sure "
+                    "that aggregations have unique output names."
+                )
+                raise ValueError(msg)
+            result_aggs = result_simple_aggs.merge(result_nunique_aggs, on=keys)
+        elif nunique_aggs and not simple_aggs:
+            result_aggs = result_nunique_aggs
+        elif simple_aggs and not nunique_aggs:
+            result_aggs = result_simple_aggs
+        else:  # pragma: no cover
+            msg = "Congrats, you entered unreachable code. Please report a bug to https://github.com/narwhals-dev/narwhals/issues."
+            raise RuntimeError(msg)
+        return from_dataframe(result_aggs.loc[:, output_names])
 
     if dataframe_is_empty:
         # Don't even attempt this, it's way too inconsistent across pandas versions.


Test patch:
diff --git a/tests/test_group_by.py b/tests/test_group_by.py
index 2bb8d435..4bd3427a 100644
--- a/tests/test_group_by.py
+++ b/tests/test_group_by.py
@@ -102,6 +102,57 @@ def test_group_by_len(constructor: Any) -> None:
     compare_dicts(result, expected)
 
 
+def test_group_by_n_unique(constructor: Any) -> None:
+    result = (
+        nw.from_native(constructor(data))
+        .group_by("a")
+        .agg(nw.col("b").n_unique())
+        .sort("a")
+    )
+    expected = {"a": [1, 3], "b": [1, 1]}
+    compare_dicts(result, expected)
+
+
+def test_group_by_std(constructor: Any) -> None:
+    data = {"a": [1, 1, 2, 2], "b": [5, 4, 3, 2]}
+    result = (
+        nw.from_native(constructor(data)).group_by("a").agg(nw.col("b").std()).sort("a")
+    )
+    expected = {"a": [1, 2], "b": [0.707107] * 2}
+    compare_dicts(result, expected)
+
+
+def test_group_by_n_unique_w_missing(constructor: Any) -> None:
+    data = {"a": [1, 1, 2], "b": [4, None, 5], "c": [None, None, 7], "d": [1, 1, 3]}
+    result = (
+        nw.from_native(constructor(data))
+        .group_by("a")
+        .agg(
+            nw.col("b").n_unique(),
+            c_n_unique=nw.col("c").n_unique(),
+            c_n_min=nw.col("b").min(),
+            d_n_unique=nw.col("d").n_unique(),
+        )
+        .sort("a")
+    )
+    expected = {
+        "a": [1, 2],
+        "b": [2, 1],
+        "c_n_unique": [1, 1],
+        "c_n_min": [4, 5],
+        "d_n_unique": [1, 1],
+    }
+    compare_dicts(result, expected)
+
+
+def test_group_by_same_name_twice() -> None:
+    import pandas as pd
+
+    df = pd.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]})
+    with pytest.raises(ValueError, match="two aggregations with the same"):
+        nw.from_native(df).group_by("a").agg(nw.col("b").sum(), nw.col("b").n_unique())
+
+
 def test_group_by_empty_result_pandas() -> None:
     df_any = pd.DataFrame({"a": [1, 2, 3], "b": [4, 3, 2]})
     df = nw.from_native(df_any, eager_only=True)