Repository: dask/dask

Issue:
`DataFrame.map_partitions` fails when passing in auxiliary Dask collections on distributed only
**Minimal Complete Verifiable Example**:

```python
import dask
import dask.dataframe as dd
import dask.array as da

df = dask.datasets.timeseries(freq="1d")
arr = da.ones((8, 9), chunks=2)

mapped = df.map_partitions(lambda df, arr: df.assign(x=df.x + arr.mean()), arr)

mapped.compute(scheduler="threads")  # works fine

import distributed
with distributed.Client() as client:
    mapped.compute()

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
...
AttributeError: 'str' object has no attribute 'mean'
```
If we make the mapper function print out `arr` with something like:
```python
def mapper(df, arr):
    print(repr(arr))
    return df.assign(x=df.x + arr.mean())
```
we see that what's passed in for `arr` is `'finalize-bca1def2-0025-456e-a35c-00b59d569556'`—the name of the key that should be getting passed in, not its actual value.

<details><summary>Some findings from debugging</summary>

I've been poking at this for a bit, but haven't been able to figure out what the root problem is yet, so I'm just leaving all these notes here in case they're helpful.

It feels like something is going wrong in Blockwise pack/unpack. But I'm also not entirely sure the Blockwise layer is reasonably constructed to begin with either. Here are a few suspicious things:

```python
def get_key(mapping, prefix: str):
    for key, v in mapping.items():
        if key.startswith(prefix):
            return v
    raise KeyError(prefix)

bw = get_key(mapped.dask.layers, "lambda-")
bw.indices
# (('make-timeseries-caa15528836dd6a39af7b7f7cdf00c52', ('.0',)),
#  ('finalize-6504dc1d-f43a-40ef-8f35-442977d9353e', None),
#  ('finalize-6504dc1d-f43a-40ef-8f35-442977d9353e', None),
#  ('make-timeseries-caa15528836dd6a39af7b7f7cdf00c52', None))
```
Why does the Blockwise layer appear to have 4 inputs? But the `finalize-` is duplicated? And the `make-timeseries` is given as both a thing with indices (`('.0',)`) and a literal (`None`)?

