scala - Spark : Add column to dataframe conditionally

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 之后的版本.