Scala + Spark：过滤包含列表元素的数据集

Question

我有一个数据集，我想根据列进行过滤。

val test = Seq(
("1", "r2_test"),
("2", "some_other_value"),
("3", "hs_2_card"),
("4", "vsx_np_v2"),
("5", "r2_test"),
("2", "some_other_value2")
).toDF("id", "my_column")

我想创建一个函数来根据此列表的元素过滤我的数据框，使用“my_column”上的包含（如果包含字符串的一部分，则必须应用过滤器）

import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions._

def filteredElements(df: DataFrame): DataFrame = {
   val elements = List("r2", "hs", "np")
   df.filter($"my_column".contains(elements))
}

但是像这样，不适用于列表，只适用于单个元素。 我该怎么做才能适应使用我的列表而不必执行多个过滤器？

下面是应用函数时预期的输出必须如何

val output = test.transform(filteredElements)

expected =
("1", "r2_test"), // contains "rs"
("3", "hs_2_card"), // contains "hs"
("4", "vsx_np_v2"), // contains "np"
("5", "r2_test"), // contains "r2"

Answer 1

解决这个问题的一种方法是使用UDF 。 我认为应该有一些方法可以用我不知道的 spark sql 函数来解决这个问题。 无论如何，您可以定义一个 udf 来告诉天气 String 是否包含元素列表中的任何值：

import org.apache.sql.functions._
val elements = List("r2", "hs", "np")

val isContainedInList = udf { (value: String) => 
  elements.exists(e => value.indexOf(e) != -1)
}

您可以在选择、筛选中使用这个 udf，基本上在任何您想要的地方：

def filteredElements(df: DataFrame): DataFrame = {
   df.filter(isContainedInList($"my_column"))
}

结果如预期的那样：

+---+---------+
| id|my_column|
+---+---------+
|  1|  r2_test|
|  3|hs_2_card|
|  4|vsx_np_v2|
|  5|  r2_test|
+---+---------+

Answer 2

您可以在没有 udf 的情况下在一行中完成（性能更好且更简单）：

df.filter(col("my_column").isNotNull).filter(row => elements.exists(row.getAs[String]("my_column").contains)).show()