Pyspark DataFrame - 如何將一列從分類值轉換為整數？

Question

我有一個 pyspark dataframe，我想將其中一列從字符串轉換為 int。 例子：

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

所以，在那種情況下：[red=1, green=2, yellow=3 and black=4]。

但我不知道所有的 colors 以便手動分配。 所以，我需要一種方法來自動進行歸因。

有人可以幫我嗎？

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

Answer 2

如果您想要一個代碼更少的解決方案並且您的類別不需要以特殊方式排序，您可以使用dense_rank函數中的 dense_rank。

import pyspark.sql.functions as F
from pyspark.sql.window import Window

df.withColumn("categ_num", F.dense_rank().over(Window.orderBy("categories")))

請記住，window 函數會導致運行時間更長。

Answer 3

SparkML中有一個StringIndexer 。