Pandas-only P2P shuffling

.github/workflows/tests.yaml

@@ -83,12 +83,12 @@ jobs:
           - os: ubuntu-latest
             environment: mindeps
             label: pandas
-            extra_packages: [numpy=1.21, pandas=1.3, pyarrow=7, pyarrow-hotfix]
+            extra_packages: [numpy=1.21, pandas=1.3]
             partition: "ci1"
           - os: ubuntu-latest
             environment: mindeps
             label: pandas
-            extra_packages: [numpy=1.21, pandas=1.3, pyarrow=7, pyarrow-hotfix]
+            extra_packages: [numpy=1.21, pandas=1.3]
             partition: "not ci1"
 
           - os: ubuntu-latest

continuous_integration/environment-3.10.yaml

@@ -45,8 +45,8 @@ dependencies:
   - zict  # overridden by git tip below
   - zstandard
   - pip:
-      - git+https://github.com/dask/dask
-      - git+https://github.com/dask-contrib/dask-expr
+      - git+https://github.com/hendrikmakait/dask@p2p-pandas
+      - git+https://github.com/hendrikmakait/dask-expr@p2p-pandas
       - git+https://github.com/dask/zict
       # Revert after https://github.com/dask/distributed/issues/8614 is fixed
       # - git+https://github.com/dask/s3fs

continuous_integration/environment-3.11.yaml

@@ -45,8 +45,8 @@ dependencies:
   - zict  # overridden by git tip below
   - zstandard
   - pip:
-      - git+https://github.com/dask/dask
-      - git+https://github.com/dask-contrib/dask-expr
+      - git+https://github.com/hendrikmakait/dask@p2p-pandas
+      - git+https://github.com/hendrikmakait/dask-expr@p2p-pandas
       - git+https://github.com/dask/zict
       # Revert after https://github.com/dask/distributed/issues/8614 is fixed
       # - git+https://github.com/dask/s3fs

continuous_integration/environment-3.12.yaml

@@ -45,8 +45,8 @@ dependencies:
   - zict  # overridden by git tip below
   - zstandard
   - pip:
-      - git+https://github.com/dask/dask
-      - git+https://github.com/dask-contrib/dask-expr
+      - git+https://github.com/hendrikmakait/dask@p2p-pandas
+      - git+https://github.com/hendrikmakait/dask-expr@p2p-pandas
       - git+https://github.com/dask/zict
       # Revert after https://github.com/dask/distributed/issues/8614 is fixed
       # - git+https://github.com/dask/s3fs

continuous_integration/environment-3.9.yaml

@@ -52,7 +52,7 @@ dependencies:
   - zict
   - zstandard
   - pip:
-      - git+https://github.com/dask/dask
-      - git+https://github.com/dask-contrib/dask-expr
+      - git+https://github.com/hendrikmakait/dask@p2p-pandas
+      - git+https://github.com/hendrikmakait/dask-expr@p2p-pandas
       - git+https://github.com/dask/crick  # Only tested here
       - keras

continuous_integration/environment-mindeps.yaml

@@ -22,7 +22,7 @@ dependencies:
   # Distributed depends on the latest version of Dask
   - pip
   - pip:
-    - git+https://github.com/dask/dask
+    - git+https://github.com/hendrikmakait/dask@p2p-pandas
   # test dependencies
   - pytest
   - pytest-cov

distributed/protocol/serialize.py

@@ -854,6 +854,17 @@ def _deserialize_memoryview(header, frames):
     return out
 
 
+@dask_serialize.register(PickleBuffer)
+def _serialize_picklebuffer(obj):
+    return _serialize_memoryview(obj.raw())
+
+
+@dask_deserialize.register(PickleBuffer)
+def _deserialize_picklebuffer(header, frames):
+    out = _deserialize_memoryview(header, frames)
+    return PickleBuffer(out)
+
+
 #########################
 # Descend into __dict__ #
 #########################

distributed/shuffle/__init__.py

@@ -1,14 +1,12 @@
 from __future__ import annotations
 
-from distributed.shuffle._arrow import check_minimal_arrow_version
 from distributed.shuffle._merge import HashJoinP2PLayer, hash_join_p2p
 from distributed.shuffle._rechunk import rechunk_p2p
 from distributed.shuffle._scheduler_plugin import ShuffleSchedulerPlugin
 from distributed.shuffle._shuffle import P2PShuffleLayer, rearrange_by_column_p2p
 from distributed.shuffle._worker_plugin import ShuffleWorkerPlugin
 
 __all__ = [
-    "check_minimal_arrow_version",
     "hash_join_p2p",
     "HashJoinP2PLayer",
     "P2PShuffleLayer",