Pyspark下训练测试集唯一值不同时One hot encoding列对齐实现

方案1：使用Spark ML Pipeline（最推荐，生产环境首选） #

核心思路：所有特征编码的转换器仅在训练集上拟合，固定分类列的编码规则，测试集直接套用训练好的转换器即可，无需单独编码。

• 第一步：针对分类列拟合StringIndexer，仅输入训练集数据，得到的标签字典仅包含训练集的唯一值，设置handleInvalid="keep"参数，测试集中未出现的类别会被统一标记为一个额外的索引值。
• 第二步：使用训练集拟合OneHotEncoder，设置dropLast=false保证训练集所有类别都生成对应的编码位。
• 第三步：将上述两个组件组装为Pipeline，仅在训练集上拟合得到PipelineModel，直接用该模型转换测试集，最后将OHE输出的向量列截取前N位（N为训练集该列的唯一值数量），即可得到和训练集完全对齐的独热编码，未知类别自动对应全0向量。

代码示例（PySpark）：

from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorSlicer

# 假设训练集为train_df，测试集为test_df，分类列名为col
# 1. 定义StringIndexer，仅用训练集拟合
indexer = StringIndexer(inputCol="col", outputCol="col_index", handleInvalid="keep")
# 2. 定义OneHotEncoder
encoder = OneHotEncoder(inputCol="col_index", outputCol="col_ohe_raw", dropLast=False)
# 3. 提前获取训练集分类列的唯一值数量
indexer_model = indexer.fit(train_df)
train_label_count = len(indexer_model.labels)
# 4. 定义向量切片器，仅保留训练集对应类别的编码位，过滤未知类别的额外编码位
slicer = VectorSlicer(inputCol="col_ohe_raw", outputCol="col_ohe", indices=list(range(train_label_count)))
# 5. 组装Pipeline并拟合训练集
pipeline = Pipeline(stages=[indexer, encoder, slicer])
pipeline_model = pipeline.fit(train_df)

# 转换训练集和测试集，得到的col_ohe列完全对齐
train_ohe_df = pipeline_model.transform(train_df)
test_ohe_df = pipeline_model.transform(test_df)

按你的示例，转换后的测试集col_ohe列对应规则为：a→[1,0,0]，e→[0,0,0]，f→[0,0,0]，完全符合预期。

方案2：自定义UDF生成显式列（适合需要单独OHE列而非向量列的场景） #

核心思路：先提取训练集分类列的所有唯一值列表，固定顺序和列名，直接遍历列表为每个值生成对应的0/1列，测试集不在列表中的值所有列都为0。

代码示例（PySpark）：

from pyspark.sql import functions as F

# 提取训练集分类列的唯一值列表，顺序固定
train_unique_vals = [row[0] for row in train_df.select("col").distinct().collect()]
# 遍历生成每列的OHE结果
for val in train_unique_vals:
    test_df = test_df.withColumn(f"col_{val}", F.when(F.col("col") == val, 1).otherwise(0))
# 最终测试集仅保留col_a、col_b、col_c三列，和训练集完全对齐

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