Move the change into value_counts.

dask · Feb 10, 2023 · dfc4c4b · dfc4c4b
1 parent 19eab76
commit dfc4c4b
Showing 1 changed file with 9 additions and 12 deletions.
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
@@ -1543,16 +1543,10 @@ def _single_agg(
             with check_numeric_only_deprecation():
                 meta = func(self._meta_nonempty, **chunk_kwargs)
 
-        if is_series_like(meta):
-            # in pandas 2.0, Series returned from value_counts have a name
-            # different from original object, but here, column name should
-            # still reflect the original object name
-            if func is _value_counts:
-                columns = self._meta.apply(pd.Series).name
-            else:
-                columns = meta.name
-        else:
-            columns = meta.columns
+        if "columns" not in chunk_kwargs:
+            chunk_kwargs["columns"] = (
+                meta.name if is_series_like(meta) else meta.columns
+            )
 
         args = [self.obj] + (self.by if isinstance(self.by, list) else [self.by])
 
@@ -1565,7 +1559,6 @@ def _single_agg(
                 chunk=_apply_chunk,
                 chunk_kwargs={
                     "chunk": func,
-                    "columns": columns,
                     **self.observed,
                     **self.dropna,
                     **chunk_kwargs,
@@ -1590,7 +1583,6 @@ def _single_agg(
             chunk=_apply_chunk,
             chunk_kwargs=dict(
                 chunk=func,
-                columns=columns,
                 **self.observed,
                 **chunk_kwargs,
                 **self.dropna,
@@ -2957,13 +2949,18 @@ def agg(self, arg=None, split_every=None, split_out=1, shuffle=None, **kwargs):
 
     @derived_from(pd.core.groupby.SeriesGroupBy)
     def value_counts(self, split_every=None, split_out=1, shuffle=None):
+        # in pandas 2.0, Series returned from value_counts have a name
+        # different from original object, but here, column name should
+        # still reflect the original object name
+        chunk_kwargs = {"columns": self._meta.apply(pd.Series).name}
         return self._single_agg(
             func=_value_counts,
             token="value_counts",
             aggfunc=_value_counts_aggregate,
             split_every=split_every,
             split_out=split_out,
             shuffle=shuffle,
+            chunk_kwargs=chunk_kwargs,
         )
 
     @derived_from(pd.core.groupby.SeriesGroupBy)