feat: Allow expressions as group_by keys

narwhals/_arrow/dataframe.py

@@ -473,7 +473,7 @@ def with_columns(self: ArrowDataFrame, *exprs: ArrowExpr) -> ArrowDataFrame:
 
         return self._with_native(native_frame, validate_column_names=False)
 
-    def group_by(self: Self, *keys: str, drop_null_keys: bool) -> ArrowGroupBy:
+    def group_by(self: Self, *keys: ArrowExpr, drop_null_keys: bool) -> ArrowGroupBy:
         from narwhals._arrow.group_by import ArrowGroupBy
 
         return ArrowGroupBy(self, keys, drop_null_keys=drop_null_keys)

narwhals/_arrow/expr.py

@@ -178,7 +178,8 @@ def func(df: ArrowDataFrame) -> Sequence[ArrowSeries]:
                     )
                     raise NotImplementedError(msg)
 
-                tmp = df.group_by(*partition_by, drop_null_keys=False).agg(self)
+                ns = self.__narwhals_namespace__()
+                tmp = df.group_by(ns.col(*partition_by), drop_null_keys=False).agg(self)
                 tmp = df.simple_select(*partition_by).join(
                     tmp,
                     how="left",

narwhals/_arrow/group_by.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import collections
+from itertools import chain
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import ClassVar
@@ -41,18 +42,26 @@ class ArrowGroupBy(EagerGroupBy["ArrowDataFrame", "ArrowExpr"]):
     }
 
     def __init__(
-        self,
+        self: Self,
         compliant_frame: ArrowDataFrame,
-        keys: Sequence[str],
+        keys: Sequence[ArrowExpr],
         /,
         *,
         drop_null_keys: bool,
     ) -> None:
-        if drop_null_keys:
-            self._compliant_frame = compliant_frame.drop_nulls(keys)
-        else:
-            self._compliant_frame = compliant_frame
-        self._keys: list[str] = list(keys)
+        compliant_frame = compliant_frame.with_columns(*keys)
+
+        self._keys: list[str] = list(
+            chain.from_iterable(
+                evaluate_output_names_and_aliases(
+                    expr=key, df=compliant_frame, exclude=[]
+                )[1]
+                for key in keys
+            )
+        )
+        self._compliant_frame = (
+            compliant_frame.drop_nulls(self._keys) if drop_null_keys else compliant_frame
+        )
         self._grouped = pa.TableGroupBy(self.compliant.native, self._keys)
 
     def agg(self: Self, *exprs: ArrowExpr) -> ArrowDataFrame:

narwhals/_compliant/dataframe.py

@@ -12,6 +12,7 @@
 from typing import TypeVar
 from typing import overload
 
+from narwhals._compliant.typing import CompliantExprAny
 from narwhals._compliant.typing import CompliantExprT_contra
 from narwhals._compliant.typing import CompliantSeriesT
 from narwhals._compliant.typing import EagerExprT_contra
@@ -127,7 +128,7 @@ def filter(self, predicate: CompliantExprT_contra | Incomplete) -> Self: ...
     def gather_every(self, n: int, offset: int) -> Self: ...
     def get_column(self, name: str) -> CompliantSeriesT: ...
     def group_by(
-        self, *keys: str, drop_null_keys: bool
+        self, *keys: CompliantExprAny, drop_null_keys: bool
     ) -> CompliantGroupBy[Self, Any]: ...
     def head(self, n: int) -> Self: ...
     def item(self, row: int | None, column: int | str | None) -> Any: ...
