Skip to content

Commit

Permalink
fix: support groupby.iter for cudf (#1265)
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli authored Oct 28, 2024
1 parent 33c9527 commit 425dbe4
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 14 deletions.
17 changes: 9 additions & 8 deletions narwhals/_pandas_like/group_by.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,15 +103,16 @@ def _from_native_frame(self, df: PandasLikeDataFrame) -> PandasLikeDataFrame:

def __iter__(self) -> Iterator[tuple[Any, PandasLikeDataFrame]]:
indices = self._grouped.indices
for key in indices:
if (
self._df._implementation is Implementation.PANDAS
and self._df._backend_version < (2, 2)
): # pragma: no cover
pass
else: # pragma: no cover
if (
self._df._implementation is Implementation.PANDAS
and self._df._backend_version < (2, 2)
) or (self._df._implementation is Implementation.CUDF): # pragma: no cover
for key in indices:
yield (key, self._from_native_frame(self._grouped.get_group(key)))
else:
for key in indices:
key = tupleify(key) # noqa: PLW2901
yield (key, self._from_native_frame(self._grouped.get_group(key)))
yield (key, self._from_native_frame(self._grouped.get_group(key)))


def agg_pandas( # noqa: PLR0915
Expand Down
7 changes: 1 addition & 6 deletions tests/group_by_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,12 +74,7 @@ def test_invalid_group_by() -> None:
)


def test_group_by_iter(
constructor_eager: ConstructorEager, request: pytest.FixtureRequest
) -> None:
if "cudf" in str(constructor_eager):
# https://github.com/rapidsai/cudf/issues/17187
request.applymarker(pytest.mark.xfail)
def test_group_by_iter(constructor_eager: ConstructorEager) -> None:
df = nw.from_native(constructor_eager(data), eager_only=True)
expected_keys = [(1,), (3,)]
keys = []
Expand Down

0 comments on commit 425dbe4

Please sign in to comment.