Pyspark：将分类列转换为二进制数组

【问题标题】：Pyspark: transform categorial column into a binary arrayPyspark：将分类列转换为二进制数组
【发布时间】：2020-06-01 10:08:23
【问题描述】：

我有一个数据框，其中有一列，其中每个值都是一个字符串，我希望对其进行转换，以便特定 id 的字符串的每个外观（表示在不同的列中）都将作为一个阵列在适当的位置。例如，如果我有这个数据集：

category   id
a          1
b          1
c          1
a          2
d          2

我想把它变成：

id   result
1    [1,1,1,0]
2    [1,0,0,1]

你知道怎么做吗？

编辑：我尝试使用字符串索引器，但它所做的只是在我真正想要获取数组本身时计算出现次数。

【问题讨论】：

标签： python pyspark azure-databricks

【解决方案1】：

您可以先旋转并用 0 填充 null，然后将除 id 列之外的列转换为数组：

import pyspark.sql.functions as F

res = df.groupBy('id').pivot("category").agg(F.count("category")).fillna(0)
res.select("id",F.array(*[F.col(i) for i in res.columns[1:]]).alias("result")).show()

+---+------------+
| id|      result|
+---+------------+
|  1|[1, 1, 1, 0]|
|  2|[1, 0, 0, 1]|
+---+------------+

请注意，如果元素可以重复，则计数可以大于 1，在这种情况下，您可以添加 when+otherwise：

(res.select("id",F.array(*[F.when(F.col(i)>=1,1).otherwise(0) 
                       for i in res.columns[1:]]).alias("result"))).show()

【讨论】：