code: def _test_write_large_data(
      tmp_path, ext, write_fn, read_fn, use_bytes, write_kwargs=None
  ):
      # Test 2G input with single task
      num_blocks_per_task = 200
      block_size = 10 * 1024 * 1024
  
      ds = ray.data.read_datasource(
          RandomBytesDatasource(),
          parallelism=1,
          num_batches_per_task=num_blocks_per_task,
          row_size=block_size,
          use_bytes=use_bytes,
      )
  
      # This should succeed without OOM.
      # https://github.com/ray-project/ray/pull/37966.
      out_dir = os.path.join(tmp_path, ext)
      write_kwargs = {} if write_kwargs is None else write_kwargs
      write_fn(ds, out_dir, **write_kwargs)
  
      max_heap_memory = (
          ds._write_ds._get_stats_summary().get_max_heap_memory()
      )
      assert max_heap_memory < (num_blocks_per_task * block_size / 2), (
          max_heap_memory,
          ext,
      )
  
      # Make sure we can read out a record.
      if read_fn is not None:
          assert read_fn(out_dir).count() == num_blocks_per_task
  
imports: import os
  import time
  from dataclasses import astuple, dataclass
  import numpy
  import pandas
  import pyarrow
  import pytest
  import ray
  from ray.data import Dataset
  from ray.data._internal.arrow_block import ArrowBlockAccessor
  from ray.data._internal.lazy_block_list import LazyBlockList
  from ray.data.block import BlockMetadata
  from ray.data.datasource import Datasource
  from ray.data.datasource.csv_datasource import CSVDatasource
  from ray.data.datasource.datasource import Reader, ReadTask
  from ray.data.tests.conftest import CoreExecutionMetrics, assert_blocks_expected_in_plasma, assert_core_execution_metrics_equals, get_initial_core_execution_metrics_snapshot
  from ray.tests.conftest import *
label_desc: Use image node_id 6 for calling a global function with example usage: _test_write_large_data(tmp_path, ext, write_fn, read_fn, use_bytes, write_kwargs) without return types