如何使用PySpark创建自定义Transformer并集成至Pipeline以避免数据泄露？

问题分析 #

你当前的实现直接基于全量数据集计算概率，导致训练集引入了测试集的未来时间窗口数据，引发数据泄露。通过自定义PySpark Transformer并集成到Pipeline，可确保训练阶段仅使用训练集数据计算，测试阶段独立处理测试集，从根源避免泄露。

自定义Transformer实现 #

继承pyspark.ml.Transformer，将贝叶斯概率估算逻辑封装到_transform方法中，并支持平滑系数rho的参数配置：

from pyspark.ml import Transformer
from pyspark.ml.param.shared import HasInputCols, HasOutputCol, Param, Params, TypeConverters
from pyspark.sql import DataFrame
import pyspark.sql.functions as F
from pyspark.sql.window import Window

class BayesianPassProbabilityTransformer(Transformer, HasInputCols, HasOutputCol):
    # 定义平滑系数rho参数
    rho = Param(Params._dummy(), "rho", "Smoothing parameter for Bayesian estimation", TypeConverters.toFloat)
    
    def __init__(self, inputCols=None, outputCol=None, rho=0.3):
        super().__init__()
        self._setDefault(rho=0.3)
        self.setInputCols(inputCols)
        self.setOutputCol(outputCol)
        self.setRho(rho)
        
    def setRho(self, value):
        return self._set(rho=value)
    
    def getRho(self):
        return self.getOrDefault(self.rho)
    
    def _transform(self, df: DataFrame) -> DataFrame:
        # 1. 按softbin和时间窗口聚合统计PASS/FAIL总数
        softbins_df = df.groupBy('softbin_first_test', 'time_window') \
            .agg(
                F.sum('PASS').alias('total_pass'),
                F.sum('FAIL').alias('total_fail')
            )
        
        # 2. 计算基础通过率/失败率
        softbins_total_df = softbins_df.withColumn('total', F.col('total_pass') + F.col('total_fail')) \
            .withColumn('prior_pass_prob', F.col('total_pass') / F.col('total')) \
            .withColumn('prior_fail_prob', F.col('total_fail') / F.col('total'))
        
        # 3. 基于滚动窗口的贝叶斯估算
        rho_val = self.getRho()
        prior_window = Window.partitionBy('softbin_first_test') \
            .orderBy(F.col('time_window')) \
            .rowsBetween(Window.unboundedPreceding, Window.currentRow)
        
        result_df = softbins_total_df.withColumn('prior_pass', F.mean('prior_pass_prob').over(prior_window)) \
            .withColumn('prior_fail', F.mean('prior_fail_prob').over(prior_window)) \
            .withColumn('alpha', F.col('prior_pass') * (1 - rho_val) / rho_val) \
            .withColumn('beta', F.col('prior_fail') * (1 - rho_val) / rho_val) \
            .withColumn(self.getOutputCol(), F.col('alpha') / (F.col('alpha') + F.col('beta')))
        
        # 将计算结果关联回原始数据集（如需保留原始所有列）
        return df.join(result_df, on=['softbin_first_test', 'time_window'], how='left')

集成到Pipeline并解决数据泄露 #

将自定义Transformer加入Pipeline后，训练阶段仅针对训练集计算滚动概率，测试阶段独立处理测试集，完全隔离数据：

# 1. 按时间拆分训练/测试集（示例按日期拆分，根据实际时间格式调整）
train_df = counts_df.filter(F.col('time_window') <= '2023-06-01')
test_df = counts_df.filter(F.col('time_window') > '2023-06-01')

# 2. 初始化Transformer实例
prob_transformer = BayesianPassProbabilityTransformer(
    inputCols=['softbin_first_test', 'time_window', 'PASS', 'FAIL'],
    outputCol='pass_prob',
    rho=0.3
)

# 3. 构建Pipeline（可添加其他步骤如特征工程、模型训练）
from pyspark.ml import Pipeline
pipeline = Pipeline(stages=[prob_transformer])

# 4. 训练Pipeline（Transformer无拟合逻辑，仅验证参数）
pipeline_model = pipeline.fit(train_df)

# 5. 分别转换训练集和测试集
train_result = pipeline_model.transform(train_df)
test_result = pipeline_model.transform(test_df)

关键说明 #

• 数据隔离：训练集转换时，滚动窗口仅包含训练集内的历史时间数据，不会触及测试集的未来数据；测试集转换时同理，仅使用自身的时间序列数据，彻底避免泄露。
• 灵活性：可通过setRho方法调整平滑系数，也可修改_transform方法适配已聚合的输入数据（跳过第一步groupBy）。
• 兼容性：该Transformer可与PySpark Pipeline的其他组件（如VectorAssembler、LogisticRegression）无缝集成，构建完整的机器学习工作流。

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