Question

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

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

数据

ID，组
   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 API在Case-otherwise语句上应用多个条件

1 个答案: