Spark Dataset使用flatMap转换时报TaskNotSerializable错误如何修复？

问题根源 #

TaskNotSerializable是Spark闭包序列化的典型报错，本次触发原因如下：

• 你在flatMap算子中传入的匿名函数隐式捕获了它所属的外部类job.DataSink的实例，而DataSink没有实现序列化接口，Spark需要将闭包序列化分发到Executor节点时校验失败。
• 你写的DataFrame版本逻辑没有报错，是因为DataFrame高阶API的算子逻辑会被Spark内部Catalyst引擎解析为逻辑计划再序列化分发，不需要序列化用户定义的外部类实例，自然不会触发序列化校验。

修复方案 #

方案1：切断闭包对外部类的引用 #

将需要传入闭包的变量提前声明为方法内局部变量，避免匿名函数捕获外部类实例即可：

def filterData(input: Dataset[CustomType], idsList: List[Int]): Dataset[CustomType] = {
  // 声明为方法局部变量，切断闭包对外部类的引用
  val validIds = idsList
  input.flatMap { record =>
    val filtered = record.data.filter(rec => validIds.contains(rec.id))
    if (filtered.nonEmpty) {
      Seq(record.withFields(filtered))
    } else {
      Iterable.empty
    }
  }
}

方案2：使用DataFrame API实现（推荐） #

你已经写过对应的DataFrame版本，该版本不仅没有序列化风险，性能还比自定义flatMap更高，Spark会自动优化执行计划，直接复用即可：

def filterData(input: Dataset[CustomType], idsList: List[Int]): Dataset[CustomType] = {
  import org.apache.spark.sql.functions._
  input.withColumn("arr", explode($"data"))
    .filter($"arr.id".isin(idsList: _*))
    .drop("arr")
    .as[CustomType]
}

如果idsList规模较大（超过1000条），可以进一步优化为array_exists+广播变量的实现，避免explode带来的数据膨胀：

def filterData(input: Dataset[CustomType], idsList: List[Int]): Dataset[CustomType] = {
  import org.apache.spark.sql.functions._
  // 广播id列表，避免重复分发
  val bcValidIds = input.sparkSession.sparkContext.broadcast(idsList.toSet)
  input.filter(array_exists($"data", item => bcValidIds.value.contains(item.getAs[Int]("id"))))
    .as[CustomType]
}

方案3：外部类实现序列化接口 #

如果必须使用Dataset的函数式API，也可以直接让job.DataSink类实现Serializable接口，不过该方案不推荐，后续外部类新增不可序列化字段时会再次触发相同问题。

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