Shuffle use pyarrow more broadly

distributed/shuffle/_disk.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
 import contextlib
+import io
 import pathlib
 import shutil
 import threading
 from collections.abc import Callable, Generator, Iterable
 from contextlib import contextmanager
 from typing import Any
 
+from packaging.version import parse as parse_version
 from toolz import concat
 
 from distributed.metrics import context_meter, thread_time
@@ -18,6 +20,19 @@
 from distributed.utils import Deadline, empty_context, log_errors, nbytes
 
 
+def open_file(*args: Any, **kwargs: Any) -> io.IOBase:
+    open_file = open
+    try:
+        import pyarrow as pa
+
+        if parse_version(pa.__version__) >= parse_version("15.0.0"):
+            # Only >15 support append mode
+            open_file = pa.OSFile
+    except ModuleNotFoundError:
+        pass
+    return open_file(*args, **kwargs)
+
+
 class ReadWriteLock:
     _condition: threading.Condition
     _n_reads: int
@@ -177,7 +192,7 @@
             if self._closed:
                 raise RuntimeError("Already closed")
 
-            with open(self.directory / str(id), mode="ab") as f:
+            with open_file(str(self.directory / str(id)), mode="ab") as f:
                 f.writelines(frames)
 
         context_meter.digest_metric("disk-write", 1, "count")

distributed/shuffle/_shuffle.py

@@ -367,9 +367,9 @@ def split_by_partition(
     Split data into many arrow batches, partitioned by final partition
     """
     import numpy as np
+    import pyarrow as pa
 
-    partitions = t.select([column]).to_pandas()[column].unique()
-    partitions.sort()
+    partitions = np.array(pa.compute.unique(t[column]).sort())
     t = t.sort_by(column)
 
     partition = np.asarray(t[column])