首页 > 解决方案 > Pyspark 数据帧分组

问题描述

我有如下数据框

|123  |124  |125 |
+-----+-----+-----+
|    1|    2|    3|
|    9|    9|    4|
|    4|   12|    1|
|    2|    4|    8|
|    7|    6|    3|
|   19|   11|    2|
|   21|   10|   10

我需要数据

1:[123,125]
2:[123,124,125]
3:[125]

订单不需要排序。我是 pyspark 中的数据框的新手,任何帮助将不胜感激

标签: dataframepyspark

解决方案


pyspark中没有melt或API 可以直接完成此操作。pivot相反,将 RDD 中的平面映射到一个新的数据框并聚合:

df.show()                                                                                                                                                                                           

+---+---+---+
|123|124|125|
+---+---+---+
|  1|  2|  3|
|  9|  9|  4|
|  4| 12|  1|
|  2|  4|  8|
|  7|  6|  3|
| 19| 11|  2|
| 21| 10| 10|
+---+---+---+

对于 RDD 中的每一列或每一行,输出一行有两列:列的值和列名:

cols = df.columns
(df.rdd
 .flatMap(lambda row: [(row[c], c) for c in cols]).toDF(["value", "column_name"])
 .show())

+-----+-----------+
|value|column_name|
+-----+-----------+
|    1|        123|
|    2|        124|
|    3|        125|
|    9|        123|
|    9|        124|
|    4|        125|
|    4|        123|
|   12|        124|
|    1|        125|
|    2|        123|
|    4|        124|
|    8|        125|
|    7|        123|
|    6|        124|
|    3|        125|
|   19|        123|
|   11|        124|
|    2|        125|
|   21|        123|
|   10|        124|
+-----+-----------+

然后,按值分组并将列名聚合到一个列表中:

from pyspark.sql import functions as f 

(df.rdd
 .flatMap(lambda row: [(row[c], c) for c in cols]).toDF(["value", "column_name"])
 .groupby("value").agg(f.collect_list("column_name"))
 .show())

+-----+-------------------------+
|value|collect_list(column_name)|
+-----+-------------------------+
|   19|                    [123]|
|    7|                    [123]|
|    6|                    [124]|
|    9|               [123, 124]|
|    1|               [123, 125]|
|   10|               [124, 125]|
|    3|               [125, 125]|
|   12|                    [124]|
|    8|                    [125]|
|   11|                    [124]|
|    2|          [124, 123, 125]|
|    4|          [125, 123, 124]|
|   21|                    [123]|
+-----+-------------------------+

推荐阅读