带虚拟变量的Pypark矩阵

Question

有两列：

ID  Text
1    a
2    b
3    c

我怎样才能用如下虚拟变量创建矩阵：

ID a b c
1  1 0 0
2  0 1 0
3  0 0 1

使用pyspark库及其功能？

Answer 1

如果虚拟变量用于建模，则可以使用OneHotEncoder/OneHotEncoderRestimator（取决于pyspark的版本）。

Spark MLlib指南中给出了一个示例。

请注意，如果您需要删除一个列来防止线性依赖（就像回归一样），您应该使用dropLast=True。

Answer 2

另一种解决方案是使用Spark的pivot方法，该方法自Spark 1.6.0以来一直存在。

示例：

from pyspark.sql import functions as F

df = sqlContext.createDataFrame([
    (1, "a"),
    (2, "b"),
    (3, "c"),],
    ["ID", "Text"])    

pivoted = df.groupBy("ID").pivot("Text").agg(F.lit(1))
pivoted.show()
# +---+----+----+----+
# | ID|   a|   b|   c|
# +---+----+----+----+
# |  1|   1|null|null|
# |  3|null|null|   1|
# |  2|null|   1|null|
# +---+----+----+----+

要删除缺失的值，只需使用na方法：

pivoted.na.fill(0).show()
# +---+---+---+---+
# | ID|  a|  b|  c|
# +---+---+---+---+
# |  1|  1|  0|  0|
# |  3|  0|  0|  1|
# |  2|  0|  1|  0|
# +---+---+---+---+

数据透视比ksindi提出的解决方案更为普遍，因为它可以聚合数字。也就是说，ksindi提出的解决方案在这种特殊情况下更有效，因为它只需要对数据进行一次传递，如果您考虑传递以获得类别，则需要两次传递。对于透视，还可以将类别作为第二个位置参数添加到pivot，以提高效率。但是，groupBy调用已经会导致洗牌，这使得这种方法变得更慢。

注意：groupBy静默调用假定示例中的ID列包含唯一的值以获得所需的输出。如果示例数据帧看起来像：

df = sqlContext.createDataFrame([
    (1, "a"),
    (2, "b"),
    (3, "c"),
    (3, "a")],
    ["ID", "Text"]) 

这个解决方案的结果是

df.groupBy("ID").pivot("Text").agg(F.lit(1)).na.fill(0).show()
# +---+---+---+---+
# | ID|  a|  b|  c|
# +---+---+---+---+
# |  1|  1|  0|  0|
# |  3|  1|  0|  1|
# |  2|  0|  1|  0|
# +---+---+---+---+

而映射解决方案最终会成为

df.select("ID", *exprs).show()
# +---+---+---+---+
# | ID|  c|  b|  a|
# +---+---+---+---+
# |  1|  0|  0|  1|
# |  2|  0|  1|  0|
# |  3|  1|  0|  0|
# |  3|  0|  0|  1|
# +---+---+---+---+

Answer 3

from pyspark.sql import functions as F

df = sqlContext.createDataFrame([
    (1, "a"),
    (2, "b"),
    (3, "c"),
], ["ID", "Text"])

categories = df.select("Text").distinct().rdd.flatMap(lambda x: x).collect()

exprs = [F.when(F.col("Text") == category, 1).otherwise(0).alias(category)
         for category in categories]

df.select("ID", *exprs).show()

输出

+---+---+---+---+
| ID|  a|  b|  c|
+---+---+---+---+
|  1|  1|  0|  0|
|  2|  0|  1|  0|
|  3|  0|  0|  1|
+---+---+---+---+