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)

    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:
    
petastorm/test_util/shuffling_analysis.py