如果集合中存在列，则过滤 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

This 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

如果集合中存在列，则过滤 spark/scala 数据框

Filter spark/scala dataframe if column is present in set

scala

filter

apache-spark

spark-dataframe