chore: import upstream snapshot with attribution

python/ray/util/dask/tests/test_dask_optimization.py

@@ -0,0 +1,83 @@
+import sys
+from unittest import mock
+
+import dask
+import dask.dataframe as dd
+import numpy as np
+import pandas as pd
+import pytest
+from packaging.version import Version
+
+from ray.tests.conftest import *  # noqa
+from ray.util.dask import dataframe_optimize
+
+try:
+    import dask_expr  # noqa: F401
+
+    DASK_EXPR_INSTALLED = True
+except ImportError:
+    DASK_EXPR_INSTALLED = False
+    pass
+
+if Version(dask.__version__) < Version("2025.1") and not DASK_EXPR_INSTALLED:
+    from dask.dataframe.shuffle import SimpleShuffleLayer
+
+    from ray.util.dask.optimizations import (
+        MultipleReturnSimpleShuffleLayer,
+        rewrite_simple_shuffle_layer,
+    )
+
+pytestmark = pytest.mark.skipif(
+    Version(dask.__version__) >= Version("2025.1") or DASK_EXPR_INSTALLED,
+    reason="Skip dask tests for Dask 2025.1+",
+)
+
+
+def test_rewrite_simple_shuffle_layer(ray_start_regular_shared):
+    npartitions = 10
+    df = dd.from_pandas(
+        pd.DataFrame(
+            np.random.randint(0, 100, size=(100, 2)), columns=["age", "grade"]
+        ),
+        npartitions=npartitions,
+    )
+    # We set max_branch=npartitions in order to ensure that the task-based
+    # shuffle happens in a single stage, which is required in order for our
+    # optimization to work.
+    a = df.set_index(["age"], shuffle="tasks", max_branch=npartitions)
+
+    dsk = a.__dask_graph__()
+    keys = a.__dask_keys__()
+    assert any(type(v) is SimpleShuffleLayer for k, v in dsk.layers.items())
+    dsk = rewrite_simple_shuffle_layer(dsk, keys)
+    assert all(type(v) is not SimpleShuffleLayer for k, v in dsk.layers.items())
+    assert any(
+        type(v) is MultipleReturnSimpleShuffleLayer for k, v in dsk.layers.items()
+    )
+
+
+@mock.patch("ray.util.dask.optimizations.rewrite_simple_shuffle_layer")
+def test_dataframe_optimize(mock_rewrite, ray_start_regular_shared):
+    def side_effect(dsk, keys):
+        return rewrite_simple_shuffle_layer(dsk, keys)
+
+    mock_rewrite.side_effect = side_effect
+    with dask.config.set(dataframe_optimize=dataframe_optimize):
+        npartitions = 10
+        df = dd.from_pandas(
+            pd.DataFrame(
+                np.random.randint(0, 100, size=(100, 2)), columns=["age", "grade"]
+            ),
+            npartitions=npartitions,
+        )
+        # We set max_branch=npartitions in order to ensure that the task-based
+        # shuffle happens in a single stage, which is required in order for our
+        # optimization to work.
+        a = df.set_index(["age"], shuffle="tasks", max_branch=npartitions).compute()
+
+    assert mock_rewrite.call_count == 2
+    assert a.index.is_monotonic_increasing
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))