Repository: dask/dask

Issue:
Dataframe merge fails with KeyError
**What happened**:

Dask optimization fails when trying to merge two dataframes

**What you expected to happen**:

The two dataframes are merged

**Minimal Complete Verifiable Example**:

Load a reasonably-large dataset to the cluster
```
base_uri = "http://datasets.imdbws.com"
file_list = [
    "name.basics.tsv.gz",
    "title.basics.tsv.gz",
    "title.principals.tsv.gz",
    "title.ratings.tsv.gz"
]
for file_name in file_list:
    (dd.read_csv("/".join((base_uri, file_name)),
                 sep="\t", 
                 blocksize=None, 
                 dtype={
                     'isOriginalTitle': 'object', 
                     'isAdult': 'object',
                     'startYear': 'object'
                 }
                )
     .repartition(partition_size="100MB")
     .to_csv(f"s3://imdb-datasets/{'.'.join(file_name.split('.')[:2])}.csv.gz",
             compression="gzip",
             storage_options=storage_options())
    )
```

Load the files as Dask dataframes

```
people = dd.read_csv("s3://imdb-datasets/name.basics.csv.gz/*", 
                     compression="gzip",
                     blocksize=None,
                     storage_options=storage_options())

movies = dd.read_csv("s3://imdb-datasets/title.basics.csv.gz/*",
                     compression="gzip",
                     blocksize=None,
                     dtype={'startYear': 'object', 'isAdult': 'object'},
                     storage_options=storage_options())

ratings = dd.read_csv("s3://imdb-datasets/title.ratings.csv.gz/*",
                      compression="gzip",
                      blocksize=None,
                      storage_options=storage_options())

principals = dd.read_csv("s3://imdb-datasets/title.principals.csv.gz/*", 
                         compression="gzip",
                         blocksize=None,
                         storage_options=storage_options())[["tconst", "nconst", "category", "job"]]
```

Do some merges

```
filtered_ratings = ratings[ratings.numVotes > 1000]

relevant_movies = filtered_ratings[["tconst"]]

filtered_movies = (
    movies
    [['tconst', 'primaryTitle', 'originalTitle', 'startYear', 'runtimeMinutes']]
    .merge(relevant_movies, on='tconst', how='inner')
)

category_index = principals.category.unique().reset_index()
job_index = principals.job.unique().reset_index()

filtered_crew = (
    principals
    .merge(category_index, on='category', how='left')
    .merge(job_index, on='job', how='left', suffixes=('_category', '_job'))
    .merge(relevant_movies, on='tconst', how='inner')
    [["tconst", "nconst", "index_category", "index_job"]]
) 
```

And if I try to merge `filtered_crew` with `people``

```
people.merge(filtered_crew, on="nconst", how="inner").head()
```

I get the following error

```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-6-4b2a05156842> in <module>
----> 1 people.merge(filtered_crew, on="nconst", how="inner").head()

~/.local/lib/python3.8/site-packages/dask/dataframe/core.py in head(self, n, npartitions, compute)
   1034             Whether to compute the result, default is True.
   1035         """
-> 1036         return self._head(n=n, npartitions=npartitions, compute=compute, safe=True)
   1037 
   1038     def _head(self, n, npartitions, compute, safe):

~/.local/lib/python3.8/site-packages/dask/dataframe/core.py in _head(self, n, npartitions, compute, safe)
   1067 
   1068         if compute:
-> 1069             result = result.compute()
   1070         return result
   1071 

~/.local/lib/python3.8/site-packages/dask/base.py in compute(self, **kwargs)
    277         dask.base.compute
    278         """
--> 279         (result,) = compute(self, traverse=False, **kwargs)
    280         return result
    281 

~/.local/lib/python3.8/site-packages/dask/base.py in compute(*args, **kwargs)
    553     )
    554 
--> 555     dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
    556     keys, postcomputes = [], []
    557     for x in collections:

~/.local/lib/python3.8/site-packages/dask/base.py in collections_to_dsk(collections, optimize_graph, optimizations, **kwargs)
    329         for opt, val in groups.items():
    330             dsk, keys = _extract_graph_and_keys(val)
--> 331             dsk = opt(dsk, keys, **kwargs)
    332 
    333             for opt in optimizations:

~/.local/lib/python3.8/site-packages/dask/dataframe/optimize.py in optimize(dsk, keys, **kwargs)
     32     if fuse_subgraphs is None:
     33         fuse_subgraphs = True
