Repository: narwhals-dev/narwhals
Pull request #122 (2024-05-09 12:14:41)

allow iteration over groups


Patch:
diff --git a/narwhals/_pandas_like/group_by.py b/narwhals/_pandas_like/group_by.py
index 8f9917f6..54c63311 100644
--- a/narwhals/_pandas_like/group_by.py
+++ b/narwhals/_pandas_like/group_by.py
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
 from typing import Iterable
+from typing import Iterator
 
 from narwhals._pandas_like.utils import is_simple_aggregation
 from narwhals._pandas_like.utils import item
@@ -24,23 +25,22 @@ class PandasGroupBy:
     def __init__(self, df: PandasDataFrame, keys: list[str]) -> None:
         self._df = df
         self._keys = list(keys)
+        self._grouped = self._df._dataframe.groupby(
+            list(self._keys),
+            sort=False,
+            as_index=True,
+        )
 
     def agg(
         self,
         *aggs: IntoPandasExpr | Iterable[IntoPandasExpr],
         **named_aggs: IntoPandasExpr,
     ) -> PandasDataFrame:
-        df = self._df._dataframe
         exprs = parse_into_exprs(
             self._df._implementation,
             *aggs,
             **named_aggs,
         )
-        grouped = df.groupby(
-            list(self._keys),
-            sort=False,
-            as_index=True,
-        )
         implementation: str = self._df._implementation
         output_names: list[str] = copy(self._keys)
         for expr in exprs:
@@ -54,7 +54,7 @@ class PandasGroupBy:
             output_names.extend(expr._output_names)
 
         return agg_pandas(
-            grouped,
+            self._grouped,
             exprs,
             self._keys,
             output_names,
@@ -70,6 +70,9 @@ class PandasGroupBy:
             implementation=self._df._implementation,
         )
 
+    def __iter__(self) -> Iterator[tuple[Any, PandasDataFrame]]:
+        return self._grouped.__iter__()  # type: ignore[no-any-return]
+
 
 def agg_pandas(  # noqa: PLR0913
     grouped: Any,
diff --git a/narwhals/group_by.py b/narwhals/group_by.py
index 91d290b9..5cf776ac 100644
--- a/narwhals/group_by.py
+++ b/narwhals/group_by.py
@@ -1,7 +1,12 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
+from typing import Any
 from typing import Iterable
+from typing import Iterator
+
+from narwhals.utils import flatten
+from narwhals.utils import tupleify
 
 if TYPE_CHECKING:
     from narwhals.dataframe import DataFrame
@@ -12,14 +17,23 @@ if TYPE_CHECKING:
 class GroupBy:
     def __init__(self, df: DataFrame, *keys: str | Iterable[str]) -> None:
         self._df = df
-        self._keys = keys
+        self._keys = flatten(keys)
+        self._grouped = self._df._dataframe.group_by(self._keys)
 
     def agg(
         self, *aggs: IntoExpr | Iterable[IntoExpr], **named_aggs: IntoExpr
     ) -> DataFrame:
         aggs, named_aggs = self._df._flatten_and_extract(*aggs, **named_aggs)
         return self._df.__class__(
-            self._df._dataframe.group_by(*self._keys).agg(*aggs, **named_aggs),
+            self._grouped.agg(*aggs, **named_aggs),
+        )
+
+    def __iter__(self) -> Iterator[tuple[Any, DataFrame]]:
+        import narwhals as nw
+
+        yield from (
+            (tupleify(key), nw.from_native(df, eager_only=True))
+            for (key, df) in self._grouped.__iter__()
         )
 
 
@@ -27,11 +41,12 @@ class LazyGroupBy:
     def __init__(self, df: LazyFrame, *keys: str | Iterable[str]) -> None:
         self._df = df
         self._keys = keys
+        self._grouped = self._df._dataframe.group_by(*self._keys)
 
     def agg(
         self, *aggs: IntoExpr | Iterable[IntoExpr], **named_aggs: IntoExpr
     ) -> LazyFrame:
         aggs, named_aggs = self._df._flatten_and_extract(*aggs, **named_aggs)
         return self._df.__class__(
-            self._df._dataframe.group_by(*self._keys).agg(*aggs, **named_aggs),
+            self._grouped.agg(*aggs, **named_aggs),
         )
diff --git a/narwhals/utils.py b/narwhals/utils.py
index c21ef40b..cf8afaf2 100644
--- a/narwhals/utils.py
+++ b/narwhals/utils.py
@@ -4,10 +4,13 @@ import re
 from typing import Any
 from typing import Iterable
 from typing import Sequence
+from typing import TypeVar
 
 from narwhals.dependencies import get_pandas
 from narwhals.dependencies import get_polars
 
+T = TypeVar("T")
+
 
 def remove_prefix(text: str, prefix: str) -> str:
     if text.startswith(prefix):
@@ -29,6 +32,12 @@ def flatten(args: Any) -> list[Any]:
     return args  # type: ignore[no-any-return]
 
 
+def tupleify(arg: Any) -> Any:
+    if not isinstance(arg, (list, tuple)):  # pragma: no cover
+        return (arg,)
+    return arg
+
+
 def _is_iterable(arg: Any | Iterable[Any]) -> bool:
     from narwhals.series import Series
 


Test patch:
diff --git a/tests/group_by_test.py b/tests/group_by_test.py
index e69de29b..ec0b85ef 100644
--- a/tests/group_by_test.py
+++ b/tests/group_by_test.py
@@ -0,0 +1,37 @@
+from typing import Any
+
+import pandas as pd
+import polars as pl
+import pytest
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+data = {
+    "a": [1, 1, 3],
+    "b": [4, 4, 6],
+    "c": [7, 8, 9],
+}
+
+
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+def test_group_by_iter(constructor: Any) -> None:
+    df = nw.from_native(constructor(data), eager_only=True)
+    expected_keys = [(1,), (3,)]
+    keys = []
+    for key, sub_df in df.group_by("a"):
+        if key == (1,):
+            expected = {"a": [1, 1], "b": [4, 4], "c": [7, 8]}
+            compare_dicts(sub_df, expected)
+            assert isinstance(sub_df, nw.DataFrame)
+        keys.append(key)
+    assert sorted(keys) == sorted(expected_keys)
+    expected_keys = [(1, 4), (3, 6)]  # type: ignore[list-item]
+    keys = []
+    for key, _df in df.group_by("a", "b"):
+        keys.append(key)
+    assert sorted(keys) == sorted(expected_keys)
+    keys = []
+    for key, _df in df.group_by(["a", "b"]):
+        keys.append(key)
+    assert sorted(keys) == sorted(expected_keys)