code: def test_streaming_split(
      ray_start_regular_shared,
      restore_data_context,
  ):
      # Tests that custom DataContext can be properly propagated
      # when using `streaming_split()`.
      block_size = 123 * 1024 * 1024
      data_context = DataContext.get_current()
      data_context.target_max_block_size = block_size
      data_context.set_config("foo", "bar")
  
      def f(x):
          assert (
              DataContext.get_current().target_max_block_size
              == block_size
          )
          assert DataContext.get_current().get_config("foo") == "bar"
          return x
  
      num_splits = 2
      splits = (
          ray.data.range(10, parallelism=10)
          .map(f)
          .streaming_split(num_splits)
      )
  
      @ray.remote(num_cpus=0)
      def consume(split):
          for _ in split.iter_rows():
              pass
  
      assert (
          ray.get([consume.remote(split) for split in splits])
          == [None] * num_splits
      )
  
file_path: ray/python/ray/data/tests/test_context_propagation.py
imports: import pandas
  import pytest
  import ray
  from ray._private.test_utils import run_string_as_driver
  from ray.data.block import BlockMetadata
  from ray.data.context import DataContext
  from ray.data.datasource import Datasource, ReadTask
  from ray.tests.conftest import *
label_desc: Use image node_id 7 for calling a global function with example usage: test_streaming_split(ray_start_regular_shared, restore_data_context) and returns: x