Pyspark-根据列表或字典创建列

时间:2020-09-24 09:01:42

标签: pyspark pyspark-dataframes

我是Pyspark的新手,我有一个包含客户ID和文本以及相关值的数据框

+------+-----+------+
|id    |text |value |
+------+-----+------+
|   1  |  Cat|     5|
|   1  |  Dog|     4|
|   2  |  Oil|     1|

我想根据关键字列表来解析text列,并创建一列来告诉我是否在文本字段中包含关键字并提取关联的值,预期结果是

List_keywords = ["Dog", Cat"]

+------+-----+------+--------+---------+--------+---------+
|id    |text |value |bool_Dog|value_Dog|bool_cat|value_cat|
+------+-----+------+--------+---------+--------+---------+
|   1  |  Cat|     5|0       |        0|       1|        5|
|   1  |  Dog|     4|1       |        4|       0|        0|
|   2  |  Oil|     1|0       |        0|       0|        0|

做到这一点的最佳方法是什么?我当时正在考虑创建一个包含关键字的列表或词典,并使用for循环对其进行解析,但是我敢肯定有一种更好的方法可以做到这一点。

非常感谢您

1 个答案:

答案 0 :(得分:1)

请参阅解决方案

import pyspark.sql.functions as F
data = [[1, 'Cat', 5], [1, 'Dog', 4], [2, 'Oil', 3]]
df = spark.createDataFrame(data, ['id', 'text', 'value'])
df.show()

+---+----+-----+
| id|text|value|
+---+----+-----+
|  1| Cat|    5|
|  1| Dog|    4|
|  2| Oil|    1|
+---+----+-----+
keywords = ['Dog','Cat']

(
    df
    .groupby('id', 'text', 'value')
    .pivot('text', keywords)
    .agg(
        F.count('value').alias('bool'),
        F.max('value').alias('value')
    )
    .fillna(0)
    .sort('text')
).show()

+---+----+-----+--------+---------+--------+---------+
| id|text|value|Dog_bool|Dog_value|Cat_bool|Cat_value|
+---+----+-----+--------+---------+--------+---------+
|  1| Cat|    5|       0|        0|       1|        5|
|  1| Dog|    4|       1|        4|       0|        0|
|  2| Oil|    1|       0|        0|       0|        0|
+---+----+-----+--------+---------+--------+---------+