diff --git a/narwhals/_compliant/group_by.py b/narwhals/_compliant/group_by.py index d9271170be..ead3ee9674 100644 --- a/narwhals/_compliant/group_by.py +++ b/narwhals/_compliant/group_by.py @@ -108,20 +108,19 @@ def _temporary_name(key: str) -> str: key_str = str(key) # pandas allows non-string column names :sob: return f"_{key_str}_tmp{'_' * (tmp_name_length - len(key_str) - 5)}" - output_names = _evaluate_aliases(compliant_frame, keys) - + keys_aliases = [expr._evaluate_aliases(compliant_frame) for expr in keys] safe_keys = [ # multi-output expression cannot have duplicate names, hence it's safe to suffix key.name.map(_temporary_name) if (metadata := key._metadata) and metadata.expansion_kind.is_multi_output() # otherwise it's single named and we can use Expr.alias - else key.alias(_temporary_name(new_name)) - for key, new_name in zip(keys, output_names) + else key.alias(_temporary_name(new_names[0])) + for key, new_names in zip(keys, keys_aliases) ] return ( compliant_frame.with_columns(*safe_keys), _evaluate_aliases(compliant_frame, safe_keys), - output_names, + list(chain.from_iterable(keys_aliases)), ) diff --git a/tests/frame/group_by_test.py b/tests/frame/group_by_test.py index ea30ff90be..d8e22e471c 100644 --- a/tests/frame/group_by_test.py +++ b/tests/frame/group_by_test.py @@ -558,14 +558,19 @@ def test_group_by_raise_drop_null_keys_with_exprs( def test_group_by_selector(constructor: Constructor) -> None: - data = {"a": [1, 1, 1], "b": [4, 4, 6], "c": [7.5, 8.5, 9.0]} + data = { + "a": [1, 1, 1], + "b": [4, 4, 6], + "c": ["foo", "foo", "bar"], + "x": [7.5, 8.5, 9.0], + } result = ( nw.from_native(constructor(data)) - .group_by(nw.selectors.by_dtype(nw.Int64)) - .agg(nw.col("c").mean()) + .group_by(nw.selectors.by_dtype(nw.Int64), "c") + .agg(nw.col("x").mean()) .sort("a", "b") ) - expected = {"a": [1, 1], "b": [4, 6], "c": [8.0, 9.0]} + expected = {"a": [1, 1], "b": [4, 6], "c": ["foo", "bar"], "x": [8.0, 9.0]} assert_equal_data(result, expected)