我需要在Spark DataFrame
中实现以下SQL逻辑
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|
+---+---+---+
我已经将嵌套案例 when 语句实现为:
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 语句更进一步,则带有"否则"语句的上述逻辑的可读性有点混乱。
当 Spark 中的语句DataFrames
时,有没有更好的方法来实现嵌套案例?
这里没有嵌套,因此不需要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")
对于更复杂的逻辑,我更喜欢使用 UDF 以获得更好的可读性:
val selectCase = udf((tc: String, amt: String) =>
if (Seq("a", "b").contains(tc)) "Y"
else if (tc == "a" && amt.toInt <= 0) "N"
else null
)
dataset1.withColumn("REASON", selectCase(col("tc"), col("amt")))
.show
简单地在数据集上使用selectExpr
。dataset1.selectExpr("*", "CASE WHEN tc in ('a') AND amt > 0 THEN 'N' WHEN tc in ('a','b') THEN 'Y' ELSE NULL END
REASON").show()
+---+---+---+------+
|KEY| tc|amt|REASON|
+---+---+---+------+
| 66| a| 4| N|
| 67| a| 0| Y|
| 70| b| 4| Y|
| 71| d| 4| null|
+---+---+---+------+
第二个条件应该放在第一个条件之前,因为第一个条件更通用。
当 tc 在 ('a') 和 amt> 0
然后是 'N'