Repository: dask/dask

Issue:
Implement shuffle-based-groupby for simpler aggregations
In #9406 @rjzamora did some nice work in implementing a shuffle-based groupby/agg operation for dask dataframe. However, it is currently only implemented for `groupby(...).agg({...})`, and some of the simpler aggregation functions which might benefit from it are un-implemented. This includes `groupby(...).min`, `groupby(...).sum and similar:

https://github.com/dask/dask/blob/024df342f77ec6d7b41227b6619461732825c7c6/dask/dataframe/groupby.py#L1239-L1250

So, the following snippet would use a shuffle-based approach:
```python
import dask.datasets

ddf = dask.datasets.timeseries()
ddf.groupby("id").agg({"x": "mean"}, split_out=2)
```

While the following semantically identical one would not:

```python
import dask.datasets

ddf = dask.datasets.timeseries()
ddf.groupby("id").x.mean(split_out=2)
```


It should be relatively straightforward to retrofit the simpler aggregate functions to also take advantage of the new shuffle-base approach where appropriate.



Discussion:

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -562,7 +562,7 @@ def index(self):
         return self.map_partitions(
             getattr,
             "index",
-            token=self._name + "-index",
+            token=key_split(self._name) + "-index",
             meta=self._meta.index,
             enforce_metadata=False,
         )
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -85,6 +85,16 @@ def _determine_levels(by):
         return 0
 
 
+def _determine_shuffle(shuffle, split_out):
+    """Determine the default shuffle behavior based on split_out"""
+    if shuffle is None:
+        if split_out > 1:
+            return shuffle or config.get("shuffle", None) or "tasks"
+        else:
+            return False
+    return shuffle
+
+
 def _normalize_by(df, by):
     """Replace series with column names wherever possible."""
     if not isinstance(df, DataFrame):
@@ -1245,7 +1255,7 @@ def _meta_nonempty(self):
         )
         return _maybe_slice(grouped, self._slice)
 
