from __future__ import annotations

from io import BytesIO
from typing import TYPE_CHECKING

from packaging.version import parse

if TYPE_CHECKING:
    import pandas as pd
    import pyarrow as pa


def check_dtype_support(meta_input: pd.DataFrame) -> None:
    import pandas as pd

    for name in meta_input:
        column = meta_input[name]
        # FIXME: PyArrow does not support complex numbers: https://issues.apache.org/jira/browse/ARROW-638
        if pd.api.types.is_complex_dtype(column):
            raise TypeError(
                f"p2p does not support data of type '{column.dtype}' found in column '{name}'."
            )
        # FIXME: PyArrow does not support sparse data: https://issues.apache.org/jira/browse/ARROW-8679
        if isinstance(column.dtype, pd.SparseDtype):
            raise TypeError("p2p does not support sparse data found in column '{name}'")


def check_minimal_arrow_version() -> None:
    """Verify that the the correct version of pyarrow is installed to support
    the P2P extension.

    Raises a RuntimeError in case pyarrow is not installed or installed version
    is not recent enough.
    """
    # First version to introduce Table.sort_by
    minversion = "7.0.0"
    try:
        import pyarrow as pa
    except ImportError:
        raise RuntimeError(f"P2P shuffling requires pyarrow>={minversion}")

    if parse(pa.__version__) < parse(minversion):
        raise RuntimeError(
            f"P2P shuffling requires pyarrow>={minversion} but only found {pa.__version__}"
        )


def convert_partition(data: bytes, meta: pd.DataFrame) -> pd.DataFrame:
    import pyarrow as pa

    from dask.dataframe.dispatch import from_pyarrow_table_dispatch

    file = BytesIO(data)
    end = len(data)
    shards = []
    while file.tell() < end:
        sr = pa.RecordBatchStreamReader(file)
        shards.append(sr.read_all())
    table = pa.concat_tables(shards, promote=True)

    df = from_pyarrow_table_dispatch(meta, table, self_destruct=True)
    return df.astype(meta.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)


def serialize_table(table: pa.Table) -> bytes:
    import pyarrow as pa

    stream = pa.BufferOutputStream()
    with pa.ipc.new_stream(stream, table.schema) as writer:
        writer.write_table(table)
    return stream.getvalue().to_pybytes()


def deserialize_table(buffer: bytes) -> pa.Table:
    import pyarrow as pa

    with pa.ipc.open_stream(pa.py_buffer(buffer)) as reader:
        return reader.read_all()

distributed/shuffle/_arrow.py
from __future__ import annotations

from io import BytesIO
from typing import TYPE_CHECKING

from packaging.version import parse

if TYPE_CHECKING:
    import pandas as pd
    import pyarrow as pa


def check_dtype_support(meta_input: pd.DataFrame) -> None:
    import pandas as pd

    for name in meta_input:
        column = meta_input[name]
        # FIXME: PyArrow does not support complex numbers: https://issues.apache.org/jira/browse/ARROW-638
        if pd.api.types.is_complex_dtype(column):
            raise TypeError(
                f"p2p does not support data of type '{column.dtype}' found in column '{name}'."
            )
        # FIXME: PyArrow does not support sparse data: https://issues.apache.org/jira/browse/ARROW-8679
        if isinstance(column.dtype, pd.SparseDtype):
            raise TypeError("p2p does not support sparse data found in column '{name}'")


def check_minimal_arrow_version() -> None:
    """Verify that the the correct version of pyarrow is installed to support
    the P2P extension.

    Raises a RuntimeError in case pyarrow is not installed or installed version
    is not recent enough.
    """
    # First version to introduce Table.sort_by
    minversion = "7.0.0"
    try:
        import pyarrow as pa
    except ImportError:
        raise RuntimeError(f"P2P shuffling requires pyarrow>={minversion}")

    if parse(pa.__version__) < parse(minversion):
        raise RuntimeError(
            f"P2P shuffling requires pyarrow>={minversion} but only found {pa.__version__}"
        )


def convert_partition(data: bytes, meta: pd.DataFrame) -> pd.DataFrame:
    import pandas as pd
    import pyarrow as pa

    from dask.dataframe.dispatch import from_pyarrow_table_dispatch

    file = BytesIO(data)
    end = len(data)
    shards = []
    while file.tell() < end:
        sr = pa.RecordBatchStreamReader(file)
        shards.append(sr.read_all())
    table = pa.concat_tables(shards, promote=True)

    def default_types_mapper(pyarrow_dtype: pa.DataType) -> object:
        # Avoid converting strings from `string[pyarrow]` to `string[python]`
        # if we have *some* `string[pyarrow]`
        if (
            pyarrow_dtype in {pa.large_string(), pa.string()}
            and pd.StringDtype("pyarrow") in meta.dtypes.values
        ):
            return pd.StringDtype("pyarrow")
        return None

    df = from_pyarrow_table_dispatch(
        meta, table, self_destruct=True, types_mapper=default_types_mapper
    )
    return df.astype(meta.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), promote=True
    )


def serialize_table(table: pa.Table) -> bytes:
    import pyarrow as pa

    stream = pa.BufferOutputStream()
    with pa.ipc.new_stream(stream, table.schema) as writer:
        writer.write_table(table)
    return stream.getvalue().to_pybytes()


