如果列存在于set中，则过滤spark/scala数据框架

Question

我使用的是Spark 1.4.0，这是我目前的版本:

data.filter($"myColumn".in(lit("A"), lit("B"), lit("C"), ...))

函数lit将字面值转换为列。

理想情况下，我会把我的A, B, C放在一个集合中，并像这样检查:

val validValues = Set("A", "B", "C", ...)
data.filter($"myColumn".in(validValues))

正确的语法是什么?是否有其他简明的解决方案?

Answer 1

Spark 1.4及以上版本:

val validValues = Set("A", "B", "C").map(lit(_))
data.filter($"myColumn".in(validValues.toSeq: _*))

Spark 1.5及以上版本:

val validValues = Set("A", "B", "C")
data.filter($"myColumn".isin(validValues.toSeq: _*))

Answer 2

此PR已合并到Spark 2.4中。现在可以输入

val profileDF = Seq(
  Some(1), Some(2), Some(3), Some(4),
  Some(5), Some(6), Some(7), None
).toDF("profileID")
val validUsers: Set[Any] = Set(6, 7.toShort, 8L, "3")
val result = profileDF.withColumn("isValid", $"profileID".isInCollection(validUsers))
result.show(10)
"""
+---------+-------+
|profileID|isValid|
+---------+-------+
|        1|  false|
|        2|  false|
|        3|   true|
|        4|  false|
|        5|  false|
|        6|   true|
|        7|   true|
|     null|   null|
+---------+-------+
 """.stripMargin