refactor: Simplify ArrowGroupBy.__iter__ (#2133)

Author: dangotbanned

narwhals/_arrow/group_by.py

@@ -10,6 +10,7 @@
 import pyarrow as pa
 import pyarrow.compute as pc
 
+from narwhals._arrow.utils import extract_py_scalar
 from narwhals._expression_parsing import evaluate_output_names_and_aliases
 from narwhals._expression_parsing import is_simple_aggregation
 from narwhals.utils import generate_temporary_column_name
@@ -156,20 +157,9 @@ def __iter__(self: Self) -> Iterator[tuple[Any, ArrowDataFrame]]:
             *it, "", null_handling="replace", null_replacement=null_token
         )
         table = table.add_column(i=0, field_=col_token, column=key_values)
-
-        yield from (
-            (
-                next(
-                    (
-                        t := self._df._from_native_frame(
-                            table.filter(pc.equal(table[col_token], v)).drop([col_token])
-                        )
-                    )
-                    .simple_select(*self._keys)
-                    .head(1)
-                    .iter_rows(named=False, buffer_size=512)
-                ),
-                t,
+        for v in pc.unique(key_values):
+            t = self._df._from_native_frame(
+                table.filter(pc.equal(table[col_token], v)).drop([col_token])
             )
-            for v in pc.unique(key_values)
-        )
+            row = t.simple_select(*self._keys).row(0)
+            yield tuple(extract_py_scalar(el) for el in row), t

narwhals/_arrow/series.py

@@ -56,6 +56,8 @@
     from narwhals.utils import Version
 
 
+# TODO @dangotbanned: move into `_arrow.utils`
+# Lots of modules are importing inline
 @overload
 def maybe_extract_py_scalar(
     value: pa.Scalar[_BasicDataType[_AsPyType]],

narwhals/_arrow/utils.py

@@ -65,6 +65,12 @@ def extract_regex(
 """Alias for `pyarrow.scalar`."""
 
 
+def extract_py_scalar(value: Any, /) -> Any:
+    from narwhals._arrow.series import maybe_extract_py_scalar
+
+    return maybe_extract_py_scalar(value, return_py_scalar=True)
+
+
 def chunked_array(
     arr: ArrowArray | list[Iterable[pa.Scalar[Any]]] | ArrowChunkedArray,
 ) -> ArrowChunkedArray: