Pyspark 利用Pandas UDF 进行模型预测

当我们训练好了一个模型，想在大规模的数据上预测时，pyspark的Pandas UDF在加载模型并使用模型对数据进行分布式预测很有用。

但在spark 3.0 以前，当模型很大时，因为在同一个python工作进程中，每个批次Pandas UDF都会反复加载同一个模型，会造成很高的额外开销。

Spark 3.0 以后， 因为Pandas UDF支持Series to Series， Iterator of Series to Iterator of Series 等，所以就只用加载一次模型。

下面的示例代码给出了一个使用pyspark 调用模型训练的示例

#  一种模型预测示例
import os
from pyspark.sql import functions as F
from pyspark.sql.functions import pandas_udf

# 这里的定义变量是因为模型存储在s3上
filename = '/tmp/catboost_model_1'
dest_path = 'remote_filesystem/catboost_model_1'

@pandas_udf("float")
def predict(features):
    # 定义模型
    model_cb = CatBoostClassifier()
    # 如果不能保证spark的每个节点上都已经有模型了，一定要作下面的判断，不然会报错类似 Model file doesn't exist: /tmp/catboost_model_1 的错误
    if not os.path.exists(filename):
        cmd = '/usr/lib/hadoop/bin/hadoop' + '-get -f {} {}'.format(dest_path, filename)
        os.system(cmd)
    # 加载模型
    model_cb.load_model(filename, format='cbm')
    logger.debug('finished load model')
    # 预测
    # return pd.Series(model_cb.predict(features)) # 直接预测结果
    return pd.Series(model_cb.predict_proba(features)[:, 1]) # 预测概率

# 将输入特征列组成结构体
feature_col = F.struct(F.col('feature1'), F.col('feature2'), F.col('feature3'))
df = df.withColumn('prediction', predict(feature_col))