Repository: dask/dask

Issue:
Possible regression introduced by PR 9302
Over in https://github.com/coiled/coiled-runtime/issues/272  and https://github.com/coiled/coiled-runtime/issues/271 we reported a potential regression, that after running a git-bisect we were able to identify it was caused by PR https://github.com/dask/dask/pull/9302 even though it's not supposed to affect this workflow, as the option is off by default. 

The test we run consists of this query https://github.com/coiled/coiled-runtime/blob/f858057b8810bce1a397aa391a1f188c83d47cf7/tests/benchmarks/h2o/test_h2o_benchmarks.py#L67-L73

```
dask-version               |     duration [s]
------------------------------------------------------

2022.8.0+18.gc1bc90aa5     |     70.4672958850861  
2022.8.0+19.g097decd5f     |     87.210559129715       (PR https://github.com/dask/dask/pull/9302)
```
This PR shouldn't have affected the behavior as the flag was not turned on. Not quite sure what could be happening here, but this seems like a legitimate regression. A couple of things to know is that `id3` has high cardinality but the total amount of data, in this case, is not that big.

cc: @ian-r-rose @rjzamora Thoughts?


Discussion:
Thanks for the report @ncclementi! This is probably due to the change to the `sort` arg in #9428 . It is now overriding the defaults for some of the aca functions (which previously had `False`, and now have `None`).
>Thanks for the report @ncclementi! This is probably due to the change to the sort arg in https://github.com/dask/dask/issues/9428 . It is now overriding the defaults for some of the aca functions (which previously had False, and now have None).

