code: def test_shuffle(shutdown_only, restore_data_context, shuffle_op):
      ray.init(
          _system_config={
              "max_direct_call_object_size": 250,
          },
          num_cpus=2,
          object_store_memory=int(100e6),
      )
  
      # Test AllToAll and Map -> AllToAll Datasets. Check that Map inherits
      # AllToAll's target block size.
      ctx = DataContext.get_current()
      ctx.min_parallelism = 1
      ctx.target_min_block_size = 1
      mem_size = 800_000
      shuffle_fn, kwargs, fusion_supported = shuffle_op
  
      ctx.target_shuffle_max_block_size = 10_000 * 8
      num_blocks_expected = (
          mem_size // ctx.target_shuffle_max_block_size
      )
      block_size_expected = ctx.target_shuffle_max_block_size
      last_snapshot = get_initial_core_execution_metrics_snapshot()
  
      ds = shuffle_fn(ray.data.range(100_000), **kwargs).materialize()
      assert (
          num_blocks_expected
          <= ds.num_blocks()
          <= num_blocks_expected * 1.5
      )
      # map * reduce intermediate blocks + 1 metadata ref per map/reduce task.
      # If fusion is not supported, the un-fused map stage produces 1 data and 1
      # metadata per task.
      num_intermediate_blocks = (
          num_blocks_expected**2
          + num_blocks_expected * (2 if fusion_supported else 4)
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          # Dataset.sort produces some empty intermediate blocks because the
          # input range is already partially sorted.
          num_intermediate_blocks,
          # Data is written out once before map phase if fusion is disabled, once
          # during map phase, once during reduce phase.
          total_bytes_expected=mem_size * 2
          + (0 if fusion_supported else mem_size),
      )
  
      ds = shuffle_fn(
          ray.data.range(100_000).map(lambda x: x), **kwargs
      ).materialize()
      if not fusion_supported:
          # TODO(swang): For some reason BlockBuilder's estimated
          # memory usage for range(1000)->map is 2x the actual memory usage.
          # Remove once https://github.com/ray-project/ray/issues/40246 is fixed.
          num_blocks_expected = int(num_blocks_expected * 2.2)
      assert (
          num_blocks_expected
          <= ds.num_blocks()
          <= num_blocks_expected * 1.5
      )
      num_intermediate_blocks = (
          num_blocks_expected**2
          + num_blocks_expected * (2 if fusion_supported else 4)
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          # Dataset.sort produces some empty intermediate blocks because the
          # input range is already partially sorted.
          num_intermediate_blocks,
          # Data is written out once before map phase if fusion is disabled, once
          # during map phase, once during reduce phase.
          total_bytes_expected=mem_size * 2
          + (0 if fusion_supported else mem_size),
      )
  
      ctx.target_shuffle_max_block_size //= 2
      num_blocks_expected = (
          mem_size // ctx.target_shuffle_max_block_size
      )
      block_size_expected = ctx.target_shuffle_max_block_size
  
      ds = shuffle_fn(ray.data.range(100_000), **kwargs).materialize()
      assert (
          num_blocks_expected
          <= ds.num_blocks()
          <= num_blocks_expected * 1.5
      )
      num_intermediate_blocks = (
          num_blocks_expected**2
          + num_blocks_expected * (2 if fusion_supported else 4)
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          num_intermediate_blocks,
          total_bytes_expected=mem_size * 2
          + (0 if fusion_supported else mem_size),
      )
  
      ds = shuffle_fn(
          ray.data.range(100_000).map(lambda x: x), **kwargs
      ).materialize()
      if not fusion_supported:
          num_blocks_expected = int(num_blocks_expected * 2.2)
          block_size_expected //= 2.2
      assert (
          num_blocks_expected
          <= ds.num_blocks()
          <= num_blocks_expected * 1.5
      )
      num_intermediate_blocks = (
          num_blocks_expected**2
          + num_blocks_expected * (2 if fusion_supported else 4)
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          num_intermediate_blocks,
          total_bytes_expected=mem_size * 2
          + (0 if fusion_supported else mem_size),
      )
  
      # Setting target max block size does not affect map ops when there is a
      # shuffle downstream.
      ctx.target_max_block_size = ctx.target_shuffle_max_block_size * 2
      ds = shuffle_fn(
          ray.data.range(100_000).map(lambda x: x), **kwargs
      ).materialize()
      assert (
          num_blocks_expected
          <= ds.num_blocks()
          <= num_blocks_expected * 1.5
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          num_intermediate_blocks,
          total_bytes_expected=mem_size * 2
          + (0 if fusion_supported else mem_size),
      )
  
file_path: ray/python/ray/data/tests/test_block_sizing.py
imports: import pytest
  import ray
  from ray.data import Dataset
  from ray.data.context import DataContext
  from ray.data.tests.conftest import *
  from ray.data.tests.conftest import assert_blocks_expected_in_plasma, get_initial_core_execution_metrics_snapshot
  from ray.tests.conftest import *
label_desc: Use image node_id 2 for calling a global function with example usage: test_shuffle(shutdown_only, restore_data_context, shuffle_op) without return types