将列中的列表拆分为 pyspark 中的一个热编码特征

Question

我有一个如下所示的 pyspark 数据框：

id	tmp_list	other features
1	['Spain', 'Italy']	xxx
2	['Spain', 'France', 'USA', 'India']	yyy
3	['Spain', 'Germany']	zzz

以及如下国家列表：

EU_countries = ['Spain', 'Italy', 'France', 'Germany']

我想做以下事情：

1. 从tmp_list 列中提取所有唯一值
2. 为EU_countries 中存在的所有值创建新列。对于EU_countries 中不存在的值，创建一个名为other_countries 的列。本质上，为 EU_countries 列表中的每个条目创建列 + 一个名为 other_countries 的额外列。
3. 如果id 包含EU_countries 列表中的任何国家/地区，则新列Spain 应具有1 作为值否则0。同样适用于EU_countries 列表中的其他国家/地区。
4. 如果id 包含EU_countries 列表中不存在的任何国家/地区，则应填写other_countries 列1 否则0。

这是我正在寻找的最终输出：

id	Spain	Italy	France	Germany	other_countries	other features
1	1	1	0	0	0	xxx
2	1	0	1	0	1	yyy
3	1	0	0	1	0	zzz

我对此感到头疼。有人可以帮我解决这个问题吗？

非常感谢任何帮助！非常感谢！

Answer 1

第 1 步：将tmp_list 中的所有非欧盟国家/地区替换为常量字符串：

from pyspark.sql import functions as F

df = ...
EU_countries = ['Spain', 'Italy', 'France', 'Germany']

def replaceNonEU(c):
    cond = c == EU_countries[0]
    for country in EU_countries[1:]:
        cond |= (c == country)
    return F.when(cond, c).otherwise(F.lit("other_countries"))

df = df.withColumn("tmp_list", F.array_distinct(F.transform("tmp_list", replaceNonEU)))

#+---+--------------------------------+--------------+
#|id |tmp_list                        |other features|
#+---+--------------------------------+--------------+
#|1  |[Spain, Italy]                  |xxx           |
#|2  |[Spain, France, other_countries]|yyy           |
#|3  |[Spain, Germany]                |zzz           |
#+---+--------------------------------+--------------+

第 2 步：为tmp_list 中的每个可能值创建一个新列，指示该值是否是tmp_list 的元素：

for c in EU_countries + ['other_countries']:
    df = df.withColumn(c, F.array_contains("tmp_list", c).cast("int"))
df = df.drop("tmp_list")

#+---+--------------+-----+-----+------+-------+---------------+
#| id|other features|Spain|Italy|France|Germany|other_countries|
#+---+--------------+-----+-----+------+-------+---------------+
#|  1|           xxx|    1|    1|     0|      0|              0|
#|  2|           yyy|    1|    0|     1|      0|              1|
#|  3|           zzz|    1|    0|     0|      1|              0|
#+---+--------------+-----+-----+------+-------+---------------+

Answer 2

推理和工作就像我在 pandas 中一样。

1. 爆炸
2. 创建将非 EU_countries 归为 other_countries 的类别
3. get_dummies。在这一点上，我感谢this post

代码如下；

df=df.select('*').withColumn('tmp_list1', F.explode(col('tmp_list')))#Create new column with exploded list
df=df.select('*').withColumn('Cat', when(col('tmp_list1').isin(EU_countries),df.tmp_list1).otherwise('other_countries'))#Create another column Cat
df.groupBy("tmp_list",'other features').pivot("Cat").agg(F.lit(1)).na.fill(0).show()#Get dummies


+---------------------------+--------------+------+-------+-----+-----+---------------+
|tmp_list                   |other features|France|Germany|Italy|Spain|other_countries|
+---------------------------+--------------+------+-------+-----+-----+---------------+
|[Spain, Germany]           |zzz           |0     |1      |0    |1    |0              |
|[Spain, Italy]             |xxx           |0     |0      |1    |1    |0              |
|[Spain, France, USA, India]|yyy           |1     |0      |0    |1    |1              |
+---------------------------+--------------+------+-------+-----+-----+---------------+