如何使用 Spark Dataframe API 对 Case-Otherwise 语句应用多个条件答案

【问题标题】：How To Apply Multiple Conditions on Case-Otherwise Statement Using Spark Dataframe API如何使用 Spark Dataframe API 对 Case-Otherwise 语句应用多个条件
【发布时间】：2015-11-04 11:28:09
【问题描述】：

我正在尝试使用Spark Dataframe API 中的withColumn 语句向现有数据框添加新列。下面的代码有效，但我想知道是否有一种方法可以选择多个组。比方说第 1、2、3、4 组，而不仅仅是第 1 组。我想我可以写四次when 语句。我在一些帖子中看到人们这样做。但是，在R 中，有一个%in% 运算符可以指定一个变量是否包含向量中的值，但我不知道Spark 中是否有这样的东西。我查看了 Spark API 文档，但大多数函数不包含任何示例。

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

选择第 1 组的 Spark Dataframe 示例代码：

 val df1 = df.withColumn("Selected_Group", when($"Group" === 1, 1).otherwise(0))

数据

ID、组
1, 0
2、1
3、2
. .
. .
100, 99

【问题讨论】：

标签： r apache-spark dataframe apache-spark-sql

【解决方案1】：

使用 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"))

使用原始 SQL：

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

用Column.in方法：

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

或when函数：

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

【讨论】：