Spark ML

最后更新: 2026-08-06 17:32:00 +0800 +0800
原文以英文撰写。您正在阅读AI翻译版本。

Spark 包含一个名为 MLlib 的机器学习库。“Spark ML”通常用于指代 MLlib 基于 DataFrame 的 API,因为其 Java 和 Scala 类位于 org.apache.spark.ml 下,而 Python 类位于 pyspark.ml 下。

以下名称容易混淆:

名称含义
MLlibSpark 的完整机器学习库
spark.ml主要的、基于 DataFrame 的 MLlib API
spark.mllib较旧的、基于 RDD 的 API,目前处于维护模式
Spark MLspark.ml API 的非正式名称

spark.mllib 并未被移除,但新应用代码通常应使用 spark.ml。DataFrame API 与 Spark SQL 集成,具有一致的流水线模型,并在 Java、Scala 和 Python 之间提供相似的 API。

Spark ML 提供分布式特征准备、模型拟合、评估、超参数调优和批量推理。当训练数据或评分 数据已经足够大、能够从 Spark 中获益时,Spark ML 最有用。它不是通用的深度学习框架, 也不会仅仅因为在集群上运行,就自动使每种算法都变得高效。

Spark ML 提供的功能 #

该库涵盖多个类别:

  • 分类: 逻辑回归、决策树、随机森林、梯度提升树、朴素贝叶斯、线性支持向量机和 多层感知器
  • 回归: 线性回归和广义线性回归、决策树、随机森林、梯度提升树、生存回归、保序 回归和因子分解机
  • 聚类: K-means、二分 K-means、高斯混合、潜在狄利克雷分配和幂迭代聚类
  • 推荐: 交替最小二乘法
  • 频繁模式: FP-growth 和 prefix-span
  • 特征工程: 分词、类别编码、特征组装、缩放、哈希、TF-IDF、离散化、插补、多项式 扩展、降维和特征选择
  • 统计: 相关性、假设检验和向量摘要
  • 工作流: 流水线、评估器、交叉验证、训练验证拆分、参数网格和模型持久化

Java、Scala、Python 和 R 之间的算法支持和参数存在差异。不要假设所有绑定都具有完全 相同的覆盖范围,应查阅集群所使用 Spark 版本的 API 文档。

作为 ML 数据集的 DataFrame #

Spark ML 消费并生成 DataFrame。监督学习数据集通常包含以下结构:

text
+-----+-------------------+
|label|features           |
+-----+-------------------+
|1.0  |[2.1, 0.0, 3.4]    |
|0.0  |(3,[0,2],[1.2,4.1])|
+-----+-------------------+

label 通常是 Double,而 features 是 org.apache.spark.ml.linalg.Vector。向量可以是:

  • Dense: 存储每个值。
  • Sparse: 存储向量大小,以及非零索引和对应的值。

对于独热编码类别和哈希文本等高维数据,稀疏向量不可或缺。将其展开为稠密数组可能会 显著增加内存和网络使用量。

默认列名只是约定。大多数算法允许调用方修改这些名称:

java
new LogisticRegression()
    .setLabelCol("isFraud")
    .setFeaturesCol("modelFeatures")
    .setPredictionCol("predictedFraud");

其他常见输出列包括:

  • rawPrediction:未归一化的模型分数
  • probability:模型提供时的类别概率
  • prediction:选定的类别或回归值

DataFrame schema 使流水线阶段可以组合,但检查发生在运行时。Spark 会在阶段拟合或转换 时验证所需列是否存在,并且类型是否兼容。

Transformer 与 Estimator #

Spark ML 围绕两个抽象统一算法。

Transformer 将一个 DataFrame 转换为另一个 DataFrame:

text
input DataFrame --transform()--> output DataFrame

它通常会保留输入列,同时追加新列。特征 Transformer 和训练好的模型都是 Transformer。 例如,Tokenizer 将文本列转换为单词数组,而 LogisticRegressionModel 将特征向量 转换为预测列。

Estimator 从 DataFrame 中学习并生成模型:

text
training DataFrame --fit()--> Model

Model 是 Transformer:

