feat: Add maintain_order to DataFrame.group_by

src/narwhals/_arrow/dataframe.py

@@ -408,11 +408,17 @@ def with_columns(self, *exprs: ArrowExpr) -> Self:
         return self._with_native(native_frame, validate_column_names=False)
 
     def group_by(
-        self, keys: Sequence[str] | Sequence[ArrowExpr], *, drop_null_keys: bool
+        self,
+        keys: Sequence[str] | Sequence[ArrowExpr],
+        *,
+        drop_null_keys: bool,
+        maintain_order: bool = False,
     ) -> ArrowGroupBy:
         from narwhals._arrow.group_by import ArrowGroupBy
 
-        return ArrowGroupBy(self, keys, drop_null_keys=drop_null_keys)
+        return ArrowGroupBy(
+            self, keys, drop_null_keys=drop_null_keys, maintain_order=maintain_order
+        )
 
     def join(
         self,

src/narwhals/_arrow/group_by.py

@@ -67,11 +67,17 @@ def __init__(
         /,
         *,
         drop_null_keys: bool,
+        maintain_order: bool = False,
     ) -> None:
         self._df = df
         frame, self._keys, self._output_key_names = self._parse_keys(df, keys=keys)
         self._compliant_frame = frame.drop_nulls(self._keys) if drop_null_keys else frame
         self._drop_null_keys = drop_null_keys
+        self._row_index = (
+            generate_temporary_column_name(n_bytes=8, columns=frame.columns)
+            if maintain_order
+            else None
+        )
 
     def _configure_agg(
         self, expr: ArrowExpr, /
@@ -94,6 +100,9 @@ def _configure_agg(
         return self._remap_expr_name(function_name), option
 
     def _configure_grouped(self, *exprs: ArrowExpr) -> pa.TableGroupBy:
+        frame = self.compliant
+        if self._row_index is not None:
+            frame = frame.with_row_index(self._row_index, order_by=None)
         order_by = ()
         use_threads = True
         for expr in exprs:
@@ -116,14 +125,14 @@ def _configure_grouped(self, *exprs: ArrowExpr) -> pa.TableGroupBy:
             raise NotImplementedError(msg)
         if order_by:
             return pa.TableGroupBy(
-                self.compliant.sort(*order_by, descending=False, nulls_last=False).native,
+                frame.sort(*order_by, descending=False, nulls_last=False).native,
                 self._keys,
                 use_threads=use_threads,
             )
         if not use_threads:
-            return pa.TableGroupBy(self.compliant.native, self._keys, use_threads=False)
+            return pa.TableGroupBy(frame.native, self._keys, use_threads=False)
         # TODO(unassigned): combine with `return` above once PyArrow 15 is the minimum.
-        return pa.TableGroupBy(self.compliant.native, self._keys)
+        return pa.TableGroupBy(frame.native, self._keys)
 
     def agg(self, *exprs: ArrowExpr) -> ArrowDataFrame:
         self._ensure_all_simple(exprs)
@@ -160,6 +169,12 @@ def agg(self, *exprs: ArrowExpr) -> ArrowDataFrame:
                 [(output_name, function_name, option) for output_name in output_names]
             )
 
+        if self._row_index is not None:
+            # The smallest row index in each group is where that group first appears.
+            new_column_names.append(self._row_index)
+            expected_pyarrow_column_names.append(f"{self._row_index}_min")
+            aggs.append((self._row_index, "min", None))
+
         result_simple = grouped.aggregate(aggs)
 
         # Rename columns, being very careful
@@ -181,6 +196,8 @@ def agg(self, *exprs: ArrowExpr) -> ArrowDataFrame:
         ]
         new_column_names = [new_column_names[i] for i in index_map]
         result_simple = result_simple.rename_columns(new_column_names)
+        if self._row_index is not None:
+            result_simple = result_simple.sort_by(self._row_index).drop([self._row_index])
         return self.compliant._with_native(result_simple).rename(
             dict(zip(self._keys, self._output_key_names, strict=False))
         )

src/narwhals/_compliant/dataframe.py

@@ -241,6 +241,7 @@ def group_by(
         keys: Sequence[str] | Sequence[CompliantExprT_contra],
         *,
         drop_null_keys: bool,
+        maintain_order: bool = False,
     ) -> DataFrameGroupBy[Self, Any]: ...
     def item(self, row: int | None, column: int | str | None) -> Any: ...
     def iter_columns(self) -> Iterator[CompliantSeriesT]: ...

src/narwhals/_pandas_like/dataframe.py

@@ -601,11 +601,17 @@ def collect(
 
     # --- actions ---
     def group_by(
-        self, keys: Sequence[str] | Sequence[PandasLikeExpr], *, drop_null_keys: bool
+        self,
+        keys: Sequence[str] | Sequence[PandasLikeExpr],
+        *,
+        drop_null_keys: bool,
+        maintain_order: bool = False,
     ) -> PandasLikeGroupBy:
         from narwhals._pandas_like.group_by import PandasLikeGroupBy
 
-        return PandasLikeGroupBy(self, keys, drop_null_keys=drop_null_keys)
+        return PandasLikeGroupBy(
+            self, keys, drop_null_keys=drop_null_keys, maintain_order=maintain_order
+        )
 
     def _join_inner(
         self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str], suffix: str

src/narwhals/_pandas_like/group_by.py

@@ -10,6 +10,7 @@
 from narwhals._exceptions import issue_warning
 from narwhals._expression_parsing import evaluate_output_names_and_aliases
 from narwhals._pandas_like.utils import make_group_by_kwargs
+from narwhals._utils import generate_temporary_column_name
 from narwhals.dependencies import is_pandas_like_dataframe
 
 if TYPE_CHECKING:
@@ -262,6 +263,7 @@ def __init__(
         /,
         *,
         drop_null_keys: bool,
+        maintain_order: bool = False,
     ) -> None:
         self._original_columns = tuple(df.columns)
         self._drop_null_keys = drop_null_keys
@@ -270,14 +272,19 @@ def __init__(
         )
         self._exclude: tuple[str, ...] = (*self._keys, *self._output_key_names)
         self._group_by_kwargs = make_group_by_kwargs(drop_null_keys=drop_null_keys)
+        self._row_index = (
+            generate_temporary_column_name(n_bytes=8, columns=self.compliant.columns)
+            if maintain_order
+            else None
+        )
 
         # Drop index to avoid potential collisions:
         # https://github.com/narwhals-dev/narwhals/issues/1907.
         self._native = self.compliant.native
         if set(self._native.index.names).intersection(self.compliant.columns):
             self._native = self._native.reset_index(drop=True)
 
-    def agg(self, *exprs: PandasLikeExpr) -> PandasLikeDataFrame:  # noqa: PLR0912
+    def agg(self, *exprs: PandasLikeExpr) -> PandasLikeDataFrame:
         all_aggs_are_simple = True
         agg_exprs: list[AggExpr] = []
         order_by = ()
@@ -292,13 +299,7 @@ def agg(self, *exprs: PandasLikeExpr) -> PandasLikeDataFrame:  # noqa: PLR0912
                     raise NotImplementedError(msg)
                 order_by = _current_order_by
 
-        if order_by:
-            grouped: NativeGroupBy = self._native.sort_values(
-                list(order_by), na_position="first"
-            ).groupby(self._keys.copy(), **self._group_by_kwargs)
-        else:
-            grouped = self._native.groupby(self._keys.copy(), **self._group_by_kwargs)
-        self._grouped = grouped
+        grouped = self._grouped = self._configure_grouped(order_by)
 
         if all_aggs_are_simple:
             result: pd.DataFrame
@@ -314,6 +315,9 @@ def agg(self, *exprs: PandasLikeExpr) -> PandasLikeDataFrame:  # noqa: PLR0912
         else:
             result = self._apply_aggs(grouped, exprs)
 
+        if order_by and self._row_index is not None:
+            result = self._restore_order(result, grouped, self._row_index)
+
         impl = self.compliant._implementation
         backend_version = impl._backend_version()
         if impl.is_pandas() and backend_version < (3, 0):  # pragma: no cover
@@ -324,6 +328,32 @@ def agg(self, *exprs: PandasLikeExpr) -> PandasLikeDataFrame:  # noqa: PLR0912
 
         return self._select_results(result, agg_exprs)
 
+    def _configure_grouped(self, order_by: Sequence[str]) -> NativeGroupBy:
+        if not order_by:
+            return self._native.groupby(self._keys.copy(), **self._group_by_kwargs)
+        native = self._native
+        if self._row_index is not None:
+            # With `sort=False`, pandas returns groups in order of first appearance,
+            # but sorting by `order_by` changes which rows appear first.
+            native = (
+                self.compliant._with_native(native, validate_column_names=False)
+                .with_row_index(self._row_index, order_by=None)
+                .native
+            )
+        return native.sort_values(list(order_by), na_position="first").groupby(
+            self._keys.copy(), **self._group_by_kwargs
+        )
+
+    def _restore_order(
+        self, result: pd.DataFrame, grouped: NativeGroupBy, row_index: str
+    ) -> pd.DataFrame:
+        # Align on the group keys rather than position: e.g. `last` returns its
+        # rows in order of last appearance.
+        ns = self.compliant.__narwhals_namespace__()
+        first_seen = grouped[row_index].min()
+        columns = list(result.columns)
+        return ns._concat_by_index([result, first_seen]).sort_values(row_index)[columns]
+
     def _select_results(
         self, df: pd.DataFrame, /, agg_exprs: Sequence[AggExpr]
     ) -> PandasLikeDataFrame:

src/narwhals/_polars/dataframe.py

@@ -630,11 +630,17 @@ def to_dict(
         return self.native.to_dict(as_series=False)
 
     def group_by(
-        self, keys: Sequence[str] | Sequence[PolarsExpr], *, drop_null_keys: bool
+        self,
+        keys: Sequence[str] | Sequence[PolarsExpr],
+        *,
+        drop_null_keys: bool,
+        maintain_order: bool = False,
     ) -> PolarsGroupBy:
         from narwhals._polars.group_by import PolarsGroupBy
 
-        return PolarsGroupBy(self, keys, drop_null_keys=drop_null_keys)
+        return PolarsGroupBy(
+            self, keys, drop_null_keys=drop_null_keys, maintain_order=maintain_order
+        )
 
     def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
         to_drop = parse_columns_to_drop(self, columns, strict=strict)

src/narwhals/_polars/group_by.py

@@ -29,13 +29,16 @@ def __init__(
         /,
         *,
         drop_null_keys: bool,
+        maintain_order: bool = False,
     ) -> None:
         self._keys = list(keys)
         self._compliant_frame = df.drop_nulls(keys) if drop_null_keys else df
         self._grouped = (
-            self.compliant.native.group_by(keys)
+            self.compliant.native.group_by(keys, maintain_order=maintain_order)
             if is_sequence_of(keys, str)
-            else self.compliant.native.group_by(arg.native for arg in keys)
+            else self.compliant.native.group_by(
+                (arg.native for arg in keys), maintain_order=maintain_order
+            )
         )
 
     def agg(self, *aggs: PolarsExpr) -> PolarsDataFrame:

src/narwhals/dataframe.py

@@ -1792,16 +1792,25 @@ def filter(
 
     @overload
     def group_by(
-        self, *keys: IntoExpr | Iterable[IntoExpr], drop_null_keys: Literal[False] = ...
+        self,
+        *keys: IntoExpr | Iterable[IntoExpr],
+        drop_null_keys: Literal[False] = ...,
+        maintain_order: bool = ...,
     ) -> GroupBy[Self]: ...
 
     @overload
     def group_by(
-        self, *keys: str | Iterable[str], drop_null_keys: Literal[True]
+        self,
+        *keys: str | Iterable[str],
+        drop_null_keys: Literal[True],
+        maintain_order: bool = ...,
     ) -> GroupBy[Self]: ...
 
     def group_by(
-        self, *keys: IntoExpr | Iterable[IntoExpr], drop_null_keys: bool = False
+        self,
+        *keys: IntoExpr | Iterable[IntoExpr],
+        drop_null_keys: bool = False,
+        maintain_order: bool = False,
     ) -> GroupBy[Self]:
         r"""Start a group by operation.
 
@@ -1810,6 +1819,9 @@ def group_by(
                 column names.
             drop_null_keys: if True, then groups where any key is null won't be included
                 in the result.
+            maintain_order: if True, groups are returned in the order in which they
+                first appear in the data, both by `agg` and when iterating over the
+                groups. If False (default), no order is guaranteed.
 
         Examples:
             >>> import pandas as pd
@@ -1856,13 +1868,28 @@ def group_by(
             0  a  0  4.0
             1  b  1  3.0
             2  c  1  1.0
+
+            Use `maintain_order=True` to get the groups in order of first appearance.
+
+            >>> nw.from_native(df_native, eager_only=True).group_by(
+            ...     "a", maintain_order=True
+            ... ).agg(nw.col("c").sum()).to_native()
+               a  c
+            0  a  8
+            1  b  6
+            2  c  1
         """
         from narwhals.group_by import GroupBy
 
         flat_keys = flatten(keys)
 
         if all(isinstance(key, str) for key in flat_keys):
-            return GroupBy(self, flat_keys, drop_null_keys=drop_null_keys)
+            return GroupBy(
+                self,
+                flat_keys,
+                drop_null_keys=drop_null_keys,
+                maintain_order=maintain_order,
+            )
 
         from narwhals import col
         from narwhals.expr import Expr
@@ -1882,7 +1909,12 @@ def group_by(
         check_expressions_preserve_length(
             *expr_flat_keys, function_name="DataFrame.group_by"
         )
-        return GroupBy(self, expr_flat_keys, drop_null_keys=drop_null_keys)
+        return GroupBy(
+            self,
+            expr_flat_keys,
+            drop_null_keys=drop_null_keys,
+            maintain_order=maintain_order,
+        )
 
     def sort(
         self,

src/narwhals/group_by.py

@@ -25,11 +25,12 @@ def __init__(
         /,
         *,
         drop_null_keys: bool,
+        maintain_order: bool = False,
     ) -> None:
         self._df: DataFrameT = df
         self._keys = keys
         self._grouped = self._df._compliant_frame.group_by(
-            self._keys, drop_null_keys=drop_null_keys
+            self._keys, drop_null_keys=drop_null_keys, maintain_order=maintain_order
         )
 
     def agg(self, *aggs: Expr | Iterable[Expr], **named_aggs: Expr) -> DataFrameT: