diff --git a/dask_expr/_reductions.py b/dask_expr/_reductions.py index 357b4849..ce4a0de7 100644 --- a/dask_expr/_reductions.py +++ b/dask_expr/_reductions.py @@ -244,7 +244,8 @@ def _lower(self): # Reset the index if we we used it for shuffling if split_by_index: - shuffled = SetIndexBlockwise(shuffled, split_by, True, None) + idx = list(self._meta.index.names) if split_by != ["index"] else split_by + shuffled = SetIndexBlockwise(shuffled, idx, True, None) # Convert back to Series if necessary if self.shuffle_by_index is not False: diff --git a/dask_expr/tests/test_groupby.py b/dask_expr/tests/test_groupby.py index e5b45e6d..05ca7f3c 100644 --- a/dask_expr/tests/test_groupby.py +++ b/dask_expr/tests/test_groupby.py @@ -157,6 +157,13 @@ def test_groupby_no_numeric_only(pdf, func): assert_eq(agg, expect) +def test_value_counts_split_out(pdf): + df = from_pandas(pdf, npartitions=10) + result = df.groupby("x").y.value_counts(split_out=True) + expected = pdf.groupby("x").y.value_counts() + assert_eq(result, expected) + + def test_unique(df, pdf): result = df.groupby("x")["y"].unique() expected = pdf.groupby("x")["y"].unique()