def create_test_dataset(tmp_url, rows, num_files=2, spark=None, use_summary_metadata=False):
    """
    Creates a test dataset under tmp_dir, with rows and num_files that has TestSchema.
    :param tmp_url: The URL of the temp directory to store the test dataset in.
    :param rows: The number of rows for the dataset.
    :param num_files: The number of files to partition the data between.
    :param spark: An optional spark session to use
    :param use_summary_metadata: If True, _metadata file will be created.
    :return: A list of the dataset dictionary.
    """

    shutdown = False
    if not spark:
        spark_session = SparkSession \
            .builder \
            .appName('petastorm_end_to_end_test') \
            .master('local[*]')

        spark = spark_session.getOrCreate()
        shutdown = True
    spark_context = spark.sparkContext

    with materialize_dataset(spark, tmp_url, TestSchema, use_summary_metadata=use_summary_metadata):
        id_rdd = spark_context.parallelize(rows, numSlices=40)

        # Make up some random data and store it for referencing in the tests
        random_dicts_rdd = id_rdd.map(_randomize_row).cache()
        dataset_dicts = random_dicts_rdd.collect()

        def _partition_key_to_str(row):
            row['partition_key'] = str(row['partition_key'])
            return row

        random_dicts_rdd = random_dicts_rdd.map(_partition_key_to_str)

        random_rows_rdd = random_dicts_rdd.map(partial(dict_to_spark_row, TestSchema))

        # Create a spark dataframe with the random rows
        dataframe = spark. \
            createDataFrame(random_rows_rdd, TestSchema.as_spark_schema()).sort('id')

        # Save a parquet
        dataframe. \
            coalesce(num_files). \
            write.option('compression', 'none'). \
            partitionBy('partition_key'). \
            mode('overwrite'). \
            parquet(tmp_url)

    # Create list of objects to build row group indexes
    indexers = [
        SingleFieldIndexer(TestSchema.id.name, TestSchema.id.name),
        SingleFieldIndexer(TestSchema.sensor_name.name, TestSchema.sensor_name.name),
        SingleFieldIndexer(TestSchema.string_array_nullable.name, TestSchema.string_array_nullable.name),
        SingleFieldIndexer(TestSchema.partition_key.name, TestSchema.partition_key.name),
    ]
    build_rowgroup_index(tmp_url, spark_context, indexers)

    if shutdown:
        spark.stop()

    return dataset_dicts

    Creates a test dataset under tmp_dir, with rows and num_files that has TestSchema.
    :param tmp_url: The URL of the temp directory to store the test dataset in.
    :param rows: The number of rows for the dataset.
    :param num_files: The number of files to partition the data between.
    :param spark: An optional spark session to use
    :param use_summary_metadata: If True, _metadata file will be created.
    :return: A list of the dataset dictionary.
    
petastorm/tests/test_common.py