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

Question

我有一个pyspark数据帧，如下所示：

以及以下国家的名单：

EU_countries=['西班牙'，'意大利'，'法国'，'德国']

我想做以下工作：

1. 从列tmp\u列表中提取所有唯一值
2. 为EU\u国家/地区中存在的所有值创建新列。对于EU\u国家中不存在的值，创建一个名为other\u countries的列。基本上，为EU\u国家中的每个条目创建列列出一个名为other\u countries的额外列
3. 如果id包含EU_国家列表中的任何国家，则新列西班牙应具有1作为值，否则0。这同样适用于EU_国家列表中的其他国家
4. 如果id包含EU\u国家列表中不存在的任何国家，则应填写other\u countries列1else0

这是我想要的最终输出：

我为这件事伤了脑筋。谁能帮我一下吗？

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




Answer 1

推理和工作，就像我在熊猫里一样。

1. 爆炸

下面的代码；

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

Answer 2

步骤1：用常量字符串替换tmp_列表中的所有非欧盟国家：

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