Pyspark 数据框 - 将字符串映射到数字

在数据处理中，经常需要将字符串映射为数字进行进一步的数据分析。在Pyspark中，我们可以使用StringIndexer将字符串映射为数字。

1. 导入必要的模块

from pyspark.ml.feature import StringIndexer
from pyspark.sql.functions import col

2. 创建一个数据框

data = [("John", "Smith", "London"),
        ("Smith", "David", "Paris"),
        ("Jane", "Doe", "New York"),
        ("Tom", "Smith", "Tokyo")]

df = spark.createDataFrame(data, ["first_name", "last_name", "city"])
df.show()

"""
+----------+---------+--------+
|first_name|last_name|    city|
+----------+---------+--------+
|      John|    Smith|  London|
|     Smith|    David|   Paris|
|      Jane|      Doe|New York|
|       Tom|    Smith|   Tokyo|
+----------+---------+--------+
"""

3. 使用 StringIndexer 将字符串映射为数字

stringIndexer = StringIndexer(inputCol="city", outputCol="city_index")
model = stringIndexer.fit(df)
indexed = model.transform(df)
indexed.show()

"""
+----------+---------+--------+----------+
|first_name|last_name|    city|city_index|
+----------+---------+--------+----------+
|      John|    Smith|  London|       0.0|
|     Smith|    David|   Paris|       2.0|
|      Jane|      Doe|New York|       1.0|
|       Tom|    Smith|   Tokyo|       3.0|
+----------+---------+--------+----------+
"""

我们可以看到，“London”被映射为了0.0，“New York”被映射为了1.0，“Paris”被映射为了2.0，“Tokyo”被映射为了3.0。

注意：

如果要将多列字符串映射为数字，我们可以使用 StringIndexer的 setInputCols 和 setOutputCols 方法。

例如：

stringIndexer = StringIndexer(inputCols=["first_name", "last_name", "city"], outputCols=["first_name_index", "last_name_index", "city_index"])
model = stringIndexer.fit(df)
indexed = model.transform(df)
indexed.show()

以上就是如何在 Pyspark 中将字符串映射为数字。