text
test DataFrame --model.transform()--> predictions

例如:

java
LogisticRegression trainer = new LogisticRegression()
    .setMaxIter(50)
    .setRegParam(0.01);

LogisticRegressionModel model = trainer.fit(training);
Dataset<Row> predictions = model.transform(test);

fit() 和 transform() 返回描述分布式 DataFrame 工作的对象。Spark SQL 在可能的情况 下保持惰性;write、evaluate、count 或 collect 等 action 会触发作业执行。

参数 #

每个 Estimator 和 Transformer 都实现通用参数 API。参数属于特定的阶段实例:

java
LogisticRegression lr = new LogisticRegression()
    .setMaxIter(50)
    .setRegParam(0.01)
    .setElasticNetParam(0.5);

可以使用以下代码检查支持的参数及其默认值:

java
System.out.println(lr.explainParams());
System.out.println(lr.extractParamMap());

传递给 fit() 或 transform() 的 ParamMap 会覆盖之前在实例上设置的值:

java
ParamMap override = new ParamMap()
    .put(lr.maxIter(), 100)
    .put(lr.regParam(), 0.05);

LogisticRegressionModel model = lr.fit(training, override);

由于参数绑定到阶段实例,因此一个参数网格中的两个逻辑回归阶段可以拥有相互独立的值。

特征工程 #

大多数算法要求一个数值向量列,但源数据通常包含数值、类别、布尔值和文本列。特征 Transformer 将这些列转换为可用向量。

数值特征 #

VectorAssembler 组合标量列和向量列:

java
VectorAssembler assembler = new VectorAssembler()
    .setInputCols(new String[] {"age", "income", "accountAgeDays"})
    .setOutputCol("rawFeatures");

缺失值必须被有意处理。Imputer 可以使用均值、中位数或众数替换缺失数值。它不会让 缺失本身在语义上无害;所选策略必须适合数据。

StandardScaler、MinMaxScaler 和 MaxAbsScaler 会重新缩放向量。对于基于距离和基于 梯度的算法,缩放很重要;而决策树通常不需要缩放。

对稀疏向量进行中心化时需要谨慎。减去非零均值会使隐式零值变为非零值,从而可能使稀疏 数据集变为稠密数据集。

类别特征 #

StringIndexer 将类别字符串转换为数值索引:

java
StringIndexer countryIndexer = new StringIndexer()
    .setInputCol("country")
    .setOutputCol("countryIndex")
    .setHandleInvalid("keep");

OneHotEncoder 将索引转换为稀疏指示向量:

java
OneHotEncoder countryEncoder = new OneHotEncoder()
    .setInputCol("countryIndex")
    .setOutputCol("countryVector");

Indexer 必须只在训练数据上拟合。在测试数据或生产数据上独立拟合,可能会为相同类别 分配不同索引。将两个阶段放入已拟合的 Pipeline 中,可以保留训练时的映射。

handleInvalid 控制收到 null 或未见过的值时的行为。静默选择错误策略可能导致生产评分 失败,或掩盖数据质量问题。

文本特征 #

基本的文本流水线可以使用:

text
text
  -> RegexTokenizer
  -> StopWordsRemover
  -> HashingTF or CountVectorizer
  -> IDF
  -> feature vector

HashingTF 不需要学习词汇表,但允许发生哈希冲突。CountVectorizer 学习显式词汇表, 更容易检查,但需要存储和分发该词汇表。

Spark ML 内置的文本特征适用于经典模型。它们不能替代现代语言模型分词器、嵌入或面向 GPU 的深度学习框架。

流水线 #

Pipeline 将特征准备和训练记录为一个有序工作流:

text
raw data
  -> StringIndexer
  -> OneHotEncoder
  -> VectorAssembler
  -> StandardScaler
  -> LogisticRegression

Transformer 阶段调用 transform()。Estimator 阶段调用 fit() 生成 Model,然后该 Model 转换后续阶段的数据。

java
StringIndexer indexer = new StringIndexer()
    .setInputCol("country")
    .setOutputCol("countryIndex")
    .setHandleInvalid("keep");

