Repository: narwhals-dev/narwhals

Issue:
bug: pyarrow backend raises with `.group_by('a').agg(c=nw.col('b').sum(), d=nw.col('b').sum())`
```python
import pyarrow as pa
import narwhals as nw
import pandas as pd
import polars as pl

data = {"a": [1, 1, 1, 1, 1, 2, 2, 2, 2, 2], "b": list(range(10))}

nw.from_native(pa.table(data)).group_by('a').agg(c=nw.col('b').sum(), d=nw.col('b').sum())
```
```
$ python t.py 
Traceback (most recent call last):
  File "/home/marcogorelli/polars-api-compat-dev/t.py", line 13, in <module>
    nw.from_native(pa.table(data)).group_by('a').agg(c=nw.col('b').sum(), d=nw.col('b').sum())
  File "/home/marcogorelli/polars-api-compat-dev/narwhals/group_by.py", line 114, in agg
    self._grouped.agg(*aggs, **named_aggs),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/marcogorelli/polars-api-compat-dev/narwhals/_arrow/group_by.py", line 83, in agg
    return agg_arrow(
           ^^^^^^^^^^
  File "/home/marcogorelli/polars-api-compat-dev/narwhals/_arrow/group_by.py", line 191, in agg_arrow
    ).select(output_names)
      ^^^^^^^^^^^^^^^^^^^^
  File "pyarrow/table.pxi", line 4330, in pyarrow.lib.Table.select
  File "pyarrow/table.pxi", line 1715, in pyarrow.lib._Tabular._ensure_integer_index
KeyError: 'Field "c" does not exist in schema'
```

Discussion:
MarcoGorelli: thanks for your review 🙏 !

Patch:
diff --git a/.github/workflows/extremes.yml b/.github/workflows/extremes.yml
index b1939004..84a0cb2c 100644
--- a/.github/workflows/extremes.yml
+++ b/.github/workflows/extremes.yml
@@ -25,12 +25,10 @@ jobs:
           cache-suffix: ${{ matrix.python-version }}
           cache-dependency-glob: "pyproject.toml"
       - name: install-minimum-versions
-        run: uv pip install tox virtualenv setuptools pandas==0.25.3 polars==0.20.3 numpy==1.17.5 pyarrow==11.0.0 scipy==1.5.0 scikit-learn==1.1.0 tzdata --system
+        run: uv pip install pipdeptree tox virtualenv setuptools pandas==0.25.3 polars==0.20.3 numpy==1.17.5 pyarrow==11.0.0 "pyarrow-stubs<17" scipy==1.5.0 scikit-learn==1.1.0 tzdata --system
       - name: install-reqs
         run: |
           uv pip install -e ".[dev]" --system
-          : # pyspark >= 3.3.0 is not compatible with pandas==0.25.3
-          uv pip uninstall pyspark --system
       - name: show-deps
         run: uv pip freeze
       - name: Run pytest
@@ -55,11 +53,13 @@ jobs:
           cache-suffix: ${{ matrix.python-version }}
           cache-dependency-glob: "pyproject.toml"
       - name: install-minimum-versions
-        run: uv pip install tox virtualenv setuptools pandas==1.1.5 polars==0.20.3 numpy==1.17.5 pyarrow==11.0.0 pyspark==3.3.0 scipy==1.5.0 scikit-learn==1.1.0 tzdata --system
+        run: uv pip install pipdeptree tox virtualenv setuptools pandas==1.1.5 polars==0.20.3 numpy==1.17.5 pyarrow==11.0.0 "pyarrow-stubs<17" pyspark==3.3.0 scipy==1.5.0 scikit-learn==1.1.0 tzdata --system
       - name: install-reqs
         run: uv pip install -e ".[dev]" --system
       - name: show-deps
         run: uv pip freeze
+      - name: show-deptree
+        run: pipdeptree
       - name: Run pytest
         run: pytest tests --cov=narwhals --cov=tests --cov-fail-under=50 --runslow
 
