Skip to content

Commit

Permalink
Fix groupby(as_index=False).size not reseting index (#17499)
Browse files Browse the repository at this point in the history
closes #17478

Also fixes a bug where the `Series.name` attribute wasn't preserved with `size`

Authors:
  - Matthew Roeschke (https://github.com/mroeschke)
  - GALI PREM SAGAR (https://github.com/galipremsagar)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Matthew Murray (https://github.com/Matt711)
  - GALI PREM SAGAR (https://github.com/galipremsagar)

URL: #17499
  • Loading branch information
mroeschke authored Dec 4, 2024
1 parent cd3e352 commit 47e49d0
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 2 deletions.
7 changes: 5 additions & 2 deletions python/cudf/cudf/core/groupby/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -497,11 +497,14 @@ def size(self):
col = cudf.core.column.column_empty(
len(self.obj), "int8", masked=False
)
return (
cudf.Series._from_column(col)
result = (
cudf.Series._from_column(col, name=getattr(self.obj, "name", None))
.groupby(self.grouping, sort=self._sort, dropna=self._dropna)
.agg("size")
)
if not self._as_index:
result = result.rename("size").reset_index()
return result

@_performance_tracking
def cumcount(self, ascending: bool = True):
Expand Down
14 changes: 14 additions & 0 deletions python/cudf/cudf/tests/test_groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -4074,3 +4074,17 @@ def test_get_group_list_like():

with pytest.raises(KeyError):
df.groupby(["a"]).get_group([1])


def test_size_as_index_false():
df = pd.DataFrame({"a": [1, 2, 1], "b": [1, 2, 3]}, columns=["a", "b"])
expected = df.groupby("a", as_index=False).size()
result = cudf.from_pandas(df).groupby("a", as_index=False).size()
assert_eq(result, expected)


def test_size_series_with_name():
ser = pd.Series(range(3), name="foo")
expected = ser.groupby(ser).size()
result = cudf.from_pandas(ser).groupby(ser).size()
assert_eq(result, expected)

0 comments on commit 47e49d0

Please sign in to comment.