Repository: dask/dask

Issue:
latest dask `map_partitions` doesn't pass list as expected.
# Reproducible example
I used the `HIGGS` dataset from UCI machine learning reprository for demonstraction purose.  But any other dataset should suffice.

The failure happens only on dask 2021.05.0, older versions work fine.  Also using tuple instead of list works fine.

``` python
from distributed import LocalCluster, Client, wait
from dask import dataframe as dd


def map_partition_fn(block, meta_list):
    # Works with  dask-2.22.0, assert error with 2021.05.0
    #
    # distributed.worker - WARNING - Compute Failed
    # Function:  execute_task
    # args:      ((subgraph_callable, [(<function read_block_from_file at 0x7f84259231f0>, <OpenFile 'higgs/HIGGS.csv'>, 640000000, 64000000, b'\n'), None, False]))
    # kwargs:    {}
    # Exception: AssertionError(<class 'str'>)
    assert isinstance(meta_list, list), type(meta_list)
    return block


def main(client):
    df: dd.DataFrame = dd.read_csv("higgs/HIGGS.csv")
    meta = [i for i in range(len(df.columns))]

    mapped = df.map_partitions(map_partition_fn, meta_list=meta)
    mapped = client.persist(mapped)
    wait(mapped)


if __name__ == "__main__":
    with LocalCluster() as cluster:
        with Client(cluster) as client:
            main(client)

```

- Dask version: 2021.05.0
- Python version: 3.8.5
- Operating System: Ubuntu 20.04.2
- Install method (conda, pip, source): pip

cc @jakirkham 


Discussion:
Thanks for raising an issue @trivialfis! Since I don't have the `HIGGS.csv` file locally on my machine, I wrote our demo timeseries dataset using `to_csv` to generate local CSV files to use

```python
    from dask.datasets import timeseries
    df = timeseries().to_csv("test_data/*.csv")
    df = dd.read_csv("test_data/*.csv")
```

Unfortunately when I run the full example with `2021.05.0` (included below), I'm not able to reproduce the `AssertionError` and the code snippet runs successfully. Do you see the same assertion error when you use the timeseries dataset? 

<details>
<summary>Details:</summary>

```python
from distributed import LocalCluster, Client, wait
from dask import dataframe as dd


def map_partition_fn(block, meta_list):
    assert isinstance(meta_list, list), type(meta_list)
    return block


def main(client):
    # df: dd.DataFrame = dd.read_csv("higgs/HIGGS.csv")
    from dask.datasets import timeseries
    df = timeseries().to_csv("test_data/*.csv")
    df = dd.read_csv("test_data/*.csv")
    meta = [i for i in range(len(df.columns))]

    mapped = df.map_partitions(map_partition_fn, meta_list=meta)
    mapped = client.persist(mapped)
    wait(mapped)


if __name__ == "__main__":
    import dask
    import distributed
    print(f"{dask.__version__ = }")
    print(f"{distributed.__version__ = }")

    with LocalCluster() as cluster:
        with Client(cluster) as client:
            main(client)
```

</details>
Hmm, your example works in my environment.  I cut out the first 10 rows from HIGGS:
[minor.csv](https://github.com/dask/dask/files/6517611/minor.csv)


And I can reproduce it with this.
Thanks @trivialfis, I'm able to reproduce using `minor.csv`. I think this is related to Dask recently disabling low-level task fusion by default for DataFrames (xref https://github.com/dask/dask/pull/7620). If I turn task fusion back on with `dask.config.set({"optimization.fuse.active": True})` then `meta_list` is a list as expected and the code snippet runs successfully

cc'ing @rjzamora in case you have any thoughts on this 
Thanks for raising this @trivialfis !  As James mentioned, the task fusion default was recently changed.  I'll look into this asap.
Thanks Rick! 😄
This does seem to be a (somewhat) long-standing Blockwise bug that is only now being exposed by the fact that we are shipping the HLG to the scheduler by default. I haven't figured out a fix just yet, but it looks like "long" `map_partitions` list arguments are [automatically converted to Delayed](https://github.com/dask/dask/blob/d6bbbb08c92652eae2820e93edc2f3fe502391d3/dask/array/core.py#L426), and we are failing to serialize/deserialize them correctly.

I'll try to poke around a bit now, but it may take me a bit of time to figure out what is going wrong (cc @madsbk - In case you have thoughts)

Resolving patch:
diff --git a/dask/blockwise.py b/dask/blockwise.py
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -259,12 +259,9 @@ def blockwise(
     new_axes = {index_subs((k,), sub)[0]: v for k, v in new_axes.items()}
 
     # Unpack dask values in non-array arguments
-    argpairs = toolz.partition(2, arrind_pairs)
-
-    # separate argpairs into two separate tuples
     inputs = []
     inputs_indices = []
-    for name, index in argpairs:
+    for name, index in toolz.partition(2, arrind_pairs):
         inputs.append(name)
         inputs_indices.append(index)
 
@@ -498,9 +495,19 @@ def __dask_distributed_pack__(
         # Embed literals in `dsk`
         keys2 = []
         indices2 = []
+        global_dependencies = set()
         for key, (val, index) in zip(keys, self.indices):
-            if index is None:  # Literal
-                dsk[key] = val
+            if index is None:
+                try:
+                    val_is_a_key = val in all_hlg_keys
+                except TypeError:  # not hashable
+                    val_is_a_key = False
+                if val_is_a_key:
+                    keys2.append(key)
+                    indices2.append((val, index))
+                    global_dependencies.add(val)
+                else:
+                    dsk[key] = val  # Literal
             else:
                 keys2.append(key)
                 indices2.append((val, index))
@@ -558,7 +565,7 @@ def __dask_distributed_pack__(
                 raise CancelledError(stringify(future.key))
 
         # All blockwise tasks will depend on the futures in `indices`
-        global_dependencies = {stringify(f.key) for f in indices_unpacked_futures}
+        global_dependencies |= {stringify(f.key) for f in indices_unpacked_futures}
 
         return {
             "output": self.output,


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
@@ -486,6 +486,24 @@ async def test_futures_in_subgraphs(c, s, a, b):
     ddf = await c.submit(dd.categorical.categorize, ddf, columns=["day"], index=False)
 
 
+@gen_cluster(client=True)
+async def test_map_partitions_da_input(c, s, a, b):
+    """Check that map_partitions can handle a dask array input"""
+    np = pytest.importorskip("numpy")
+    pd = pytest.importorskip("pandas")
+    da = pytest.importorskip("dask.array")
+    datasets = pytest.importorskip("dask.datasets")
+
+    def f(d, a):
+        assert isinstance(d, pd.DataFrame)
+        assert isinstance(a, np.ndarray)
+        return d
+
+    df = datasets.timeseries(freq="1d").persist()
+    arr = da.ones((1,), chunks=1).persist()
+    await c.compute(df.map_partitions(f, arr, meta=df._meta))
+
+
 @gen_cluster(client=True)
 async def test_annotation_pack_unpack(c, s, a, b):
     hlg = HighLevelGraph({"l1": MaterializedLayer({"n": 42})}, {"l1": set()})