2026/10/21 7:37:12

Spark.ml机器学习指南:Pipeline、Transformer与Estimator实战解析

Spark.ml机器学习指南:Pipeline、Transformer与Estimator实战解析 很多人第一次看到Spark.ml这个写法会以为它是 Spark MLlib 的简写其实它指的是 Apache Spark 官方机器学习库里“基于 DataFrame 的那一套 API”。另一套老接口基于 RDD现在官方更推荐用 DataFrame API。如果你正准备用 Spark 做大规模机器学习想把单机 sklearn 的脚本往集群上搬或者想彻底搞清楚 Pipeline、Transformer、Estimator 这些概念之间到底怎么配合这篇文章值得你花十分钟看完。我会从“为什么选 Spark.ml”讲起再到一份完整能跑通的示例代码最后把我在实际调试中踩过的坑和排查思路一并列出来尽量做到拿来就能用。1. 先搞清楚Spark.ml 到底是一套什么样的机器学习 API1.1 它和 MLlib 的关系Spark.ml 并不是一个独立于 MLlib 的东西它其实就是 MLlib 的一部分对应的是“基于 DataFrame 的机器学习流水线 API”。你可以在org.apache.spark.mlScala 里或pyspark.mlPySpark 里引入这些类。而老的 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再变成 TransformerEstimator 的接口里有一个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用户IDageint年龄genderstring性别income_levelstring收入等级比如 low/mid/highvisit_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( inputColgender, outputColgender_idx, handleInvalidkeep ) gender_encoder OneHotEncoder( inputColgender_idx, outputColgender_vec ) income_indexer StringIndexer( inputColincome_level, outputColincome_idx, handleInvalidkeep ) income_encoder OneHotEncoder( inputColincome_idx, outputColincome_vec ) assembler VectorAssembler( inputCols[age, gender_vec, income_vec, visit_cnt, purchase_amt], outputColfeatures )handleInvalidkeep这个参数我强烈建议你加上。默认情况下如果推理阶段的数据里出现了训练阶段没见过的字符串值StringIndexer会直接抛异常。生产环境经常有新客群、新城市出现keep会让未知类别保留到专门的索引里模型至少能跑完虽然效果可能差一点但不会中断线上任务。3.3 用随机森林训练一个模型随机森林在大数据场景下表现稳定不用做太多特征缩放而且能输出特征重要性适合当第一个模型。我把整个流程包成一个 Pipelinefrom pyspark.ml import Pipeline from pyspark.ml.classification import RandomForestClassifier rf RandomForestClassifier( labelCollabel, featuresColfeatures, numTrees100, maxDepth10, seed42 ) 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( labelCollabel, metricNameareaUnderROC ) 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( estimatorpipeline, estimatorParamMapsparam_grid, evaluatorevaluator, numFolds3, parallelism2 ) 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如果训练数据只有几百 MB200 个分区完全多余每个 task 处理的数据很少调度开销反而变大。反过来如果数据跑到 TB 级分区数太少又会导致个别 Executor 内存爆掉。我的判断标准是控制每个分区 128MB256MB可以先数一下输入文件大小再用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默认行为会让作业直接挂掉。遇到这个问题你已经知道答案了使用handleInvalidkeep或者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 异常新类别未处理设handleInvalidkeep作业 Executor OOM特征向量过于稠密或 join 数据膨胀检查分区大小用稀疏向量调整 join 策略模型加载失败版本不兼容或缺少自定义类保持同版本环境检查 classpath任务整体运行很慢分区数不合理调整spark.sql.shuffle.partitionsAUC 忽高忽低交叉验证折数太少随机种子不固定固定 seed增加 fold最后再分享一个小技巧。我调参时经常只盯着 AUC 看后来发现交叉验证折数少的时候指标波动很大。现在我会在随机森林里固定seed42并且把numTrees和maxDepth的搜索组合先控制在 4 组以内跑完以后用best_model.stages[-1].featureImportances对一遍特征顺序确认它不是靠某个异常特征冲高的。这个习惯帮我少走了很多弯路你也不妨试试。