Internal cleanup of P2P code (#1154)

dask · Oct 25, 2024 · 9992c63 · 9992c63
1 parent da840e8
commit 9992c63
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 11 deletions.
diff --git a/dask_expr/_merge.py b/dask_expr/_merge.py
@@ -607,9 +607,8 @@ def _lower(self):
         return None
 
     def _layer(self) -> dict:
-        from distributed.shuffle._core import ShuffleId, barrier_key
+        from distributed.shuffle._core import ShuffleId, barrier_key, p2p_barrier
         from distributed.shuffle._merge import merge_unpack
-        from distributed.shuffle._shuffle import shuffle_barrier
 
         dsk = {}
         token_left = _tokenize_deterministic(
@@ -667,9 +666,9 @@ def _layer(self) -> dict:
                 self.right_index,
             )
 
-        dsk[_barrier_key_left] = (shuffle_barrier, token_left, transfer_keys_left)
+        dsk[_barrier_key_left] = (p2p_barrier, token_left, transfer_keys_left)
         dsk[_barrier_key_right] = (
-            shuffle_barrier,
+            p2p_barrier,
             token_right,
             transfer_keys_right,
         )

diff --git a/dask_expr/_shuffle.py b/dask_expr/_shuffle.py
@@ -568,12 +568,8 @@ def _meta(self):
         return self.frame._meta.drop(columns=self.partitioning_index)
 
     def _layer(self):
-        from distributed.shuffle._shuffle import (
-            ShuffleId,
-            barrier_key,
-            shuffle_barrier,
-            shuffle_unpack,
-        )
+        from distributed.shuffle._core import p2p_barrier
+        from distributed.shuffle._shuffle import ShuffleId, barrier_key, shuffle_unpack
 
         dsk = {}
         token = self._name.split("-")[-1]
@@ -604,7 +600,7 @@ def _layer(self):
                 True,
             )
 
-        dsk[_barrier_key] = (shuffle_barrier, token, transfer_keys)
+        dsk[_barrier_key] = (p2p_barrier, token, transfer_keys)
 
         # TODO: Decompose p2p Into transfer/barrier + unpack
         name = self._name