chore: import upstream snapshot with attribution

python/ray/util/dask/optimizations.py

@@ -0,0 +1,149 @@
+import warnings
+
+import dask
+from dask import core
+from dask.dataframe.core import _concat
+from dask.highlevelgraph import HighLevelGraph
+
+from .scheduler import MultipleReturnFunc, multiple_return_get
+
+try:
+    from dask.dataframe.optimize import optimize
+    from dask.dataframe.shuffle import SimpleShuffleLayer, shuffle_group
+except ImportError:
+    # SimpleShuffleLayer doesn't exist in this version of Dask.
+    # This is the case for dask>=2025.1.0.
+    SimpleShuffleLayer = None
+try:
+    import dask_expr  # noqa: F401
+
+    SimpleShuffleLayer = None
+except ImportError:
+    pass
+
+
+if SimpleShuffleLayer is not None:
+
+    class MultipleReturnSimpleShuffleLayer(SimpleShuffleLayer):
+        @classmethod
+        def clone(cls, layer: SimpleShuffleLayer):
+            # TODO(Clark): Probably don't need this since SimpleShuffleLayer
+            # implements __copy__() and the shallow clone should be enough?
+            return cls(
+                name=layer.name,
+                column=layer.column,
+                npartitions=layer.npartitions,
+                npartitions_input=layer.npartitions_input,
+                ignore_index=layer.ignore_index,
+                name_input=layer.name_input,
+                meta_input=layer.meta_input,
+                parts_out=layer.parts_out,
+                annotations=layer.annotations,
+            )
+
+        def __repr__(self):
+            return (
+                f"MultipleReturnSimpleShuffleLayer<name='{self.name}', "
+                f"npartitions={self.npartitions}>"
+            )
+
+        def __reduce__(self):
+            attrs = [
+                "name",
+                "column",
+                "npartitions",
+                "npartitions_input",
+                "ignore_index",
+                "name_input",
+                "meta_input",
+                "parts_out",
+                "annotations",
+            ]
+            return (
+                MultipleReturnSimpleShuffleLayer,
+                tuple(getattr(self, attr) for attr in attrs),
+            )
+
+        def _cull(self, parts_out):
+            return MultipleReturnSimpleShuffleLayer(
+                self.name,
+                self.column,
+                self.npartitions,
+                self.npartitions_input,
+                self.ignore_index,
+                self.name_input,
+                self.meta_input,
+                parts_out=parts_out,
+            )
+
+        def _construct_graph(self):
+            """Construct graph for a simple shuffle operation."""
+
+            shuffle_group_name = "group-" + self.name
+            shuffle_split_name = "split-" + self.name
+
+            dsk = {}
+            n_parts_out = len(self.parts_out)
+            for part_out in self.parts_out:
+                # TODO(Clark): Find better pattern than in-scheduler concat.
+                _concat_list = [
+                    (shuffle_split_name, part_out, part_in)
+                    for part_in in range(self.npartitions_input)
+                ]
+                dsk[(self.name, part_out)] = (_concat, _concat_list, self.ignore_index)
+                for _, _part_out, _part_in in _concat_list:
+                    dsk[(shuffle_split_name, _part_out, _part_in)] = (
+                        multiple_return_get,
+                        (shuffle_group_name, _part_in),
+                        _part_out,
+                    )
+                    if (shuffle_group_name, _part_in) not in dsk:
+                        dsk[(shuffle_group_name, _part_in)] = (
+                            MultipleReturnFunc(
+                                shuffle_group,
+                                n_parts_out,
+                            ),
+                            (self.name_input, _part_in),
+                            self.column,
+                            0,
+                            self.npartitions,
+                            self.npartitions,
+                            self.ignore_index,
+                            self.npartitions,
+                        )
+
+            return dsk
+
+    def rewrite_simple_shuffle_layer(dsk, keys):
+        if not isinstance(dsk, HighLevelGraph):
+            dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
+        else:
+            dsk = dsk.copy()
+
+        layers = dsk.layers.copy()
+        for key, layer in layers.items():
+            if type(layer) is SimpleShuffleLayer:
+                dsk.layers[key] = MultipleReturnSimpleShuffleLayer.clone(layer)
+        return dsk
+
+    def dataframe_optimize(dsk, keys, **kwargs):
+        if not isinstance(keys, (list, set)):
+            keys = [keys]
+        keys = list(core.flatten(keys))
+
+        if not isinstance(dsk, HighLevelGraph):
+            dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
+
+        dsk = rewrite_simple_shuffle_layer(dsk, keys=keys)
+        return optimize(dsk, keys, **kwargs)
+
+else:
+
+    def dataframe_optimize(dsk, keys, **kwargs):
+        warnings.warn(
+            "Custom dataframe shuffle optimization only works on "
+            "dask>=2024.11.0,<2025.1.0, you are on version "
+            f"{dask.__version__}."
+            "Doing no additional optimization aside from the default one."
+        )
+        return None