Repository: dask/dask

Issue:
Msgpack TypeError when performing dask-geopandas sjoin
**What happened**:
Performing a spatial join with `dask-geopandas` fails with a `msgpack` TypeError.

**Minimal Complete Verifiable Example**:
Not quite an MCVE, but the code below produces the error. Notebook + environment.yml [here](https://github.com/coiled/coiled-resources/blob/rrp-spatialjoin/spatial-join/debugging.ipynb).

```python
from coiled.v2 import Cluster

cluster = Cluster(
    name="spatial-join",
    software="coiled-examples/spatial-join",
    n_workers=50,
    worker_memory="16Gib",
)

from distributed import Client
client = Client(cluster)

import dask.dataframe as dd

ddf = dd.read_parquet(
    "s3://coiled-datasets/dask-book/nyc-tlc/2009-2013/*",
    engine="pyarrow",
    storage_options={"anon": True},
)

import dask_geopandas

ddf = ddf.set_geometry(
    dask_geopandas.points_from_xy(ddf, "pickup_longitude", "pickup_latitude"),
)

ddf = ddf.set_crs(4326)

import geopandas as gpd

ngbhoods = gpd.read_file(
    "CHS_2009_DOHMH_2010B/CHS_2009_DOHMH_2010B.shp"
)
ngbhoods = ngbhoods[["FIRST_UHF_", "UHF_CODE", "geometry"]]
ngbhoods = ngbhoods[:34]
ngbhoods = ngbhoods.to_crs(epsg=4326)

joined = ddf.sjoin(ngbhoods, predicate="within")

joined1.head()
```

TRACEBACK:
```
distributed.protocol.core - CRITICAL - Failed to Serialize
Traceback (most recent call last):
  File "/Users/rpelgrim/mambaforge/envs/dask-dataframe/lib/python3.9/site-packages/distributed/protocol/core.py", line 76, in dumps
    frames[0] = msgpack.dumps(msg, default=_encode_default, use_bin_type=True)
  File "/Users/rpelgrim/mambaforge/envs/dask-dataframe/lib/python3.9/site-packages/msgpack/__init__.py", line 35, in packb
    return Packer(**kwargs).pack(o)
  File "msgpack/_packer.pyx", line 294, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 300, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 297, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 291, in msgpack._cmsgpack.Packer._pack
TypeError: can not serialize 'numpy.int64' object
distributed.comm.utils - INFO - Unserializable Message: [{'op': 'update-graph-hlg', 'hlg': {'layers': [{'__module__': 'dask.blockwise', '__name__': 'Blockwise', 'state': {'output': '_set_crs-6ca0659b532951365eb2f3c74d7d5dda', 'output_indices': ('.0',), 'func': b'\x80\x04\x95\x03\x00\x01\x00\x00\x00\x00\x00\x8c\x11dask.optimization\x94\x8c\x10SubgraphCallable\x94\x93\x94(}\x94(\x8c)_set_crs-6ca0659b532951365eb2f3c74d7d5dda\x94\x8cEpoints_from_xy-set_geometry-_set_crs-6ca0659b532951365eb2f3c74d7d5dda\x94\x8c0read-parquet-520df5ddd4d8ac310c949ff50485dcbb_.
```

**Anything else we need to know?**:
- Paired with @crusaderky briefly. We confirmed that `dd.from_pandas(ngbhoods, npartitions=1).compute()` works so it can make a roundtrip fine. His first instinct is that something is likely off in `dask.dataframe` where it's letting `numpy.int64` slip to `msgpack`.

**Environment**:

- Dask version: 2022.01.0
- Python version: 3.9.7
- Operating System: MacOS
- Install method (conda, pip, source): conda





Discussion:
cc @rjzamora @madsbk, looks like an HLG serialization issue
I ran into this as well in the past (also with dask-geopandas). I did a fix on the dask-geopandas side (ensuring we always convert the np.int64 to a python int already on our side), https://github.com/geopandas/dask-geopandas/pull/113, but in the end I didn't merge it because I could not reproduce the issue ...

Resolving patch:
diff --git a/dask/blockwise.py b/dask/blockwise.py
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -791,10 +791,10 @@ def cull(
         # collect a set of required output blocks (tuples), and
         # only construct graph for these blocks in `make_blockwise_graph`
 
-        output_blocks = set()
+        output_blocks: set[tuple[int, ...]] = set()
         for key in keys:
             if key[0] == self.output:
-                output_blocks.add(key[1:])
+                output_blocks.add(tuple(map(int, key[1:])))
         culled_deps = self._cull_dependencies(all_hlg_keys, output_blocks)
         out_size_iter = (self.dims[i] for i in self.output_indices)
         if prod(out_size_iter) != len(culled_deps):


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
@@ -491,6 +491,28 @@ def test_blockwise_different_optimization(c):
     np.testing.assert_equal(y_value, expected)
 
 
+def test_blockwise_cull_allows_numpy_dtype_keys(c):
+    # Regression test for https://github.com/dask/dask/issues/9072
+    da = pytest.importorskip("dask.array")
+    np = pytest.importorskip("numpy")
+
+    # Create a multi-block array.
+    x = da.ones((100, 100), chunks=(10, 10))
+
+    # Make a layer that pulls a block out of the array, but
+    # refers to that block using a numpy.int64 for the key rather
+    # than a python int.
+    name = next(iter(x.dask.layers))
+    block = {("block", 0, 0): (name, np.int64(0), np.int64(1))}
+    dsk = HighLevelGraph.from_collections("block", block, [x])
+    arr = da.Array(dsk, "block", ((10,), (10,)), dtype=x.dtype)
+
+    # Stick with high-level optimizations to force serialization of
+    # the blockwise layer.
+    with dask.config.set({"optimization.fuse.active": False}):
+        da.assert_eq(np.ones((10, 10)), arr, scheduler=c)
+
+
 @gen_cluster(client=True)
 async def test_combo_of_layer_types(c, s, a, b):
     """Check pack/unpack of a HLG that has every type of Layers!"""