如何从PySpark ML Pipeline对象中获取转换后的数据集？

嘿，刚入门Spark ML的时候我也被Pipeline这个点绕晕过！别担心，咱们一步步来搞定用Pipeline做链式ETL转换，并且顺利拿到最终的DataFrame。

核心逻辑先搞懂 #

Spark ML的Pipeline本质是把一系列数据处理步骤串起来，但它的fit()方法不是直接输出转换后的数据——而是生成一个PipelineModel（包含了所有步骤的执行逻辑，要是有需要训练的组件也会训练好）。真正生成转换后DataFrame的是PipelineModel的transform()方法！

用代码示例一步步来 #

假设我们要做的ETL是：输入包含姓名、年龄、城市的DataFrame，依次新增年龄分组列、城市索引列、城市独热向量列，最终输出所有列都保留的结果。

1. 导入依赖并创建基础数据

from pyspark.sql import SparkSession
from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, OneHotEncoder
from pyspark.ml.pipeline import Transformer
from pyspark.sql.functions import col, udf
from pyspark.sql.types import StringType

# 初始化SparkSession
spark = SparkSession.builder.appName("ETLPipelineDemo").getOrCreate()

# 示例输入DataFrame
raw_df = spark.createDataFrame(
    [("Alice", 25, "New York"), ("Bob", 30, "London"), ("Charlie", 35, "Paris")],
    ["name", "age", "city"]
)

2. 定义你的链式转换步骤

这里既有Spark自带的Transformer，也可以自定义符合你需求的ETL转换（比如新增年龄分组列）：

# 自定义Transformer：根据年龄新增age_group列
class AgeGroupTransformer(Transformer):
    def __init__(self):
        super(AgeGroupTransformer, self).__init__()
    
    def _transform(self, df):
        # 定义UDF来生成年龄分组
        def get_group(age):
            return "Young" if age < 30 else "Adult"
        age_group_udf = udf(get_group, StringType())
        return df.withColumn("age_group", age_group_udf(col("age")))

# 步骤1：新增年龄分组列
age_stage = AgeGroupTransformer()
# 步骤2：将城市字符串转为索引列（新增city_index）
city_index_stage = StringIndexer(inputCol="city", outputCol="city_index", handleInvalid="keep")
# 步骤3：将城市索引转为独热向量列（新增city_onehot）
city_onehot_stage = OneHotEncoder(inputCol="city_index", outputCol="city_onehot")

3. 组装Pipeline并执行转换

# 把所有步骤按顺序组装成Pipeline
pipeline = Pipeline(stages=[age_stage, city_index_stage, city_onehot_stage])

# 第一步：fit生成PipelineModel（对于纯ETL步骤，这里主要是验证流程，不会做训练）
pipeline_model = pipeline.fit(raw_df)

# 第二步：用transform执行所有转换，得到最终的DataFrame！
transformed_df = pipeline_model.transform(raw_df)

4. 查看结果

# 打印所有数据
transformed_df.show(truncate=False)

# 查看Schema确认新增的列
transformed_df.printSchema()

关键细节提醒 #

• 如果你的Pipeline里包含Estimator（比如训练机器学习模型的步骤），fit()会先训练这些组件生成对应的Model，再把所有步骤整合成PipelineModel；而纯ETL的Pipeline（全是Transformer），fit()只是做流程验证，真正的转换都在transform()里。
• transform()会保留原始DataFrame的所有列，同时新增每个步骤定义的输出列，完全符合你“每次新增一列”的需求。
• 要是后续需要重复用这套转换逻辑处理新数据，直接用已经生成的pipeline_model.transform(new_df)就行，不用重新fit。

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