如何使用Scala过滤Spark DataFrame中的数组列

当然可行！而且你想在展开数组前过滤的思路特别棒——这样能避免不必要的数据膨胀，处理效率会高很多，毕竟你提到大部分单词都不带#。下面给你两种实用的实现方式，优先推荐第一种内置函数的方案，性能更优：

1. 使用Spark内置数组过滤函数（Spark 3.1+ 推荐） #

Spark 3.1及以上版本提供了专门的filter函数来直接操作数组列，完美匹配你的需求。不需要展开数组，直接在列层面过滤出以#开头的元素：

import org.apache.spark.sql.functions.{col, filter}

// 假设你的原始DataFrame名为df
val filteredDf = df.withColumn(
  "filtered_words", // 生成的过滤后数组列名
  filter(col("words"), word => word.startsWith("#"))
)

说明： #

• filter函数接收两个参数：要处理的数组列，以及一个判断每个元素是否保留的lambda表达式。
• 处理后会生成新列filtered_words，仅包含原数组中以#开头的单词；如果原数组没有符合条件的元素，新列会是空数组。
• 这种方式完全基于Spark内置优化，执行效率远高于先展开再过滤的方式，也比自定义UDF更高效。

2. 低版本Spark兼容方案（自定义UDF） #

如果你的Spark版本低于3.1，没有内置的数组filter函数，可以用自定义UDF来实现相同逻辑：

import org.apache.spark.sql.functions.udf

// 定义UDF：过滤数组中以#开头的元素
val filterHashtagsUdf = udf((words: Seq[String]) => {
  words.filter(_.startsWith("#"))
})

// 将UDF应用到DataFrame
val filteredDf = df.withColumn(
  "filtered_words",
  filterHashtagsUdf(col("words"))
)

说明： #

• 这个UDF接收字符串序列作为输入，返回过滤后的序列。
• 注意：UDF的性能通常不如Spark内置函数，因为内置函数是经过Spark执行计划优化的，所以如果条件允许，尽量升级到Spark 3.1+使用第一种方案。

测试示例 #

你可以用下面的测试数据验证效果：

// 创建测试DataFrame
val testData = Seq(
  ("2024-01-01", Array("hello", "#spark", "#scala", "world")),
  ("2024-01-02", Array("foo", "bar")),
  ("2024-01-03", Array("#bigdata", "dataengineering"))
).toDF("created_at", "words")

// 应用过滤逻辑后，filtered_words列的结果为：
// ["#spark", "#scala"], [], ["#bigdata"]

如果不需要保留原words列，你可以直接覆盖它（把"filtered_words"改成"words"），或者用filteredDf.drop("words")删除原列。

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