Memory spike and CancelledError when calling `dataframe.set_index()` on single machine

Hi all, I'm trying to use Dask's distributed scheduler to work with a large dataframe on a single machine, but it's crashing when I try to modify the index. Any ideas what might be causing this?

**What happened**:

I'm using the distributed scheduler on a single machine (a 2019 MacBook Pro, 6 cores, 32GB RAM). When I call `set_index` on a large Dask DataFrame and then try to write the result to Parquet, the main process's memory usage blows up to about 40GB and then fails with a `CancelledError`:

```
/Users/janek/Development/cms/.venv/lib/python3.7/site-packages/distributed/worker.py:3693: UserWarning: Large object of size 2.10 MiB detected in task graph:
[["('to-parquet-a3f2e5760f0b5dfd7d318e5894cc656f', ... _data.pq.dask']
Consider scattering large objects ahead of time
with client.scatter to reduce scheduler burden and
keep data on workers

future = client.submit(func, big_data) # bad

big_future = client.scatter(big_data) # good
future = client.submit(func, big_future) # good
% (format_bytes(len(b)), s)
distributed.batched - INFO - Batched Comm Closed:
Traceback (most recent call last):
File "/Users/janek/.pyenv/versions/3.7.3/lib/python3.7/runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "/Users/janek/.pyenv/versions/3.7.3/lib/python3.7/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/Users/janek/Development/cms/fail.py", line 18, in
dd.to_parquet(df, './indexed_data.pq.dask', write_index=True)
File "/Users/janek/Development/cms/.venv/lib/python3.7/site-packages/dask/dataframe/io/parquet/core.py", line 648, in to_parquet
out = out.compute(**compute_kwargs)
File "/Users/janek/Development/cms/.venv/lib/python3.7/site-packages/dask/base.py", line 285, in compute
(result,) = compute(self, traverse=False, **kwargs)
File "/Users/janek/Development/cms/.venv/lib/python3.7/site-packages/dask/base.py", line 567, in compute
results = schedule(dsk, keys, **kwargs)
File "/Users/janek/Development/cms/.venv/lib/python3.7/site-packages/distributed/client.py", line 2673, in get
results = self.gather(packed, asynchronous=asynchronous, direct=direct)
File "/Users/janek/Development/cms/.venv/lib/python3.7/site-packages/distributed/client.py", line 1988, in gather
asynchronous=asynchronous,
File "/Users/janek/Development/cms/.venv/lib/python3.7/site-packages/distributed/client.py", line 854, in sync
self.loop, func, *args, callback_timeout=callback_timeout, **kwargs
File "/Users/janek/Development/cms/.venv/lib/python3.7/site-packages/distributed/utils.py", line 354, in sync
raise exc.with_traceback(tb)
File "/Users/janek/Development/cms/.venv/lib/python3.7/site-packages/distributed/utils.py", line 337, in f
result[0] = yield future
File "/Users/janek/Development/cms/.venv/lib/python3.7/site-packages/tornado/gen.py", line 735, in run
value = future.result()
concurrent.futures._base.CancelledError
```

**What you expected to happen**:

I expected the re-indexed frame to be written to a file, without using a lot of memory.

**Minimal Complete Example**:

Unfortunately it's not easy to share my dataset, but my failing code looks like this:

```python
import dask.dataframe as dd
from dask.distributed import Client

dask_client = Client('tcp://192.168.7.66:8766')

df = dd.read_parquet('./bad_data.pq.dask')
df = df.set_index('claim_id', shuffle='disk')
dd.to_parquet(df, './indexed_data.pq.dask', write_index=True)
```

This consistently produces a memory blowup and crash. I'm running the cluster separately like this:

```bash
dask-scheduler --dashboard-address :8777 --port 8766
dask-worker tcp://192.168.7.66:8766 --nprocs 4 --nthreads 3 --memory-limit 8GB
```

**Anything else we need to know?**:

My dataset is somewhat sensitive and I can't put it in the cloud, which is why I'm stuck with a single machine.

The frame in question has about 1 billion rows. For the minimal example above, I've reduced the dataframe to just a single column plus an index, both `int64` and hence both taking about 8GB (according to `df.memory_usage(deep=True)`), and I still get the same failure. However, my full dataset has more columns and uses 300 GB or so, which is why I'm using Dask.

To isolate the problem, I looked into what `set_index()` is doing. It seems to boil down to calling `_calculate_divisions()` and then `set_partition()`, which itself does a `map_partitions(set_partitions_pre, ...)` followed by some other stuff. So, I saved the result of `_calculate_divisions` and then made a test script that just calls `map_partitions(set_partitions_pre, ...)` in the same way and tries to write the result to Parquet. This produces the same failure.

```python
import pickle

import dask.dataframe as dd
from dask.distributed import Client

dask_client = Client('tcp://192.168.7.66:8766')

df = dd.read_parquet('./bad_data.pq.dask')

with open ('./divisions.pkl', 'rb') as f:
divisions_ = pickle.load(f)
assert divisions_ == tuple(sorted(divisions_))
assert len(divisions_) == df.npartitions + 1

index = 'claim_id'
dtype = df[index].dtype
meta = df._meta._constructor_sliced([0])
divisions = df._meta._constructor_sliced(divisions_, dtype=dtype)

partitions = df[index].map_partitions(
dd.shuffle.set_partitions_pre, divisions=divisions, meta=meta,
)
df2 = df.assign(_partitions=partitions)

dd.to_parquet(df2, '/Volumes/JS CMS Data/scratch/intermediate_data.pq.dask', write_index=True)
```

**Environment**:

- Dask version: 2021.06.0
- Distributed version: 2021.06.0
- Python version: 3.7.3 (default, Dec 8 2019, 12:38:49) \n[Clang 11.0.0 (clang-1100.0.33.12)]
- Operating System: Mac OS 10.15.7
- Install method: pip

`pyarrow` is installed; `fastparquet` is not.

Thanks!