Question

有两列：

ID  Text
1    a
2    b
3    c

如何使用虚拟变量创建矩阵，如下所示：

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

数据透视比solution proposed by 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

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

可以在 Spark MLlib Guide 中显示一个示例。

请注意，如果您需要删除一列以防止线性依赖（如回归的情况），则应使用 dropLast=True。

具有虚拟变量的pyspark矩阵

3 个答案: