☰
Spark.ml机器学习指南:Pipeline、Transformer与Estimator实战解析
2026/10/21 7:37:03 网站建设 项目流程

很多人第一次看到Spark.ml这个写法,会以为它是 Spark MLlib 的简写,其实它指的是 Apache Spark 官方机器学习库里“基于 DataFrame 的那一套 API”。另一套老接口基于 RDD,现在官方更推荐用 DataFrame API。如果你正准备用 Spark 做大规模机器学习,想把单机 sklearn 的脚本往集群上搬,或者想彻底搞清楚 Pipeline、Transformer、Estimator 这些概念之间到底怎么配合,这篇文章值得你花十分钟看完。我会从“为什么选 Spark.ml”讲起,再到一份完整能跑通的示例代码,最后把我在实际调试中踩过的坑和排查思路一并列出来,尽量做到拿来就能用。

1. 先搞清楚:Spark.ml 到底是一套什么样的机器学习 API

1.1 它和 MLlib 的关系

Spark.ml 并不是一个独立于 MLlib 的东西,它其实就是 MLlib 的一部分,对应的是“基于 DataFrame 的机器学习流水线 API”。你可以在org.apache.spark.ml(Scala 里)或pyspark.ml(PySpark 里)引入这些类。而老的 RDD API 则在org.apache.spark.mllib下,比如以前常见到的LogisticRegressionWithSGD、SVMWithSGD这些类。

简单记忆方式:spark.ml是“新接口”,spark.mllib是“旧接口”。虽然mllib这个包名看起来像总称,但新特性基本都不会往它里面加了。我见过不少迁移项目,最大的工作量往往不是换算法,而是要把老的 RDD 格式的 LabeledPoint 改成 DataFrame 的列式输入,这一改,代码结构基本都重写了。

1.2 为什么官方后来把重心放到了 DataFrame 上

最核心的原因有三个:统一数据源、统一编程模型、统一发布上线。

DataFrame 本身带 schema,每一列都知道自己是什么类型,这能让 Catalyst 优化器在下游做很多自动优化,比如列裁剪、谓词下推。用老 RDD API 时,特征工程和算法训练之间经常要写一堆 map、reduce 操作,数据在不同阶段怎么组织全靠人肉记忆,代码一长就很容易失控。而 DataFrame 这套接口把特征表、样本表、预测结果表全部统一成一张表,特征工程用 Spark SQL 就能顺手完成,训练前只需要把特征合并成一个 Vector 列。

统一编程模型这点对我来说最重要。以前离线训练是一套代码,在线推理时再重写一套逻辑,训练时忘了某个特征归一化、上线时又加错字段,这种问题特别多。Spark.ml 用 Pipeline 把特征处理和模型训练串成一条链,训练时 fit,预测时 transform,同一份 Pipeline 代码可以直接用于批量推理,不用为上线再维护一套。

1.3 用 Spark.ml 做机器学习适合什么场景

Spark.ml 里内置的算法虽然不算特别多,但 LR、GBDT、随机森林、FM、ALS 这些工业界常用算法都有,覆盖了分类、回归、聚类、推荐、频繁模式挖掘这些主流场景。它真正舒服的领域是:数据量已经超过单机内存、特征工程比较复杂、并且希望训练和特征处理用同一条链路管理起来的大数据机器学习任务。

如果数据只有几万行,我会继续用 pandas + sklearn,没必要引入 Spark 的调度和序列化开销。如果是需要大规模 GPU 深度学习的任务,Spark.ml 也不是主要选择,它更偏向传统机器学习算法。认清这个边界,后面选型才不会跑偏。

2. 核心组件拆解:Transformer、Estimator 和 Pipeline 是怎么配合的

2.1 Transformer:把一张表变成另一张表

Transformer 的概念很直白:接口里有一个transform()方法,输入一个 DataFrame,输出一个新的 DataFrame。典型的例子是VectorAssembler,它把多列数值拼成一个向量列,本身不需要学习任何参数,输入一列,输出增加一个features列。

但并不是所有 Transformer 都像 VectorAssembler 一样无状态。以StandardScalerModel为例,它保存了训练阶段学到的均值和方差,预测阶段用这些保存的参数去归一化新数据。所以虽然它也是一个 Transformer,但它的产生依赖一次 fit 过程,这就要引出 Estimator 了。

2.2 Estimator:先 fit,再变成 Transformer

Estimator 的接口里有一个fit()方法,传入 DataFrame,返回一个 Model。这个 Model 本身就是一个 Transformer。把逻辑回归算法LogisticRegression看成就是一个 Estimator,调用fit()以后得到的LogisticRegressionModel,就是可以用来做预测的 Transformer。

这里有个容易混淆的地方:特征处理里也有 Estimator。比如StringIndexer,它需要先扫描数据,统计字符串到数字索引的映射,训练时 fit 出一张映射表,然后才能 transform 新数据,所以它也是 Estimator。而VectorAssembler不需要学任何东西,所以它是纯 Transformer。在 Pipeline 里,这两种东西会被统一处理。

提示:判断一个组件是 Transformer 还是 Estimator,就看它需要不需要 fit。需要 fit 才能用的,在 Pipeline 里就是 Estimator;不需要 fit、直接 transform 的,就是 Transformer。

2.3 Pipeline 的生命周期和持久化

Pipeline 就是一个 Stage 数组,每个 Stage 要么是 Transformer,要么是 Estimator。调用pipeline.fit(trainingData)时,Spark 会从头开始按顺序执行:如果当前 Stage 是 Transformer,就直接 transform;如果是 Estimator,就先 fit 得到 Model,然后让这个 Model 继续 transform,输出作为下一个 Stage 的输入。

训练完成后,你得到一个PipelineModel,这也就是最终的预测流水线。预测阶段不需要重复 fit,直接把新数据丢给pipelineModel.transform(testData),所有转换器都会用训练时学到的参数来处理。最大的好处就是训练代码和推理代码结构一致,不会出现“训练时做了归一化,推理时漏掉”的情况。

Pipeline 整条链路支持save和load,可以把模型和所有特征处理逻辑一起打包存到 HDFS、本地文件系统或对象存储上。上线时直接加载,不需要自己记着一串特征处理步骤。

2.4 Param 和 ParamMap 是调参的地基

每个 Estimator 和 Transformer 都有若干Param,比如maxDepth、numTrees、handleInvalid。你可以用 set 方法直接设值,比如rf.setMaxDepth(10);更灵活的方式是用ParamGridBuilder构造一组ParamMap,把不同参数组合传给CrossValidator去做网格搜索。

这里要注意,Pipeline 里某个 Estimator 的参数在参数网格里引用的是原始对象,比如保存在变量里的rf。用CrossValidator包住 Pipeline 以后,CrossValidator 会克隆 Pipeline 和里面的 Stage,所以不同参数组合之间不会互相污染,这是它内部已经处理好的事,不需要你手动 clone。

3. 完整实操:从原始业务表到一个可评估的分类模型

3.1 先造一份能跑通的最小场景

为了讲原理,我准备了一个简化版电商用户行为预测场景,目标是根据用户的访问特征预测是否产生复购。假设原始订单表经过 ETL 后长这样:

列名类型说明
user_idstring用户ID
ageint年龄
genderstring性别
income_levelstring收入等级,比如 low/mid/high
visit_cntdouble最近30天访问次数
purchase_amtdouble最近30天消费金额
labelint是否复购,1是0否

这里 gender 和 income_level 是典型的分类型变量,不能直接塞进向量里。age、visit_cnt、purchase_amt 是连续值。实际操作中,原始表里可能还有一堆 id 列、时间列和评论字段,我在示例里都提前去掉了,这一步省下来的资源很可观。

3.2 特征处理:StringIndexer + OneHotEncoder + VectorAssembler

在 Spark.ml 里,所有算法都只认数值类型,但并不是说把字符串直接转成整数就行。gender 的 male/female 如果用 0/1 表示没问题,但像 city 这种高基数类别,直接映射成 0、1、2、3,会让算法误以为城市之间有大小远近关系。所以标准的做法是先StringIndexer把字符串转成类别索引,再用OneHotEncoder把类别索引转成稀疏向量。

我用 PySpark 写一段核心代码:

from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler gender_indexer = StringIndexer( inputCol="gender", outputCol="gender_idx", handleInvalid="keep" ) gender_encoder = OneHotEncoder( inputCol="gender_idx", outputCol="gender_vec" ) income_indexer = StringIndexer( inputCol="income_level", outputCol="income_idx", handleInvalid="keep" ) income_encoder = OneHotEncoder( inputCol="income_idx", outputCol="income_vec" ) assembler = VectorAssembler( inputCols=["age", "gender_vec", "income_vec", "visit_cnt", "purchase_amt"], outputCol="features" )

handleInvalid="keep"这个参数我强烈建议你加上。默认情况下,如果推理阶段的数据里出现了训练阶段没见过的字符串值,StringIndexer会直接抛异常。生产环境经常有新客群、新城市出现,keep会让未知类别保留到专门的索引里,模型至少能跑完,虽然效果可能差一点,但不会中断线上任务。

3.3 用随机森林训练一个模型

随机森林在大数据场景下表现稳定,不用做太多特征缩放,而且能输出特征重要性,适合当第一个模型。我把整个流程包成一个 Pipeline:

