具有虚拟变量的pyspark矩阵

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

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

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