code: 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
  
file_path: ray/python/ray/util/dask/tests/test_dask_optimization.py
imports: import dask
  import dask.dataframe
  from dask.dataframe.shuffle import SimpleShuffleLayer
  from unittest import mock
  import numpy
  import pandas
  import pytest
  from ray.tests.conftest import *
  from ray.util.dask import dataframe_optimize
  from ray.util.dask.optimizations import rewrite_simple_shuffle_layer, MultipleReturnSimpleShuffleLayer
label_desc: Use image node_id 2 for calling a global function with example usage: test_dataframe_optimize(mock_rewrite, ray_start_regular_shared) and returns: rewrite_simple_shuffle_layer