Тестирование PySpark: создание тестовых данных, состоящих из массива структур

Question

Я хотел бы сгенерировать некоторые тестовые данные для моих модульных тестов в PySpark. Одним из полей входной строки является массив структур: basket: array<struct<price:bigint,product_id:string>>. Какой лучший способ добиться этого?

Answer 1

Вот один из способов использования python и две вспомогательные функции, отвечающие за генерацию случайных данных:

from pyspark.sql.types import *
from random import randrange, uniform

array_size = 2

def create_row(array_size):
  return ([{"price" : uniform(1.0, 100.0), "product_id" : randrange(10) + 1} for _ in range(array_size)],)

def generate_data(data_size):
  return [create_row(array_size) for _ in range(data_size)] 

# create 5 rows
rows = generate_data(5)

# string schema
schema = "basket: array<struct<price:double,product_id:string>>"

# static typing schema
# schema = StructType([
#             StructField('basket',
#                 ArrayType(
#                     StructType(
#                           [
#                               StructField('price', DoubleType()),
#                               StructField('product_id', StringType()),
#                           ]
#                     )
#                 )
#             )])

df = spark.createDataFrame(rows, schema)

df.show(10, False)

# +--------------------------------------------------+
# |basket                                            |
# +--------------------------------------------------+
# |[[61.40674765573896, 9], [5.994467505720648, 7]]  |
# |[[1.1388272509974906, 10], [47.32070824053193, 3]]|
# |[[42.423106687845795, 2], [70.99107361888588, 4]] |
# |[[50.019594333009806, 8], [63.51239439900147, 4]] |
# |[[68.15711374321089, 9], [70.06617125228864, 10]] |
# +--------------------------------------------------+

• create_row : создаст новую строку (представлен здесь как кортеж) с array_size items. price будет иметь значение в диапазоне от 1,0 до 100,0 и product_id в диапазоне от 1 до 10, пожалуйста, не стесняйтесь изменять границы соответствующим образом. Кроме того, мы обрабатываем каждый элемент массива (пары product_id-price) с помощью словаря python.

• generate_data : вызовы create_row data_size раз и возврат случайные сгенерированные строки в список.