Spark集群使用H2O MOJO模型预测时的并行化问题排查

这个错误的核心原因是你在并行Task中引用了Driver端的非序列化对象，导致Spark无法正确将这些对象传递到Executor节点执行。具体来说：

• easyModel（以及底层的MojoModel）没有实现Java序列化接口，无法被Spark序列化后分发到Executor；
• 你在rowToRowData中直接使用了Driver端的df对象，而DataFrame内部包含RDD的依赖关系，这些依赖无法被正确序列化到Executor，从而触发了ClassCastException。

当你用collect()把数据拉到Driver端串行执行时，所有操作都在本地JVM中完成，不需要序列化这些对象，所以不会报错；但并行运行时，Spark需要把Task中引用的所有对象序列化后发送到Executor，这就暴露了序列化问题。

可行的解决方案 #

我们需要做两个关键调整：

1. 在每个Executor的Task中本地加载MOJO模型（而不是从Driver传递序列化后的模型）；
2. 避免在Task中引用Driver端的DataFrame对象，改用广播字段名来完成Row到RowData的转换。

方案1：使用mapPartitions + 广播字段名 + 本地加载模型

这个方案的优势是每个Partition只加载一次模型（比map中每条数据加载一次性能高很多），同时避免序列化Driver端的模型对象。

import hex.genmodel.easy.EasyPredictModelWrapper
import hex.genmodel.easy.RowData
import hex.genmodel.MojoModel

// 1. 广播DataFrame的字段名，避免在Task中引用Driver端的df对象
val fieldNames = df.schema.fieldNames
val broadcastFieldNames = spark.sparkContext.broadcast(fieldNames)

// 2. 使用mapPartitions，在每个Partition中本地加载模型并处理数据
val predictions = df.rdd.mapPartitions { iter =>
    // 每个Partition加载一次MOJO模型（确保所有Executor节点都能访问到该路径，比如HDFS或本地共享路径）
    val mojo = MojoModel.load("/path/to/mojo.zip")
    val easyModel = new EasyPredictModelWrapper(mojo)
    val fields = broadcastFieldNames.value

    iter.map { r =>
        // 转换Row到RowData，使用广播的字段名
        val rowAsMap = r.getValuesMap[Any](fields)
        val rowData = rowAsMap.foldLeft(new RowData()) {
            case (rd, (k, v)) =>
                if (v != null) {
                    rd.put(k, v.toString)
                }
                rd
        }
        // 执行预测
        val prediction = easyModel.predictBinomial(rowData).label
        (r.getAs[String]("id"), prediction.toInt)
    }
}.toDF("id", "prediction")

方案2：优化MOJO模型分发（无需本地文件）

如果你的Executor节点无法直接访问MOJO的本地路径，可以在Driver端把MOJO文件读成字节数组，通过广播变量分发到Executor，再从字节流加载模型：

import hex.genmodel.easy.EasyPredictModelWrapper
import hex.genmodel.easy.RowData
import hex.genmodel.MojoModel
import java.io.{ByteArrayInputStream, ByteArrayOutputStream, FileInputStream}

// 1. 在Driver端读取MOJO文件为字节数组
val baos = new ByteArrayOutputStream()
val fis = new FileInputStream("/path/to/mojo.zip")
val buffer = new Array[Byte](1024)
var len = fis.read(buffer)
while (len != -1) {
    baos.write(buffer, 0, len)
    len = fis.read(buffer)
}
fis.close()
val mojoBytes = baos.toByteArray
val broadcastMojoBytes = spark.sparkContext.broadcast(mojoBytes)

// 2. 广播字段名
val fieldNames = df.schema.fieldNames
val broadcastFieldNames = spark.sparkContext.broadcast(fieldNames)

// 3. mapPartitions中从广播的字节流加载模型
val predictions = df.rdd.mapPartitions { iter =>
    val bais = new ByteArrayInputStream(broadcastMojoBytes.value)
    val mojo = MojoModel.load(bais)
    val easyModel = new EasyPredictModelWrapper(mojo)
    val fields = broadcastFieldNames.value

    iter.map { r =>
        val rowAsMap = r.getValuesMap[Any](fields)
        val rowData = rowAsMap.foldLeft(new RowData()) {
            case (rd, (k, v)) =>
                if (v != null) {
                    rd.put(k, v.toString)
                }
                rd
        }
        val prediction = easyModel.predictBinomial(rowData).label
        (r.getAs[String]("id"), prediction.toInt)
    }
}.toDF("id", "prediction")

关键注意事项 #

• 确保MOJO模型的特征名和你的DataFrame字段名完全一致，否则预测会失败；
• mapPartitions中加载模型的逻辑会在每个Partition的Task中执行一次，同一个Executor上的多个Partition会各自加载模型（如果需要更高效的模型复用，可以考虑用Executor级别的单例，比如用lazy val封装模型加载）；
• 如果你使用的是H2O的最新版本，也可以检查是否有官方提供的Spark集成工具类，避免手动转换Row到RowData。

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