Repository: dask/distributed

Issue:
Categorical column turned to NaN after P2P Shuffle
**Describe the issue**:
If you P2P shuffle while your dataframe has a categorical column that column will be turned into all NaNs.

**Minimal Complete Verifiable Example**:

```python
from dask.distributed import Client
from dask.distributed import LocalCluster
import dask.dataframe as dd

cluster = LocalCluster()
client = Client(cluster)
client.cluster.scale(1)


dask_df = dd.from_dict(
    {
        "b": ["a", "shuffle", "column"],
        "test": ["apple", "pear", "citrus"],
    },
    npartitions=2,
)
dask_df.test = dask_df.test.astype('category')  
print(list(dask_df.test.unique().compute())) # ['apple', 'pear', 'citrus']
dask_df = dask_df.shuffle(on="b").persist() 
print(list(dask_df.test.unique().compute())) # [nan]

```

**Anything else we need to know?**:
Probably due to the p2p shuffling. 
`dask_df = dask_df.shuffle(on="b", shuffle="tasks").persist() `
does not have the issue. 

If the column is not categorical it also does not happen.

Might be related to
#8183 and #8165 


**Environment**:

Dask version: 2023.9.1
Python version: 3.10
Operating System: Linux
Install method (conda, pip, source): conda



Discussion:
Thanks for the issue @NickSchouten. I'm able to reproduce. 

