From a1abbb6d4239cb629bacb4f758c5016589630f22 Mon Sep 17 00:00:00 2001 From: Marco Edward Gorelli Date: Sat, 22 Jun 2024 10:49:46 +0100 Subject: [PATCH] fix: group-by len (#330) --- narwhals/_pandas_like/group_by.py | 7 +++++-- tests/group_by_test.py | 9 +++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/narwhals/_pandas_like/group_by.py b/narwhals/_pandas_like/group_by.py index a5c57a857..203b6f840 100644 --- a/narwhals/_pandas_like/group_by.py +++ b/narwhals/_pandas_like/group_by.py @@ -118,9 +118,12 @@ def agg_pandas( # noqa: PLR0913 assert expr._depth == 1 assert expr._root_names is not None assert expr._output_names is not None + function_name = remove_prefix(expr._function_name, "col->") + function_name = POLARS_TO_PANDAS_AGGREGATIONS.get( + function_name, function_name + ) for root_name, output_name in zip(expr._root_names, expr._output_names): - name = remove_prefix(expr._function_name, "col->") - simple_aggregations[output_name] = (root_name, name) + simple_aggregations[output_name] = (root_name, function_name) aggs = collections.defaultdict(list) name_mapping = {} diff --git a/tests/group_by_test.py b/tests/group_by_test.py index ec0b85efa..52cf96dfd 100644 --- a/tests/group_by_test.py +++ b/tests/group_by_test.py @@ -35,3 +35,12 @@ def test_group_by_iter(constructor: Any) -> None: for key, _df in df.group_by(["a", "b"]): keys.append(key) assert sorted(keys) == sorted(expected_keys) + + +@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame]) +def test_group_by_len(constructor: Any) -> None: + result = ( + nw.from_native(constructor(data)).group_by("a").agg(nw.col("b").len()).sort("a") + ) + expected = {"a": [1, 3], "b": [2, 1]} + compare_dicts(result, expected)