Repository: dask/dask

Issue:
Shuffle-based groupby aggregation by default
In https://github.com/dask/dask/pull/9302, @rjzamora added a new shuffle-based groupby algorithm which @ian-r-rose demonstrated has significant performance improvements in certain cases https://github.com/dask/dask/pull/9302#issuecomment-1216782830. Today this new shuffle-based groupby is gated behind an optional `shuffle=` keyword argument (off by default). We should probably turn it on by default in cases where it makes sense. 

Here's a brief summary of @ian-r-rose's findings (see https://github.com/dask/dask/pull/9302#issuecomment-1216782830 for more details)

> TL;DR
> 
> This largely confirms @rjzamora's points above and in conversations I've had with him:
> 
>     ACA-based groupby/aggs do fine for low-cardinality cases (say, where the number of groups is less than a few percent of the dataframe size).
>     Once the cardinality gets larger, it does indeed make sense to start using the shuffle-based algorithm. To me, the range where split_out > 1 and where it's still better to use ACA looks pretty darn narrow. A rule of thumb where shuffle is on-by-default when split_out>1 doesn't sound like the worst idea in the world to me.
>     Memory-usage is indeed higher for task-based shuffling than ACA in the lower-cardinality cases. But that advantage is wiped out in higher-cardinality cases, so it actually doesn't worry me much. And task-based shuffling is less likely to knock over a worker in the last few aggs of the ACA-based approach.
>     p2p shuffling has much better memory performance. I ran into a couple of problems when using it here that I'll open as separate issues, but the results are encouraging.

Opening this issue so we don't loose track of this topic 



Discussion:
Just checking in here.  Based on the

> Once the cardinality gets larger, it does indeed make sense to start using the shuffle-based algorithm. To me, the range where split_out > 1 and where it's still better to use ACA looks pretty darn narrow. A rule of thumb where shuffle is on-by-default when split_out>1 doesn't sound like the worst idea in the world to me.

comment from @ian-r-rose in https://github.com/dask/dask/pull/9302#issuecomment-1216782830 it sounds like we should move to turn on `shuffle=` by default when `split_out > 1`. @ian-r-rose @rjzamora does that sound right to you? Are there other times when it makes sense to default to shuffling? 
> shuffle= by default when split_out > 1

I'm going to try to run some more targeted tests today and tomorrow, but broadly speaking, I think this is correct. It is likely *not* an improvement when `split_out > 1` and we are low-cardinality, but in that case I also wouldn't recommend doing `split_out > 1`, so I'm not sure that is a big deal.
>does that sound right to you? Are there other times when it makes sense to default to shuffling?

As @ian-r-rose mentioned, there are `split_out>1` cases where the cardinality is too low to benefit from the shuffle, but this simplification does seem reasonable to me.

Another alternative may be to introduce a more intuitive knob for users to turn (and to discourage the explicit use of `split_out`, `split_every`, and `shuffle`).  For example, we *could* introduce a `cardinality` kwarg, where the user can (optionally) specify the estimated ratio of output records to input records (i.e. `len(ddf.drop_duplicates()) / len(ddf)`). This ratio argument could be used within dask set appropriate values for `split_out`, `split_every`, and `shuffle`.  We could also implement something like `cardinality="calculate"` to let the user opt into an eager (but approximate) ratio calculation, and/or `cardinality="infer"` to estimate the ratio from the first non-empty partition.

Resolving patch:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -7,6 +7,7 @@
 import numpy as np
 import pandas as pd
 
+from dask import config
 from dask.base import tokenize
 from dask.dataframe._compat import PANDAS_GT_150
 from dask.dataframe.core import (
@@ -1077,15 +1078,20 @@ def wrapper(func):
             - dict of column names -> function, function name or list of such.
         split_every : int, optional
             Number of intermediate partitions that may be aggregated at once.
-            Default is 8.
+            This defaults to 8. If your intermediate partitions are likely to
+            be small (either due to a small number of groups or a small initial
+            partition size), consider increasing this number for better performance.
         split_out : int, optional
             Number of output partitions. Default is 1.
         shuffle : bool or str, optional
             Whether a shuffle-based algorithm should be used. A specific
-            algorithm name may also be specified (e.g. `"tasks"` or `"p2p"`).
+            algorithm name may also be specified (e.g. ``"tasks"`` or ``"p2p"``).
             The shuffle-based algorithm is likely to be more efficient than
             ``shuffle=False`` when ``split_out>1`` and the number of unique
-            groups is large (high cardinality). Default is ``False``.
+            groups is large (high cardinality). Default is ``False`` when
+            ``split_out = 1``. When ``split_out > 1``, it chooses the algorithm
+            set by the ``shuffle`` option in the dask config system, or ``"tasks"``
+            if nothing is set.
         """
         return func
 
@@ -1654,6 +1660,12 @@ def get_group(self, key):
 
     @_aggregate_docstring()
     def aggregate(self, arg, split_every=None, split_out=1, shuffle=None):
+        if shuffle is None:
+            if split_out > 1:
+                shuffle = shuffle or config.get("shuffle", None) or "tasks"
+            else:
+                shuffle = False
+
         column_projection = None
         if isinstance(self.obj, DataFrame):
             if isinstance(self.by, tuple) or np.isscalar(self.by):
@@ -2449,8 +2461,12 @@ def _shuffle_aggregate(
     aggregate_kwargs : dict, optional
         Keywords for the aggregate function only.
     split_every : int, optional
-        Number of intermediate partitions that may be aggregated at once.
-        Default is 8.
+        Number of partitions to aggregate into a shuffle partition.
+        Defaults to eight, meaning that the initial partitions are repartitioned
+        into groups of eight before the shuffle. Shuffling scales with the number
+        of partitions, so it may be helpful to increase this number as a performance
+        optimization, but only when the aggregated partition can comfortably
+        fit in worker memory.
     split_out : int, optional
         Number of output partitions.
     ignore_index : bool, default False
@@ -2480,8 +2496,8 @@ def _shuffle_aggregate(
         split_every = 8
     elif split_every is False:
         split_every = npartitions
-    elif split_every < 2 or not isinstance(split_every, Integral):
-        raise ValueError("split_every must be an integer >= 2")
+    elif split_every < 1 or not isinstance(split_every, Integral):
+        raise ValueError("split_every must be an integer >= 1")
 
     # Shuffle-based groupby aggregation
     chunk_name = f"{token or funcname(chunk)}-chunk"


Test patch:
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
@@ -1106,8 +1106,9 @@ def test_aggregate_dask():
                 assert len(other.dask) == len(result2.dask)
 
 
+@pytest.mark.parametrize("split_every", [1, 8])
 @pytest.mark.parametrize("split_out", [2, 32])
-def test_shuffle_aggregate(shuffle_method, split_out):
+def test_shuffle_aggregate(shuffle_method, split_out, split_every):
 
     pdf = pd.DataFrame(
         {
@@ -1122,7 +1123,7 @@ def test_shuffle_aggregate(shuffle_method, split_out):
 
     spec = {"b": "mean", "c": ["min", "max"]}
     result = ddf.groupby(["a", "b"]).agg(
-        spec, split_out=split_out, shuffle=shuffle_method
+        spec, split_out=split_out, split_every=split_every, shuffle=shuffle_method
     )
     expect = pdf.groupby(["a", "b"]).agg(spec)
 
@@ -1163,6 +1164,33 @@ def test_shuffle_aggregate_sort(shuffle_method, sort):
             )
 
 
+def test_shuffle_aggregate_defaults(shuffle_method):
+    pdf = pd.DataFrame(
+        {
+            "a": [1, 2, 3, 1, 1, 2, 4, 3, 7] * 100,
+            "b": [4, 2, 7, 3, 3, 1, 1, 1, 2] * 100,
+            "c": [0, 1, 2, 3, 4, 5, 6, 7, 8] * 100,
+            "d": [3, 2, 1, 3, 2, 1, 2, 6, 4] * 100,
+        },
+        columns=["c", "b", "a", "d"],
+    )
+    ddf = dd.from_pandas(pdf, npartitions=100)
+
+    spec = {"b": "mean", "c": ["min", "max"]}
+
+    # No shuffle layer when  split_out = 1
+    dsk = ddf.groupby("a").agg(spec, split_out=1).dask
+    assert not any("shuffle" in l for l in dsk.layers)
+
+    # split_every=1 is invalid for tree reduction
+    with pytest.raises(ValueError):
+        ddf.groupby("a").agg(spec, split_out=1, split_every=1)
+
+    # If split_out > 1, default to shuffling.
+    dsk = ddf.groupby("a").agg(spec, split_out=2, split_every=1).dask
+    assert any("shuffle" in l for l in dsk.layers)
+
+
 @pytest.mark.parametrize("axis", [0, 1])
 @pytest.mark.parametrize("group_keys", [True, False, None])
 @pytest.mark.parametrize("method", ["ffill", "bfill"])