Repository: dask/dask

Issue:
Memory performance regresssion in Dask Dataframe when shuffling on disk
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:
I've been using Dask locally on my laptop primarily just as "Pandas but for more data than can fit in memory." But I've run into issues with memory consumption scaling to more than I have on my machine when sorting data. I posted about this to [Stack Overflow](https://stackoverflow.com/questions/64903267/using-set-index-on-a-dask-dataframe-and-writing-to-parquet-causes-memory-explo) some weeks ago.

To summarize my SO question, I found that memory consumption was scaling with input size and spiking around one part of the computation in a suspicious way.

I've since done some more investigation and it seems that this spike in memory usage is a regression introduced in release 2.14.0. Here's the resource profiler output from running my sort operation using Dask **2.14.0**:

![image](https://user-images.githubusercontent.com/6401947/108763212-20396380-751f-11eb-97ce-47288fcae849.png)

Note the memory spike to ~6000mb.

**What you expected to happen**:
Here is the resource profiler output from running the same sort operation on the same data using Dask **2.13.0**:

![image](https://user-images.githubusercontent.com/6401947/108763538-78706580-751f-11eb-8a66-bf01e6162d6c.png)

Note that there is no spike.

**Minimal Complete Verifiable Example**:
Here is my sort operation.

```python
import os
import sys
import dask
import dask.dataframe as dd
from dask.diagnostics import (
    Profiler,
    ResourceProfiler,
    visualize
)


def run(input_path, output_path):
    dask.config.set(scheduler="synchronous")

    filenames = os.listdir(input_path)
    full_filenames = [os.path.join(input_path, f) for f in filenames]

    prof = Profiler()
    rprof = ResourceProfiler()
    with prof, rprof:
        df = dd.read_parquet(full_filenames)
        df = df.set_index("follower_id")
        df.to_parquet(output_path)

    visualize([prof, rprof])


if __name__ == "__main__":
    run(sys.argv[1], sys.argv[2])
```

**Anything else we need to know?**:
I suspect this regression was introduced by https://github.com/dask/dask/pull/5977. That PR added a cleanup task to the task graph implementing the shuffle operation. It looks to me like the cleanup task forces the scheduler to pull all the partitions into memory. Using Dask **2.14.0** (which includes #5977), the ordering of the tasks in the graph (produced by `visualize(color="order")` shows that all the "collect-1" tasks that read from partd are executed before any of the write tasks, which I assume explains the memory spike.

Using Dask **2.13.0**, I can see that the scheduler has an opportunity to order the tasks such that the reads from partd are paired with writes to the output directory so that only one partition is in memory at any time.

**Environment**:

- Python version: 3.8.0
- Operating System: Ubuntu
- Install method (conda, pip, source): pip



Discussion:

Resolving patch:
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -855,22 +855,14 @@ def rearrange_by_column_disk(df, column, npartitions=None, compute=False):
     dsk3 = {barrier_token: (barrier, list(dsk2))}
 
     # Collect groups
-    name1 = "shuffle-collect-1" + token
+    name = "shuffle-collect-" + token
     dsk4 = {
-        (name1, i): (collect, p, i, df._meta, barrier_token) for i in range(npartitions)
+        (name, i): (collect, p, i, df._meta, barrier_token) for i in range(npartitions)
     }
-    cleanup_token = "cleanup-" + always_new_token
-    barrier_token2 = "barrier2-" + always_new_token
-    # A task that depends on `cleanup-`, but has a small output
-    dsk5 = {(barrier_token2, i): (barrier, part) for i, part in enumerate(dsk4)}
-    # This indirectly depends on `cleanup-` and so runs after we're done using the disk
-    dsk6 = {cleanup_token: (cleanup_partd_files, p, list(dsk5))}
-
-    name = "shuffle-collect-2" + token
-    dsk7 = {(name, i): (_noop, (name1, i), cleanup_token) for i in range(npartitions)}
+
     divisions = (None,) * (npartitions + 1)
 
-    layer = toolz.merge(dsk1, dsk2, dsk3, dsk4, dsk5, dsk6, dsk7)
+    layer = toolz.merge(dsk1, dsk2, dsk3, dsk4)
     graph = HighLevelGraph.from_collections(name, layer, dependencies=dependencies)
     return DataFrame(graph, name, df._meta, divisions)
 


Test patch: