如何基于查找数据框创建数据框并在特定列中的动态和映射值上创建多列

Posted 2023-04-17

【中文标题】如何基于查找数据框创建数据框并在特定列中的动态和映射值上创建多列

【英文标题】：How to create a dataframe based on the lookup dataframe and create mulitple column on dynamic and map values in the specific columns

【发布时间】：2021-06-16 17:20:49

【问题描述】：

我有两个数据框，一个是主要的，另一个是查找数据框。我需要使用 pyspark 以自定义形式实现第三个。我需要检查 list_ids 列中的值并检查查找数据框中的匹配并在最终数据框中标记计数。我尝试过数组相交和数组查找，但它不起作用。

主要数据框：

df = spark.createDataFrame([(123, [75319, 75317]), (212, [136438, 25274]), (215, [136438, 75317])], ("ID", "list_IDs"))
df.show()

+---+---------------+
| ID|       list_IDs|
+---+---------------+
|123| [75319, 75317]|
|212|[136438, 25274]|
|215|[136438, 75317]|
+---+---------------+

查找数据框：

df_2 = spark.createDataFrame([(75319, "Wheat", 20), (75317, "Rice", 10), (136438, "Jowar", 30), (25274, "Rajma", 40)], ("ID", "Material", "Count"))
df_2.show()

+------+--------+-----+
|    ID|Material|Count|
+------+--------+-----+
| 75319|   Wheat|   A|
| 75317|    Rice|   B|
|136438|   Jowar|   C|
| 25274|   Rajma|   D|
+------+--------+-----+

需要结果数据框

+---+---------------+------+------+-------+------+
| ID|       list_IDs|Wheat | Rice | Jowar | Rajma|
+---+---------------+------+------+-------+------+
|123| [75319, 75317]|    A|    B|    0 |     0|
|212|[136438, 25274]|    0|    0|    C |     D|
|215|[136438, 75317]|    0|    B|    C |    0 |
+---+---------------+------+------+-------+------+

【参考方案1】：

您可以连接两个数据框，然后进行旋转：

import pyspark.sql.functions as F

df2 = df.join(
    df_2, 
    F.array_contains(df.list_IDs, df_2.ID)
).groupBy(df.ID, 'list_IDs').pivot('Material').agg(F.first('Count')).fillna(0)

result.show()
+---+---------------+-----+-----+----+-----+
| ID|       list_IDs|Jowar|Rajma|Rice|Wheat|
+---+---------------+-----+-----+----+-----+
|212|[136438, 25274]|   30|   40|   0|    0|
|215|[136438, 75317]|   30|    0|  10|    0|
|123| [75319, 75317]|    0|    0|  10|   20|
+---+---------------+-----+-----+----+-----+