在PySpark DataFrame中运行H2O模型预测遇ModuleNotFoundError问题求助

错误根源 #

ModuleNotFoundError: No module named 'h2o'的核心原因是Spark所有Executor节点未安装h2o库。Pandas UDF会分发到各个Executor节点执行，仅在Driver端安装h2o无法满足需求，必须保证所有Worker节点都有对应版本的h2o依赖。

此外你的代码存在语法与逻辑问题，一并修正如下：

解决方案步骤 #

1. 全集群安装h2o依赖 #

在所有Spark Driver和Executor节点上执行：

pip install h2o

如果是集群环境，可通过Spark提交参数--py-files上传h2o的whl包，或配置spark.executorEnv.PYSPARK_PYTHON指定统一的、已安装h2o的Python环境，确保依赖一致性。

2. 修正Pandas UDF代码 #

原代码存在函数定义缺冒号、pd.series大小写错误、h2o未初始化、模型未正确分发等问题，修正后代码示例：

import pandas as pd
import h2o
from pyspark.sql.functions import pandas_udf, col

# Driver端提前加载模型并广播，避免每个Task重复加载
h2o.init()
model = h2o.load_model("/path/to/your/h2o/model")
broadcast_model = spark.sparkContext.broadcast(model)

# 指定UDF返回值类型，根据模型类型调整（如分类模型用"string"）
@pandas_udf("double")
def predict_h2o_model(*cols):
    # Executor进程独立，需单独初始化h2o
    h2o.init(enable_assertions=False, nthreads=-1)
    # 获取广播的模型实例
    model = broadcast_model.value
    
    # 拼接输入特征列
    x = pd.concat(cols, axis=1)
    # 转换为H2OFrame
    h2o_df = h2o.H2OFrame(x)
    # 执行预测
    scores = model.predict(h2o_df)
    # 转换为Pandas Series并返回（注意大写S，且指定预测列名）
    return pd.Series(scores.as_data_frame()['predict'])

# 正确调用UDF：用col()包装特征列
feature_cols = ["feat1", "feat2", "feat3"]  # 替换为你的实际特征列名
df_scores = spark_df.select(
    col("cust_id"),
    predict_h2o_model(*[col(c) for c in feature_cols]).alias('model_score')
)
df_scores.show()

关键注意事项 #

• Executor端h2o初始化：每个Executor是独立进程，必须单独初始化h2o，无法复用Driver端的实例。
• 模型广播：通过广播变量传递已加载的模型，减少重复加载开销，提升执行效率。
• 返回值类型匹配：必须根据模型输出指定正确的UDF返回类型，否则会出现类型不匹配错误。
• 版本一致性：所有节点的h2o版本需与Driver端完全一致，避免版本兼容问题。

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