chore: import upstream snapshot with attribution

python/ray/data/tests/test_randomize_block_order.py

@@ -0,0 +1,37 @@
+import pytest
+
+from ray.data._internal.execution.operators.base_physical_operator import (
+    AllToAllOperator,
+)
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.logical.interfaces import LogicalPlan
+from ray.data._internal.logical.operators import RandomizeBlocks
+from ray.data._internal.planner import create_planner
+from ray.data._internal.random_config import RandomSeedConfig
+from ray.data.context import DataContext
+from ray.data.tests.test_util import get_parquet_read_logical_op
+
+
+def test_randomize_blocks_operator(ray_start_regular_shared):
+    ctx = DataContext.get_current()
+
+    planner = create_planner()
+    read_op = get_parquet_read_logical_op()
+    op = RandomizeBlocks(
+        seed_config=RandomSeedConfig(seed=0),
+        input_dependencies=[read_op],
+    )
+    plan = LogicalPlan(op, ctx)
+    physical_plan, _ = planner.plan(plan)
+    physical_op = physical_plan.dag
+
+    assert op.name == "RandomizeBlockOrder"
+    assert isinstance(physical_op, AllToAllOperator)
+    assert len(physical_op.input_dependencies) == 1
+    assert isinstance(physical_op.input_dependencies[0], MapOperator)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))