from pyspark.ml import Pipeline from pyspark.ml.classification import RandomForestClassifier rf = RandomForestClassifier( labelCol="label", featuresCol="features", numTrees=100, maxDepth=10, seed=42 ) pipeline = Pipeline(stages=[ gender_indexer, gender_encoder, income_indexer, income_encoder, assembler, rf ]) model = pipeline.fit(train_df)

fit 完之后,这个model是 PipelineModel。预测代码非常短:

pred_df = model.transform(test_df) pred_df.select("user_id", "label", "prediction", "probability").show(10)

评估指标我用 AUC,因为二分类场景下正负样本经常不平衡,AUC 比准确率稳定得多:

from pyspark.ml.evaluation import BinaryClassificationEvaluator evaluator = BinaryClassificationEvaluator( labelCol="label", metricName="areaUnderROC" ) auc = evaluator.evaluate(pred_df) print("test AUC:", auc)

3.4 用 CrossValidator 做网格调参

单次训练的模型参数基本靠经验拍脑袋,我一般会再用 CrossValidator 搜一轮参数。这里有个细节:CrossValidator 的 estimator 直接传 Pipeline 即可,参数网格里引用的仍然是之前那个rf对象:

from pyspark.ml.tuning import CrossValidator, ParamGridBuilder param_grid = ParamGridBuilder() \ .addGrid(rf.numTrees, [50, 100]) \ .addGrid(rf.maxDepth, [5, 10]) \ .build() cv = CrossValidator( estimator=pipeline, estimatorParamMaps=param_grid, evaluator=evaluator, numFolds=3, parallelism=2 ) cv_model = cv.fit(train_df) best_model = cv_model.bestModel

训练完成后,bestModel还是一个 PipelineModel,你可以从 stages 里取出随机森林模型看最佳参数:

best_rf = best_model.stages[-1] print("best numTrees:", best_rf.getNumTrees) print("best maxDepth:", best_rf.getMaxDepth)

CrossValidator 会光是训练多个参数组合,计算量乘 3 折再乘组合数。如果集群资源不宽裕,我会先跑一个很小的网格,比如两个参数各 2 个值,等确认数据没问题再扩大范围。

4. 别让训练任务卡死:性能优化与资源经验

4.1 缓存和分区的取舍

如果同一份训练数据要被 CrossValidator 反复读取,最好在 fit 之前做一次缓存:

train_df.cache()

缓存能避免每个 Stage 都重新从源数据算一遍。但缓存不是无脑用,缓存之后如果数据被多次 shuffle,缓存可能因为 Executor 内存不足反而触发频繁 GC。我一般只在数据经过清洗后、进入 Pipeline 前缓存一次,之后就不动它了。

分区数量也是老生常谈但容易踩坑的事。Spark 默认 shuffle 分区数是 200,如果训练数据只有几百 MB,200 个分区完全多余,每个 task 处理的数据很少,调度开销反而变大。反过来,如果数据跑到 TB 级,分区数太少又会导致个别 Executor 内存爆掉。我的判断标准是控制每个分区 128MB~256MB,可以先数一下输入文件大小,再用repartition或coalesce微调。

4.2 特征向量的存储占用

类别列做 OneHotEncoder 以后,结果是一个稀疏向量。比如收入等级列可能只有 5 个值,转成 one-hot 后长度为 5,但每个样本里只有一个位置是 1,如果不压缩存储,内存会无谓上涨。Spark.ml 内部用 SparseVector 保存稀疏稠密结构,只要你没有手动把它转成稠密数组,一般不需要担心。

但要注意VectorAssembler的输出顺序。它拼接列的顺序就是你传入inputCols的顺序,这个顺序在模型保存后是不能变的。我见过有人中途给输入表加了列,又没改inputCols,结果组装出的向量整体错位,模型预测效果直接崩了。调的时候先用printSchema确认列顺序,再对照featureImportances去解释特征重要性。

4.3 资源估算和 OOM 处理

如果任务在特征工程阶段 OOM,常见原因是某个 join 把一方的数据放大太多。如果 Executor 内存比较紧张,可以先查 Spark UI 看每个 Stage 的 shuffle 读写量。若是 join 一侧是小表,Spark 默认会自动做广播 join,但如果小表超过广播阈值,也会变成排序 merge join,数据热点一多就容易炸。

如果 OOM 发生在模型训练阶段,先想是不是树类模型深度太大,或者特征维度太高。随机森林在 Spark 上的并行方式是分布式建树,每个 Executor 只处理自己那一部分样本,但每棵树的内存占用和maxDepth强相关。我习惯先在本地用小数据把 train/test 流程跑通,再上集群放大参数,这样能区分到底是代码 bug 还是资源不够。

