From a682551d479aec1141af9ad9908330bda0bb6919 Mon Sep 17 00:00:00 2001 From: Stewyn Chaudhary Date: Tue, 6 Oct 2026 23:06:46 -0500 Subject: [PATCH] feat: Add `maintain_order` to `DataFrame.group_by` --- src/narwhals/_arrow/dataframe.py | 10 +++- src/narwhals/_arrow/group_by.py | 23 +++++++-- src/narwhals/_compliant/dataframe.py | 1 + src/narwhals/_pandas_like/dataframe.py | 10 +++- src/narwhals/_pandas_like/group_by.py | 46 ++++++++++++++---- src/narwhals/_polars/dataframe.py | 10 +++- src/narwhals/_polars/group_by.py | 7 ++- src/narwhals/dataframe.py | 42 +++++++++++++++-- src/narwhals/group_by.py | 3 +- tests/frame/group_by_test.py | 64 ++++++++++++++++++++++++++ 10 files changed, 191 insertions(+), 25 deletions(-) diff --git a/src/narwhals/_arrow/dataframe.py b/src/narwhals/_arrow/dataframe.py index cd0b643f7a..bd52c3eed6 100644 --- a/src/narwhals/_arrow/dataframe.py +++ b/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, diff --git a/src/narwhals/_arrow/group_by.py b/src/narwhals/_arrow/group_by.py index f9f5b18dce..eb0ac4f9ed 100644 --- a/src/narwhals/_arrow/group_by.py +++ b/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)) ) diff --git a/src/narwhals/_compliant/dataframe.py b/src/narwhals/_compliant/dataframe.py index b1d28620f1..0cc62835c6 100644 --- a/src/narwhals/_compliant/dataframe.py +++ b/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]: ... diff --git a/src/narwhals/_pandas_like/dataframe.py b/src/narwhals/_pandas_like/dataframe.py index 0844eea431..a1034da4f1 100644 --- a/src/narwhals/_pandas_like/dataframe.py +++ b/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 diff --git a/src/narwhals/_pandas_like/group_by.py b/src/narwhals/_pandas_like/group_by.py index 0de26c8ee8..6f5fcbfd1d 100644 --- a/src/narwhals/_pandas_like/group_by.py +++ b/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,6 +272,11 @@ 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. @@ -277,7 +284,7 @@ def __init__( 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: diff --git a/src/narwhals/_polars/dataframe.py b/src/narwhals/_polars/dataframe.py index ee27771603..96e551fdc0 100644 --- a/src/narwhals/_polars/dataframe.py +++ b/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) diff --git a/src/narwhals/_polars/group_by.py b/src/narwhals/_polars/group_by.py index e0722565bb..2aab01c877 100644 --- a/src/narwhals/_polars/group_by.py +++ b/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: diff --git a/src/narwhals/dataframe.py b/src/narwhals/dataframe.py index 78abb282eb..922da0dc96 100644 --- a/src/narwhals/dataframe.py +++ b/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, diff --git a/src/narwhals/group_by.py b/src/narwhals/group_by.py index 6b469637d5..325fb75300 100644 --- a/src/narwhals/group_by.py +++ b/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: diff --git a/tests/frame/group_by_test.py b/tests/frame/group_by_test.py index 2c03a1f890..ba59a780b3 100644 --- a/tests/frame/group_by_test.py +++ b/tests/frame/group_by_test.py @@ -806,3 +806,67 @@ def test_multi_column_expansion(constructor: Constructor) -> None: .sort("a", descending=True) ) assert_equal_data(result, expected) + + +@pytest.mark.parametrize( + "agg", + [ + nw.col("b").sum(), + nw.len(), + nw.all().mean(), + nw.col("b").first(order_by="c"), + nw.all().last(order_by="c"), + ], + ids=["sum", "len", "multi-output", "first-order-by", "last-order-by"], +) +def test_group_by_maintain_order( + constructor_eager: ConstructorEager, agg: nw.Expr, request: pytest.FixtureRequest +) -> None: + request.applymarker( + pytest.mark.xfail( + "pyarrow_table" in str(constructor_eager) and (PYARROW_VERSION < (14, 0)), + reason="https://github.com/apache/arrow/issues/36709", + raises=NotImplementedError, + ) + ) + # Enough rows and groups for multi-threaded backends to shuffle the groups + # when the order isn't maintained. + n_rows, n_groups = 20_000, 500 + keys = [(i * 7919) % n_groups for i in range(n_rows)] + data = {"a": keys, "b": list(range(n_rows)), "c": list(range(n_rows, 0, -1))} + df = nw.from_native(constructor_eager(data), eager_only=True) + result = df.group_by("a", maintain_order=True).agg(agg) + assert result["a"].to_list() == list(dict.fromkeys(keys)) + + +def test_group_by_maintain_order_null_keys(constructor_eager: ConstructorEager) -> None: + data = {"a": ["y", None, "x", "y", None, "z"], "b": [1, 2, 3, 4, 5, 6]} + df = nw.from_native(constructor_eager(data), eager_only=True) + result = df.group_by("a", maintain_order=True).agg(nw.col("b").sum()) + expected = {"a": ["y", None, "x", "z"], "b": [5, 7, 3, 6]} + assert_equal_data(result, expected) + result = df.group_by("a", drop_null_keys=True, maintain_order=True).agg( + nw.col("b").sum() + ) + expected = {"a": ["y", "x", "z"], "b": [5, 3, 6]} + assert_equal_data(result, expected) + + +def test_group_by_maintain_order_multiple_and_expr_keys( + constructor_eager: ConstructorEager, +) -> None: + data = {"a": [3, 1, 3, 2, 1], "b": ["q", "p", "q", "p", "r"], "c": [1, 2, 3, 4, 5]} + df = nw.from_native(constructor_eager(data), eager_only=True) + result = df.group_by("a", "b", maintain_order=True).agg(nw.col("c").sum()) + expected = {"a": [3, 1, 2, 1], "b": ["q", "p", "p", "r"], "c": [4, 2, 4, 5]} + assert_equal_data(result, expected) + result = df.group_by(nw.col("a") * 10, maintain_order=True).agg(nw.col("c").max()) + expected = {"a": [30, 10, 20], "c": [3, 5, 4]} + assert_equal_data(result, expected) + + +def test_group_by_maintain_order_iter(constructor_eager: ConstructorEager) -> None: + keys = [(i * 7919) % 500 for i in range(20_000)] + df = nw.from_native(constructor_eager({"a": keys}), eager_only=True) + result = [key for key, _ in df.group_by("a", maintain_order=True)] + assert result == [(key,) for key in dict.fromkeys(keys)]