Repository: dask/distributed

Issue:
Object columns cause P2P shuffling errors after v2023.9.2
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**:

Since v2023.9.2, using object columns with `dataframe.convert-string` disabled causes a `P2P shuffling <ID> failed during unpack phase` error stemming from a PyArrow error `Unsupported cast from string to null using function cast_null`. I have automatic string conversions disabled because I sometimes have need to nest arrays as object columns and I don't want them converted to strings.


**Minimal Complete Verifiable Example**:

```python
import dask
import dask.distributed as dd
import dask.dataframe as ddf
import numpy as np
import pandas as pd


if __name__ == '__main__':
    dask.config.set({"dataframe.convert-string": False})
    df = pd.DataFrame({
        'ints': [1, 2, 3],
        'arrays': [
            np.asarray([1, 2, 3]),
            np.asarray([4, 5, 6]),
            np.asarray([7, 8, 9]),
        ],
    })

    with dd.Client(processes=True) as client:
        ddf.from_pandas(
            df,
            npartitions=1,
        ).groupby(
            'ints',
        ).apply(
            lambda x: x,
            meta=df.iloc[:0],
        ).compute()

```

**Anything else we need to know?**:

I bisected the issue and found https://github.com/dask/distributed/commit/e57d1c55 to be the cause, specifically this line: https://github.com/dask/distributed/blob/e57d1c55c0a3ccf96785b75987a2fb87c96ddd27/distributed/shuffle/_arrow.py#L119
At this point it has arrays of the correct PyArrow types but the schema has `null` for the object columns. If I remove the `_copy_table()` function and instead add `promote=True` to the calls to `pa.concat_tables()`, the example above no longer errors. 

I suspect the goal of using a single schema object was to avoid any schema collision issues, but maybe in the cases where the schema contains nulls promote should be used instead like the old `convert_partition()` function did?

I have also attempted to update my code to use nested arrays declared with a `pd.ArrowDtype`, but that ironically causes different arrows in PyArrow itself when it attempts to parse the datatype back in after serialization.

**Environment**:

- Dask version: 2023.9.2
- Python version: Python 3.11.5
- Operating System: Debian Bookworm (Docker Container)
- Install method (conda, pip, source): pip & source for bisect



Discussion:
We ran into the same issue. For us it only occurs if we use strings in the columns. This is the minimum example I came up with:
```
import dask
import dask.dataframe as dd
import pandas as pd
from dask.distributed import Client


if __name__ == '__main__':
    dask.config.set({"dataframe.convert-string": False})
    client = Client("tcp://192.168.179.28:9003")

    d1 = dd.from_pandas(pd.DataFrame({'x':[1,2], 'y':[3, 4]}), chunksize=1000)
    d2 = dd.from_pandas(pd.DataFrame({'x':[1,2], 'z':['hi', 'hey']}), chunksize=1000)

    d1.merge(d2, how="outer", on="x", shuffle='p2p').compute()
```

In my case the error only happens if I use how="outer", with how="inner" there is no crash.

**Environment**

- Dask version: 2023.9.2 and 2023.9.3
- Python version: Python 3.11.5
- Operating System: Debian Bookworm (Docker Container)
- Install method (conda, pip, source): pip

P2P shuffling is breaking multiple things. Earlier, the merge operation with left_index=True and right_index=True was not working, and now this error is occurring.
@aiudirog 
Is there any way to have nested array in a column without disabling dataframe.convert-string?

Resolving patch:
diff --git a/distributed/shuffle/_arrow.py b/distributed/shuffle/_arrow.py
--- a/distributed/shuffle/_arrow.py
+++ b/distributed/shuffle/_arrow.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from packaging.version import parse
 
@@ -80,15 +80,12 @@ def deserialize_table(buffer: bytes) -> pa.Table:
         return reader.read_all()
 
 
-def read_from_disk(path: Path, meta: pd.DataFrame) -> tuple[Any, int]:
+def read_from_disk(path: Path) -> tuple[list[pa.Table], int]:
     import pyarrow as pa
 
-    from dask.dataframe.dispatch import pyarrow_schema_dispatch
-
     batch_size = parse_bytes("1 MiB")
     batch = []
     shards = []
