如何在pyspark.ml管道中创建自定义常量值填充器并优化缺失值填充？

要实现支持常量值填充的自定义转换器，你需要继承PySpark ML的Transformer基类，并实现核心的参数定义、转换逻辑和Schema校验方法。以下是具体实现步骤和代码示例：

1. 导入依赖模块 #

from pyspark.ml import Transformer
from pyspark.ml.param.shared import HasInputCols, HasOutputCols, Param, TypeConverters
from pyspark.sql import DataFrame
from pyspark.sql.functions import col, lit

2. 自定义常量填充转换器类 #

class ConstantImputer(Transformer, HasInputCols, HasOutputCols):
    # 定义填充常量的参数，支持类型转换
    fill_value = Param(
        Param._dummy(), 
        "fill_value", 
        "用于填充缺失值的常量", 
        typeConverter=TypeConverters.toFloat
    )

    def __init__(self, inputCols=None, outputCols=None, fill_value=0.0):
        super().__init__()
        self._setDefault(fill_value=0.0)
        self.setInputCols(inputCols)
        self.setOutputCols(outputCols)
        self.setFillValue(fill_value)

    # 填充值的设置/获取方法
    def setFillValue(self, value):
        return self._set(fill_value=value)
    
    def getFillValue(self):
        return self.getOrDefault(self.fill_value)

    # 核心转换逻辑
    def _transform(self, df: DataFrame) -> DataFrame:
        input_cols = self.getInputCols()
        output_cols = self.getOutputCols()
        fill_val = self.getFillValue()

        transformed_df = df
        # 遍历输入列，对每个列执行缺失值填充
        for in_col, out_col in zip(input_cols, output_cols):
            transformed_df = transformed_df.withColumn(out_col, col(in_col).fillna(fill_val))
        
        return transformed_df

    # Schema校验与输出Schema定义
    def _transformSchema(self, schema):
        input_cols = self.getInputCols()
        output_cols = self.getOutputCols()

        # 验证输入列存在
        for col_name in input_cols:
            if col_name not in schema.names:
                raise ValueError(f"输入列 {col_name} 不存在于数据Schema中")
        
        # 定义输出列的Schema，类型与输入列一致且非空
        for in_col, out_col in zip(input_cols, output_cols):
            col_type = schema[in_col].dataType
            schema = schema.add(out_col, col_type, nullable=False)
        
        return schema

3. 使用示例 #

# 创建测试数据集
data = [(1.0, None), (None, 2.0), (3.0, 4.0)]
df = spark.createDataFrame(data, ["col1", "col2"])

# 初始化转换器，指定输入列、输出列和填充常量
imputer = ConstantImputer(
    inputCols=["col1", "col2"], 
    outputCols=["col1_filled", "col2_filled"], 
    fill_value=0.0
)

# 执行转换
result_df = imputer.transform(df)
result_df.show()

执行后输出结果：

+----+----+-----------+-----------+
|col1|col2|col1_filled|col2_filled|
+----+----+-----------+-----------+
| 1.0|null|        1.0|        0.0|
|null| 2.0|        0.0|        2.0|
| 3.0| 4.0|        3.0|        4.0|
+----+----+-----------+-----------+

额外说明 #

• 若需要填充字符串类型常量，可修改fill_value的类型转换器为TypeConverters.toString，并调整默认值。
• 该转换器可直接集成到PySpark ML的Pipeline中，和其他Estimator/Transformer配合使用。
• 如果不需要保留原列，可将outputCols设置为与inputCols同名，直接覆盖原列数据。

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