web-dev-qa-db-fra.com

Cas imbriqué Spark Dataframe When Statement

J'ai besoin d'implémenter la logique SQL ci-dessous dans Spark 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;

Ma saisie DataFrame est la suivante:

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|
+---+---+---+

J'ai implémenté le cas imbriqué lorsque déclaration en tant que:

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|
+---+---+---+------+

La lisibilité de la logique ci-dessus avec l'instruction "sinon" n'est pas très compliquée si les instructions imbriquées quand vont plus loin.

Existe-t-il un meilleur moyen d'implémenter des cas imbriqués lorsque des déclarations dans Spark DataFrames? 

sqlapache-sparkdataframeapache-spark-sql
5
RaAm

Il n'y a pas de nidification ici, donc il n'y a pas besoin de otherwise. Tout ce dont vous avez besoin est chaîné when:

import spark.implicits._

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

ELSE NULL est implicite, vous pouvez donc l'omettre complètement.

Le modèle que vous utilisez est plus applicable pour folding sur une structure de données:

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

où when - otherwise suit naturellement le modèle de récursivité et null fournit le scénario de base.

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

Veuillez noter qu'il est impossible d'atteindre le résultat "N" avec cette chaîne de conditions. Si tc est égal à "a", il sera capturé par la première clause. Si ce n'est pas le cas, les deux prédicats ne seront pas satisfaits et la valeur par défaut sera NULL. Vous devriez plutôt:

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

Pour une logique plus complexe, je préfère utiliser les fonctions UDF pour une meilleure lisibilité:

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
2
Raphael Roth