Even weirder: when we pack the layer, the `indices` change:
```python
>>> packed = mapped.dask.__dask_distributed_pack__(None, {})
>>> packed["layers"][-1]
{'__module__': 'dask.blockwise',
 '__name__': 'Blockwise',
 'state': {'output': 'lambda-a76ecb36114baf96c010d10066849013',
  'output_indices': ('.0',),
  'func': b"\x80\x04\x95*\x08\x00\x00\x00\x00\x00\x00\x8c\x11dask.optimization\x94\x8c\x10SubgraphCallable\x94\x93\x94(}\x94(\x8c'lambda-a76ecb36114baf96c010d10066849013\x94(\x8c\ndask.utils\x94\x8c\x05apply\x94\x93\x94\x8c\x13dask.dataframe.core\x94\x8c\x11apply_and_enforce\x94\x93\x94]\x94(\x8c\x02_0\x94\x8c\x02_1\x94e\x8c\x08builtins\x94\x8c\x04dict\x94\x93\x94]\x94(]\x94(\x8c\x05_func\x94\x8c\x17cloudpickle.cloudpickle\x94\x8c\r_builtin_type\x94\x93\x94\x8c\nLambdaType\x94\x85\x94R\x94(h\x16\x8c\x08CodeType\x94\x85\x94R\x94(K\x02K\x00K\x00K\x02K\x04KCC\x16|\x00j\x00|\x00j\x01|\x01\xa0\x02\xa1\x00\x17\x00d\x01\x8d\x01S\x00\x94N\x8c\x01x\x94\x85\x94\x86\x94\x8c\x06assign\x94h\x1e\x8c\x04mean\x94\x87\x94\x8c\x02df\x94\x8c\x03arr\x94\x86\x94\x8c\x1e<ipython-input-3-8592a9ca823d>\x94\x8c\x08<lambda>\x94K\x01C\x00\x94))t\x94R\x94}\x94(\x8c\x0b__package__\x94N\x8c\x08__name__\x94\x8c\x08__main__\x94uNNNt\x94R\x94\x8c\x1ccloudpickle.cloudpickle_fast\x94\x8c\x12_function_setstate\x94\x93\x94h1}\x94}\x94(h.h(\x8c\x0c__qualname__\x94h(\x8c\x0f__annotations__\x94}\x94\x8c\x0e__kwdefaults__\x94N\x8c\x0c__defaults__\x94N\x8c\n__module__\x94h/\x8c\x07__doc__\x94N\x8c\x0b__closure__\x94N\x8c\x17_cloudpickle_submodules\x94]\x94\x8c\x0b__globals__\x94}\x94u\x86\x94\x86R0e]\x94(\x8c\x05_meta\x94\x8c\x11pandas.core.frame\x94\x8c\tDataFrame\x94\x93\x94)\x81\x94}\x94(\x8c\x04_mgr\x94\x8c\x1epandas.core.internals.managers\x94\x8c\x0cBlockManager\x94\x93\x94)\x81\x94(]\x94(\x8c\x18pandas.core.indexes.base\x94\x8c\n_new_Index\x94\x93\x94hQ\x8c\x05Index\x94\x93\x94}\x94(\x8c\x04data\x94\x8c\x15numpy.core.multiarray\x94\x8c\x0c_reconstruct\x94\x93\x94\x8c\x05numpy\x94\x8c\x07ndarray\x94\x93\x94K\x00\x85\x94C\x01b\x94\x87\x94R\x94(K\x01K\x04\x85\x94h[\x8c\x05dtype\x94\x93\x94\x8c\x02O8\x94\x89\x88\x87\x94R\x94(K\x03\x8c\x01|\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK?t\x94b\x89]\x94(\x8c\x02id\x94\x8c\x04name\x94h\x1e\x8c\x01y\x94et\x94bhlNu\x86\x94R\x94\x8c\x1dpandas.core.indexes.datetimes\x94\x8c\x12_new_DatetimeIndex\x94\x93\x94hq\x8c\rDatetimeIndex\x94\x93\x94}\x94(hW\x8c\x1cpandas.core.arrays.datetimes\x94\x8c\rDatetimeArray\x94\x93\x94)\x81\x94}\x94(\x8c\x05_data\x94hZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x00\x85\x94hd\x8c\x02M8\x94\x89\x88\x87\x94R\x94(K\x04\x8c\x01<\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00}\x94(C\x02ns\x94K\x01K\x01K\x01t\x94\x86\x94t\x94b\x89h)t\x94b\x8c\x05_freq\x94\x8c\x1bpandas._libs.tslibs.offsets\x94\x8c\x04Hour\x94\x93\x94K\x01\x89\x86\x94R\x94\x8c\x06_dtype\x94h\x83\x8c\x06_cache\x94}\x94(\x8c\x08_ndarray\x94h\x7f\x8c\x04ndim\x94K\x01uubhl\x8c\ttimestamp\x94\x8c\x02tz\x94N\x8c\x04freq\x94h\x90u\x86\x94R\x94e]\x94(hZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x02K\x00\x86\x94hd\x8c\x02f8\x94\x89\x88\x87\x94R\x94(K\x03h\x84NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94b\x89h)t\x94bhZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x01K\x00\x86\x94hd\x8c\x02i8\x94\x89\x88\x87\x94R\x94(K\x03h\x84NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94b\x89h)t\x94bhZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x01K\x00\x86\x94hg\x89]\x94t\x94be]\x94(hShU}\x94(hWhZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x02\x85\x94hg\x89]\x94(h\x1ehmet\x94bhlNu\x86\x94R\x94hShU}\x94(hWhZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x01\x85\x94hg\x89]\x94hkat\x94bhlNu\x86\x94R\x94hShU}\x94(hWhZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x01\x85\x94hg\x89]\x94hlat\x94bhlNu\x86\x94R\x94e}\x94\x8c\x060.14.1\x94}\x94(\x8c\x04axes\x94hP\x8c\x06blocks\x94]\x94(}\x94(\x8c\x06values\x94h\x9e\x8c\x08mgr_locs\x94h\x0e\x8c\x05slice\x94\x93\x94K\x02K\x04K\x01\x87\x94R\x94u}\x94(h\xd7h\xa7h\xd8h\xdaK\x00K\x01K\x01\x87\x94R\x94u}\x94(h\xd7h\xb0h\xd8h\xdaK\x01K\x02K\x01\x87\x94R\x94ueust\x94b\x8c\x04_typ\x94\x8c\tdataframe\x94\x8c\t_metadata\x94]\x94\x8c\x05attrs\x94}\x94\x8c\x06_flags\x94}\x94\x8c\x17allows_duplicate_labels\x94\x88subee\x86\x94t\x94\x8c\x02_1\x94\x8c-finalize-6504dc1d-f43a-40ef-8f35-442977d9353e\x94\x8c\x02_2\x94h\xf0\x8c\x02_3\x94\x8c0make-timeseries-caa15528836dd6a39af7b7f7cdf00c52\x94uh\x04\x8c\x02_0\x94\x85\x94\x8c\x11subgraph_callable\x94t\x94R\x94.",
  'func_future_args': (),
  'global_dependencies': set(),
  'indices': ['make-timeseries-caa15528836dd6a39af7b7f7cdf00c52', ('.0',)],
  'is_list': [False, False],
  'numblocks': {'make-timeseries-caa15528836dd6a39af7b7f7cdf00c52': (30,)},
  'concatenate': None,
  'new_axes': {},
  'output_blocks': None,
  'dims': {'.0': 30},
  'io_deps': {}},
 'annotations': {}}
```
Notice the `'indices': ['make-timeseries-caa15528836dd6a39af7b7f7cdf00c52', ('.0',)],`. It seems that all of the literal inputs got lost.

