Spark MLlib 是 Apache Spark 的机器学习库，提供分布式机器学习算法和工具，适用于大规模数据处理。以下是 Spark MLlib（基于 DataFrame 的 ML Pipeline API） 的使用方法及建模流程，包含关键步骤和示例代码。
1.建模流程概述
-
数据准备 ：加载数据并转换为
DataFrame格式。 -
特征工程 ：使用转换器（
Transformer）处理特征。 -
模型训练 ：选择算法（
Estimator）训练模型。 -
模型评估 ：使用评估器（
Evaluator）验证性能。 -
调优与部署：超参数调优、保存/加载模型。
2.详细步骤与python代码示例
1. 数据准备
from pyspark.sql import SparkSession
# 创建 SparkSession
spark = SparkSession.builder \
.appName("MLlibExample") \
.getOrCreate()
# 加载数据（示例：CSV 文件）
data = spark.read.csv("path/to/data.csv", header=True, inferSchema=True)
data.show(5)
2. 特征工程
from pyspark.ml.feature import VectorAssembler, StringIndexer, StandardScaler
**# 示例：处理类别特征和数值特征
(1) 将字符串类别列转换为数值索引
indexer = StringIndexer(inputCol="category_column", outputCol="category_index")
indexed_data = indexer.fit(data).transform(data)**
# (2) 合并特征为向量列（Spark ML 要求输入为特征向量）
feature_cols = "numerical_feature1", "numerical_feature2", "category_index"
assembler = VectorAssembler(inputCols=feature_cols, outputCol="raw_features")
assembled_data = assembler.transform(indexed_data)
# (3) 标准化特征（可选）
scaler = StandardScaler(inputCol="raw_features", outputCol="scaled_features")
scaler_model = scaler.fit(assembled_data)
scaled_data = scaler_model.transform(assembled_data)
3. 划分训练集和测试集
train_data, test_data = scaled_data.randomSplit(0.7, 0.3, seed=42)
4. 模型训练
from pyspark.ml.classification import LogisticRegression
# 初始化模型
lr = LogisticRegression(
featuresCol="scaled_features", # 指定特征列
labelCol="label", # 指定标签列
maxIter=100,
regParam=0.01
)
# 训练模型
model = lr.fit(train_data)
5. 模型预测与评估
# 预测测试集
predictions = model.transform(test_data)
# 使用评估器计算指标（如 AUC）
from pyspark.ml.evaluation import BinaryClassificationEvaluator
evaluator = BinaryClassificationEvaluator(
labelCol="label",
rawPredictionCol="rawPrediction",
metricName="areaUnderROC"
)
auc = evaluator.evaluate(predictions)
print(f"AUC = {auc}")
6. 超参数调优（交叉验证）
from pyspark.ml.tuning import ParamGridBuilder, CrossValidator
# 定义参数网格
param_grid = ParamGridBuilder() \
.addGrid(lr.regParam, 0.01, 0.1, 1.0) \
.addGrid(lr.maxIter, 50, 100) \
.build()
# 创建交叉验证器
cv = CrossValidator(
estimator=lr,
estimatorParamMaps=param_grid,
evaluator=evaluator,
numFolds=3
)
# 运行交叉验证
cv_model = cv.fit(train_data)
# 获取最佳模型
best_model = cv_model.bestModel
7. 模型保存与加载
# 保存模型
best_model.save("path/to/save/model")
# 加载模型
from pyspark.ml.classification import LogisticRegressionModel
loaded_model = LogisticRegressionModel.load("path/to/save/model")
3.关键概念
-
Transformer ：用于数据转换（如
VectorAssembler,StandardScaler）。 -
Estimator ：用于训练模型（如
LogisticRegression,RandomForestClassifier）。 -
Pipeline：将多个转换器和估计器串联为工作流。
python 代码：
from pyspark.ml import Pipeline
pipeline = Pipeline(stages=indexer, assembler, scaler, lr)
pipeline_model = pipeline.fit(train_data)
4.常见算法
-
分类：逻辑回归、决策树、随机森林、梯度提升树。
-
回归：线性回归、广义线性回归。
-
聚类：K-means、高斯混合模型。
-
推荐：交替最小二乘法（ALS）。
5.注意事项
-
数据分布式存储：确保数据已分区并适合并行处理。
-
特征向量列 ：所有特征必须合并为
Vector类型列（通过VectorAssembler）。 -
内存管理 ：调整
executor memory和driver memory避免 OOM。 -
稀疏数据 ：使用
SparseVector优化存储。
通过以上流程，可以快速使用 Spark MLlib 构建分布式机器学习模型，适用于 TB 级数据处理。