P2P rechunking (#9939)

hendrikmakait · web-flow · commit a58a2a74228d · 2023-02-24T19:02:20.000+01:00
diff --git a/dask/array/core.py b/dask/array/core.py
@@ -2747,7 +2747,12 @@ def squeeze(self, axis=None):
         return squeeze(self, axis)
 
     def rechunk(
-        self, chunks="auto", threshold=None, block_size_limit=None, balance=False
+        self,
+        chunks="auto",
+        threshold=None,
+        block_size_limit=None,
+        balance=False,
+        algorithm=None,
     ):
         """Convert blocks in dask array x for new chunks.
 
@@ -2759,7 +2764,7 @@ def rechunk(
         """
         from dask.array.rechunk import rechunk  # avoid circular import
 
-        return rechunk(self, chunks, threshold, block_size_limit, balance)
+        return rechunk(self, chunks, threshold, block_size_limit, balance, algorithm)
 
     @property
     def real(self):
diff --git a/dask/array/rechunk.py b/dask/array/rechunk.py
@@ -225,7 +225,14 @@ def intersect_chunks(old_chunks, new_chunks):
     return cross
 
 
-def rechunk(x, chunks="auto", threshold=None, block_size_limit=None, balance=False):
+def rechunk(
+    x,
+    chunks="auto",
+    threshold=None,
+    block_size_limit=None,
+    balance=False,
+    algorithm=None,
+):
     """
     Convert blocks in dask array x for new chunks.
 
@@ -249,6 +256,9 @@ def rechunk(x, chunks="auto", threshold=None, block_size_limit=None, balance=Fal
         This means ``balance=True`` will remove any small leftover chunks, so
         using ``x.rechunk(chunks=len(x) // N, balance=True)``
         will almost certainly result in ``N`` chunks.
+    algorithm: {'tasks', 'p2p'}, optional.
+        Algorithm to use.
+
 
     Examples
     --------
@@ -314,13 +324,24 @@ def rechunk(x, chunks="auto", threshold=None, block_size_limit=None, balance=Fal
         if new != old and not math.isnan(old) and not math.isnan(new):
             raise ValueError("Provided chunks are not consistent with shape")
 
-    steps = plan_rechunk(
-        x.chunks, chunks, x.dtype.itemsize, threshold, block_size_limit
-    )
-    for c in steps:
-        x = _compute_rechunk(x, c)
+    algorithm = algorithm or config.get("array.rechunk.algorithm")
 
-    return x
+    if algorithm == "tasks":
+        steps = plan_rechunk(
+            x.chunks, chunks, x.dtype.itemsize, threshold, block_size_limit
+        )
+        for c in steps:
+            x = _compute_rechunk(x, c)
+
+        return x
+
+    elif algorithm == "p2p":
+        from distributed.shuffle import rechunk_p2p
+
+        return rechunk_p2p(x, chunks)
+
+    else:
+        raise NotImplementedError(f"Unknown rechunking algorithm '{algorithm}'")
 
 
 def _number_of_blocks(chunks):
@@ -542,15 +563,13 @@ def plan_rechunk(
     if isinstance(block_size_limit, str):
         block_size_limit = parse_bytes(block_size_limit)
 
-    ndim = len(new_chunks)
-    steps = []
-    has_nans = [any(math.isnan(y) for y in x) for x in old_chunks]
+    has_nans = (any(math.isnan(y) for y in x) for x in old_chunks)
 
-    if ndim <= 1 or not all(new_chunks) or any(has_nans):
+    if len(new_chunks) <= 1 or not all(new_chunks) or any(has_nans):
         # Trivial array / unknown dim => no need / ability for an intermediate
-        return steps + [new_chunks]
+        return [new_chunks]
 
-    # Make it a number ef elements
+    # Make it a number of elements
     block_size_limit /= itemsize
 
     # Fix block_size_limit if too small for either old_chunks or new_chunks
@@ -565,6 +584,7 @@ def plan_rechunk(
 
     current_chunks = old_chunks
     first_pass = True
+    steps = []
 
     while True:
         graph_size = estimate_graph_size(current_chunks, new_chunks)
diff --git a/dask/dask-schema.yaml b/dask/dask-schema.yaml
@@ -97,6 +97,15 @@ properties:
           Backend to use for supported array-creation functions.
           Default is "numpy".
 
+      rechunk:
+        type: object
+        properties:
+          algorithm:
+            type: string
+            description: |
+              The algorithm to use for rechunking. Must be either "tasks" or "p2p";
+              default is "tasks". Using "p2p" requires a distributed cluster.
+
       svg:
         type: object
         properties:
diff --git a/dask/dask.yaml b/dask/dask.yaml
@@ -17,6 +17,8 @@ dataframe:
 
 array:
   backend: "numpy"  # Backend array library for input IO and data creation
+  rechunk:
+    algorithm: "tasks"  # Rechunking algorithm to use
   svg:
     size: 120  # pixels
   slicing:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
@@ -115,7 +115,7 @@ def _determine_shuffle(shuffle, split_out):
             # For more context, see
             # https://github.com/dask/dask/pull/9826/files#r1072395307
             # https://github.com/dask/distributed/issues/5502
-            return shuffle or config.get("shuffle", None) or "tasks"
+            return config.get("shuffle", None) or "tasks"
         else:
             return False
     return shuffle