diff --git a/narwhals/_arrow/group_by.py b/narwhals/_arrow/group_by.py
index 7e2422a1..290118f0 100644
--- a/narwhals/_arrow/group_by.py
+++ b/narwhals/_arrow/group_by.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from copy import copy
+import collections
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
@@ -69,7 +69,6 @@ class ArrowGroupBy:
             namespace=self._df.__narwhals_namespace__(),
             **named_aggs,
         )
-        output_names: list[str] = copy(self._keys)
         for expr in exprs:
             if expr._output_names is None:
                 msg = (
@@ -78,13 +77,11 @@ class ArrowGroupBy:
                     "`nw.col('a', 'b')`\n"
                 )
                 raise ValueError(msg)
-            output_names.extend(expr._output_names)
 
         return agg_arrow(
             self._grouped,
             exprs,
             self._keys,
-            output_names,
             self._df._from_native_frame,
         )
 
@@ -126,7 +123,6 @@ def agg_arrow(
     grouped: pa.TableGroupBy,
     exprs: Sequence[CompliantExpr[ArrowSeries]],
     keys: list[str],
-    output_names: list[str],
     from_dataframe: Callable[[Any], ArrowDataFrame],
 ) -> ArrowDataFrame:
     import pyarrow.compute as pc
@@ -178,17 +174,38 @@ def agg_arrow(
                 )
 
         aggs: list[Any] = []
-        name_mapping = {}
+        expected_pyarrow_column_names = keys.copy()
+        new_column_names = keys.copy()
         for output_name, (
             aggregation_args,
             pyarrow_output_name,
         ) in simple_aggregations.items():
             aggs.append(aggregation_args)
-            name_mapping[pyarrow_output_name] = output_name
+            expected_pyarrow_column_names.append(pyarrow_output_name)
+            new_column_names.append(output_name)
+
         result_simple = grouped.aggregate(aggs)