-    schema = pyarrow_schema_dispatch(meta, preserve_index=True)
 
     with pa.OSFile(str(path), mode="rb") as f:
         size = f.seek(0, whence=2)
@@ -103,17 +100,17 @@ def read_from_disk(path: Path, meta: pd.DataFrame) -> tuple[Any, int]:
 
             if offset - prev >= batch_size:
                 table = pa.concat_tables(batch)
-                shards.append(_copy_table(table, schema))
+                shards.append(_copy_table(table))
                 batch = []
                 prev = offset
     if batch:
         table = pa.concat_tables(batch)
-        shards.append(_copy_table(table, schema))
+        shards.append(_copy_table(table))
     return shards, size
 
 
-def _copy_table(table: pa.Table, schema: pa.Schema) -> pa.Table:
+def _copy_table(table: pa.Table) -> pa.Table:
     import pyarrow as pa
 
     arrs = [pa.concat_arrays(column.chunks) for column in table.columns]
-    return pa.table(data=arrs, schema=schema)
+    return pa.table(data=arrs, schema=table.schema)
diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py
--- a/distributed/shuffle/_shuffle.py
+++ b/distributed/shuffle/_shuffle.py
@@ -499,8 +499,8 @@ def _(partition_id: int, meta: pd.DataFrame) -> pd.DataFrame:
     def _get_assigned_worker(self, id: int) -> str:
         return self.worker_for[id]
 
-    def read(self, path: Path) -> tuple[Any, int]:
-        return read_from_disk(path, self.meta)
+    def read(self, path: Path) -> tuple[pa.Table, int]:
+        return read_from_disk(path)
 
 
 @dataclass(frozen=True)


Test patch:
diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
--- a/distributed/shuffle/tests/test_shuffle.py
+++ b/distributed/shuffle/tests/test_shuffle.py
@@ -1129,7 +1129,7 @@ def __init__(self, value: int) -> None:
 
     out = {}
     for k in range(npartitions):
-        shards, _ = read_from_disk(tmp_path / str(k), meta)
+        shards, _ = read_from_disk(tmp_path / str(k))
         out[k] = convert_shards(shards, meta)
 
     shuffled_df = pd.concat(df for df in out.values())
@@ -2100,7 +2100,7 @@ async def test_replace_stale_shuffle(c, s, a, b):
 
 
 @gen_cluster(client=True)
-async def test_handle_null_partitions_p2p_shuffling(c, s, *workers):
+async def test_handle_null_partitions_p2p_shuffling(c, s, a, b):
     data = [
         {"companies": [], "id": "a", "x": None},
         {"companies": [{"id": 3}, {"id": 5}], "id": "b", "x": None},
@@ -2113,8 +2113,8 @@ async def test_handle_null_partitions_p2p_shuffling(c, s, *workers):
     result = await c.compute(ddf)
     dd.assert_eq(result, df)
 
-    await c.close()
-    await asyncio.gather(*[check_worker_cleanup(w) for w in workers])
+    await check_worker_cleanup(a)
+    await check_worker_cleanup(b)
     await check_scheduler_cleanup(s)
 
 
@@ -2133,7 +2133,35 @@ def make_partition(i):
     result = await result
     expected = await expected
     dd.assert_eq(result, expected)
-    del result
+
+    await check_worker_cleanup(a)
+    await check_worker_cleanup(b)
+    await check_scheduler_cleanup(s)
+
+
+@gen_cluster(client=True)
+async def test_handle_object_columns_p2p(c, s, a, b):
+    with dask.config.set({"dataframe.convert-string": False}):
+        df = pd.DataFrame(
+            {
+                "a": [1, 2, 3],
+                "b": [
+                    np.asarray([1, 2, 3]),
+                    np.asarray([4, 5, 6]),
+                    np.asarray([7, 8, 9]),
+                ],
+                "c": ["foo", "bar", "baz"],
+            }
+        )
+
+        ddf = dd.from_pandas(
+            df,
+            npartitions=2,
+        )
+        shuffled = ddf.shuffle(on="a")
+
+        result = await c.compute(shuffled)
+    dd.assert_eq(result, df)
 
     await check_worker_cleanup(a)
     await check_worker_cleanup(b)