Repository: dask/distributed

Issue:
DataFrame P2P Shuffling with stable ordering
The P2P algorithm as is currently does not strictly guarantee ordering. This can be problematic for some order sensitive operations like groupby + first (https://github.com/dask/dask/issues/10034) or for a `drop_duplicates` with `keep` (https://github.com/dask/dask/issues/10708)

It's a little work but should be possible to get P2P to be stable


Discussion:
For the sake of documenting and being precise: When talking about _stable_ ordering, the only thing we can guarantee with P2P is stable ordering between rows of the same shuffle key (i.e., the combination of values of the rows/index we shuffle on). With shuffling as a hashing-based operation, any ordering between keys is impossible.

Resolving patch:
diff --git a/distributed/shuffle/_arrow.py b/distributed/shuffle/_arrow.py
--- a/distributed/shuffle/_arrow.py
+++ b/distributed/shuffle/_arrow.py
@@ -13,6 +13,9 @@
     import pyarrow as pa
 
 
+_INPUT_PARTITION_ID_COLUMN = "__input_partition_id__"
+
+
 def check_dtype_support(meta_input: pd.DataFrame) -> None:
     import pandas as pd
 
@@ -66,6 +69,8 @@ def convert_shards(shards: list[pa.Table], meta: pd.DataFrame) -> pd.DataFrame:
     from dask.dataframe.dispatch import from_pyarrow_table_dispatch
 
     table = concat_tables(shards)
+    table = table.sort_by(_INPUT_PARTITION_ID_COLUMN)
+    table = table.drop([_INPUT_PARTITION_ID_COLUMN])
 
     df = from_pyarrow_table_dispatch(meta, table, self_destruct=True)
     reconciled_dtypes = {}
@@ -87,10 +92,33 @@ def convert_shards(shards: list[pa.Table], meta: pd.DataFrame) -> pd.DataFrame:
     return df.astype(reconciled_dtypes, copy=False)
 
 
-def list_of_buffers_to_table(data: list[bytes]) -> pa.Table:
+def buffers_to_table(data: list[tuple[int, bytes]]) -> pa.Table:
+    import numpy as np
+    import pyarrow as pa
+
     """Convert a list of arrow buffers and a schema to an Arrow Table"""
 
-    tables = (deserialize_table(buffer) for buffer in data)
+    def _create_input_partition_id_array(
+        table: pa.Table, input_partition_id: int
+    ) -> pa.ChunkedArray:
+        arrays = (
+            np.full((batch.num_rows,), input_partition_id)
+            for batch in table.to_batches()
+        )
+        return pa.chunked_array(arrays)
+
+    tables = (
+        (input_partition_id, deserialize_table(buffer))
+        for input_partition_id, buffer in data
+    )
+    tables = (
+        table.append_column(
+            _INPUT_PARTITION_ID_COLUMN,
+            _create_input_partition_id_array(table, input_partition_id),
+        )
+        for input_partition_id, table in tables
+    )
+
     return concat_tables(tables)
 
 
diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py
--- a/distributed/shuffle/_shuffle.py
+++ b/distributed/shuffle/_shuffle.py
@@ -30,11 +30,11 @@
 from distributed.exceptions import Reschedule
 from distributed.metrics import context_meter
 from distributed.shuffle._arrow import (
+    buffers_to_table,
     check_dtype_support,
     check_minimal_arrow_version,
     convert_shards,
     deserialize_table,
-    list_of_buffers_to_table,
     read_from_disk,
     serialize_table,
 )
@@ -475,7 +475,7 @@ async def _receive(self, data: list[tuple[int, bytes]]) -> None:
         filtered = []
         for d in data:
             if d[0] not in self.received:
-                filtered.append(d[1])
+                filtered.append(d)
                 self.received.add(d[0])
                 self.total_recvd += sizeof(d)
         del data
@@ -489,8 +489,10 @@ async def _receive(self, data: list[tuple[int, bytes]]) -> None:
             self._exception = e
             raise
 
-    def _repartition_buffers(self, data: list[bytes]) -> dict[NDIndex, bytes]:
-        table = list_of_buffers_to_table(data)
+    def _repartition_buffers(
+        self, data: list[tuple[int, bytes]]
+    ) -> dict[NDIndex, bytes]:
+        table = buffers_to_table(data)
         groups = split_by_partition(table, self.column)
         assert len(table) == sum(map(len, groups.values()))
         del data


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
@@ -44,8 +44,8 @@
 from distributed.core import ConnectionPool
 from distributed.scheduler import TaskState as SchedulerTaskState
 from distributed.shuffle._arrow import (
+    buffers_to_table,
     convert_shards,
-    list_of_buffers_to_table,
     read_from_disk,
     serialize_table,
 )
@@ -210,6 +210,32 @@ async def test_basic_integration(c, s, a, b, npartitions, disk):
     await check_scheduler_cleanup(s)
 
 
+@pytest.mark.parametrize("disk", [True, False])
+@gen_cluster(client=True)
+async def test_stable_ordering(c, s, a, b, disk):
+    df = dask.datasets.timeseries(
+        start="2000-01-01",
+        end="2000-02-01",
+        dtypes={"x": int, "y": int},
+        freq="10 s",
+    )
+    df["x"] = df["x"] % 19
+    df["y"] = df["y"] % 23
+    with dask.config.set(
+        {"dataframe.shuffle.method": "p2p", "distributed.p2p.disk": disk}
+    ):
+        shuffled = dd.shuffle.shuffle(df, "x")
+    result, expected = await c.compute([shuffled, df], sync=True)
+    dd.assert_eq(
+        result.drop_duplicates("x", keep="first"),
+        expected.drop_duplicates("x", keep="first"),
+    )
+
+    await check_worker_cleanup(a)
+    await check_worker_cleanup(b)
+    await check_scheduler_cleanup(s)
+
+
 @pytest.mark.parametrize("processes", [True, False])
 @gen_test()
 async def test_basic_integration_local_cluster(processes):
@@ -1145,13 +1171,13 @@ def __init__(self, value: int) -> None:
     assert set(data) == set(worker_for.cat.categories)
     assert sum(map(len, data.values())) == len(df)
 
-    batches = {worker: [serialize_table(t)] for worker, t in data.items()}
+    batches = {worker: [(0, serialize_table(t))] for worker, t in data.items()}
 
     # Typically we communicate to different workers at this stage
     # We then receive them back and reconstute them
 
     by_worker = {
-        worker: list_of_buffers_to_table(list_of_batches)
+        worker: buffers_to_table(list_of_batches)
         for worker, list_of_batches in batches.items()
     }
     assert sum(map(len, by_worker.values())) == len(df)