4.4 模型保存和上线的常见方式

训练完以后,PipelineModel 可以整体保存:

model.write().overwrite().save("/path/to/spark_ml_model")

加载时用:

from pyspark.ml import PipelineModel loaded_model = PipelineModel.load("/path/to/spark_ml_model")

SavedModel 整个链路都带上了,所以预测时不用重新写特征处理逻辑,直接:

result = loaded_model.transform(new_data)

不过要注意,这个保存方式对于 PySpark 模型来说依赖 Python 的序列化格式,Spark 小版本升级后可能会出现兼容问题。我的习惯是在上线前写一个小的回放测试,用原来训练集里的样本跑一遍 transform,对比加载后的预测结果和内存模型预测结果。如果结果一致,才敢把模型推到生产。

5. 常见问题与排查笔记

5.1 DataFrame 类型不一致导致的报错

症状是调用 Pipeline 时抛异常,提示某个列的期待类型和实际类型不一致。这一类报错大部分来自VectorAssembler要求所有输入列必须是数值型或向量型,如果你把字符串列直接塞进去,它一定会拒绝。

排查思路很简单:先执行data.printSchema(),看清楚每个列的类型。如果是 StringType,就先用cast("double")转成数值。如果是字符串类别,不要硬塞进 Vector,先走 StringIndexer,再接 OneHotEncoder。

提示:VectorAssembler的输入列中如果混入 null,也会让后续算法直接报错或者跑出异常结果。清洗阶段最好统一处理 null,比如用na.fill填充,或用drop过滤掉空行,别把脏数据留到训练。

5.2 StringIndexer 遇到新类别

训练数据里性别只有男、女两种,预测数据里突然出现“未知”,StringIndexer默认行为会让作业直接挂掉。遇到这个问题,你已经知道答案了:

使用handleInvalid="keep"或者"skip"。keep 会把未知类别映射到单独的索引;skip 会在 transform 时丢弃包含未知类别的行。生产环境我一般选 keep,至少不会因为个别异常值中断整批预测。但是测试结果的时候要注意,unknown 类别可能会给模型带来额外偏置,必要时可以用IndexToString配合检查。

5.3 模型保存后再加载失败

模型报错可能是版本不兼容。不同 Spark 版本之间,序列化结构有变化,尤其是使用 PySpark 保存的模型。另一个常见问题是自定义 UDF 或自定义 Transformer 没有打包进 Spark 的 classpath。如果你写了一个自定义 Stage,保存和加载时务必要把包含这个类的 jar 或 Python 文件一并提交。

排查方法是:换一台和训练环境相同版本的 Spark 来加载模型;如果找不到完全一致的环境,就用 MLeap 这种跨平台模型导出方案。不过跨平台方案需要额外处理 Pipeline 里所有的特征转换器,成本也不低。

5.4 训练任务某一个 Stage 特别慢

Spark UI 打开后,如果看到某一个 Stage 内某个 Task 运行时间明显长于其他 Task,而其他 Task 都在等它,大概率是数据倾斜。几个 Task 分到的数据块大小严重不均衡,通常是源表分区不均匀,或者某个 join key 是热点。

可以先通过 groupBy 统计一下热门 key,比如复购用户的 id 是否只集中在某几个值。解决思路是加盐打散 join key,或者用repartition重新调整分区。还有一种更不影响业务的用法:在 Spark 3.0 以后开启spark.sql.adaptive.skewJoin.enabled,它能在 join 过程中自动拆分倾斜分区,减少人工处理成本。

5.5 常见问题速查表

现象可能原因解决思路
训练报 Column 类型错误字符串列直接进 VectorAssembler先转数值或走 StringIndexer
预测阶段抛 Invalid 异常新类别未处理设handleInvalid="keep"
作业 Executor OOM特征向量过于稠密或 join 数据膨胀检查分区大小,用稀疏向量,调整 join 策略
模型加载失败版本不兼容或缺少自定义类保持同版本环境,检查 classpath
任务整体运行很慢分区数不合理调整spark.sql.shuffle.partitions
AUC 忽高忽低交叉验证折数太少,随机种子不固定固定 seed,增加 fold

最后再分享一个小技巧。我调参时经常只盯着 AUC 看,后来发现交叉验证折数少的时候,指标波动很大。现在我会在随机森林里固定seed=42,并且把numTrees和maxDepth的搜索组合先控制在 4 组以内,跑完以后用best_model.stages[-1].featureImportances对一遍特征顺序,确认它不是靠某个异常特征冲高的。这个习惯帮我少走了很多弯路,你也不妨试试。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询