Pyspark Dataframe - 文字列を数値にマップする

Question

データの特定の列（この場合は文字列）を変換し、数値表現に変換する方法を探しています。例えば、私は値を持つ文字列のデータフレームがあります

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

をこのようにデータフレーム上でforeachを実行すると、次のようになります。

def f(x): 
    if(x == 'Medium'): 
     return 2 
    elif(x == "Low"): 
     return 3 
    else: 
     return 1 

a = df.select("level").rdd.foreach(f) 

しかし、これは "None"型を返します。思考？いつものように助けてくれてありがとう！

Answer 1

これは確かにあなたが試しているラインに沿って行うことができます - foreachの代わりにmapの操作が必要です。

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を使用して、一時的な中間のRDDを呼び出さずにあなたのf(x)：

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