Don't switch order of randomize_block_order and write operators (ray-…

…project#32555)

It's incorrect to the switch the order of randomize_block_order and write operators: the write operator will always have to be the last operator.

python/ray/data/_internal/plan.py

@@ -1165,7 +1165,7 @@ def _reorder_stages(stages: List[Stage]) -> List[Stage]:
  reorder_buf.append(s)
  else:
  # Barrier: flush the reorder buffer.
- if isinstance(s, AllToAllStage):
+ if isinstance(s, AllToAllStage) or s.name == "write":
  output.extend(reorder_buf)
  reorder_buf = []
  output.append(s)

python/ray/data/tests/test_optimize.py

@@ -374,6 +374,19 @@ def test_write_fusion(ray_start_regular_shared, tmp_path):
  assert "MapBatches(<lambda>)->write" in stats, stats
 
 
+def test_write_doesnt_reorder_randomize_block(ray_start_regular_shared, tmp_path):
+ path = os.path.join(tmp_path, "out")
+ ds = ray.data.range(100).randomize_block_order().map_batches(lambda x: x)
+ ds.write_csv(path)
+ stats = ds._write_ds.stats()
+
+ # The randomize_block_order will switch order with the following map_batches,
+ # but not the tailing write operator.
+ assert "read->MapBatches(<lambda>)" in stats, stats
+ assert "randomize_block_order" in stats, stats
+ assert "write" in stats, stats
+
+
 def test_optimize_fuse(ray_start_regular_shared):
  context = DatasetContext.get_current()