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