如何将PySpark ML预测结果与erp_number标识数据关联

核心方案是保留erp_number列但不将其纳入特征，无需事后拼接——PySpark的Pipeline不会自动移除未被用作特征的原始列，只要确保它不被包含在特征处理流程中即可。

具体修改步骤 #

• 移除开头删除erp_number的代码，保留该标识列
• 在筛选分类列、数值列时，明确排除erp_number
• 确保特征列列表中不包含任何与erp_number相关的处理结果

修改后的完整代码 #

def create_predictions(data, module):
    # 保留erp_number列，不做删除操作
    # data = data.drop("erp_number")

    # 识别分类列，排除erp_number
    categorical_columns = [
        field.name for field in data.schema.fields 
        if isinstance(field.dataType, StringType) and field.name != "erp_number"
    ]

    # 识别数值列，排除分类列、目标列和erp_number
    numerical_columns = [
        field.name for field in data.schema.fields 
        if field.name not in categorical_columns 
        and field.name != module 
        and field.name != "erp_number"
    ]

    # 创建StringIndexer和OneHotEncoder处理流程
    stages = []
    for categorical_col in categorical_columns:
        string_indexer = StringIndexer(inputCol=categorical_col, outputCol=categorical_col + "_index", handleInvalid="keep")
        encoder = OneHotEncoder(inputCols=[string_indexer.getOutputCol()], outputCols=[categorical_col + "_vec"])
        stages += [string_indexer, encoder]

    # 组装特征列，仅包含处理后的分类特征和数值特征
    feature_columns = [c + "_vec" for c in categorical_columns] + numerical_columns
    assembler = VectorAssembler(inputCols=feature_columns, outputCol="features")
    stages += [assembler]

    # 添加GBT分类器到流程
    gbt = GBTClassifier(labelCol=module, featuresCol="features", predictionCol="prediction")
    stages += [gbt]

    # 构建并训练Pipeline
    pipeline = Pipeline(stages=stages)
    model = pipeline.fit(data)

    # 生成预测结果，erp_number列会保留在输出DataFrame中
    predictions = model.transform(data)

    return predictions

补充说明 #

如果之前尝试保留erp_number却报错，大概率是因为该列被误归类到了categorical_columns中，导致StringIndexer对其进行了处理，但后续未将其纳入特征列——直接在分类列筛选时排除它，就能避免这类不必要的处理和报错。

另外，PySpark其实支持类似bind_cols的列拼接操作，比如通过join（需保证行顺序一致，可临时添加索引列），但在这个场景下完全没必要，保留原始列的方式更高效简洁。

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