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

Question

我正在尝试使用 Spark Dataframe API 中的 withColumn 语句向现有数据框添加新列。下面的代码有效，但我想知道是否有一种方法可以 select 多个组。让我们说第 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))

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

How To Apply Multiple Conditions on Case-Otherwise Statement Using Spark Dataframe API

r

dataframe

apache-spark

apache-spark-sql