如何基于条件高效过滤Spark DataFrame中的行

示例Spark DataFrame数据如下：

product     type
table     Furniture
chair     Furniture
TV        Electronic
...

需求规则：若数据集中存在任意type为Furniture的行，则删除所有type为Electronic的行。
实际场景数据规模达百万行级别，常规实现思路为先统计type为Furniture的行数，若统计值大于0则删除type为Electronic的行，该方案执行效率较低。

原有方案效率低的核心原因是触发了两次全表扫描：第一次扫描统计Furniture行数，第二次扫描执行过滤，额外IO开销很高。下面两种方案都可以大幅提升执行效率：

方案1：优化存在性判断逻辑（改造成本最低） #

绝大多数人写存在性判断时会误用count()，这个算子会遍历所有匹配行做计数，实际上判断是否存在只要找到第一条匹配记录就可以终止计算，完全不需要统计总数，优化后代码如下（以Scala API为例，PySpark逻辑一致）：

import org.apache.spark.sql.functions.col

// 找到第一条Furniture记录就返回，不会遍历全表所有匹配项
val hasFurniture = df.filter(col("type") === "Furniture").take(1).nonEmpty

val result = if (hasFurniture) {
  df.filter(col("type") =!= "Electronic")
} else {
  df
}

这个方案的存在性判断阶段不会做全表扫描，只要命中第一条Furniture记录就会终止计算，仅在最终过滤阶段做一次全表扫描，改造成本极低，性能比原count方案高数倍。

方案2：单次全表扫描实现（性能最优，超大数据量首选） #

如果要完全消除额外扫描开销，可以用累加器配合缓存实现单次遍历完成判断+过滤，全程只触发一次全表IO：

import org.apache.spark.sql.functions._
import org.apache.spark.util.LongAccumulator

// 初始化累加器，用来标记是否存在Furniture类型数据
val hasFurnitureAcc: LongAccumulator = spark.sparkContext.longAccumulator("hasFurniture")

// 遍历数据时，只要碰到Furniture类型就把累加器标记为1
val markedDf = df.withColumn("tmp_mark", when(col("type") === "Furniture", {
  hasFurnitureAcc.add(1)
  lit(true)
}).otherwise(lit(false)))

// 触发缓存和全表计算，填充累加器值
markedDf.cache()
markedDf.foreachPartition(_ => {})

val result = if (hasFurnitureAcc.value > 0) {
  markedDf.filter(col("type") =!= "Electronic").drop("tmp_mark")
} else {
  markedDf.drop("tmp_mark")
}

这个方案在第一次全表遍历读数据时，就同步完成了Furniture存在性的标记，后续过滤直接基于缓存的数据计算，没有额外的磁盘IO开销，百万级数据下执行速度比原方案快一倍左右。

优化提示：所有Spark场景下判断“是否存在至少N条匹配数据”的需求，都不要用count()，优先用take(N)，可以避免大量无效计算。

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