我正在寻找一种方法来转换给定的数据列,在本例中为字符串,并将它们转换为数字表示。例如,我有一个带有值的字符串数据框:
+------------+
| level |
+------------+
| Medium|
| Medium|
| Medium|
| High|
| Medium|
| Medium|
| Low|
| Low|
| High|
| Low|
| Low|
我想创建一个新列,将这些值转换为:
"High"= 1, "Medium" = 2, "Low" = 3
+------------+
| level_num|
+------------+
| 2|
| 2|
| 2|
| 1|
| 2|
| 2|
| 3|
| 3|
| 1|
| 3|
| 3|
我尝试过定义一个函数并对数据框进行预测,如下所示:
def f(x):
if(x == 'Medium'):
return 2
elif(x == "Low"):
return 3
else:
return 1
a = df.select("level").rdd.foreach(f)
但是这会返回“无”类型。思考?一如既往地感谢您的帮助!
答案 0 :(得分:4)
您当然可以按照您尝试过的方式执行此操作 - 您需要map
操作而不是foreach
。
spark.version
# u'2.2.0'
from pyspark.sql import Row
# toy data:
df = spark.createDataFrame([Row("Medium"),
Row("High"),
Row("High"),
Row("Low")
],
["level"])
df.show()
# +------+
# | level|
# +------+
# |Medium|
# | High|
# | High|
# | Low|
# +------+
将f(x)
与这些玩具数据结合使用,我们得到:
df.select("level").rdd.map(lambda x: f(x[0])).collect()
# [2, 1, 1, 3]
还有一个map
将为您提供数据框:
df.select("level").rdd.map(lambda x: f(x[0])).map(lambda x: Row(x)).toDF(["level_num"]).show()
# +---------+
# |level_num|
# +---------+
# | 2|
# | 1|
# | 1|
# | 3|
# +---------+
但最好不要使用数据框函数when
而不是f(x)
调用临时中间RDD来执行此操作:
from pyspark.sql.functions import col, when
df.withColumn("level_num", when(col("level")=='Medium', 2).when(col("level")=='Low', 3).otherwise(1)).show()
# +------+---------+
# | level|level_num|
# +------+---------+
# |Medium| 2|
# | High| 1|
# | High| 1|
# | Low| 3|
# +------+---------+