Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle as_index in GroupBy.apply #13951

Merged
merged 2 commits into from
Aug 25, 2023
+13 −0
Merged

Handle as_index in GroupBy.apply #13951

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cab0e9e
test, fix
brandon-b-miller Aug 24, 2023
2f73ae3
parametrize over engine
brandon-b-miller Aug 24, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions python/cudf/cudf/core/groupby/groupby.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1436,6 +1436,9 @@ def mult(df):

if self._sort:
result = result.sort_index()
if self._as_index is False:
result = result.reset_index()
result[None] = result.pop(0)
Comment on lines +1440 to +1441
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
result = result.reset_index()
result[None] = result.pop(0)
result = result.reset_index(drop=True)

Could we do this instead?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When I try that, I get a Series still. We're trying to end up with a DataFrame with a column named None:

0    1.0
1    2.5
dtype: float64

vs

   y  None
0  0   1.0
1  1   2.5

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, got it 👍 Then this change LGTM

return result

@_cudf_nvtx_annotate
Expand Down
10 changes: 10 additions & 0 deletions python/cudf/cudf/tests/test_groupby.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,16 @@ def test_groupby_as_index_single_agg(pdf, gdf, as_index):
assert_groupby_results_equal(pdf, gdf)


@pytest.mark.parametrize("engine", ["cudf", "jit"])
@pytest.mark.parametrize("as_index", [True, False])
def test_groupby_as_index_apply(pdf, gdf, as_index, engine):
gdf = gdf.groupby("y", as_index=as_index).apply(
lambda df: df["x"].mean(), engine=engine
)
pdf = pdf.groupby("y", as_index=as_index).apply(lambda df: df["x"].mean())
assert_groupby_results_equal(pdf, gdf)


@pytest.mark.parametrize("as_index", [True, False])
def test_groupby_as_index_multiindex(pdf, gdf, as_index):
pdf = pd.DataFrame(
Expand Down