Scala+Spark：如何基于列表元素模糊匹配过滤DataFrame

原始数据与需求 #

我有如下Spark DataFrame：

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

需要编写一个函数，基于列表List("r2", "hs", "np")中的元素对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)

// 预期结果行：
("1", "r2_test"), // 包含"r2"
("3", "hs_2_card"), // 包含"hs"
("4", "vsx_np_v2"), // 包含"np"
("5", "r2_test"), // 包含"r2"

修改方案 #

可以通过两种方式实现一次filter完成多元素匹配：

方法1：多条件逻辑或拼接 #

遍历列表生成每个元素的contains条件，再用逻辑或连接所有条件：

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

def filteredElements(df: DataFrame): DataFrame = {
  val elements = List("r2", "hs", "np")
  // 生成所有元素的contains条件，并用OR连接
  val filterCondition = elements.map(element => $"my_column".contains(element)).reduce(_ || _)
  df.filter(filterCondition)
}

方法2：正则表达式匹配 #

将列表元素拼接成正则或模式，用rlike进行匹配：

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

def filteredElements(df: DataFrame): DataFrame = {
  val elements = List("r2", "hs", "np")
  // 拼接成正则或模式："r2|hs|np"
  val regexPattern = elements.mkString("|")
  df.filter($"my_column".rlike(regexPattern))
}

两种方法都无需多次调用filter，执行test.transform(filteredElements)即可得到预期结果。

内容的提问来源于stack exchange，提问作者Malkath