OneHotEncoder encoder = new OneHotEncoder()
    .setInputCol("countryIndex")
    .setOutputCol("countryVector");

VectorAssembler assembler = new VectorAssembler()
    .setInputCols(new String[] {"age", "income", "countryVector"})
    .setOutputCol("features");

LogisticRegression classifier = new LogisticRegression()
    .setLabelCol("label")
    .setFeaturesCol("features");

Pipeline pipeline = new Pipeline().setStages(new PipelineStage[] {
    indexer,
    encoder,
    assembler,
    classifier
});

PipelineModel model = pipeline.fit(training);
Dataset<Row> predictions = model.transform(test);

Pipeline.fit() 本身是 Estimator 操作,并返回作为 Transformer 的 PipelineModel。 原始 Pipeline 中的每个 Estimator 都会在该模型中变为已拟合的 Transformer。

这不只是语法上的便利。它保证验证数据和生产数据使用从训练数据学习的特征映射。在拆分 数据集之前拟合特征,会将验证数据的信息泄漏到训练过程中,从而产生过于乐观的评估结果。

Pipeline 阶段虽然以线性方式列出,但可以通过输入列和输出列形成 DAG。阶段必须按拓扑顺序 排列,并且应使用唯一实例,因为每个阶段都有一个参数系统使用的唯一 ID。

完整的训练流程 #

可靠的监督学习作业通常遵循以下顺序:

  1. 使用显式 schema 读取原始数据。
  2. 验证标签、范围、null 处理和行标识。
  3. 将数据拆分为训练集、验证集和最终测试集。
  4. 只使用训练数据拟合完整的特征与模型 Pipeline。
  5. 使用训练数据和验证数据选择参数。
  6. 仅在未接触的测试数据上对选定模型评估一次。
  7. 保存已拟合的 PipelineModel 及其运行元数据。

例如:

java
Dataset<Row>[] split = input.randomSplit(
    new double[] {0.8, 0.2},
    12345L
);

Dataset<Row> training = split[0].cache();
Dataset<Row> test = split[1];

PipelineModel model = pipeline.fit(training);
Dataset<Row> predictions = model.transform(test);

BinaryClassificationEvaluator evaluator =
    new BinaryClassificationEvaluator()
        .setLabelCol("label")
        .setRawPredictionCol("rawPrediction")
        .setMetricName("areaUnderROC");

double areaUnderRoc = evaluator.evaluate(predictions);

在 API 支持时设置种子,但不要假设分布式作业具有逐位确定性。分区方式、浮点归约顺序、 输入文件布局和算法实现都可能影响结果。

对于时间相关数据,随机拆分通常是错误的。应使用较早的事件进行训练、较晚的事件进行 验证,以避免从未来学习。分组实体可能还需要按组拆分,以确保同一用户或设备的事件 不会出现在两侧。

评估 #

选择与任务匹配的 Evaluator:

任务常用 Evaluator
二分类BinaryClassificationEvaluator
多分类MulticlassClassificationEvaluator
多标签分类MultilabelClassificationEvaluator
回归RegressionEvaluator
排名RankingEvaluator
聚类ClusteringEvaluator

对于类别不平衡的分类问题,仅看准确率可能产生误导。应根据业务成本检查精确率、召回率、 F-measure、精确率-召回率曲线下面积、校准结果和混淆矩阵。排名性能良好的模型仍然可能 需要单独选择决策阈值。

RMSE、MAE 和 \(R^2\) 等回归指标回答不同问题。应根据错误成本选择指标,而不是选择看起来 最好的数字。

评估是分布式 Spark 作业,但某些摘要最终会向 driver 返回小型矩阵、历史记录或标量指标。 不要仅为了计算 Spark 可以在集群中聚合的指标,就收集整个预测 DataFrame。

超参数调优 #

Spark ML 提供两个模型选择 Estimator:

  • CrossValidator 在多个 fold 上评估每个参数映射。
  • TrainValidationSplit 使用一次训练/验证拆分评估每个参数映射,成本更低,但统计稳健性 较弱。

两者都需要:

  1. 一个 Estimator,通常是完整的 Pipeline
  2. 一个参数映射数组
  3. 一个 Evaluator
