diff --git a/distributed/shuffle/_arrow.py b/distributed/shuffle/_arrow.py index 93cd8deccd..4a330422aa 100644 --- a/distributed/shuffle/_arrow.py +++ b/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 diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py index 608ac2ea4b..435a086f68 100644 --- a/distributed/shuffle/_shuffle.py +++ b/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"]) diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py index b7f18865e5..9b7df78a12 100644 --- a/distributed/shuffle/tests/test_shuffle.py +++ b/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}):