Apache Spark,将一个"CASE WHEN ... ELSE ..."计算列添加到现有的DataFrame中

Question

我正在尝试使用Scala API向现有DataFrame添加"CASE WHEN ... ELSE ..."计算列.启动数据帧:

color
Red
Green
Blue

期望的数据帧(SQL语法:CASE WHEN color == Green THEN 1 ELSE 0 END AS bool):

color bool
Red   0
Green 1
Blue  0

我该如何实现这个逻辑？

Answer 1

在即将发布的SPARK 1.4.0版本中(应在未来几天内发布).您可以使用when/otherwise语法:

// Create the dataframe
val df = Seq("Red", "Green", "Blue").map(Tuple1.apply).toDF("color")

// Use when/otherwise syntax
val df1 = df.withColumn("Green_Ind", when($"color" === "Green", 1).otherwise(0))

如果您使用的是SPARK 1.3.0,则可以选择使用UDF:

// Define the UDF
val isGreen = udf((color: String) => {
  if (color == "Green") 1
  else 0
})
val df2 = df.withColumn("Green_Ind", isGreen($"color"))

Answer 2

在Spark 1.5.0中:您还可以使用SQL语法expr函数

val df3 = df.withColumn("Green_Ind", expr("case when color = 'green' then 1 else 0 end"))

或普通的spark-sql

df.registerTempTable("data")
val df4 = sql(""" select *, case when color = 'green' then 1 else 0 end as Green_ind from data """)