java
ParamMap[] grid = new ParamGridBuilder()
    .addGrid(classifier.regParam(), new double[] {0.001, 0.01, 0.1})
    .addGrid(classifier.elasticNetParam(), new double[] {0.0, 0.5, 1.0})
    .build();

CrossValidator cv = new CrossValidator()
    .setEstimator(pipeline)
    .setEvaluator(new BinaryClassificationEvaluator())
    .setEstimatorParamMaps(grid)
    .setNumFolds(3)
    .setParallelism(3)
    .setSeed(12345L);

CrossValidatorModel selectedModel = cv.fit(training);

示例会训练:

text
3 regParam values
x 3 elasticNetParam values
x 3 folds
= 27 model fits

选择最佳参数后,交叉验证会在完整输入数据上再次拟合该配置。因此,大型网格会消耗大量 CPU、executor 内存、存储 I/O 和时间。

parallelism 允许并发评估参数配置。它不会产生免费的容量:过高的并行度可能使多个 训练作业竞争相同的 executor,导致性能下降。

当特征阶段参数也是问题的一部分时,应调优完整的 Pipeline。这样还会在每个 fold 内分别 拟合有状态的特征 Transformer,从而避免验证泄漏。

分布式执行 #

driver 创建 Pipeline 阶段,分析 DataFrame 计划,提交拟合和转换作业,并保存生成的模型 元数据。executor 扫描分区,计算特征转换,计算分布式统计量和梯度,构建部分树统计量, 并对行进行评分。

具体执行模式取决于算法:

  • 特征转换通常使用分区上的普通 DataFrame 操作。
  • 迭代算法反复运行作业并聚合更新。
  • 树算法交换直方图和分裂统计量。
  • ALS 对用户因子和物品因子进行分区,并反复交换所需的因子块。
  • 部分已学习的模型状态会被收集到 driver,然后分发以进行预测。

“分布式”并不意味着每个 executor 独立训练一个模型对象。driver 协调一次拟合操作,而 executor 执行其分区级工作。

模型大小仍然重要。系数、树、词汇表、类别映射以及其他拟合状态可能需要驻留在 driver 上、被序列化或发送到 executor。数据集能够放入集群,并不意味着任意大的模型都能放入。

如果重复拟合模型会导致训练数据被重复读取和计算,应缓存重复使用的训练数据:

java
Dataset<Row> training = preparedTraining.cache();
training.count();

该 action 会在调优开始前物化缓存。在最后一次拟合后取消持久化:

java
training.unpersist();

缓存每个中间 DataFrame 会浪费内存,并可能增加 spill。只持久化会被多个 action 重用的 高成本数据。

保存和加载模型 #

大多数 Spark ML Estimator、Model、Pipeline 和 PipelineModel 都实现 Spark 的 ML 持久化 API:

java
model.write()
    .overwrite()
    .save("hdfs:///models/fraud/2026-07-20");

稍后加载:

java
PipelineModel model = PipelineModel.load(
    "hdfs:///models/fraud/2026-07-20"
);

Dataset<Row> predictions = model.transform(input);

保存的 PipelineModel 包含已拟合的特征阶段和训练好的模型。但它不会自动包含重现训练 运行所需的全部运维事实。应额外存储以下元数据:

  • Spark 和应用程序版本
  • 源代码修订版本
  • 训练数据快照或版本
  • 特征和标签定义
  • 拆分方法和随机种子
  • 评估指标和选定阈值
  • 参数映射
  • 预期输入和输出 schema

Spark 致力于保证小版本和补丁版本之间的持久化兼容性。跨主要 Spark 版本的兼容性仅为 尽力而为,并不保证。升级生产推理作业前,应验证加载和预测行为。

批量推理 #

最简单的推理作业是加载模型并转换 DataFrame:

java
PipelineModel model = PipelineModel.load(modelPath);

Dataset<Row> predictions = model.transform(input)
    .select("recordId", "probability", "prediction");

predictions.write()
    .mode(SaveMode.Overwrite)
    .parquet(outputPath);

转换是惰性的。输出 action 运行时,executor 执行评分。

