如何在PySpark Feature Store中使用score_batch获取概率值？

核心思路 #

要让score_batch返回概率值而非仅预测标签，需在日志模型阶段配置模型输出逻辑，让模型同时输出标签和概率，而非仅默认的预测标签。以下是两种实用方案：

方案一：用PyFunc封装模型（通用所有模型类型） #

通过自定义PyFunc模型封装原分类器，让它在预测时同时返回标签和概率列：

1. 定义自定义PyFunc模型类 #

from databricks.feature_engineering import FeatureEngineeringClient
import mlflow.pyfunc
import pandas as pd

class ProbabilityOutputModel(mlflow.pyfunc.PythonModel):
    def __init__(self, base_model):
        self.base_model = base_model
    
    def predict(self, context, input_df):
        # 生成预测标签
        pred_labels = self.base_model.predict(input_df)
        # 二分类场景取正类概率，多分类可保留全概率矩阵
        pred_probs = self.base_model.predict_proba(input_df)[:, 1]
        # 多分类场景可替换为：pred_probs = self.base_model.predict_proba(input_df)
        # 将结果合并到输入DataFrame返回
        return input_df.assign(prediction=pred_labels, probability=pred_probs)

2. 日志封装后的模型 #

# 初始化封装模型
wrapped_model = ProbabilityOutputModel(base_model=model)

# 日志到特征存储
fe = FeatureEngineeringClient()
fe.log_model(
    model=wrapped_model,
    artifact_path=artifact_path,
    flavor=mlflow.pyfunc,
    training_set=training_set,
    registered_model_name=model_name
)

方案二：针对Sklearn模型直接配置（更简洁） #

如果使用Sklearn分类器（如RandomForestClassifier），可在日志模型时指定自定义预测函数和模型签名，直接输出标签+概率：

1. 定义模型签名（指定输出列） #

from mlflow.models.signature import ModelSignature
from mlflow.types import DataType, Schema, ColSpec

# 根据你的特征定义输入Schema
input_schema = Schema([
    ColSpec(DataType.double, "feature_col1"),
    ColSpec(DataType.double, "feature_col2"),
    # 补充所有特征列
])

# 定义输出Schema：包含预测标签和概率
output_schema = Schema([
    ColSpec(DataType.integer, "prediction"),  # 标签列
    ColSpec(DataType.double, "probability")   # 概率列
])
model_signature = ModelSignature(inputs=input_schema, outputs=output_schema)

2. 日志模型时配置预测函数 #

fe = FeatureEngineeringClient()

# 自定义预测函数：同时返回标签和概率
def custom_predict_fn(model, data):
    labels = model.predict(data)
    probs = model.predict_proba(data)[:, 1]  # 二分类取正类概率
    return pd.DataFrame({"prediction": labels, "probability": probs})

# 日志模型
fe.log_model(
    model=model,
    artifact_path=artifact_path,
    flavor=mlflow.sklearn,
    training_set=training_set,
    registered_model_name=model_name,
    predict_fn=custom_predict_fn,
    signature=model_signature
)

批量打分推理 #

完成模型日志后，直接调用score_batch即可得到包含预测标签和概率的结果：

prediction_df = fe.score_batch(
  model_uri=uc_modeling.get_lastest_model_uri(), 
  df=batch_input_df
)

# 查看结果列
print(prediction_df.columns)
# 输出包含：原输入列 + prediction（标签） + probability（概率）

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