code: 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
  
imports: from __future__ import division
  import random
  from collections import OrderedDict
  from decimal import Decimal
  from functools import partial
  import numpy
  import pytz
  from pyspark import Row
  from pyspark.sql import SparkSession
  from pyspark.sql.types import StringType, DecimalType, DoubleType, StructField, IntegerType, StructType, DateType, TimestampType, ShortType, ArrayType
  from petastorm.codecs import CompressedImageCodec, NdarrayCodec, ScalarCodec
  from petastorm.etl.dataset_metadata import materialize_dataset
  from petastorm.etl.rowgroup_indexers import SingleFieldIndexer
  from petastorm.etl.rowgroup_indexing import build_rowgroup_index
  from petastorm.unischema import Unischema, UnischemaField, dict_to_spark_row
label_desc: Use image node_id 4 for calling a global function with example usage: create_test_dataset(tmp_url, rows, num_files, spark, use_summary_metadata) and returns: dataset_dicts, row