Pyspark Dataframe - 将字符串映射到数字

时间:2017-11-30 16:09:46

标签: apache-spark pyspark apache-spark-sql spark-dataframe pyspark-sql

我正在寻找一种方法来转换给定的数据列,在本例中为字符串,并将它们转换为数字表示。例如,我有一个带有值的字符串数据框:

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

但是这会返回“无”类型。思考?一如既往地感谢您的帮助!

1 个答案:

答案 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| 
# +------+---------+