---> 34     dsk, _ = fuse(
     35         dsk,
     36         keys,

~/.local/lib/python3.8/site-packages/dask/optimization.py in fuse(dsk, keys, dependencies, ave_width, max_width, max_height, max_depth_new_edges, rename_keys, fuse_subgraphs)
    598             child = children_stack[-1]
    599             if child != parent:
--> 600                 children = reducible & deps[child]
    601                 while children:
    602                     # Depth-first search

KeyError: ('assign-6806e9502a8527bf71183e8bcc82e738', 48)
```

**Anything else we need to know?**:

I can get the same result successfully by not using the `filtered_crew` dataframe

```
relevant_people = principals.merge(relevant_movies, on="tconst", how='inner')[['nconst']].drop_duplicates()
filtered_people = people.merge(relevant_people, on="nconst", how="inner")
filtered_people.head()
```

**Environment**:

- Dask version: 2021.1.0
- Python version: 3.8.0
- Operating System: Linux
- Install method (conda, pip, source): Dask helm chart



Discussion:
What is `storage_options=storage_options()` ?
The function just returns the credentials and endpoint to read from a S3 bucket.

```
{
    "key": "AWS_ACCESS_KEY_ID",
    "secret": "AWS_SECRET_ACCESS_KEY",
    "client_kwargs": {"endpoint_url": "http://mys3endpoint"},
}
```
Could you confirm that

```
ddf1 = (
    dd.read_csv(
        "http://datasets.imdbws.com/name.basics.tsv.gz",
         sep="\t", 
         blocksize=None, 
         dtype={'isOriginalTitle': 'object', 'isAdult': 'object', 'startYear': 'object'}
    )
    .repartition(partition_size="100MB")
)

ddf2 = (
    dd.read_csv(
        "http://datasets.imdbws.com/title.principals.tsv.gz",
         sep="\t", 
         blocksize=None, 
         dtype={'isOriginalTitle': 'object', 'isAdult': 'object', 'startYear': 'object'}
    )
    .repartition(partition_size="100MB")
)

ddf1.merge(ddf2, on="nconst", how="inner").head()
```

reproduces your problem as well (at least here it does). Please try to make your minimum verifiable example as minimal as possible.
Yes i does! It triggers the same KeyError. Thank you so much @rubenvdg!
It seems like this might have to do with the blockwise work that @rjzamora is doing. 

Resolving patch:
diff --git a/dask/blockwise.py b/dask/blockwise.py
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -325,6 +325,7 @@ def __dask_distributed_pack__(self, client):
             "func_future_args": func_future_args,
             "global_dependencies": global_dependencies,
             "indices": indices,
+            "is_list": [isinstance(x, list) for x in indices],
             "numblocks": self.numblocks,
             "concatenate": self.concatenate,
             "new_axes": self.new_axes,
@@ -337,11 +338,21 @@ def __dask_distributed_pack__(self, client):
 
     @classmethod
     def __dask_distributed_unpack__(cls, state, dsk, dependencies, annotations):
+
+        # Make sure we convert list items back from tuples in `indices`.
+        # The msgpack serialization will have converted lists into
+        # tuples, and tuples may be stringified during graph
+        # materialization (bad if the item was not a key).
+        indices = [
+            list(ind) if is_list else ind
+            for ind, is_list in zip(state["indices"], state["is_list"])
+        ]
+
         raw, raw_deps = make_blockwise_graph(
             state["func"],
             state["output"],
             state["output_indices"],
-            *state["indices"],
+            *indices,
             new_axes=state["new_axes"],
             numblocks=state["numblocks"],
             concatenate=state["concatenate"],
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -359,8 +359,12 @@ def _cull_dependencies(self, keys, parts_out=None):
         for part in parts_out:
             out = self.inputs[part]
             for k in range(self.nsplits):
-                _part = inp_part_map[insert(out, self.stage, k)]
-                deps[(self.name, part)].add((self.name_input, _part))
+                _inp = insert(out, self.stage, k)
+                _part = inp_part_map[_inp]
+                if self.stage == 0 and _part >= self.npartitions_input:
+                    deps[(self.name, part)].add(("group-" + self.name, _inp, "empty"))
+                else:
+                    deps[(self.name, part)].add((self.name_input, _part))
         return deps
 
     def _cull(self, parts_out):


Test patch:
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
@@ -83,6 +83,33 @@ def test_futures_to_delayed_dataframe(c):
         ddf = dd.from_delayed([1, 2])
 
 
+@pytest.mark.parametrize("fuse", [True, False])
+def test_fused_blockwise_dataframe_merge(c, fuse):
+    pd = pytest.importorskip("pandas")
+    dd = pytest.importorskip("dask.dataframe")
+
+    # Generate two DataFrames with more partitions than
+    # the `max_branch` default used for shuffling (32).
+    # We need a multi-stage shuffle to cover #7178 fix.
+    size = 35
+    df1 = pd.DataFrame({"x": range(size), "y": range(size)})
+    df2 = pd.DataFrame({"x": range(size), "z": range(size)})
+    ddf1 = dd.from_pandas(df1, npartitions=size) + 10
+    ddf2 = dd.from_pandas(df2, npartitions=5) + 10
+    df1 += 10
+    df2 += 10
+
+    with dask.config.set({"optimization.fuse.active": fuse}):
+        ddfm = ddf1.merge(ddf2, on=["x"], how="left")
+        ddfm.head()  # https://github.com/dask/dask/issues/7178
+        dfm = ddfm.compute().sort_values("x")
+        # We call compute above since `sort_values` is not
+        # supported in `dask.dataframe`
+    dd.utils.assert_eq(
+        dfm, df1.merge(df2, on=["x"], how="left").sort_values("x"), check_index=False
+    )
+
+
 def test_futures_to_delayed_bag(c):
     db = pytest.importorskip("dask.bag")
     L = [1, 2, 3]