Pyspark объединяет столбец StructType как массив его элементов для каждой строки

Question

Я пытаюсь сделать что-то, что кажется довольно простым, но почему-то не могу понять, как это сделать с pyspark.

У меня есть df с двумя столбцами (для упрощения) 'id' и 'strcol', с возможными идентификаторами дубликатов

Я хочу сделать df.groupBy('id'), который бы возвращал для каждого id массив значений strcol

простой пример:

|--id--|--strCol--|
|   a  |  {'a':1} |
|   a  |  {'a':2} |
|   b  |  {'b':3} |
|   b  |  {'b':4} |
|------|----------|
would become
|--id--|-------aggsStr------|
|   a  |  [{'a':1},{'a':2}] |
|   b  |  [{'b':3},{'b':4}] |
|------|--------------------|

Я пытался использовать apply с pandas udf, но, похоже, отказывается возвращать массивы. (или возможно я не использовал это правильно)

Answer 1

Ты можешь использовать collect_list от pyspark.sql.functions модуль:

from pyspark.sql import functions as F
agg = df.groupby("id").agg(F.collect_list("strCol"))

Полнофункциональный пример:

import pandas as pd
from pyspark.sql import functions as F

data =  {'id': ['a', 'a', 'b', 'b'], 'strCol': [{'a':1}, {'a':2}, {'b':3}, {'b':4}]}

df_aux = pd.DataFrame(data)

# df type: DataFrame[id: string, strCol: map<string,bigint>]
df = spark.createDataFrame(df_aux) 


# agg type: # DataFrame[id: string, collect_list(strCol): array<map<string,bigint>>]
agg = df.groupby("id").agg(F.collect_list("strCol")) 

Надеюсь, это помогло!