python - 将列中的列表拆分为 pyspark 中的一个热编码功能

Question

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

<表类="s-表"><头>idtmp_list其他功能<正文>1['西类牙', '意大利']xxx2['西类牙', '法国', '美国', '印度']yyy3['西类牙', '德国']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 .

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

<表类="s-表"><头>id西类牙意大利法国德国其他国家其他功能<正文>111000xxx210101yyy310010zzz

我为此伤透了脑筋。有人可以帮我解决这个问题吗？

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

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|
#+---+--------------+-----+-----+------+-------+---------------+