-        result_simple = result_simple.rename_columns(
-            [name_mapping.get(col, col) for col in result_simple.column_names]
-        ).select(output_names)
+
+        # Rename columns, being very careful
+        expected_old_names_indices: dict[str, list[int]] = collections.defaultdict(list)
+        for idx, item in enumerate(expected_pyarrow_column_names):
+            expected_old_names_indices[item].append(idx)
+        if not (
+            set(result_simple.column_names) == set(expected_pyarrow_column_names)
+            and len(result_simple.column_names) == len(expected_pyarrow_column_names)
+        ):  # pragma: no cover
+            msg = (
+                f"Safety assertion failed, expected {expected_pyarrow_column_names} "
+                f"got {result_simple.column_names}, "
+                "please report a bug at https://github.com/narwhals-dev/narwhals/issues"
+            )
+            raise AssertionError(msg)
+        index_map: list[int] = [
+            expected_old_names_indices[item].pop(0) for item in result_simple.column_names
+        ]
+        new_column_names = [new_column_names[i] for i in index_map]
+
+        result_simple = result_simple.rename_columns(new_column_names)
         return from_dataframe(result_simple)
 
     msg = (
diff --git a/narwhals/_pandas_like/group_by.py b/narwhals/_pandas_like/group_by.py
index a1abe2eb..b014f4a9 100644
--- a/narwhals/_pandas_like/group_by.py
+++ b/narwhals/_pandas_like/group_by.py
@@ -12,7 +12,6 @@ from typing import Sequence
 from narwhals._expression_parsing import is_simple_aggregation
 from narwhals._expression_parsing import parse_into_exprs
 from narwhals._pandas_like.utils import native_series_from_iterable
-from narwhals._pandas_like.utils import rename
 from narwhals._pandas_like.utils import select_columns_by_name
 from narwhals.utils import Implementation
 from narwhals.utils import find_stacklevel
@@ -202,25 +201,46 @@ def agg_pandas(  # noqa: PLR0915
                 else:
                     simple_aggregations[output_name] = (root_name, function_name)
 
-        simple_aggs = collections.defaultdict(list)
-        name_mapping = {}
-        for output_name, named_agg in simple_aggregations.items():
-            simple_aggs[named_agg[0]].append(named_agg[1])
-            name_mapping[f"{named_agg[0]}_{named_agg[1]}"] = output_name
+        simple_aggs: dict[str, list[str]] = collections.defaultdict(list)
+        expected_old_names: list[str] = []
+        new_names: list[str] = []
+        for output_name, (col_name, function) in simple_aggregations.items():
+            simple_aggs[col_name].append(function)
+            new_names.append(output_name)
+            expected_old_names.append(f"{col_name}_{function}")
         if simple_aggs:
             result_simple_aggs = grouped.agg(simple_aggs)
             result_simple_aggs.columns = [
                 f"{a}_{b}" for a, b in result_simple_aggs.columns
             ]
-            result_simple_aggs = rename(
-                result_simple_aggs,
-                columns=name_mapping,
-                implementation=implementation,
-                backend_version=backend_version,
+            if not (
+                set(result_simple_aggs.columns) == set(expected_old_names)
+                and len(result_simple_aggs.columns) == len(expected_old_names)
+            ):  # pragma: no cover
+                msg = (
+                    f"Safety assertion failed, expected {expected_old_names} "
+                    f"got {result_simple_aggs.columns}, "
+                    "please report a bug at https://github.com/narwhals-dev/narwhals/issues"
+                )
+                raise AssertionError(msg)
+
+            # Rename columns, being very careful
+            expected_old_names_indices: dict[str, list[int]] = collections.defaultdict(
+                list
             )
+            for idx, item in enumerate(expected_old_names):
+                expected_old_names_indices[item].append(idx)
+            index_map: list[int] = [
+                expected_old_names_indices[item].pop(0)
+                for item in result_simple_aggs.columns
+            ]
+            new_names = [new_names[i] for i in index_map]
+            result_simple_aggs.columns = new_names
+
             # Keep inplace=True to avoid making a redundant copy.
             # This may need updating, depending on https://github.com/pandas-dev/pandas/pull/51466/files
             result_simple_aggs.reset_index(inplace=True)  # noqa: PD002
+
         if nunique_aggs:
             result_nunique_aggs = grouped[list(nunique_aggs.values())].nunique(
                 dropna=False
diff --git a/pyproject.toml b/pyproject.toml
index bc13ae65..aaba74e4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,7 +45,7 @@ dev = [
   "pandas",
   "polars",
   "pyarrow",
-  "pyspark; python_version < '3.12'",
+  "pyspark; python_version >= '3.9' and python_version < '3.12'",
   "pyarrow-stubs",
   "pytest",
   "pytest-cov",


Test patch:
diff --git a/tests/group_by_test.py b/tests/group_by_test.py
index 61e63e74..9efda1fa 100644
--- a/tests/group_by_test.py
+++ b/tests/group_by_test.py
@@ -364,3 +364,19 @@ def test_group_by_shift_raises(
         ValueError, match=".*(failed to aggregate|Non-trivial complex aggregation found)"
     ):
         df.group_by("b").agg(nw.col("a").shift(1))
+
+
+def test_double_same_aggregation(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if "dask" in str(constructor) or "modin" in str(constructor):
+        # bugged in dask https://github.com/dask/dask/issues/11612
+        # and modin lol https://github.com/modin-project/modin/issues/7414
+        # At least cudf gets it right
+        request.applymarker(pytest.mark.xfail)
+    if "pandas" in str(constructor) and PANDAS_VERSION < (1,):
+        request.applymarker(pytest.mark.xfail)
+    df = nw.from_native(constructor({"a": [1, 1, 2], "b": [4, 5, 6]}))
+    result = df.group_by("a").agg(c=nw.col("b").mean(), d=nw.col("b").mean()).sort("a")
+    expected = {"a": [1, 2], "c": [4.5, 6], "d": [4.5, 6]}
+    assert_equal_data(result, expected)