dask/tests/test_distributed.py::test_map_partitions_df_input failing

[jrbourbeau]

dask/tests/test_distributed.py::test_map_partitions_df_input is failing on main with

RuntimeError: shuffle_unpack failed during shuffle 53a3327765b9171a0d39308129460022

See the full traceback in this CI build

cc @hendrikmakait

xref #10334

[jrbourbeau]

Looks like we're not properly handling how we pass meta around

>   return df.astype(meta.dtypes)
E   AttributeError: 'Serialize' object has no attribute 'dtypes'

../distributed/distributed/shuffle/_arrow.py:59: AttributeError

EDIT: Something simple like

diff --git a/distributed/shuffle/_arrow.py b/distributed/shuffle/_arrow.py
index 6a0bc81e..84620c07 100644
--- a/distributed/shuffle/_arrow.py
+++ b/distributed/shuffle/_arrow.py
@@ -4,6 +4,7 @@ from io import BytesIO
 from typing import TYPE_CHECKING

 from packaging.version import parse
+from distributed.protocol import Serialize

 if TYPE_CHECKING:
     import pandas as pd
@@ -56,6 +57,8 @@ def convert_partition(data: bytes, meta: pd.DataFrame) -> pd.DataFrame:
         shards.append(sr.read_all())
     table = pa.concat_tables(shards)
     df = table.to_pandas(self_destruct=True)
+    if isinstance(meta, Serialize):
+        meta = meta.data
     return df.astype(meta.dtypes)

fixes the problem. But I've not thought deeply yet about why we're getting a Serialize instance here in the first place. There may be a better fix

[rjzamora]

But I've not thought deeply yet about why we're getting a Serialize instance here in the first place.

Pretty sure this all comes back to #8581 (I see use_processes=False in test_map_partitions_df_input). Though I expected the pickle changes to fix this, so this hunch could be wrong.