Как определить схему для Spark DataFrame со многими столбцами - PullRequest
Новая
       9

Как определить схему для Spark DataFrame со многими столбцами

1 голос
/ 09 мая 2020

У меня есть Spark DF in_df с более чем 300 столбцами с одним столбцом строк, а остальные дублируются. Мне нужно запустить на нем UDF GroupedMap Pandas и определить схему вывода перед запуском. В ситуации, когда выходные данные должны иметь одинаковое количество столбцов, но разных типов, как вы определяете эту схему? Несколько примеров Pandas UDF, которые я могу найти, обычно просто используют схему in в качестве схемы вывода.

Один метод, который я видел, использует withColumn и cast() на in_df. Это лучшая практика? Что, если я хочу, чтобы мой вывод имел совершенно другую форму, чем in_df, но слишком много столбцов для ручного кодирования? Мне не удалось найти для этого хороший ресурс.

1 Ответ

0 голосов
/ 09 мая 2020

Uisng pyspark.sql.types.StructType.fromJson() вы можете динамически построить схему из json.

В соответствии с вашим требованием я изменил тип данных, используя для «col_e», вы можете изменить DataTypes на один или несколько столбцов в зависимости от вашего варианта использования. 

df = spark.read.csv('test.csv',header=True,inferSchema=True)
fields = []
for f in json.loads(df.schema.json())["fields"]:
    if f["name"] == "col_e":
        fields.append(StructField("col_e", StringType(), True))
    else:
        fields.append(StructField.fromJson(f))

schema = StructType(fields)

@F.pandas_udf(schema, F.PandasUDFType.GROUPED_MAP)
def many_cols_data(pdf):
    pdf['col_e'] = "test"
    return pdf

df.groupBy(
    'col_a'
).apply(
    many_cols_data
).show()

тест входного файла. csv 

col_a,col_b,col_c,col_d,col_e
a,2,3,4,5
b,2,3,4,5
c,2,3,4,5

, что дает 

+-----+-----+-----+-----+-----+
|col_a|col_b|col_c|col_d|col_e|
+-----+-----+-----+-----+-----+
|    c|    2|    3|    4| test|
|    b|    2|    3|    4| test|
|    a|    2|    3|    4| test|
+-----+-----+-----+-----+-----+
...