Spark:有条件地将列添加到数据框

Question

我正在尝试获取输入数据:

A    B       C
--------------
4    blah    2
2            3
56   foo     3

并根据B是否为空来在末尾添加一列:

A    B       C     D
--------------------
4    blah    2     1
2            3     0
56   foo     3     1

我可以通过将输入数据帧注册为临时表,然后键入SQL查询来轻松完成此操作.

但我真的想知道如何使用Scala方法执行此操作,而不必在Scala中键入SQL查询.

我已经尝试过了.withColumn,但我无法做到我想做的事情.

Answer 1

尝试withColumn使用when如下功能:

val sqlContext = new SQLContext(sc)
import sqlContext.implicits._ // for `toDF` and $""
import org.apache.spark.sql.functions._ // for `when`

val df = sc.parallelize(Seq((4, "blah", 2), (2, "", 3), (56, "foo", 3), (100, null, 5)))
    .toDF("A", "B", "C")

val newDf = df.withColumn("D", when($"B".isNull or $"B" === "", 0).otherwise(1))

newDf.show() 节目

+---+----+---+---+
|  A|   B|  C|  D|
+---+----+---+---+
|  4|blah|  2|  1|
|  2|    |  3|  0|
| 56| foo|  3|  1|
|100|null|  5|  0|
+---+----+---+---+

我添加了(100, null, 5)用于测试isNull案例的行.

我试过这段代码Spark 1.6.0但是在代码中评论过when,它适用于之后的版本1.4.0.