def deserialize_table(buffer: bytes) -> pa.Table:
    import pyarrow as pa

    with pa.ipc.open_stream(pa.py_buffer(buffer)) as reader:
        return reader.read_all()

distributed/shuffle/_arrow.py
diff --git a/distributed/shuffle/_arrow.py b/distributed/shuffle/_arrow.py
index eec104573ba..794d88d0a49 100644
--- a/distributed/shuffle/_arrow.py
+++ b/distributed/shuffle/_arrow.py
@@ -79,7 +79,9 @@ 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)
+    return pa.concat_tables(
+        (deserialize_table(buffer) for buffer in data), promote=True
+    )
 
 
 def serialize_table(table: pa.Table) -> bytes:
diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
index ebaa7a05e46..0e73cd5efde 100644
--- a/distributed/shuffle/tests/test_shuffle.py
+++ b/distributed/shuffle/tests/test_shuffle.py
@@ -18,6 +18,7 @@
 from distributed.shuffle._core import ShuffleId, ShuffleRun, barrier_key
 from distributed.worker import Status
 
+np = pytest.importorskip("numpy")
 pd = pytest.importorskip("pandas")
 dd = pytest.importorskip("dask.dataframe")
 
@@ -2057,6 +2058,28 @@ async def test_handle_null_partitions_p2p_shuffling(c, s, *workers):
     await check_scheduler_cleanup(s)
 
 
+@gen_cluster(client=True)
+async def test_handle_null_partitions_p2p_shuffling_2(c, s, a, b):
+    def make_partition(i):
+        """Return null column for one partition"""
+        if i % 2 == 1:
+            return pd.DataFrame({"a": np.random.random(10), "b": [None] * 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])
+    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)
+    await check_scheduler_cleanup(s)
+
+
 @gen_cluster(client=True)
 async def test_set_index_p2p(c, s, *workers):
     df = pd.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8], "b": 1})

P2P with null partitions fails
I came across a use case where P2P shuffling fails when a column in a partition has all null values. Here's a minimal reproducer:

```python
import numpy as np
import pandas as pd
import dask.dataframe as dd
from dask.distributed import Client

def make_partition(i):
    """Return null column for one partition"""
    if i == 1:
        return pd.DataFrame({"a": np.random.random(10), "b": None})
    return pd.DataFrame({"a": np.random.random(10), "b": np.random.random(10)})


if __name__ == "__main__":
    with Client() as client:
        ddf = dd.from_map(make_partition, range(10))
        result = ddf.set_index("a", shuffle="p2p").compute()
        print(result)
```

which raises the following error:

```
Traceback (most recent call last):
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/shuffle/_shuffle.py", line 96, in shuffle_barrier
    return _get_worker_plugin().barrier(id, run_ids)
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/shuffle/_worker_plugin.py", line 925, in barrier
    result = sync(self.worker.loop, self._barrier, shuffle_id, run_ids)
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/utils.py", line 426, in sync
    raise exc.with_traceback(tb)
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/utils.py", line 399, in f
    result = yield future
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/tornado/gen.py", line 767, in run
    value = future.result()
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/shuffle/_worker_plugin.py", line 689, in _barrier
    await shuffle.barrier()
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/shuffle/_worker_plugin.py", line 116, in barrier
    await self.scheduler.shuffle_barrier(id=self.id, run_id=self.run_id)
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/core.py", line 1374, in send_recv_from_rpc
    return await send_recv(comm=comm, op=key, **kwargs)
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/core.py", line 1158, in send_recv
    raise exc.with_traceback(tb)
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/core.py", line 930, in _handle_comm
    result = await result
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/shuffle/_scheduler_plugin.py", line 139, in barrier
    await self.scheduler.broadcast(
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/scheduler.py", line 6169, in broadcast
    results = await All(
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/utils.py", line 252, in All
    result = await tasks.next()
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/scheduler.py", line 6147, in send_message
    resp = await send_recv(
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/core.py", line 1160, in send_recv
    raise Exception(response["exception_text"])
Exception: ArrowInvalid('Schema at index 1 was different: \na: double\nb: null\n_partitions: int64\n__index_level_0__: int64\nvs\na: double\nb: double\n_partitions: int64\n__index_level_0__: int64')

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/james/projects/dask/dask/test-p2p-shuffle.py", line 16, in <module>
    result = ddf.set_index("a", shuffle="p2p").compute()
  File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/distributed/shuffle/_shuffle.py", line 98, in shuffle_barrier
    raise RuntimeError(f"shuffle_barrier failed during shuffle {id}") from e
RuntimeError: shuffle_barrier failed during shuffle 2b30bc4838ba6b632ee7d432b2b31dc8
```

Interestingly this snippet _usually_ fails, but I happened to notice it sometimes runs successfully which is in itself is also interesting 

```
                 b
a
0.015788  0.677673
0.019857  0.481580
0.027898  0.564877
0.031679  0.442530
0.048167  0.990417
...            ...
0.957410  0.651139
0.969251       NaN
0.976877  0.369628
0.984942       NaN
0.999345  0.926310

[100 rows x 1 columns]
```

cc @hendrikmakait for visibility 