Now if we unpack the HLG and look at one of those `lambda-` tasks, it's a task-spec dict, but the `args` is a stringified key wrapped in `Serialize`. Is that expected?

```python
>>> unpacked = dask.highlevelgraph.HighLevelGraph.__dask_distributed_unpack__(packed)
>>> up_task_ser = get_key(unpacked["dsk"], "('lambda")
>>> up_task_ser
{'function': b"\x80\x04\x95*\x08\x00\x00\x00\x00\x00\x00\x8c\x11dask.optimization\x94\x8c\x10SubgraphCallable\x94\x93\x94(}\x94(\x8c'lambda-a76ecb36114baf96c010d10066849013\x94(\x8c\ndask.utils\x94\x8c\x05apply\x94\x93\x94\x8c\x13dask.dataframe.core\x94\x8c\x11apply_and_enforce\x94\x93\x94]\x94(\x8c\x02_0\x94\x8c\x02_1\x94e\x8c\x08builtins\x94\x8c\x04dict\x94\x93\x94]\x94(]\x94(\x8c\x05_func\x94\x8c\x17cloudpickle.cloudpickle\x94\x8c\r_builtin_type\x94\x93\x94\x8c\nLambdaType\x94\x85\x94R\x94(h\x16\x8c\x08CodeType\x94\x85\x94R\x94(K\x02K\x00K\x00K\x02K\x04KCC\x16|\x00j\x00|\x00j\x01|\x01\xa0\x02\xa1\x00\x17\x00d\x01\x8d\x01S\x00\x94N\x8c\x01x\x94\x85\x94\x86\x94\x8c\x06assign\x94h\x1e\x8c\x04mean\x94\x87\x94\x8c\x02df\x94\x8c\x03arr\x94\x86\x94\x8c\x1e<ipython-input-3-8592a9ca823d>\x94\x8c\x08<lambda>\x94K\x01C\x00\x94))t\x94R\x94}\x94(\x8c\x0b__package__\x94N\x8c\x08__name__\x94\x8c\x08__main__\x94uNNNt\x94R\x94\x8c\x1ccloudpickle.cloudpickle_fast\x94\x8c\x12_function_setstate\x94\x93\x94h1}\x94}\x94(h.h(\x8c\x0c__qualname__\x94h(\x8c\x0f__annotations__\x94}\x94\x8c\x0e__kwdefaults__\x94N\x8c\x0c__defaults__\x94N\x8c\n__module__\x94h/\x8c\x07__doc__\x94N\x8c\x0b__closure__\x94N\x8c\x17_cloudpickle_submodules\x94]\x94\x8c\x0b__globals__\x94}\x94u\x86\x94\x86R0e]\x94(\x8c\x05_meta\x94\x8c\x11pandas.core.frame\x94\x8c\tDataFrame\x94\x93\x94)\x81\x94}\x94(\x8c\x04_mgr\x94\x8c\x1epandas.core.internals.managers\x94\x8c\x0cBlockManager\x94\x93\x94)\x81\x94(]\x94(\x8c\x18pandas.core.indexes.base\x94\x8c\n_new_Index\x94\x93\x94hQ\x8c\x05Index\x94\x93\x94}\x94(\x8c\x04data\x94\x8c\x15numpy.core.multiarray\x94\x8c\x0c_reconstruct\x94\x93\x94\x8c\x05numpy\x94\x8c\x07ndarray\x94\x93\x94K\x00\x85\x94C\x01b\x94\x87\x94R\x94(K\x01K\x04\x85\x94h[\x8c\x05dtype\x94\x93\x94\x8c\x02O8\x94\x89\x88\x87\x94R\x94(K\x03\x8c\x01|\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK?t\x94b\x89]\x94(\x8c\x02id\x94\x8c\x04name\x94h\x1e\x8c\x01y\x94et\x94bhlNu\x86\x94R\x94\x8c\x1dpandas.core.indexes.datetimes\x94\x8c\x12_new_DatetimeIndex\x94\x93\x94hq\x8c\rDatetimeIndex\x94\x93\x94}\x94(hW\x8c\x1cpandas.core.arrays.datetimes\x94\x8c\rDatetimeArray\x94\x93\x94)\x81\x94}\x94(\x8c\x05_data\x94hZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x00\x85\x94hd\x8c\x02M8\x94\x89\x88\x87\x94R\x94(K\x04\x8c\x01<\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00}\x94(C\x02ns\x94K\x01K\x01K\x01t\x94\x86\x94t\x94b\x89h)t\x94b\x8c\x05_freq\x94\x8c\x1bpandas._libs.tslibs.offsets\x94\x8c\x04Hour\x94\x93\x94K\x01\x89\x86\x94R\x94\x8c\x06_dtype\x94h\x83\x8c\x06_cache\x94}\x94(\x8c\x08_ndarray\x94h\x7f\x8c\x04ndim\x94K\x01uubhl\x8c\ttimestamp\x94\x8c\x02tz\x94N\x8c\x04freq\x94h\x90u\x86\x94R\x94e]\x94(hZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x02K\x00\x86\x94hd\x8c\x02f8\x94\x89\x88\x87\x94R\x94(K\x03h\x84NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94b\x89h)t\x94bhZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x01K\x00\x86\x94hd\x8c\x02i8\x94\x89\x88\x87\x94R\x94(K\x03h\x84NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94b\x89h)t\x94bhZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x01K\x00\x86\x94hg\x89]\x94t\x94be]\x94(hShU}\x94(hWhZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x02\x85\x94hg\x89]\x94(h\x1ehmet\x94bhlNu\x86\x94R\x94hShU}\x94(hWhZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x01\x85\x94hg\x89]\x94hkat\x94bhlNu\x86\x94R\x94hShU}\x94(hWhZh]K\x00\x85\x94h_\x87\x94R\x94(K\x01K\x01\x85\x94hg\x89]\x94hlat\x94bhlNu\x86\x94R\x94e}\x94\x8c\x060.14.1\x94}\x94(\x8c\x04axes\x94hP\x8c\x06blocks\x94]\x94(}\x94(\x8c\x06values\x94h\x9e\x8c\x08mgr_locs\x94h\x0e\x8c\x05slice\x94\x93\x94K\x02K\x04K\x01\x87\x94R\x94u}\x94(h\xd7h\xa7h\xd8h\xdaK\x00K\x01K\x01\x87\x94R\x94u}\x94(h\xd7h\xb0h\xd8h\xdaK\x01K\x02K\x01\x87\x94R\x94ueust\x94b\x8c\x04_typ\x94\x8c\tdataframe\x94\x8c\t_metadata\x94]\x94\x8c\x05attrs\x94}\x94\x8c\x06_flags\x94}\x94\x8c\x17allows_duplicate_labels\x94\x88subee\x86\x94t\x94\x8c\x02_1\x94\x8c-finalize-6504dc1d-f43a-40ef-8f35-442977d9353e\x94\x8c\x02_2\x94h\xf0\x8c\x02_3\x94\x8c0make-timeseries-caa15528836dd6a39af7b7f7cdf00c52\x94uh\x04\x8c\x02_0\x94\x85\x94\x8c\x11subgraph_callable\x94t\x94R\x94.",
 'args': <Serialize: ["('make-timeseries-caa15528836dd6a39af7b7f7cdf00c52', 0)"]>}
```

