code: def test_memory_usage(
      ray_start_regular, restore_data_context, use_push_based_shuffle
  ):
      DataContext.get_current().use_push_based_shuffle = (
          use_push_based_shuffle
      )
  
      parallelism = 2
      ds = ray.data.range(int(1e8), parallelism=parallelism)
      ds = ds.random_shuffle().materialize()
  
      stats = ds._get_stats_summary()
      # TODO(swang): Sort on this dataset seems to produce significant skew, so
      # one task uses much more memory than the other.
      for op_stats in stats.operators_stats:
          assert op_stats.memory["max"] < 2000
  
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 17 for calling a global function with example usage: test_memory_usage(ray_start_regular, restore_data_context, use_push_based_shuffle) without return types