code: def test_dataset(
      shutdown_only,
      restore_data_context,
      compute,
  ):
      def identity_fn(x):
          return x
  
      def empty_fn(x):
          return {}
  
      class IdentityClass:
          def __call__(self, x):
              return x
  
      class EmptyClass:
          def __call__(self, x):
              return {}
  
      ctx = ray.data.DataContext.get_current()
      # 1MiB.
      ctx.target_max_block_size = 1024 * 1024
  
      if compute == "tasks":
          compute = ray.data._internal.compute.TaskPoolStrategy()
          identity_func = identity_fn
          empty_func = empty_fn
          func_name = "identity_fn"
      else:
          compute = ray.data.ActorPoolStrategy()
          identity_func = IdentityClass
          empty_func = EmptyClass
          func_name = "IdentityClass"
  
      ray.shutdown()
      # We need at least 2 CPUs to run a actorpool streaming
      ray.init(num_cpus=2, object_store_memory=1e9)
      # Test 10 tasks, each task returning 10 blocks, each block has 1 row and each
      # row has 1024 bytes.
      num_blocks_per_task = 10
      num_tasks = 10
  
      @ray.remote
      def warmup():
          return np.zeros(ctx.target_max_block_size, dtype=np.uint8)
  
      last_snapshot = get_initial_core_execution_metrics_snapshot()
      ds = ray.data.read_datasource(
          RandomBytesDatasource(),
          parallelism=num_tasks,
          num_batches_per_task=num_blocks_per_task,
          row_size=ctx.target_max_block_size,
      )
      # Note the following calls to ds will not fully execute it.
      assert ds.schema() is not None
      assert ds.count() == num_blocks_per_task * num_tasks
      assert ds.num_blocks() == num_tasks
      assert (
          ds.size_bytes()
          >= 0.7
          * ctx.target_max_block_size
          * num_blocks_per_task
          * num_tasks
      )
      last_snapshot = assert_core_execution_metrics_equals(
          CoreExecutionMetrics(
              task_count={
                  "_get_datasource_or_legacy_reader": 1,
                  "_execute_read_task_split": 1,
              },
              object_store_stats={
                  "cumulative_created_plasma_bytes": lambda count: True,
                  "cumulative_created_plasma_objects": lambda count: True,
              },
          ),
          last_snapshot,
      )
  
      # Too-large blocks will get split to respect target max block size.
      map_ds = ds.map_batches(identity_func, compute=compute)
      map_ds = map_ds.materialize()
      num_blocks_expected = num_tasks * num_blocks_per_task
      assert map_ds.num_blocks() == num_blocks_expected
      assert_core_execution_metrics_equals(
          CoreExecutionMetrics(
              task_count={
                  "MapWorker(ReadRandomBytes->MapBatches"
                  f"({func_name})).get_location": lambda count: True,
                  "_MapWorker.__init__": lambda count: True,
                  "_MapWorker.get_location": lambda count: True,
                  f"ReadRandomBytes->MapBatches({func_name})": num_tasks,
              },
          ),
          last_snapshot,
      )
      assert_blocks_expected_in_plasma(
          last_snapshot,
          num_blocks_expected,
          block_size_expected=ctx.target_max_block_size,
      )
  
      # Blocks smaller than requested batch size will get coalesced.
      map_ds = ds.map_batches(
          empty_func,
          batch_size=num_blocks_per_task * num_tasks,
          compute=compute,
      )
      map_ds = map_ds.materialize()
      assert map_ds.num_blocks() == 1
      map_ds = ds.map(identity_func, compute=compute)
      map_ds = map_ds.materialize()
      assert map_ds.num_blocks() == num_blocks_per_task * num_tasks
  
      ds_list = ds.split(5)
      assert len(ds_list) == 5
      for new_ds in ds_list:
          assert (
              new_ds.num_blocks() == num_blocks_per_task * num_tasks / 5
          )
  
      train, test = ds.train_test_split(test_size=0.25)
      assert (
          train.num_blocks() == num_blocks_per_task * num_tasks * 0.75
      )
      assert test.num_blocks() == num_blocks_per_task * num_tasks * 0.25
  
      new_ds = ds.union(ds, ds)
      assert new_ds.num_blocks() == num_tasks * 3
      new_ds = new_ds.materialize()
      assert new_ds.num_blocks() == num_blocks_per_task * num_tasks * 3
  
      new_ds = ds.random_shuffle()
      assert new_ds.num_blocks() == num_tasks
      new_ds = ds.randomize_block_order()
      assert new_ds.num_blocks() == num_tasks
      assert (
          ds.groupby("one").count().count()
          == num_blocks_per_task * num_tasks
      )
  
      new_ds = ds.zip(ds)
      new_ds = new_ds.materialize()
      assert new_ds.num_blocks() == num_blocks_per_task * num_tasks
  
      assert len(ds.take(5)) == 5
      assert len(ds.take_all()) == num_blocks_per_task * num_tasks
      for batch in ds.iter_batches(batch_size=10):
          assert len(batch["one"]) == 10
  
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 2 for calling a global function with example usage: test_dataset(shutdown_only, restore_data_context, compute) and returns: x, dict, np, x, dict