cc @hendrikmakait @fjetter for visibility
bisect blames #7879, i.e. first affected version was already [2023.6.0](https://github.com/dask/distributed/releases/tag/2023.6.0)

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,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Iterable
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -34,8 +35,7 @@ def check_minimal_arrow_version() -> None:
     Raises a ModuleNotFoundError if pyarrow is not installed or an
     ImportError if the installed version is not recent enough.
     """
-    # First version that supports concatenating extension arrays (apache/arrow#14463)
-    minversion = "12.0.0"
+    minversion = "7.0.0"
     try:
         import pyarrow as pa
     except ModuleNotFoundError:
@@ -46,22 +46,52 @@ def check_minimal_arrow_version() -> None:
         )
 
 
-def convert_shards(shards: list[pa.Table], meta: pd.DataFrame) -> pd.DataFrame:
+def concat_tables(tables: Iterable[pa.Table]) -> pa.Table:
     import pyarrow as pa
 
+    if parse(pa.__version__) >= parse("14.0.0"):
+        return pa.concat_tables(tables, promote_options="permissive")
+    try:
+        return pa.concat_tables(tables, promote=True)
+    except pa.ArrowNotImplementedError as e:
+        if parse(pa.__version__) >= parse("12.0.0"):
+            raise e
+        raise
+
+
+def convert_shards(shards: list[pa.Table], meta: pd.DataFrame) -> pd.DataFrame:
+    import pandas as pd
+    from pandas.core.dtypes.cast import find_common_type  # type: ignore[attr-defined]
+
     from dask.dataframe.dispatch import from_pyarrow_table_dispatch
 
-    table = pa.concat_tables(shards)
+    table = concat_tables(shards)
 
     df = from_pyarrow_table_dispatch(meta, table, self_destruct=True)
-    return df.astype(meta.dtypes, copy=False)
+    reconciled_dtypes = {}
+    for column, dtype in meta.dtypes.items():
+        actual = df[column].dtype
+        if actual == dtype:
+            continue
+        # Use the specific string dtype from meta (e.g., string[pyarrow])
+        if isinstance(actual, pd.StringDtype) and isinstance(dtype, pd.StringDtype):
+            reconciled_dtypes[column] = dtype
+            continue
+        # meta might not be aware of the actual categories so the two dtype objects are not equal
+        # Also, the categories_dtype does not properly roundtrip through Arrow
+        if isinstance(actual, pd.CategoricalDtype) and isinstance(
+            dtype, pd.CategoricalDtype
+        ):
+            continue
+        reconciled_dtypes[column] = find_common_type([actual, dtype])
+    return df.astype(reconciled_dtypes, copy=False)
 
 
 def list_of_buffers_to_table(data: list[bytes]) -> pa.Table:
     """Convert a list of arrow buffers and a schema to an Arrow Table"""
-    import pyarrow as pa
 
-    return pa.concat_tables(deserialize_table(buffer) for buffer in data)
+    tables = (deserialize_table(buffer) for buffer in data)
+    return concat_tables(tables)
 
 
 def serialize_table(table: pa.Table) -> bytes:
@@ -99,18 +129,33 @@ def read_from_disk(path: Path) -> tuple[list[pa.Table], int]:
             batch.append(shard)
 
             if offset - prev >= batch_size:
-                table = pa.concat_tables(batch)
+                table = concat_tables(batch)
                 shards.append(_copy_table(table))
                 batch = []
                 prev = offset
     if batch:
-        table = pa.concat_tables(batch)
+        table = concat_tables(batch)
         shards.append(_copy_table(table))
     return shards, size
 
 
+def concat_arrays(arrays: Iterable[pa.Array]) -> pa.Array:
+    import pyarrow as pa
+
+    try:
+        return pa.concat_arrays(arrays)
+    except pa.ArrowNotImplementedError as e:
+        if parse(pa.__version__) >= parse("12.0.0"):
+            raise
+        if e.args[0].startswith("concatenation of extension"):
+            raise RuntimeError(
+                "P2P shuffling requires pyarrow>=12.0.0 to support extension types."
+            ) from e
+        raise
+
+
 def _copy_table(table: pa.Table) -> pa.Table:
     import pyarrow as pa
 
-    arrs = [pa.concat_arrays(column.chunks) for column in table.columns]
+    arrs = [concat_arrays(column.chunks) for column in table.columns]
     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
@@ -306,8 +306,6 @@ def split_by_worker(
 
     from dask.dataframe.dispatch import to_pyarrow_table_dispatch
 
-    df = df.astype(meta.dtypes, copy=False)
-
     # (cudf support) Avoid pd.Series
     constructor = df._constructor_sliced
     assert isinstance(constructor, type)


Test patch:
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -71,12 +71,12 @@ jobs:
           - os: ubuntu-latest
             environment: mindeps
             label: pandas
-            extra_packages: [numpy=1.21, pandas=1.3, pyarrow=12]
+            extra_packages: [numpy=1.21, pandas=1.3, pyarrow=7]
             partition: "ci1"
           - os: ubuntu-latest
             environment: mindeps
             label: pandas
-            extra_packages: [numpy=1.21, pandas=1.3, pyarrow=12]
+            extra_packages: [numpy=1.21, pandas=1.3, pyarrow=7]
             partition: "not ci1"
 
           - os: ubuntu-latest
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
@@ -15,6 +15,7 @@
 from unittest import mock
 
 import pytest
+from packaging.version import parse
 from tornado.ioloop import IOLoop
 
 from dask.utils import key_split
@@ -1045,13 +1046,6 @@ def __init__(self, value: int) -> None:
             [np.datetime64("2022-01-01") + i for i in range(100)],
             dtype=pd.DatetimeTZDtype(tz="Europe/Berlin"),
         ),
-        f"col{next(counter)}": pd.array(
-            [pd.Period("2022-01-01", freq="D") + i for i in range(100)],
-            dtype="period[D]",
-        ),
-        f"col{next(counter)}": pd.array(
-            [pd.Interval(left=i, right=i + 2) for i in range(100)], dtype="Interval"
-        ),
         f"col{next(counter)}": pd.array(["x", "y"] * 50, dtype="category"),
         f"col{next(counter)}": pd.array(["lorem ipsum"] * 100, dtype="string"),
         # FIXME: PyArrow does not support sparse data:
@@ -1067,6 +1061,17 @@ def __init__(self, value: int) -> None:
         # ),
     }
 
+    if parse(pa.__version__) >= parse("12.0.0"):
+        columns.update(
+            {
+                # Extension types
+                f"col{next(counter)}": pd.period_range(
+                    "2022-01-01", periods=100, freq="D"
+                ),
+                f"col{next(counter)}": pd.interval_range(start=0, end=100, freq=1),
+            }
+        )
+
     if PANDAS_GE_150:
         columns.update(
             {
@@ -2197,7 +2202,7 @@ async def test_replace_stale_shuffle(c, s, a, b):
 
 @pytest.mark.skipif(not PANDAS_GE_200, reason="requires pandas >=2.0")
 @gen_cluster(client=True)
-async def test_handle_null_partitions_p2p_shuffling(c, s, a, b):
+async def test_handle_null_partitions(c, s, a, b):
     data = [
         {"companies": [], "id": "a", "x": None},
         {"companies": [{"id": 3}, {"id": 5}], "id": "b", "x": None},
@@ -2216,7 +2221,7 @@ async def test_handle_null_partitions_p2p_shuffling(c, s, a, b):
 
 
 @gen_cluster(client=True)
-async def test_handle_null_partitions_p2p_shuffling_2(c, s, a, b):
+async def test_handle_null_partitions_2(c, s, a, b):
     def make_partition(i):
         """Return null column for every other partition"""
         if i % 2 == 1:
@@ -2237,7 +2242,7 @@ def make_partition(i):
 
 
 @gen_cluster(client=True)
-async def test_handle_object_columns_p2p(c, s, a, b):
+async def test_handle_object_columns(c, s, a, b):
     with dask.config.set({"dataframe.convert-string": False}):
         df = pd.DataFrame(
             {
@@ -2266,21 +2271,31 @@ async def test_handle_object_columns_p2p(c, s, a, b):
 
 
 @gen_cluster(client=True)
-async def test_reconcile_mismatching_partitions_p2p_shuffling(c, s, a, b):
+async def test_reconcile_partitions(c, s, a, b):
     def make_partition(i):
         """Return mismatched column types for every other partition"""
         if i % 2 == 1:
-            return pd.DataFrame({"a": np.random.random(10), "b": [True] * 10})
+            return pd.DataFrame(
+                {"a": np.random.random(10), "b": np.random.randint(1, 10, 10)}
+            )
         return pd.DataFrame({"a": np.random.random(10), "b": np.random.random(10)})
 
     ddf = dd.from_map(make_partition, range(50))
     out = ddf.shuffle(on="a", shuffle="p2p", ignore_index=True)
-    result, expected = c.compute([ddf, out])
+
+    if parse(pa.__version__) >= parse("14.0.0"):
+        result, expected = c.compute([ddf, out])
+        result = await result
+        expected = await expected
+        dd.assert_eq(result, expected)
+        del result
+    else:
+        with raises_with_cause(
+            RuntimeError, r"shuffling \w+ failed", pa.ArrowInvalid, "incompatible types"
+        ):
+            await c.compute(out)
+        await c.close()
     del out
-    result = await result
-    expected = await expected
-    dd.assert_eq(result, expected)
-    del result
 
     await check_worker_cleanup(a)
     await check_worker_cleanup(b)
@@ -2288,7 +2303,7 @@ def make_partition(i):
 
 
 @gen_cluster(client=True)
-async def test_raise_on_incompatible_partitions_p2p_shuffling(c, s, a, b):
+async def test_raise_on_incompatible_partitions(c, s, a, b):
     def make_partition(i):
         """Return incompatible column types for every other partition"""
         if i % 2 == 1:
@@ -2297,10 +2312,19 @@ def make_partition(i):
 
     ddf = dd.from_map(make_partition, range(50))
     out = ddf.shuffle(on="a", shuffle="p2p", ignore_index=True)
-    with raises_with_cause(
-        RuntimeError, "failed during transfer", ValueError, "could not convert"
-    ):
-        await c.compute(out)
+    if parse(pa.__version__) >= parse("14.0.0"):
+        with raises_with_cause(
+            RuntimeError,
+            r"shuffling \w* failed",
+            pa.ArrowTypeError,
+            "incompatible types",
+        ):
+            await c.compute(out)
+    else:
+        with raises_with_cause(
+            RuntimeError, r"shuffling \w* failed", pa.ArrowInvalid, "incompatible types"
+        ):
+            await c.compute(out)
 
     await c.close()
     await check_worker_cleanup(a)
@@ -2309,7 +2333,55 @@ def make_partition(i):
 
 
 @gen_cluster(client=True)
-async def test_set_index_p2p(c, s, *workers):
+async def test_handle_categorical_data(c, s, a, b):
+    """Regression test for https://github.com/dask/distributed/issues/8186"""
+    df = dd.from_dict(
+        {
+            "a": [1, 2, 3, 4, 5],
+            "b": [
+                "x",
+                "y",
+                "x",
+                "y",
+                "z",
+            ],
+        },
+        npartitions=2,
+    )
+    df.b = df.b.astype("category")
+    shuffled = df.shuffle("a", shuffle="p2p")
+    result, expected = await c.compute([shuffled, df], sync=True)
+    dd.assert_eq(result, expected, check_categorical=False)
+
+    await check_worker_cleanup(a)
+    await check_worker_cleanup(b)
+    await check_scheduler_cleanup(s)
+
+
+@gen_cluster(client=True)
+async def test_handle_floats_in_int_meta(c, s, a, b):
+    """Regression test for https://github.com/dask/distributed/issues/8183"""
+    df1 = pd.DataFrame(
+        {
+            "a": [1, 2],
+        },
+    )
+    df2 = pd.DataFrame(
+        {
+            "b": [1],
+        },
+    )
+    expected = df1.join(df2, how="left")
+
+    ddf1 = dd.from_pandas(df1, npartitions=1)
+    ddf2 = dd.from_pandas(df2, npartitions=1)
+    result = ddf1.join(ddf2, how="left").shuffle(on="a")
+    result = await c.compute(result)
+    dd.assert_eq(result, expected)
+
+
+@gen_cluster(client=True)
+async def test_set_index(c, s, *workers):
     df = pd.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8], "b": 1})
     ddf = dd.from_pandas(df, npartitions=3)
     ddf = ddf.set_index("a", shuffle="p2p", divisions=(1, 3, 8))
