code: def test_push_based_shuffle_stats(ray_start_cluster):
      ctx = ray.data.context.DataContext.get_current()
      try:
          original = ctx.use_push_based_shuffle
          ctx.use_push_based_shuffle = True
  
          cluster = ray_start_cluster
          cluster.add_node(
              resources={"bar:1": 100},
              num_cpus=10,
              _system_config={"max_direct_call_object_size": 0},
          )
          cluster.add_node(resources={"bar:2": 100}, num_cpus=10)
          cluster.add_node(resources={"bar:3": 100}, num_cpus=0)
  
          ray.init(cluster.address)
  
          parallelism = 100
          ds = ray.data.range(
              1000, parallelism=parallelism
          ).random_shuffle()
          ds = ds.materialize()
          assert "RandomShuffleMerge" in ds.stats()
          # Check all nodes used.
          assert "2 nodes used" in ds.stats()
          assert "1 nodes used" not in ds.stats()
  
          # Check all merge tasks are included in stats.
          internal_stats = ds._plan.stats()
          num_merge_tasks = len(
              internal_stats.metadata["RandomShuffleMerge"]
          )
          # Merge factor is 2 for random_shuffle ops.
          merge_factor = 2
          assert (
              parallelism // (merge_factor + 1)
              <= num_merge_tasks
              <= parallelism // merge_factor
          )
  
      finally:
          ctx.use_push_based_shuffle = original
  
imports: import random
  from collections import defaultdict
  import numpy
  import pandas
  import pyarrow
  import pytest
  import ray
  from ray.data import Dataset
  from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import PushBasedShuffleTaskScheduler
  from ray.data._internal.sort import SortKey
  from ray.data.block import BlockAccessor
  from ray.data.context import DataContext
  from ray.data.tests.conftest import *
  from ray.data.tests.util import extract_values
  from ray.tests.conftest import *
label_desc: Use image node_id 11 for calling a global function with example usage: test_push_based_shuffle_stats(ray_start_cluster) without return types