如何高效将Scala Spark Dataset按规则拆分为三个独立Dataset？

问题场景 #

现有Scala Spark Dataset（命名为ds），以及两个自定义判断函数isTypeA()和isTypeB()：

• isTypeA(row)：判断该行是否属于A类
• isTypeB(row)：判断该行是否属于B类

分类规则：

1. 若行满足isTypeA，无论是否满足isTypeB，均归为A类
2. 若行不满足isTypeA但满足isTypeB，归为B类
3. 若行既不满足isTypeA也不满足isTypeB，归为C类

当前实现通过多次调用filter：

val a = ds.filter(isTypeA(_))
val b = ds.filter(row => !isTypeA(row) && isTypeB(row))
val c = ds.filter(row => !isTypeA(row) && !isTypeB(row))

但该方式会导致Spark多次扫描源Dataset，性能存在优化空间。

高效实现方案 #

核心思路是仅扫描一次源Dataset，先为每行添加分类标识，再基于标识拆分出三个Dataset，避免重复计算。

步骤1：将自定义判断函数转为UDF（若使用自定义函数） #

如果isTypeA和isTypeB是自定义Scala函数，需要先转为Spark UDF以便在Dataset操作中使用：

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

// 将自定义函数转为UDF
val isTypeAUdf = udf(isTypeA _)
val isTypeBUdf = udf(isTypeB _)

步骤2：添加分类列 #

通过withColumn结合when表达式，一次性完成所有行的分类标记：

// 构造包含所有字段的结构体，传入UDF（若函数需要整行数据）
val allColumnsStruct = struct(ds.columns.map(col): _*)

val categorizedDs = ds.withColumn("category",
  // 优先判断A类
  when(isTypeAUdf(allColumnsStruct), "A")
  // 不符合A类时判断B类
  .when(not(isTypeAUdf(allColumnsStruct)) && isTypeBUdf(allColumnsStruct), "B")
  // 其余归为C类
  .otherwise("C")
)

步骤3：拆分出三类Dataset #

基于已标记的category列过滤，得到最终的三个Dataset：

val a = categorizedDs.filter(col("category") === "A").drop("category")
val b = categorizedDs.filter(col("category") === "B").drop("category")
val c = categorizedDs.filter(col("category") === "C").drop("category")

额外优化建议 #

• 若isTypeA和isTypeB的判断逻辑可以用Spark内置SQL函数实现（比如字段比较、字符串匹配等），尽量避免使用UDF。内置函数是Spark可优化的白盒，性能远优于UDF。
  示例（假设isTypeA是判断age字段大于30）：

  val categorizedDs = ds.withColumn("category",
    when(col("age") > 30, "A")
    .when(col("age") <= 30 && col("score") > 80, "B")
    .otherwise("C")
  )
  

• 如果源Dataset数据量极大，可考虑先对categorizedDs进行缓存（categorizedDs.cache()），进一步提升后续过滤操作的速度。

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