code: def test_map(shutdown_only, restore_data_context):
      ray.init(
          _system_config={
              "max_direct_call_object_size": 10_000,
          },
          num_cpus=2,
          object_store_memory=int(100e6),
      )
  
      ctx = DataContext.get_current()
      ctx.target_min_block_size = 10_000 * 8
      ctx.target_max_block_size = 10_000 * 8
      num_blocks_expected = 10
      last_snapshot = get_initial_core_execution_metrics_snapshot()
  
      # Test read.
      ds = ray.data.range(100_000, parallelism=1).materialize()
      assert (
          num_blocks_expected
          <= ds.num_blocks()
          <= num_blocks_expected + 1
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          num_blocks_expected,
          block_size_expected=ctx.target_max_block_size,
      )
  
      # Test read -> map.
      # NOTE(swang): For some reason BlockBuilder's estimated memory usage when a
      # map fn is used is 2x the actual memory usage.
      ds = (
          ray.data.range(100_000, parallelism=1)
          .map(lambda row: row)
          .materialize()
      )
      assert (
          num_blocks_expected * 2
          <= ds.num_blocks()
          <= num_blocks_expected * 2 + 1
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          num_blocks_expected * 2,
          block_size_expected=ctx.target_max_block_size // 2,
      )
  
      # Test adjusted block size.
      ctx.target_max_block_size *= 2
      num_blocks_expected //= 2
  
      # Test read.
      ds = ray.data.range(100_000, parallelism=1).materialize()
      assert (
          num_blocks_expected
          <= ds.num_blocks()
          <= num_blocks_expected + 1
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          num_blocks_expected,
          block_size_expected=ctx.target_max_block_size,
      )
  
      # Test read -> map.
      ds = (
          ray.data.range(100_000, parallelism=1)
          .map(lambda row: row)
          .materialize()
      )
      assert (
          num_blocks_expected * 2
          <= ds.num_blocks()
          <= num_blocks_expected * 2 + 1
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          num_blocks_expected * 2,
          block_size_expected=ctx.target_max_block_size // 2,
      )
  
      # Setting the shuffle block size doesn't do anything for
      # map-only Datasets.
      ctx.target_shuffle_max_block_size = ctx.target_max_block_size / 2
  
      # Test read.
      ds = ray.data.range(100_000, parallelism=1).materialize()
      assert (
          num_blocks_expected
          <= ds.num_blocks()
          <= num_blocks_expected + 1
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          num_blocks_expected,
          block_size_expected=ctx.target_max_block_size,
      )
  
      # Test read -> map.
      ds = (
          ray.data.range(100_000, parallelism=1)
          .map(lambda row: row)
          .materialize()
      )
      assert (
          num_blocks_expected * 2
          <= ds.num_blocks()
          <= num_blocks_expected * 2 + 1
      )
      last_snapshot = assert_blocks_expected_in_plasma(
          last_snapshot,
          num_blocks_expected * 2,
          block_size_expected=ctx.target_max_block_size // 2,
      )
  
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 1 for calling a global function with example usage: test_map(shutdown_only, restore_data_context) without return types