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

时间:2016-01-20 19:10:41

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

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

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,但我无法做到我想做的事。

3 个答案:

答案 0 :(得分:73)

使用函数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之后的版本。

答案 1 :(得分:3)

我的不好,我错过了问题的一部分。

最好,最干净的方法是使用for wall in...。 代码中的解释。

UDF

答案 2 :(得分:1)

这样的事情怎么样?

val newDF = df.filter($"B" === "").take(1) match {
  case Array() => df
  case _ => df.withColumn("D", $"B" === "")
}

使用take(1)应该有最小的打击