Comment appliquer plusieurs conditions sur l'instruction Case-Sinon à l'aide de l'API Spark Dataframe

Question

J'essaie d'ajouter une nouvelle colonne à un bloc de données existant en utilisant l'instruction withColumn dans Spark Dataframe API. Le code ci-dessous fonctionne mais je me demandais s'il y avait un moyen de sélectionner plus d'un groupe. Disons les Groupes 1, 2, 3, 4 au lieu du Groupe 1 seulement. Je pense que je pourrais être capable d'écrire une déclaration une quatrième fois. J'ai peut-être vu des gens faire ça dans certains posts. Cependant, dans R, il y a un opérateur %in% qui peut spécifier si une variable contient des valeurs dans un vecteur mais je ne sais pas s'il y a une telle chose dans Spark. J'ai vérifié la documentation de l'API Spark mais la plupart des fonctions ne contiennent aucun exemple.

R Sample Code: 
> library(dplyr) 
> df1 <- df %>% mutate(Selected_Group = (Group %in% 1:4)) 

Spark Dataframe Sample Code That Selects Group 1: 
> val df1 = df.withColumn("Selected_Group", when($"Group" === 1, 1).otherwise(0)) 

données

ID, Groupe
1, 0
2, 1
3, 2
. .
. .
100, 99

Une pensée serait appréciée!

Answer 1

Avec UDF:

import org.apache.spark.sql.functions.udf 

def in(s: Set[Int]) = udf((x: Int) => if (s.contains(x)) 1 else 0) 
df.withColumn("Selected_Group", in((1 to 4).toSet)($"group")) 

Avec SQL brut:

df.registerTempTable("df") 
sqlContext.sql(
    "SELECT *, CAST(group IN (1, 2, 3, 4) AS INT) AS Selected_Group FROM df" 
) 

Avec Column.in méthode:

import org.apache.spark.sql.functions.{lit, when} 
import org.apache.spark.sql.types.IntegerType 

df.withColumn(
    "Selected_Group", 
    $"group".in((1 to 4).map(lit): _*).cast(IntegerType)) 

ou when fonction:

df 
.withColumn(
    "Selected_Group", 
    when($"group".in((1 to 4).map(lit): _*), 1).otherwise(0))