@@ -2322,7 +2394,7 @@ async def test_set_index_p2p(c, s, *workers):
     await check_scheduler_cleanup(s)
 
 
-def test_shuffle_p2p_with_existing_index(client):
+def test_shuffle_with_existing_index(client):
     df = pd.DataFrame({"a": np.random.randint(0, 3, 20)}, index=np.random.random(20))
     ddf = dd.from_pandas(
         df,
@@ -2333,7 +2405,7 @@ def test_shuffle_p2p_with_existing_index(client):
     dd.assert_eq(result, df)
 
 
-def test_set_index_p2p_with_existing_index(client):
+def test_set_index_with_existing_index(client):
     df = pd.DataFrame({"a": np.random.randint(0, 3, 20)}, index=np.random.random(20))
     ddf = dd.from_pandas(
         df,
@@ -2344,7 +2416,7 @@ def test_set_index_p2p_with_existing_index(client):
     dd.assert_eq(result, df.set_index("a"))
 
 
-def test_sort_values_p2p_with_existing_divisions(client):
+def test_sort_values_with_existing_divisions(client):
     "Regression test for #8165"
     df = pd.DataFrame(
         {"a": np.random.randint(0, 3, 20), "b": np.random.randint(0, 3, 20)}
@@ -2430,7 +2502,7 @@ async def connect(self, *args, **kwargs):
     client=True,
     config={"distributed.comm.retry.count": 0, "distributed.p2p.comm.retry.count": 0},
 )
-async def test_p2p_flaky_connect_fails_without_retry(c, s, a, b):
+async def test_flaky_connect_fails_without_retry(c, s, a, b):
     df = dask.datasets.timeseries(
         start="2000-01-01",
         end="2000-01-10",
@@ -2460,7 +2532,7 @@ async def test_p2p_flaky_connect_fails_without_retry(c, s, a, b):
     client=True,
     config={"distributed.comm.retry.count": 0, "distributed.p2p.comm.retry.count": 1},
 )
-async def test_p2p_flaky_connect_recover_with_retry(c, s, a, b):
+async def test_flaky_connect_recover_with_retry(c, s, a, b):
     df = dask.datasets.timeseries(
         start="2000-01-01",
         end="2000-01-10",