-    def _aca_agg(
+    def _single_agg(
         self,
         token,
         func,
@@ -1253,9 +1263,16 @@ def _aca_agg(
         meta=None,
         split_every=None,
         split_out=1,
+        shuffle=None,
         chunk_kwargs=None,
         aggregate_kwargs=None,
     ):
+        """
+        Aggregation with a single function/aggfunc rather than a compound spec
+        like in GroupBy.aggregate
+        """
+        shuffle = _determine_shuffle(shuffle, split_out)
+
         if self.sort is None and split_out > 1:
             warnings.warn(SORT_SPLIT_OUT_WARNING, FutureWarning)
 
@@ -1273,14 +1290,39 @@ def _aca_agg(
             aggregate_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])
 
         token = self._token_prefix + token
         levels = _determine_levels(self.by)
 
+        if shuffle:
+            return _shuffle_aggregate(
+                args,
+                chunk=_apply_chunk,
+                chunk_kwargs={
+                    "chunk": func,
+                    "columns": columns,
+                    **self.observed,
+                    **self.dropna,
+                    **chunk_kwargs,
+                },
+                aggregate=_groupby_aggregate,
+                aggregate_kwargs={
+                    "aggfunc": aggfunc,
+                    "levels": levels,
+                    **self.observed,
+                    **self.dropna,
+                    **aggregate_kwargs,
+                },
+                token=token,
+                split_every=split_every,
+                split_out=split_out,
+                shuffle=shuffle,
+                sort=self.sort,
+            )
+
         return aca(
-            [self.obj, self.by]
-            if not isinstance(self.by, list)
-            else [self.obj] + self.by,
+            args,
             chunk=_apply_chunk,
             chunk_kwargs=dict(
                 chunk=func,
@@ -1493,9 +1535,13 @@ def cumcount(self, axis=no_default):
         )
 
     @derived_from(pd.core.groupby.GroupBy)
-    def sum(self, split_every=None, split_out=1, min_count=None):
-        result = self._aca_agg(
-            token="sum", func=M.sum, split_every=split_every, split_out=split_out
+    def sum(self, split_every=None, split_out=1, shuffle=None, min_count=None):
+        result = self._single_agg(
+            func=M.sum,
+            token="sum",
+            split_every=split_every,
+            split_out=split_out,
+            shuffle=shuffle,
         )
         if min_count:
             return result.where(self.count() >= min_count, other=np.NaN)
@@ -1503,9 +1549,13 @@ def sum(self, split_every=None, split_out=1, min_count=None):
             return result
 
     @derived_from(pd.core.groupby.GroupBy)
-    def prod(self, split_every=None, split_out=1, min_count=None):
-        result = self._aca_agg(
-            token="prod", func=M.prod, split_every=split_every, split_out=split_out
+    def prod(self, split_every=None, split_out=1, shuffle=None, min_count=None):
+        result = self._single_agg(
+            func=M.prod,
+            token="prod",
+            split_every=split_every,
+            split_out=split_out,
+            shuffle=shuffle,
         )
         if min_count:
             return result.where(self.count() >= min_count, other=np.NaN)
@@ -1513,65 +1563,81 @@ def prod(self, split_every=None, split_out=1, min_count=None):
             return result
 
     @derived_from(pd.core.groupby.GroupBy)
-    def min(self, split_every=None, split_out=1):
-        return self._aca_agg(
-            token="min", func=M.min, split_every=split_every, split_out=split_out
+    def min(self, split_every=None, split_out=1, shuffle=None):
+        return self._single_agg(
+            func=M.min,
+            token="min",
+            split_every=split_every,
+            split_out=split_out,
+            shuffle=shuffle,
         )
 
     @derived_from(pd.core.groupby.GroupBy)
-    def max(self, split_every=None, split_out=1):
-        return self._aca_agg(
-            token="max", func=M.max, split_every=split_every, split_out=split_out
+    def max(self, split_every=None, split_out=1, shuffle=None):
+        return self._single_agg(
+            func=M.max,
+            token="max",
+            split_every=split_every,
+            split_out=split_out,
+            shuffle=shuffle,
         )
 
     @derived_from(pd.DataFrame)
-    def idxmin(self, split_every=None, split_out=1, axis=None, skipna=True):
-        return self._aca_agg(
-            token="idxmin",
+    def idxmin(
+        self, split_every=None, split_out=1, shuffle=None, axis=None, skipna=True
+    ):
+        return self._single_agg(
             func=M.idxmin,
+            token="idxmin",
             aggfunc=M.first,
             split_every=split_every,
             split_out=split_out,
+            shuffle=shuffle,
             chunk_kwargs=dict(skipna=skipna),
         )
 
     @derived_from(pd.DataFrame)
-    def idxmax(self, split_every=None, split_out=1, axis=None, skipna=True):
-        return self._aca_agg(
-            token="idxmax",
+    def idxmax(
+        self, split_every=None, split_out=1, shuffle=None, axis=None, skipna=True
+    ):
+        return self._single_agg(
             func=M.idxmax,
+            token="idxmax",
             aggfunc=M.first,
             split_every=split_every,
             split_out=split_out,
+            shuffle=shuffle,
             chunk_kwargs=dict(skipna=skipna),
         )
 
     @derived_from(pd.core.groupby.GroupBy)
-    def count(self, split_every=None, split_out=1):
-        return self._aca_agg(
-            token="count",
+    def count(self, split_every=None, split_out=1, shuffle=None):
+        return self._single_agg(
             func=M.count,
+            token="count",
             aggfunc=M.sum,
             split_every=split_every,
             split_out=split_out,
+            shuffle=shuffle,
         )
 
     @derived_from(pd.core.groupby.GroupBy)
-    def mean(self, split_every=None, split_out=1):
-        s = self.sum(split_every=split_every, split_out=split_out)
-        c = self.count(split_every=split_every, split_out=split_out)
+    def mean(self, split_every=None, split_out=1, shuffle=None):
+        s = self.sum(split_every=split_every, split_out=split_out, shuffle=shuffle)
+        c = self.count(split_every=split_every, split_out=split_out, shuffle=shuffle)
         if is_dataframe_like(s):
             c = c[s.columns]
         return s / c
 
     @derived_from(pd.core.groupby.GroupBy)
-    def size(self, split_every=None, split_out=1):
-        return self._aca_agg(
+    def size(self, split_every=None, split_out=1, shuffle=None):
+        return self._single_agg(
             token="size",
             func=M.size,
             aggfunc=M.sum,
             split_every=split_every,
             split_out=split_out,
+            shuffle=shuffle,
         )
 
     @derived_from(pd.core.groupby.GroupBy)
@@ -1664,15 +1730,23 @@ def cov(self, ddof=1, split_every=None, split_out=1, std=False):
         return result
 
     @derived_from(pd.core.groupby.GroupBy)
-    def first(self, split_every=None, split_out=1):
-        return self._aca_agg(
-            token="first", func=M.first, split_every=split_every, split_out=split_out
+    def first(self, split_every=None, split_out=1, shuffle=None):
+        return self._single_agg(
+            func=M.first,
+            token="first",
+            split_every=split_every,
+            split_out=split_out,
+            shuffle=shuffle,
         )
 
     @derived_from(pd.core.groupby.GroupBy)
-    def last(self, split_every=None, split_out=1):
-        return self._aca_agg(
-            token="last", func=M.last, split_every=split_every, split_out=split_out
+    def last(self, split_every=None, split_out=1, shuffle=None):
+        return self._single_agg(
+            token="last",
+            func=M.last,
+            split_every=split_every,
+            split_out=split_out,
+            shuffle=shuffle,
         )
 
     @derived_from(
@@ -1706,11 +1780,7 @@ def aggregate(self, arg, split_every=None, split_out=1, shuffle=None):
                 category=FutureWarning,
             )
             split_out = 1
-        if shuffle is None:
-            if split_out > 1:
-                shuffle = shuffle or config.get("shuffle", None) or "tasks"
-            else:
-                shuffle = False
+        shuffle = _determine_shuffle(shuffle, split_out)
 
         column_projection = None
         if isinstance(self.obj, DataFrame):
@@ -2387,53 +2457,57 @@ def agg(self, arg, split_every=None, split_out=1, shuffle=None):
         )
 
     @derived_from(pd.core.groupby.SeriesGroupBy)
-    def value_counts(self, split_every=None, split_out=1):
-        return self._aca_agg(
-            token="value_counts",
+    def value_counts(self, split_every=None, split_out=1, shuffle=None):
+        return self._single_agg(
             func=_value_counts,
+            token="value_counts",
             aggfunc=_value_counts_aggregate,
             split_every=split_every,
             split_out=split_out,
+            shuffle=shuffle,
         )
 
     @derived_from(pd.core.groupby.SeriesGroupBy)
-    def unique(self, split_every=None, split_out=1):
+    def unique(self, split_every=None, split_out=1, shuffle=None):
         name = self._meta.obj.name
-        return self._aca_agg(
-            token="unique",
+        return self._single_agg(
             func=M.unique,
+            token="unique",
             aggfunc=_unique_aggregate,
             aggregate_kwargs={"name": name},
             split_every=split_every,
             split_out=split_out,
+            shuffle=shuffle,
         )
 
     @derived_from(pd.core.groupby.SeriesGroupBy)
-    def tail(self, n=5, split_every=None, split_out=1):
+    def tail(self, n=5, split_every=None, split_out=1, shuffle=None):
         index_levels = len(self.by) if isinstance(self.by, list) else 1
-        return self._aca_agg(
-            token="tail",
+        return self._single_agg(
             func=_tail_chunk,
+            token="tail",
             aggfunc=_tail_aggregate,
             meta=M.tail(self._meta_nonempty),
             chunk_kwargs={"n": n},
             aggregate_kwargs={"n": n, "index_levels": index_levels},
             split_every=split_every,
             split_out=split_out,
+            shuffle=shuffle,
         )
 
     @derived_from(pd.core.groupby.SeriesGroupBy)
-    def head(self, n=5, split_every=None, split_out=1):
+    def head(self, n=5, split_every=None, split_out=1, shuffle=None):
         index_levels = len(self.by) if isinstance(self.by, list) else 1
-        return self._aca_agg(
-            token="head",
+        return self._single_agg(
             func=_head_chunk,
+            token="head",
             aggfunc=_head_aggregate,
             meta=M.head(self._meta_nonempty),
             chunk_kwargs={"n": n},
             aggregate_kwargs={"n": n, "index_levels": index_levels},
             split_every=split_every,
             split_out=split_out,
+            shuffle=shuffle,
         )
 
 
@@ -2563,6 +2637,14 @@ def _shuffle_aggregate(
         token=chunk_name,
         **chunk_kwargs,
     )
+    if is_series_like(chunked):
+        # Temporarily convert series to dataframe for shuffle
+        series_name = chunked._meta.name
+        chunked = chunked.to_frame("__series__")
+        convert_back_to_series = True
+    else:
+        series_name = None
+        convert_back_to_series = False
 
     shuffle_npartitions = max(
         chunked.npartitions // split_every,
@@ -2609,6 +2691,9 @@ def _shuffle_aggregate(
             shuffle=shuffle,
         ).map_partitions(aggregate, **aggregate_kwargs)
 
+    if convert_back_to_series:
+        result = result["__series__"].rename(series_name)
+
     if split_out < shuffle_npartitions:
         return result.repartition(npartitions=split_out)
     return result


Test patch:
diff --git a/dask/dataframe/tests/test_format.py b/dask/dataframe/tests/test_format.py
--- a/dask/dataframe/tests/test_format.py
+++ b/dask/dataframe/tests/test_format.py
@@ -480,7 +480,7 @@ def test_index_format():
 G       ...
 H       ...
 dtype: object
-Dask Name: from_pandas, 2 graph layers"""
+Dask Name: from_pandas-index, 2 graph layers"""
     assert repr(ds.index) == exp
     assert str(ds.index) == exp
 
@@ -498,7 +498,7 @@ def test_index_format():
     7                ...
     8                ...
     Name: YYY, dtype: category
-    Dask Name: from_pandas, 2 graph layers"""
+    Dask Name: from_pandas-index, 2 graph layers"""
     )
     assert repr(ds.index) == exp
     assert str(ds.index) == exp
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -1321,7 +1321,10 @@ def test_bfill():
         lambda df: [df["a"] > 2, df["b"] > 1],
     ],
 )
-def test_dataframe_aggregations_multilevel(grouper, agg_func):
+@pytest.mark.parametrize("split_out", [1, 2])
+def test_dataframe_aggregations_multilevel(grouper, split_out, agg_func):
+    sort = split_out == 1  # Don't sort for split_out > 1
+
     def call(g, m, **kwargs):
         return getattr(g, m)(**kwargs)
 
@@ -1340,31 +1343,53 @@ def call(g, m, **kwargs):
     # covariance only works with N+1 columns
     if agg_func not in ("cov", "corr"):
         assert_eq(
-            call(pdf.groupby(grouper(pdf))["c"], agg_func),
-            call(ddf.groupby(grouper(ddf))["c"], agg_func, split_every=2),
+            call(pdf.groupby(grouper(pdf), sort=sort)["c"], agg_func),
+            call(
+                ddf.groupby(grouper(ddf), sort=sort)["c"],
+                agg_func,
+                split_out=split_out,
+                split_every=2,
+            ),
         )
 
     # not supported by pandas
     if agg_func != "nunique":
+        if agg_func in ("cov", "corr") and split_out > 1:
+            pytest.skip("https://github.com/dask/dask/issues/9509")
         assert_eq(
-            call(pdf.groupby(grouper(pdf))[["c", "d"]], agg_func),
-            call(ddf.groupby(grouper(ddf))[["c", "d"]], agg_func, split_every=2),
+            call(pdf.groupby(grouper(pdf), sort=sort)[["c", "d"]], agg_func),
+            call(
+                ddf.groupby(grouper(ddf), sort=sort)[["c", "d"]],
+                agg_func,
+                split_out=split_out,
+                split_every=2,
+            ),
         )
 
         if agg_func in ("cov", "corr"):
             # there are sorting issues between pandas and chunk cov w/dask
-            df = call(pdf.groupby(grouper(pdf)), agg_func).sort_index()
+            df = call(pdf.groupby(grouper(pdf), sort=sort), agg_func).sort_index()
             cols = sorted(list(df.columns))
             df = df[cols]
-            dddf = call(ddf.groupby(grouper(ddf)), agg_func, split_every=2).compute()
+            dddf = call(
+                ddf.groupby(grouper(ddf), sort=sort),
+                agg_func,
+                split_out=split_out,
+                split_every=2,
+            ).compute()
             dddf = dddf.sort_index()
             cols = sorted(list(dddf.columns))
             dddf = dddf[cols]
             assert_eq(df, dddf)
         else:
             assert_eq(
-                call(pdf.groupby(grouper(pdf)), agg_func),
-                call(ddf.groupby(grouper(ddf)), agg_func, split_every=2),
+                call(pdf.groupby(grouper(pdf), sort=sort), agg_func),
+                call(
+                    ddf.groupby(grouper(ddf), sort=sort),
+                    agg_func,
+                    split_out=split_out,
+                    split_every=2,
+                ),
             )
 
 
@@ -1376,11 +1401,13 @@ def call(g, m, **kwargs):
         lambda df: [df["a"] > 2, df["b"] > 1],
     ],
 )
-def test_series_aggregations_multilevel(grouper, agg_func):
+@pytest.mark.parametrize("split_out", [1, 2])
+def test_series_aggregations_multilevel(grouper, split_out, agg_func):
     """
     similar to ``test_dataframe_aggregations_multilevel``, but series do not
     support all groupby args.
     """
+    sort = split_out == 1  # Don't sort for split_out > 1
 
     def call(g, m, **kwargs):
         return getattr(g, m)(**kwargs)
@@ -1401,8 +1428,13 @@ def call(g, m, **kwargs):
     ddf = dd.from_pandas(pdf, npartitions=10)
 
     assert_eq(
-        call(pdf["c"].groupby(grouper(pdf)), agg_func),
-        call(ddf["c"].groupby(grouper(ddf)), agg_func, split_every=2),
+        call(pdf["c"].groupby(grouper(pdf), sort=sort), agg_func),
+        call(
+            ddf["c"].groupby(grouper(ddf), sort=sort),
+            agg_func,
+            split_out=split_out,
+            split_every=2,
+        ),
         # for pandas ~ 0.18, the name is not not properly propagated for
         # the mean aggregation
         check_names=(agg_func not in {"mean", "nunique"}),
@@ -2780,8 +2812,8 @@ def test_groupby_sort_true_split_out():
         M.sum(ddf.groupby("x"), split_out=2)
 
     with pytest.raises(NotImplementedError):
-        # Cannot use sort=True with split_out>1 (for now)
-        M.sum(ddf.groupby("x", sort=True), split_out=2)
+        # Cannot use sort=True with split_out>1 using non-shuffle-based approach
+        M.sum(ddf.groupby("x", sort=True), shuffle=False, split_out=2)
 
     # Can use sort=True with split_out>1 with agg() if shuffle=True
     ddf.groupby("x", sort=True).agg("sum", split_out=2, shuffle=True)
@@ -2790,10 +2822,10 @@ def test_groupby_sort_true_split_out():
 @pytest.mark.skipif(
     not PANDAS_GT_110, reason="observed only supported for newer pandas"
 )
-@pytest.mark.parametrize("known_cats", [True, False])
-@pytest.mark.parametrize("ordered_cats", [True, False])
+@pytest.mark.parametrize("known_cats", [True, False], ids=["known", "unknown"])
+@pytest.mark.parametrize("ordered_cats", [True, False], ids=["ordered", "unordererd"])
 @pytest.mark.parametrize("groupby", ["cat_1", ["cat_1", "cat_2"]])
-@pytest.mark.parametrize("observed", [True, False])
+@pytest.mark.parametrize("observed", [True, False], ids=["observed", "unobserved"])
 def test_groupby_aggregate_categorical_observed(
     known_cats, ordered_cats, agg_func, groupby, observed
 ):