@@ -255,7 +256,7 @@ def filter(self, predicate: CompliantExprT_contra | Incomplete) -> Self: ...
     )
     def gather_every(self, n: int, offset: int) -> Self: ...
     def group_by(
-        self, *keys: str, drop_null_keys: bool
+        self, *keys: CompliantExprAny, drop_null_keys: bool
     ) -> CompliantGroupBy[Self, Any]: ...
     def head(self, n: int) -> Self: ...
     def join(

narwhals/_compliant/group_by.py

@@ -14,6 +14,7 @@
 from typing import TypeVar
 
 from narwhals._compliant.typing import CompliantDataFrameT_co
+from narwhals._compliant.typing import CompliantExprAny
 from narwhals._compliant.typing import CompliantExprT_contra
 from narwhals._compliant.typing import CompliantFrameT_co
 from narwhals._compliant.typing import CompliantLazyFrameT_co
@@ -67,7 +68,7 @@ def compliant(self) -> CompliantFrameT_co:
     def __init__(
         self,
         compliant_frame: CompliantFrameT_co,
-        keys: Sequence[str],
+        keys: Sequence[CompliantExprAny],
         /,
         *,
         drop_null_keys: bool,

narwhals/_compliant/series.py

@@ -300,7 +300,7 @@ def _with_native(self, series: Any, *, preserve_broadcast: bool = False) -> Self
         """Return a new `CompliantSeries`, wrapping the native `series`.
 
         In cases when operations are known to not affect whether a result should
-        be broadcast, we can pass `preverse_broadcast=True`.
+        be broadcast, we can pass `preserve_broadcast=True`.
         Set this with care - it should only be set for unary expressions which don't
         change length or order, such as `.alias` or `.fill_null`. If in doubt, don't
         set it, you probably don't need it.

narwhals/_dask/dataframe.py

@@ -388,7 +388,7 @@ def join_asof(
             ),
         )
 
-    def group_by(self: Self, *by: str, drop_null_keys: bool) -> DaskLazyGroupBy:
+    def group_by(self: Self, *by: DaskExpr, drop_null_keys: bool) -> DaskLazyGroupBy:
         from narwhals._dask.group_by import DaskLazyGroupBy
 
         return DaskLazyGroupBy(self, by, drop_null_keys=drop_null_keys)

narwhals/_dask/group_by.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from functools import partial
+from itertools import chain
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
@@ -72,12 +73,24 @@ class DaskLazyGroupBy(DepthTrackingGroupBy["DaskLazyFrame", "DaskExpr", Aggregat
     }
 
     def __init__(
-        self: Self, df: DaskLazyFrame, keys: Sequence[str], /, *, drop_null_keys: bool
+        self: Self,
+        compliant_frame: DaskLazyFrame,
+        keys: Sequence[DaskExpr],
+        /,
+        *,
+        drop_null_keys: bool,
     ) -> None:
-        self._compliant_frame = df
-        self._keys: list[str] = list(keys)
+        self._compliant_frame = compliant_frame.with_columns(*keys)
+        self._keys: list[str] = list(
+            chain.from_iterable(
+                evaluate_output_names_and_aliases(
+                    expr=key, df=compliant_frame, exclude=[]
+                )[1]
+                for key in keys
+            )
+        )
         self._grouped = self.compliant.native.groupby(
-            list(self._keys), dropna=drop_null_keys, observed=True
+            self._keys, dropna=drop_null_keys, observed=True
         )
 
     def agg(self: Self, *exprs: DaskExpr) -> DaskLazyFrame:

narwhals/_duckdb/dataframe.py

@@ -223,7 +223,7 @@ def _with_native(self: Self, df: duckdb.DuckDBPyRelation) -> Self:
             df, backend_version=self._backend_version, version=self._version
         )
 
-    def group_by(self: Self, *keys: str, drop_null_keys: bool) -> DuckDBGroupBy:
+    def group_by(self: Self, *keys: DuckDBExpr, drop_null_keys: bool) -> DuckDBGroupBy:
         from narwhals._duckdb.group_by import DuckDBGroupBy
 
         return DuckDBGroupBy(self, keys, drop_null_keys=drop_null_keys)

narwhals/_duckdb/group_by.py

@@ -5,6 +5,7 @@
 from typing import Sequence
 
 from narwhals._compliant import LazyGroupBy
+from narwhals._expression_parsing import evaluate_output_names_and_aliases
 
 if TYPE_CHECKING:
     from duckdb import Expression  # noqa: F401
@@ -17,14 +18,27 @@
 class DuckDBGroupBy(LazyGroupBy["DuckDBLazyFrame", "DuckDBExpr", "Expression"]):
     def __init__(
         self: Self,
-        df: DuckDBLazyFrame,
-        keys: Sequence[str],
+        compliant_frame: DuckDBLazyFrame,
+        keys: Sequence[DuckDBExpr],
         /,
         *,
         drop_null_keys: bool,
     ) -> None:
-        self._compliant_frame = df.drop_nulls(subset=None) if drop_null_keys else df
-        self._keys = list(keys)
+        compliant_frame = compliant_frame.with_columns(*keys)
+
+        self._keys: list[str] = list(
+            chain.from_iterable(
+                evaluate_output_names_and_aliases(
+                    expr=key, df=compliant_frame, exclude=[]
+                )[1]
+                for key in keys
+            )
+        )
+        self._compliant_frame = (
+            compliant_frame.drop_nulls(subset=self._keys)
+            if drop_null_keys
+            else compliant_frame
+        )
 
     def agg(self: Self, *exprs: DuckDBExpr) -> DuckDBLazyFrame:
         agg_columns = list(chain(self._keys, self._evaluate_exprs(exprs)))

narwhals/_pandas_like/dataframe.py

@@ -699,7 +699,9 @@ def collect(
         raise ValueError(msg)  # pragma: no cover
 
     # --- actions ---
-    def group_by(self: Self, *keys: str, drop_null_keys: bool) -> PandasLikeGroupBy:
+    def group_by(
+        self: Self, *keys: PandasLikeExpr, drop_null_keys: bool
+    ) -> PandasLikeGroupBy:
         from narwhals._pandas_like.group_by import PandasLikeGroupBy
 
         return PandasLikeGroupBy(self, keys, drop_null_keys=drop_null_keys)

narwhals/_pandas_like/group_by.py

@@ -2,6 +2,7 @@
 
 import collections
 import warnings
+from itertools import chain
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import ClassVar
@@ -40,20 +41,30 @@ class PandasLikeGroupBy(EagerGroupBy["PandasLikeDataFrame", "PandasLikeExpr"]):
 
     def __init__(
         self: Self,
-        df: PandasLikeDataFrame,
-        keys: Sequence[str],
+        compliant_frame: PandasLikeDataFrame,
+        keys: Sequence[PandasLikeExpr],
         /,
         *,
         drop_null_keys: bool,
     ) -> None:
-        self._compliant_frame = df
-        self._keys: list[str] = list(keys)
+        self._compliant_frame = compliant_frame.with_columns(*keys)
+        self._keys: list[str] = list(
+            chain.from_iterable(
+                evaluate_output_names_and_aliases(
+                    expr=key, df=compliant_frame, exclude=[]
+                )[1]
+                for key in keys
+            )
+        )
         # Drop index to avoid potential collisions:
         # https://github.com/narwhals-dev/narwhals/issues/1907.
-        if set(df.native.index.names).intersection(df.columns):
-            native_frame = df.native.reset_index(drop=True)
+
+        if set(self._compliant_frame.native.index.names).intersection(
+            self._compliant_frame.columns
+        ):
+            native_frame = self._compliant_frame.native.reset_index(drop=True)
         else:
-            native_frame = df.native
+            native_frame = self._compliant_frame.native
         if (
             self.compliant._implementation.is_pandas()
             and self.compliant._backend_version < (1, 1)

narwhals/_polars/dataframe.py

@@ -34,6 +34,7 @@
     from typing_extensions import Self
     from typing_extensions import TypeAlias
 
+    from narwhals._polars.expr import PolarsExpr
     from narwhals._polars.group_by import PolarsGroupBy
     from narwhals._polars.group_by import PolarsLazyGroupBy
     from narwhals._polars.series import PolarsSeries
@@ -396,7 +397,7 @@ def to_dict(
         else:
             return self.native.to_dict(as_series=False)
 
-    def group_by(self: Self, *keys: str, drop_null_keys: bool) -> PolarsGroupBy:
+    def group_by(self: Self, *keys: PolarsExpr, drop_null_keys: bool) -> PolarsGroupBy:
         from narwhals._polars.group_by import PolarsGroupBy
 
         return PolarsGroupBy(self, keys, drop_null_keys=drop_null_keys)
@@ -641,7 +642,9 @@ def collect(
         msg = f"Unsupported `backend` value: {backend}"  # pragma: no cover
         raise ValueError(msg)  # pragma: no cover
 
-    def group_by(self: Self, *keys: str, drop_null_keys: bool) -> PolarsLazyGroupBy:
+    def group_by(
+        self: Self, *keys: PolarsExpr, drop_null_keys: bool
+    ) -> PolarsLazyGroupBy:
         from narwhals._polars.group_by import PolarsLazyGroupBy
 
         return PolarsLazyGroupBy(self, keys, drop_null_keys=drop_null_keys)

narwhals/_polars/group_by.py

@@ -8,6 +8,7 @@
 from narwhals._polars.utils import extract_native
 
 if TYPE_CHECKING:
+    from polars._typing import IntoExpr
     from polars.dataframe.group_by import GroupBy as NativeGroupBy
     from polars.lazyframe.group_by import LazyGroupBy as NativeLazyGroupBy
     from typing_extensions import Self
@@ -19,23 +20,36 @@
 
 class PolarsGroupBy:
     _compliant_frame: PolarsDataFrame
-    _keys: Sequence[str]
+    _keys: Sequence[IntoExpr]
 
     @property
     def compliant(self) -> PolarsDataFrame:
         return self._compliant_frame
 
     def __init__(
-        self, df: PolarsDataFrame, keys: Sequence[str], /, *, drop_null_keys: bool
+        self,
+        compliant_frame: PolarsDataFrame,
+        keys: Sequence[PolarsExpr],
+        /,
+        *,
+        drop_null_keys: bool,
     ) -> None:
-        self._compliant_frame = df
-        self._keys = list(keys)
-        df = df.drop_nulls(keys) if drop_null_keys else df
-        self._grouped: NativeGroupBy = df._native_frame.group_by(keys)
+        if not drop_null_keys:
+            self._compliant_frame = compliant_frame
+            self._keys = [extract_native(arg) for arg in keys]
+        else:
+            compliant_frame = compliant_frame.with_columns(*keys)
+            self._keys = [extract_native(arg).meta.output_name() for arg in keys]
+            self._compliant_frame = compliant_frame.drop_nulls(self._keys)
+
+        self._grouped: NativeGroupBy = self._compliant_frame._native_frame.group_by(
+            self._keys
+        )
 
     def agg(self: Self, *aggs: PolarsExpr) -> PolarsDataFrame:
-        from_native = self.compliant._with_native
-        return from_native(self._grouped.agg(extract_native(arg) for arg in aggs))
+        return self.compliant._with_native(
+            self._grouped.agg(extract_native(arg) for arg in aggs)
+        )
 
     def __iter__(self: Self) -> Iterator[tuple[tuple[str, ...], PolarsDataFrame]]:
         for key, df in self._grouped:
@@ -44,20 +58,33 @@ def __iter__(self: Self) -> Iterator[tuple[tuple[str, ...], PolarsDataFrame]]:
 
 class PolarsLazyGroupBy:
     _compliant_frame: PolarsLazyFrame
-    _keys: Sequence[str]
+    _keys: Sequence[IntoExpr]
 
     @property
     def compliant(self) -> PolarsLazyFrame:
         return self._compliant_frame
 
     def __init__(
-        self, df: PolarsLazyFrame, keys: Sequence[str], /, *, drop_null_keys: bool
+        self,
+        compliant_frame: PolarsLazyFrame,
+        keys: Sequence[PolarsExpr],
+        /,
+        *,
+        drop_null_keys: bool,
     ) -> None:
-        self._compliant_frame = df
-        self._keys = list(keys)
-        df = df.drop_nulls(keys) if drop_null_keys else df
-        self._grouped: NativeLazyGroupBy = df._native_frame.group_by(keys)
+        if not drop_null_keys:
+            self._compliant_frame = compliant_frame
+            self._keys = [extract_native(arg) for arg in keys]
+        else:
+            compliant_frame = compliant_frame.with_columns(*keys)
+            self._keys = [extract_native(arg).meta.output_name() for arg in keys]
+            self._compliant_frame = compliant_frame.drop_nulls(self._keys)
+
+        self._grouped: NativeLazyGroupBy = self._compliant_frame._native_frame.group_by(
+            self._keys
+        )
 
     def agg(self: Self, *aggs: PolarsExpr) -> PolarsLazyFrame:
-        from_native = self.compliant._with_native
-        return from_native(self._grouped.agg(extract_native(arg) for arg in aggs))
+        return self.compliant._with_native(
+            self._grouped.agg(extract_native(arg) for arg in aggs)
+        )

narwhals/_spark_like/dataframe.py

@@ -271,7 +271,9 @@ def drop(self: Self, columns: Sequence[str], *, strict: bool) -> Self:
     def head(self: Self, n: int) -> Self:
         return self._with_native(self.native.limit(num=n))
 
-    def group_by(self: Self, *keys: str, drop_null_keys: bool) -> SparkLikeLazyGroupBy:
+    def group_by(
+        self: Self, *keys: SparkLikeExpr, drop_null_keys: bool
+    ) -> SparkLikeLazyGroupBy:
         from narwhals._spark_like.group_by import SparkLikeLazyGroupBy
 
         return SparkLikeLazyGroupBy(self, keys, drop_null_keys=drop_null_keys)