Lastly, if we deserialize the task, the `subgraph_callable` only takes one input, yet has these extra `_1`, `_2`, `_3` keys in it. I imagine that's where the `'finalize-'` string that's getting passed into our mapper function and causing the error is coming from:
```python
>>> up_task_deser = distributed.worker._deserialize(**up_task_ser)
>>> up_task_deser
(subgraph_callable,
 <Serialize: ["('make-timeseries-caa15528836dd6a39af7b7f7cdf00c52', 0)"]>,
 {})
>>> up_task_deser[0].inkeys
('_0',)
>>> up_task_deser[0].dsk
{'lambda-a76ecb36114baf96c010d10066849013': (<function dask.utils.apply(func, args, kwargs=None)>,
  <function dask.dataframe.core.apply_and_enforce(*args, **kwargs)>,
  ['_0', '_1'],
  (dict,
   [['_func', <function __main__.<lambda>(df, arr)>],
    ['_meta', Empty DataFrame
     Columns: [id, name, x, y]
     Index: []]])),
 '_1': 'finalize-6504dc1d-f43a-40ef-8f35-442977d9353e',
 '_2': 'finalize-6504dc1d-f43a-40ef-8f35-442977d9353e',
 '_3': 'make-timeseries-caa15528836dd6a39af7b7f7cdf00c52'}
```

This feels maybe like #7720, but not quite?

</details>

cc @rjzamora @madsbk would be great if you could take a look at this; I'm worried this could have a wide-reaching effect.

**Environment**:

- Dask version: 2021.5.0 / 636b0d664e78febc54434131d6f7f906a6e74f6a (reproduced on both)
- Python version: 3.8
- Operating System: macOS
- Install method (conda, pip, source): pip, source

Discussion:
madsbk: Rebased on main now that #7739  has been merged.

Patch:
diff --git a/dask/blockwise.py b/dask/blockwise.py
index 11f21ca6a..b4cb0159a 100644
--- 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 @@ class Blockwise(Layer):
         # 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 @@ class Blockwise(Layer):
                 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
index 38ca876c1..0f781f49b 100644
--- 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()})