python - 将 pyspark groupedData 转换为 pandas DataFrame

Question

我需要通过 Spark 对一个大型数据集进行分组，将其作为两列 Pandas 数据帧加载，然后重新转换为 Pandas:基本上是执行 Pandas -> 'pyspark.sql.group.GroupedData' -> Pandas。两列的元素都是整数，分组后的数据需要以列表格式存储，如下:

df.a        df.b
1            3
2            5
3            8
1            2
3            1
2            6
...
spark_df = spark.createDataFrame(df)
spark_grouped_df = spark_df.groupBy('a')
type: <class 'pyspark.sql.group.GroupedData'>

此时，我需要像 Pandas df 这样的东西(之后我需要做其他对 pandas 更友好的操作):

a        b
1    | [3,2]
2    | [5,6] 
3    | [8,1]
...

如果使用 pandas，我会这样做，但太耗时:

grouped_data = pd.DataFrame(df.groupby('a',as_index = True, sort = True)['b'].apply(list))

使用 Spark，我确信速度会快得多。

有什么提示吗？谢谢!

Answer 1

您需要聚合分组数据。要获取输出格式，您可以使用collect_list函数，

>>> from pyspark.sql.functions import collect_list
>>> pdf = spark_df.groupby('a').agg(collect_list('b').alias('b')).toPandas()
>>> pdf.head()
       a    b
    0  1  [3, 2]
    1  3  [8, 1]
    2  2  [5, 6]