code: def generate_shuffle_analysis_dataset(
      spark, output_dataset_url, num_rows=1000, row_group_size=100
  ):
      """
      Generates a small dataset useful for doing analysis on shuffling algorithms
  
      :param spark: spark session
      :param output_dataset_url: location to write dataset
      :param num_rows: how many rows should the dataset include
      :param row_group_size: how many rows in each row group (there is a minimum of 5)
      :return:
      """
      spark_context = spark.sparkContext
      with materialize_dataset(
          spark, output_dataset_url, _ShuffleAnalysisSchema
      ):
          rows_rdd = (
              spark_context.parallelize(range(num_rows), numSlices=50)
              .map(lambda i: {"id": i})
              .map(partial(dict_to_spark_row, _ShuffleAnalysisSchema))
          )
          spark.createDataFrame(
              rows_rdd, _ShuffleAnalysisSchema.as_spark_schema()
          ).sort("id").coalesce(
              max(1, int(num_rows / row_group_size))
          ).write.option(
              "compression", "none"
          ).parquet(
              output_dataset_url
          )
  
file_path: petastorm/petastorm/test_util/shuffling_analysis.py
imports: from __future__ import division
  from functools import partial
  import numpy
  from pyspark.sql.types import LongType
  from petastorm import make_reader
  from petastorm.codecs import ScalarCodec
  from petastorm.etl.dataset_metadata import materialize_dataset
  from petastorm.unischema import Unischema, UnischemaField, dict_to_spark_row
label_desc: Use image node_id 1 for calling a global function with example usage: generate_shuffle_analysis_dataset(spark, output_dataset_url, num_rows, row_group_size) without return types