code: def test_flat_map(
      ray_start_regular_shared,
      restore_data_context,
  ):
      context = DataContext.get_current()
      context.foo = 70002
      ds = ray.data.range(1).flat_map(
          lambda x: [{"id": DataContext.get_current().foo}]
      )
      assert ds.take_all()[0]["id"] == 70002
  
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 4 for calling a global function with example usage: test_flat_map(ray_start_regular_shared, restore_data_context) without return types