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

Question

我尝试使用 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))

Data

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

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