Pyspark DataFrame - 如何将一列从分类值转换为 int？

Question

我有一个 pyspark 数据框，我想将该列之一从字符串转换为整数。例子：

表1：

+------------+-----+
|categories  |value|
+------------+-----+
|         red| 0.23|
|       green| 0.34|
|      yellow| 0.56|
|       black| 0.11|
|         red| 0.67|
|         red| 0.34|
|       green| 0.45|
+------------+-----+

表2：

+------------+-----+
|categ_num   |value|
+------------+-----+
|           1| 0.23|
|           2| 0.34|
|           3| 0.56|
|           4| 0.11|
|           1| 0.67|
|           1| 0.34|
|           2| 0.45|
+------------+-----+

因此，在这种情况下：[红色=1，绿色=2，黄色=3，黑色=4]。

但我不知道所有颜色以便手动分配。所以，我需要一种方法来自动进行归因。

有人可以帮我吗？

Answer 1

这段代码对我有用：

from pyspark.ml.feature import StringIndexer

df = spark.createDataFrame(
[(0, "a"), (1, "b"), (2, "c"), (3, "a"), (4, "a"), (5, "c")],
["id", "category"])

indexer = StringIndexer(inputCol="category", outputCol="categoryIndex")
indexed = indexer.fit(df).transform(df)
indexed.show()

https://spark.apache.org/docs/latest/ml-features.html#stringindexer