Preserve Pandas 3 null-object semantics in P2P shuffle

distributed/shuffle/_arrow.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+import json
 from collections.abc import Iterable
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypedDict
 
 from dask.utils import parse_bytes
 
@@ -12,6 +13,12 @@
 
 
 _INPUT_PARTITION_ID_COLUMN = "__input_partition_id__"
+_NULL_OBJECT_COLUMNS = b"dask.null_object_columns"
+
+
+class _NullObjectInfo(TypedDict):
+    marker: str
+    object_dtype: bool
 
 
 def check_dtype_support(meta_input: pd.DataFrame) -> None:
@@ -29,10 +36,81 @@ def check_dtype_support(meta_input: pd.DataFrame) -> None:
             raise TypeError("p2p does not support sparse data found in column '{name}'")
 
 
+def preserve_null_objects(table: pa.Table, df: pd.DataFrame) -> pa.Table:
+    """Retain pandas null semantics that Arrow's null type cannot represent."""
+    import numpy as np
+    import pandas as pd
+    import pyarrow as pa
+
+    from dask.dataframe._compat import PANDAS_GE_300
+
+    if not PANDAS_GE_300 or not isinstance(df, pd.DataFrame):
+        return table
+
+    def null_kind(value: object) -> int:
+        if value is None:
+            return 0
+        if value is pd.NA:
+            return 2
+        if value is pd.NaT:
+            return 3
+        return 1  # NaN
+
+    # Use the same marker names even if input partitions reorder their columns.
+    columns: dict[str, _NullObjectInfo] = {}
+    for i, column in enumerate(sorted(df.columns.tolist())):
+        series = df[column]
+        # NumPy-backed columns already round-trip their null sentinel through
+        # Arrow. Extension arrays may need pd.NA restored after object promotion.
+        if (
+            series.dtype != object
+            and getattr(series.dtype, "na_value", None) is not pd.NA
+        ):
+            continue
+        if not table[column].null_count:
+            continue
+        marker = f"__dask_null_object_{i}__"
+        while marker in table.column_names:
+            marker += "_"
+        mask = series.isna().to_numpy()
+        # 255 denotes a non-null value; 0..3 distinguish pandas null sentinels.
+        kinds = np.full(len(series), 255, dtype=np.uint8)
+        if series.dtype == object:
+            values = series.to_numpy()[mask]
+            kinds[mask] = [null_kind(value) for value in values]
+        else:
+            kinds[mask] = 2  # Extension arrays use pd.NA, as checked above.
+        table = table.append_column(marker, pa.array(kinds))
+        columns[column] = {
+            "marker": marker,
+            "object_dtype": series.dtype == object
+            and pa.types.is_null(table[column].type),
+        }
+    if columns:
+        metadata = dict(table.schema.metadata or {})
+        metadata[_NULL_OBJECT_COLUMNS] = json.dumps(columns).encode()
+        table = table.replace_schema_metadata(metadata)
+    return table
+
+
 def concat_tables(tables: Iterable[pa.Table]) -> pa.Table:
     import pyarrow as pa
 
-    return pa.concat_tables(tables, promote_options="permissive")
+    tables = list(tables)
+    null_columns: dict[str, _NullObjectInfo] = {}
+    for table in tables:
+        for column, info in json.loads(
+            (table.schema.metadata or {}).get(_NULL_OBJECT_COLUMNS, b"{}")
+        ).items():
+            if column in null_columns:
+                info["object_dtype"] |= null_columns[column]["object_dtype"]
+            null_columns[column] = info
+    result = pa.concat_tables(tables, promote_options="permissive")
+    if null_columns:
+        metadata = dict(result.schema.metadata or {})
+        metadata[_NULL_OBJECT_COLUMNS] = json.dumps(null_columns).encode()
+        result = result.replace_schema_metadata(metadata)
+    return result
 
 
 def convert_shards(
@@ -45,11 +123,29 @@ def convert_shards(
 
     table = concat_tables(shards)
     table = table.sort_by(_INPUT_PARTITION_ID_COLUMN)
-    table = table.drop([_INPUT_PARTITION_ID_COLUMN])
+    null_columns = json.loads(
+        (table.schema.metadata or {}).get(_NULL_OBJECT_COLUMNS, b"{}")
+    )
+    # Arrow promotes null columns to the other shards' types. Pandas >=3 instead
+    # retains object dtype when concatenating all-null object and typed columns.
+    null_kinds = {
+        column: table[info["marker"]].to_numpy().copy()
+        for column, info in null_columns.items()
+        if info["object_dtype"]
+    }
+    table = table.drop(
+        [_INPUT_PARTITION_ID_COLUMN]
+        + [info["marker"] for info in null_columns.values()]
+    )
 
     if drop_column:
         meta = meta.drop(columns=partition_column)
-    df = from_pyarrow_table_dispatch(meta, table, self_destruct=True)
+    kwargs = {"integer_object_nulls": True} if null_kinds else {}
+    df = from_pyarrow_table_dispatch(meta, table, self_destruct=True, **kwargs)
+    for column, kinds in null_kinds.items():
+        df[column] = df[column].astype(object)
+        for code, value in enumerate((None, float("nan"), pd.NA, pd.NaT)):
+            df.loc[kinds == code, column] = value
     reconciled_dtypes = {}
     for column, dtype in meta.dtypes.items():
         actual = df[column].dtype

distributed/shuffle/_shuffle.py

@@ -22,6 +22,7 @@
     buffers_to_table,
     convert_shards,
     deserialize_table,
+    preserve_null_objects,
     read_from_disk,
     serialize_table,
 )
@@ -101,6 +102,7 @@ def split_by_worker(
     # FIXME: If we do not preserve the index something is corrupting the
     # bytestream such that it cannot be deserialized anymore
     t = to_pyarrow_table_dispatch(df, preserve_index=True)
+    t = preserve_null_objects(t, df)
     t = t.sort_by("_worker")
     codes = np.asarray(t["_worker"])
     t = t.drop(["_worker"])

distributed/shuffle/tests/test_shuffle.py

@@ -2430,18 +2430,22 @@ async def test_handle_null_partitions(c, s, a, b):
     await assert_scheduler_cleanup(s)
 
 
-@pytest.mark.xfail(PANDAS_GE_300, reason="dask/distributed#9184", strict=True)
+@pytest.mark.parametrize("disk", [True, False])
+@pytest.mark.parametrize("null_first", [True, False])
 @gen_cluster(client=True)
-async def test_handle_null_partitions_2(c, s, a, b):
+async def test_handle_null_partitions_2(c, s, a, b, disk, null_first):
     def make_partition(i):
         """Return null column for every other 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)})
+        values = np.arange(i * 10, (i + 1) * 10, dtype=float)
+        return pd.DataFrame(
+            {"a": values, "b": [None] * 10 if (i % 2 == 0) == null_first else values}
+        )
 
     with dask.config.set({"dataframe.convert-string": False}):
         ddf = dd.from_map(make_partition, range(5), meta={"a": float, "b": float})
-    with dask.config.set({"dataframe.shuffle.method": "p2p"}):
+    with dask.config.set(
+        {"dataframe.shuffle.method": "p2p", "distributed.p2p.storage.disk": disk}
+    ):
         out = ddf.shuffle(on="a", ignore_index=True, force=True)
     expected, result = c.compute([ddf, out])
     del out
@@ -2454,6 +2458,45 @@ def make_partition(i):
     await assert_scheduler_cleanup(s)
 
 
+@pytest.mark.parametrize("null_first", [True, False])
+@pytest.mark.parametrize(
+    "null_values", [[None, None], [np.nan, np.nan], [pd.NA, None], [pd.NaT, np.nan]]
+)
+@pytest.mark.parametrize("dtype", ["float64", "int64", "bool", "str", "Int64"])
+@pytest.mark.skipif(not PANDAS_GE_300, reason="Pandas 3 null-object concat semantics")
+def test_null_object_column_survives_arrow_concatenation(
+    tmp_path, null_first, null_values, dtype
+):
+    values = [2**60 + 1, 2**60 + 3] if dtype in ("int64", "Int64") else [1, 2]
+    if dtype == "Int64":
+        values[1] = pd.NA
+    null = pd.DataFrame({"a": [0, 1], "b": pd.Series(null_values, dtype=object)})
+    typed = pd.DataFrame({"a": [2, 3], "b": pd.Series(values, dtype=dtype)})
+    # Marker names must not collide with user columns or depend on column order.
+    null["__dask_null_object_4__"] = [4, 5]
+    typed["__dask_null_object_4__"] = [6, 7]
+    typed = typed[typed.columns[::-1]]
+    frames = [null, typed] if null_first else [typed, null]
+    # Exercise both the receive-side concat and subsequent disk-buffer concats.
+    shards = []
+    worker_for = pd.Series({0: "worker"}).astype("category")
+    for partition_id, frame in enumerate(frames):
+        frame = frame.assign(_partitions=0)
+        table = split_by_worker(frame, "_partitions", frame.iloc[:0], worker_for)[
+            "worker"
+        ]
+        shards.append(buffers_to_table([(partition_id, serialize_table(table))]))
+    path = tmp_path / "shards"
+    with path.open("ab") as f:
+        for shard in shards:
+            f.write(serialize_table(shard))
+    shards, _ = read_from_disk(path)
+    result = convert_shards(shards, typed.iloc[:0], "_partitions", False)
+    result = result.drop(columns="_partitions")
+    expected = pd.concat(frames)
+    pd.testing.assert_frame_equal(result, expected)
+
+
 @gen_cluster(client=True)
 async def test_handle_object_columns(c, s, a, b):
     with dask.config.set({"dataframe.convert-string": False}):