一尘不染

Spark Dataframe嵌套的when语句

sql

我需要在Spark中实现以下SQL逻辑 DataFrame

SELECT KEY,
    CASE WHEN tc in ('a','b') THEN 'Y'
         WHEN tc in ('a') AND amt > 0 THEN 'N'
         ELSE NULL END REASON,
FROM dataset1;

我的输入DataFrame如下:

val dataset1 = Seq((66, "a", "4"), (67, "a", "0"), (70, "b", "4"), (71, "d", "4")).toDF("KEY", "tc", "amt")

dataset1.show()



+---+---+---+
|KEY| tc|amt|
+---+---+---+
| 66|  a|  4|
| 67|  a|  0|
| 70|  b|  4|
| 71|  d|  4|
+---+---+---+

我已经实现了嵌套的情况,当声明为:

dataset1.withColumn("REASON", when(col("tc").isin("a", "b"), "Y")
  .otherwise(when(col("tc").equalTo("a") && col("amt").geq(0), "N")
    .otherwise(null))).show()



+---+---+---+------+
|KEY| tc|amt|REASON|
+---+---+---+------+
| 66|  a|  4|     Y|
| 67|  a|  0|     Y|
| 70|  b|  4|     Y|
| 71|  d|  4|  null|
+---+---+---+------+

如果嵌套的when语句更进一步,则使用“ otherwise”语句的上述逻辑的可读性会很混乱。

在Spark中的语句时,有没有更好的方法来实现嵌套大小写DataFrames


阅读 524

收藏
2021-03-10

共1个答案

一尘不染

这里没有嵌套,因此不需要otherwise。您所需的全部都是链式的when

import spark.implicits._

when($"tc" isin ("a", "b"), "Y")
  .when($"tc" === "a" && $"amt" >= 0, "N")

ELSE NULL 是隐式的,因此您可以完全省略它。

您使用的模式更适用于folding数据结构:

val cases = Seq(
  ($"tc" isin ("a", "b"), "Y"),
  ($"tc" === "a" && $"amt" >= 0, "N")
)

其中when-otherwise自然遵循递归模式并null提供基本情况。

cases.foldLeft(lit(null)) {
  case (acc, (expr, value)) => when(expr, value).otherwise(acc)
}

请注意,在这种情况下,不可能达到“ N”个结果。如果tc等于“
a”,它将被第一个子句捕获。如果不是,它将无法同时满足两个谓词,并且默认为NULL。您应该宁愿:

when($"tc" === "a" && $"amt" >= 0, "N")
 .when($"tc" isin ("a", "b"), "Y")
2021-03-10