I was just about to say the same thing (sort of). The PR now passes the sort argument to `chunk` and `combine` (while it didn't pass anything at all before).  I was under the impression that passing `None` would be the same as passing nothing in this case, but I'm probably missing something.
**Side Note**: There is no requirement for the output groups to be sorted in query-3.  The pandas code even uses an explicit `sort=False` for the [published results](https://h2oai.github.io/db-benchmark/).  Therefore, I'd suggest that the benchmark you linked also be changed to `sort=False`.
> I was under the impression that passing None would be the same as passing nothing in this case

I was also under that impression -- but the `None` overrides the default of `False` in some of the chunk/agg functions, which then probably gets defaulted to `True` in [pandas](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.groupby.html#pandas-dataframe-groupby).



> **Side Note**: There is no requirement for the output groups to be sorted in query-3. The pandas code even uses an explicit `sort=False` for the [published results](https://h2oai.github.io/db-benchmark/). Therefore, I'd suggest that the benchmark you linked also be changed to `sort=False`.

Good point! Though I'm glad it was unset so as to catch this!
A bit of supporting information: shown are boxplots for duration and average memory. The first is the commit before #9302, the second is #9302, and the third is the same, but with the `sort=self.sort` args removed
![visualization(28)](https://user-images.githubusercontent.com/5728311/186559066-b4f68b56-64bd-4e7b-be73-e65d236b8cfb.png)
:
Thanks for running those experiments @ian-r-rose - Do you think the best path is to keep defaults, or is this more motivation for something like #9386 (assuming that would help this case)?
From my perspective, going forward with #9386 seems like it would be a fine choice.

Over there you said:

> [it] does not improve performance when the aggregation fits comfortably in memory (and may even introduce a slight regression)

I don't quite understand how it could introduce a regression, seems like it should be doing less work overall. Unless you mean that it performs something like a merge sort on the tree reduction?
>I don't quite understand how it could introduce a regression

I was confused by my observations, and decided not to devote time to the issue in the short term. My assumption was that the final aggregation node(s) (where `sort=True` is the default) were benefiting from the input aggregations having "pre-sorted" groups. However, I'm far from sure about this, and I don't think the regressions were significant.
Hmm, I'm still seeing `sort=False` have some worse behavior than just not passing it in -- I guess I'll need to dig a bit deeper.

Resolving patch:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -69,6 +69,13 @@
 #
 # #############################################
 
+SORT_SPLIT_OUT_WARNING = (
+    "In the future, `sort` for groupby operations will default to `True`"
+    " to match the behavior of pandas. However, `sort=True` does not work"
+    " with `split_out>1`. To retain the current behavior for multiple"
+    " output partitions, set `sort=False`."
+)
+
 
 def _determine_levels(by):
     """Determine the correct levels argument to groupby."""
@@ -1116,7 +1123,7 @@ class _GroupBy:
         Passed to pandas.DataFrame.groupby()
     dropna: bool
         Whether to drop null values from groupby index
-    sort: bool, defult None
+    sort: bool
         Passed along to aggregation methods. If allowed,
         the output aggregation will have sorted keys.
     observed: bool, default False
@@ -1132,8 +1139,8 @@ def __init__(
         slice=None,
         group_keys=GROUP_KEYS_DEFAULT,
         dropna=None,
-        sort=None,
-        observed=None,
+        sort=True,
+        observed=False,
     ):
 
         by_ = by if isinstance(by, (tuple, list)) else [by]
@@ -1249,6 +1256,9 @@ def _aca_agg(
         chunk_kwargs=None,
         aggregate_kwargs=None,
     ):
+        if self.sort is None and split_out > 1:
+            warnings.warn(SORT_SPLIT_OUT_WARNING, FutureWarning)
+
         if aggfunc is None:
             aggfunc = func
 
@@ -1566,6 +1576,9 @@ def size(self, split_every=None, split_out=1):
 
     @derived_from(pd.core.groupby.GroupBy)
     def var(self, ddof=1, split_every=None, split_out=1):
+        if self.sort is None and split_out > 1:
+            warnings.warn(SORT_SPLIT_OUT_WARNING, FutureWarning)
+
         levels = _determine_levels(self.by)
         result = aca(
             [self.obj, self.by]
@@ -1615,6 +1628,8 @@ def cov(self, ddof=1, split_every=None, split_out=1, std=False):
 
         When `std` is True calculate Correlation
         """
+        if self.sort is None and split_out > 1:
+            warnings.warn(SORT_SPLIT_OUT_WARNING, FutureWarning)
 
         levels = _determine_levels(self.by)
 
@@ -1787,7 +1802,7 @@ def aggregate(self, arg, split_every=None, split_out=1, shuffle=None):
                 chunk=_groupby_apply_funcs,
                 chunk_kwargs=dict(
                     funcs=chunk_funcs,
-                    sort=self.sort,
+                    sort=False,
                     **self.observed,
                     **self.dropna,
                 ),
@@ -1806,6 +1821,9 @@ def aggregate(self, arg, split_every=None, split_out=1, shuffle=None):
                 sort=self.sort,
             )
 
+        if self.sort is None and split_out > 1:
+            warnings.warn(SORT_SPLIT_OUT_WARNING, FutureWarning)
+
         # Check sort behavior
         if self.sort and split_out > 1:
             raise NotImplementedError(
@@ -1819,7 +1837,7 @@ def aggregate(self, arg, split_every=None, split_out=1, shuffle=None):
             chunk=_groupby_apply_funcs,
             chunk_kwargs=dict(
                 funcs=chunk_funcs,
-                sort=self.sort,
+                sort=False,
                 **self.observed,
                 **self.dropna,
             ),
@@ -1827,7 +1845,7 @@ def aggregate(self, arg, split_every=None, split_out=1, shuffle=None):
             combine_kwargs=dict(
                 funcs=aggregate_funcs,
                 level=levels,
-                sort=self.sort,
+                sort=False,
                 **self.observed,
                 **self.dropna,
             ),
@@ -2329,6 +2347,9 @@ def nunique(self, split_every=None, split_out=1):
         else:
             chunk = _nunique_series_chunk
 
+        if self.sort is None and split_out > 1:
+            warnings.warn(SORT_SPLIT_OUT_WARNING, FutureWarning)
+
         return aca(
             [self.obj, self.by]
             if not isinstance(self.by, list)
@@ -2466,7 +2487,7 @@ def _shuffle_aggregate(
     aggregate_kwargs=None,
     split_every=None,
     split_out=1,
-    sort=None,
+    sort=True,
     ignore_index=False,
     shuffle="tasks",
 ):
@@ -2502,7 +2523,7 @@ def _shuffle_aggregate(
         Number of output partitions.
     ignore_index : bool, default False
         Whether the index can be ignored during the shuffle.
-    sort : bool, default None
+    sort : bool
         If allowed, sort the keys of the output aggregation.
     shuffle : str, default "tasks"
         Shuffle option to be used by ``DataFrame.shuffle``.
@@ -2553,6 +2574,17 @@ def _shuffle_aggregate(
         aggregate_kwargs = aggregate_kwargs or {}
         aggregate_kwargs["sort"] = sort
 
+    if sort is None and split_out > 1:
+        idx = set(chunked._meta.columns) - set(chunked._meta.reset_index().columns)
+        if len(idx) > 1:
+            warnings.warn(
+                "In the future, `sort` for groupby operations will default to `True`"
+                " to match the behavior of pandas. However, `sort=True` does not work"
+                " with `split_out>1` when grouping by multiple columns. To retain the"
+                " current behavior for multiple output partitions, set `sort=False`.",
+                FutureWarning,
+            )
+
     # Perform global sort or shuffle
     if sort and split_out > 1:
         cols = set(chunked.columns)


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
@@ -784,12 +784,12 @@ def call(g, m, **kwargs):
         # covariance/correlation is not a series aggregation
         if m in ("nunique", "cov", "corr"):
             continue
-        res = call(ddf.groupby("b"), m, **{keyword: 2})
+        res = call(ddf.groupby("b", sort=False), m, **{keyword: 2})
         sol = call(pdf.groupby("b"), m)
         assert_eq(res, sol)
         assert call(ddf.groupby("b"), m)._name != res._name
 
-    res = call(ddf.groupby("b"), "var", ddof=2, **{keyword: 2})
+    res = call(ddf.groupby("b", sort=False), "var", ddof=2, **{keyword: 2})
     sol = call(pdf.groupby("b"), "var", ddof=2)
     assert_eq(res, sol)
     assert call(ddf.groupby("b"), "var", ddof=2)._name != res._name
@@ -799,12 +799,12 @@ def call(g, m, **kwargs):
         # covariance/correlation is not a series aggregation
         if m in ("cov", "corr"):
             continue
-        res = call(ddf.groupby("b").a, m, **{keyword: 2})
+        res = call(ddf.groupby("b", sort=False).a, m, **{keyword: 2})
         sol = call(pdf.groupby("b").a, m)
         assert_eq(res, sol)
         assert call(ddf.groupby("b").a, m)._name != res._name
 
-    res = call(ddf.groupby("b").a, "var", ddof=2, **{keyword: 2})
+    res = call(ddf.groupby("b", sort=False).a, "var", ddof=2, **{keyword: 2})
     sol = call(pdf.groupby("b").a, "var", ddof=2)
     assert_eq(res, sol)
     assert call(ddf.groupby("b").a, "var", ddof=2)._name != res._name
@@ -814,14 +814,14 @@ def call(g, m, **kwargs):
         # covariance/correlation is not a series aggregation
         if m in ("cov", "corr"):
             continue
-        res = call(ddf.a.groupby(ddf.b), m, **{keyword: 2})
+        res = call(ddf.a.groupby(ddf.b, sort=False), m, **{keyword: 2})
         sol = call(pdf.a.groupby(pdf.b), m)
         # There's a bug in pandas 0.18.0 with `pdf.a.groupby(pdf.b).count()`
         # not forwarding the series name. Skip name checks here for now.
         assert_eq(res, sol, check_names=False)
         assert call(ddf.a.groupby(ddf.b), m)._name != res._name
 
-    res = call(ddf.a.groupby(ddf.b), "var", ddof=2, **{keyword: 2})
+    res = call(ddf.a.groupby(ddf.b, sort=False), "var", ddof=2, **{keyword: 2})
     sol = call(pdf.a.groupby(pdf.b), "var", ddof=2)
 
     assert_eq(res, sol)
@@ -1138,7 +1138,7 @@ def test_shuffle_aggregate(shuffle_method, split_out, split_every):
     ddf = dd.from_pandas(pdf, npartitions=100)
 
     spec = {"b": "mean", "c": ["min", "max"]}
-    result = ddf.groupby(["a", "b"]).agg(
+    result = ddf.groupby(["a", "b"], sort=False).agg(
         spec, split_out=split_out, split_every=split_every, shuffle=shuffle_method
     )
     expect = pdf.groupby(["a", "b"]).agg(spec)
@@ -1203,7 +1203,7 @@ def test_shuffle_aggregate_defaults(shuffle_method):
         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
+    dsk = ddf.groupby("a", sort=False).agg(spec, split_out=2, split_every=1).dask
     assert any("shuffle" in l for l in dsk.layers)
 
 
@@ -1520,19 +1520,21 @@ def test_groupy_series_wrong_grouper():
 
 @pytest.mark.parametrize("npartitions", [1, 4, 20])
 @pytest.mark.parametrize("split_every", [2, 5])
-@pytest.mark.parametrize("split_out", [None, 1, 5, 20])
+@pytest.mark.parametrize("split_out", [1, 5, 20])
 def test_hash_groupby_aggregate(npartitions, split_every, split_out):
     df = pd.DataFrame({"x": np.arange(100) % 10, "y": np.ones(100)})
     ddf = dd.from_pandas(df, npartitions)
 
-    result = ddf.groupby("x").y.var(split_every=split_every, split_out=split_out)
+    result = ddf.groupby("x", sort=(split_out == 1)).y.var(
+        split_every=split_every, split_out=split_out
+    )
 
     dsk = result.__dask_optimize__(result.dask, result.__dask_keys__())
     from dask.core import get_deps
 
     dependencies, dependents = get_deps(dsk)
 
-    assert result.npartitions == (split_out or 1)
+    assert result.npartitions == split_out
     assert len([k for k, v in dependencies.items() if not v]) == npartitions
 
     assert_eq(result, df.groupby("x").y.var())
@@ -1545,7 +1547,7 @@ def test_split_out_multi_column_groupby():
 
     ddf = dd.from_pandas(df, npartitions=10)
 
-    result = ddf.groupby(["x", "y"]).z.mean(split_out=4)
+    result = ddf.groupby(["x", "y"], sort=False).z.mean(split_out=4)
     expected = df.groupby(["x", "y"]).z.mean()
 
     assert_eq(result, expected, check_dtype=False)
@@ -1557,8 +1559,8 @@ def test_groupby_split_out_num():
         pd.DataFrame({"A": [1, 1, 2, 2], "B": [1, 2, 3, 4]}), npartitions=2
     )
     assert ddf.groupby("A").sum().npartitions == 1
-    assert ddf.groupby("A").sum(split_out=2).npartitions == 2
-    assert ddf.groupby("A").sum(split_out=3).npartitions == 3
+    assert ddf.groupby("A", sort=False).sum(split_out=2).npartitions == 2
+    assert ddf.groupby("A", sort=False).sum(split_out=3).npartitions == 3
 
     with pytest.raises(TypeError):
         # groupby doesn't accept split_out
@@ -2669,14 +2671,14 @@ def test_groupby_split_out_multiindex(split_out, column):
     ddf = dd.from_pandas(df, npartitions=3)
 
     ddf_result_so1 = (
-        ddf.groupby(column).a.mean(split_out=1).compute().sort_values().dropna()
+        ddf.groupby(column, sort=False).a.mean(split_out=1).compute().dropna()
     )
 
     ddf_result = (
-        ddf.groupby(column).a.mean(split_out=split_out).compute().sort_values().dropna()
+        ddf.groupby(column, sort=False).a.mean(split_out=split_out).compute().dropna()
     )
 
-    assert_eq(ddf_result, ddf_result_so1, check_index=False)
+    assert_eq(ddf_result, ddf_result_so1)
 
 
 @pytest.mark.parametrize(
@@ -2772,7 +2774,10 @@ def test_groupby_sort_true_split_out():
     # Works fine for split_out==1 or sort=False/None
     M.sum(ddf.groupby("x", sort=True), split_out=1)
     M.sum(ddf.groupby("x", sort=False), split_out=2)
-    M.sum(ddf.groupby("x"), split_out=2)
+
+    # Warns for sort=None
+    with pytest.warns(FutureWarning, match="split_out>1"):
+        M.sum(ddf.groupby("x"), split_out=2)
 
     with pytest.raises(NotImplementedError):
         # Cannot use sort=True with split_out>1 (for now)
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -229,7 +229,7 @@ async def test_serializable_groupby_agg(c, s, a, b):
     df = pd.DataFrame({"x": [1, 2, 3, 4], "y": [1, 0, 1, 0]})
     ddf = dd.from_pandas(df, npartitions=2)
 
-    result = ddf.groupby("y").agg("count", split_out=2)
+    result = ddf.groupby("y", sort=False).agg("count", split_out=2)
 
     # Check Culling and Compute
     agg0 = await c.compute(result.partitions[0])