perf: streamline pyarrow and pandas group by agg (#1621)

Author: FBruzzesi

narwhals/_arrow/group_by.py

@@ -150,9 +150,10 @@ def agg_arrow(
         )
         raise ValueError(msg)
 
-    # Mapping from output name to
-    # (aggregation_args, pyarrow_output_name)  # noqa: ERA001
-    simple_aggregations: dict[str, tuple[tuple[Any, ...], str]] = {}
+    aggs: list[tuple[str, str, pc.FunctionOptions | None]] = []
+    expected_pyarrow_column_names: list[str] = keys.copy()
+    new_column_names: list[str] = keys.copy()
+
     for expr in exprs:
         if expr._depth == 0:
             # e.g. agg(nw.len()) # noqa: ERA001
@@ -161,10 +162,11 @@ def agg_arrow(
             ):  # pragma: no cover
                 msg = "Safety assertion failed, please report a bug to https://github.com/narwhals-dev/narwhals/issues"
                 raise AssertionError(msg)
-            simple_aggregations[expr._output_names[0]] = (
-                (keys[0], "count", pc.CountOptions(mode="all")),
-                f"{keys[0]}_count",
-            )
+
+            new_column_names.append(expr._output_names[0])
+            expected_pyarrow_column_names.append(f"{keys[0]}_count")
+            aggs.append((keys[0], "count", pc.CountOptions(mode="all")))
+
             continue
 
         # e.g. agg(nw.mean('a')) # noqa: ERA001
@@ -179,22 +181,13 @@ def agg_arrow(
             function_name, (function_name, None)
         )
 
-        for root_name, output_name in zip(expr._root_names, expr._output_names):
-            simple_aggregations[output_name] = (
-                (root_name, function_name, option),
-                f"{root_name}_{function_name}",
-            )
-
-    aggs: list[Any] = []
-    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)
-        expected_pyarrow_column_names.append(pyarrow_output_name)
-        new_column_names.append(output_name)
+        new_column_names.extend(expr._output_names)
+        expected_pyarrow_column_names.extend(
+            [f"{root_name}_{function_name}" for root_name in expr._root_names]
+        )
+        aggs.extend(
+            [(root_name, function_name, option) for root_name in expr._root_names]
+        )
 
     result_simple = grouped.aggregate(aggs)
 

narwhals/_pandas_like/group_by.py

@@ -166,9 +166,11 @@ def agg_pandas(  # noqa: PLR0915
     # We need to do this separately from the rest so that we
     # can pass the `dropna` kwargs.
     nunique_aggs: dict[str, str] = {}
+    simple_aggs: dict[str, list[str]] = collections.defaultdict(list)
+    expected_old_names: list[str] = []
+    new_names: list[str] = []
 
     if all_aggs_are_simple:
-        simple_aggregations: dict[str, tuple[str, str]] = {}
         for expr in exprs:
             if expr._depth == 0:
                 # e.g. agg(nw.len()) # noqa: ERA001
@@ -180,7 +182,9 @@ def agg_pandas(  # noqa: PLR0915
                     expr._function_name, expr._function_name
                 )
                 for output_name in expr._output_names:
-                    simple_aggregations[output_name] = (keys[0], function_name)
+                    new_names.append(output_name)
+                    expected_old_names.append(f"{keys[0]}_{function_name}")
+                    simple_aggs[keys[0]].append(function_name)
                 continue
 
             # e.g. agg(nw.mean('a')) # noqa: ERA001
@@ -199,15 +203,10 @@ def agg_pandas(  # noqa: PLR0915
                 if is_n_unique:
                     nunique_aggs[output_name] = root_name
                 else:
-                    simple_aggregations[output_name] = (root_name, function_name)
+                    new_names.append(output_name)
+                    expected_old_names.append(f"{root_name}_{function_name}")
+                    simple_aggs[root_name].append(function_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 = [