Spark ML 天然适合批量推理。低延迟在线服务是不同的工作负载。为每个请求启动 Spark 作业并不合适,在每个 Web 服务实例中嵌入 Spark driver 通常也会带来较高运维成本。

在线服务可以采用以下方法:

  • 在面向服务的运行时中复现兼容模型;
  • 将支持的模型类型导出为可互操作格式;
  • 使用专用模型服务器;或
  • 在延迟要求允许时,以微批次对请求进行评分。

并非每个 Pipeline 阶段都有可移植表示。在部署独立的服务代码前,应测试离线 Spark 特征与在线特征之间的一致性。

自定义 Transformer 和 Estimator #

自定义 Transformer 可以将应用特定的特征逻辑封装为 Pipeline 阶段。它必须定义参数、 输入/输出 schema 行为、transform() 以及复制/持久化语义。

优先使用原生 Spark SQL 表达式,而不是逐行 UDF:

java
Dataset<Row> prepared = input.withColumn(
    "amountLog",
    functions.log1p(functions.col("amount"))
);

原生表达式可以参与 Catalyst 优化和代码生成。Java 或 Scala UDF 会引入不透明的调用边界; Python UDF 除此之外还会跨越 JVM/Python 边界,除非存在向量化执行路径。

如果需要从数据中学习自定义状态,应实现一个返回对应 Model 的 Estimator。除非需要持久化、 调优和流水线复用,否则,一次性的 DataFrame 转换序列可能比自定义阶段更简单。

Spark ML 与其他 ML 系统的比较 #

Spark ML 在以下情况下最具优势:

  • 训练数据存储在分布式数据湖中;
  • 特征准备需要大型连接、聚合或扫描;
  • 经典算法能够满足建模要求;
  • 批量训练和批量推理符合产品需求;并且
  • 组织已经能够可靠地运行 Spark。

以下情况下吸引力较低:

  • 准备好的数据集可以轻松放入一台机器;
  • 需要基于 GPU 的深度学习;
  • 所需算法尚未实现;
  • 极低延迟在线推理是主要工作负载;或
  • 实验迭代主要受集群启动和调度支配。

对于小型和中型表格数据集,scikit-learn、XGBoost、LightGBM 或单机统计软件包可能更 简单、更快。PyTorch 和 TensorFlow 等深度学习框架提供更丰富的神经网络和 GPU 支持。 Spark 仍然可以准备它们的数据集,但使用 Spark 进行特征工程并不要求最终模型也使用 Spark ML。

常见错误 #

  1. 新代码使用 spark.mllib。 除非维护现有 RDD 应用,否则优先使用基于 DataFrame 的 spark.ml API。
  2. 在拆分数据前拟合特征阶段。 这会将信息泄漏到验证集和测试集。
  3. 将训练数据或预测数据收集到 driver。 将分布式数据保留为 DataFrame,并在 executor 上聚合。
  4. 从稀疏数据创建稠密向量。 内存和 shuffle 大小可能大幅增长。
  5. 忽略未知类别和 null。 通过阶段参数和数据质量检查定义明确的生产行为。
  6. 为类别不平衡问题选择准确率。 评估与业务成本一致的指标。
  7. 启动巨大的参数网格。 模型数量是参数值数量与 fold 数量的笛卡尔积。
  8. 假设每个阶段都能线性扩展。 模型状态、driver 聚合、迭代、数据倾斜和 shuffle 仍可能成为瓶颈。
  9. 逐行记录日志或执行远程调用。 这些副作用会在 executor 上并行执行,且难以安全重试。
  10. 将保存的模型视为完整的血缘信息。 单独记录数据、代码、参数、schema 和指标。

Spark ML 提供了一种在分布式 DataFrame 上表达经典机器学习工作流的一致方式。其核心抽象 不是单个算法,而是一个同时携带特征准备和预测能力的已拟合 PipelineModel。只有在分布式 数据处理确实属于问题的一部分时才使用它,并在整个工作流中关注评估、模型大小、服务需求 和数据泄漏。

评论

在 GitHub 上阅读或评论(发表评论需要登录)