Spark MLlib 实现线性回归算法

Spark MLlib 是 Apache Spark 提供的一个机器学习库，其中包括了许多常见的机器学习算法，包括线性回归。

下面是使用 Spark MLlib 实现线性回归算法的一般步骤：

1. 加载数据

首先需要加载训练数据，这里我们假设数据已经存储在 HDFS 上。可以使用 `SparkSession` 创建一个 `DataFrame` 对象，然后使用 `read` 方法将数据加载进来。

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("LinearRegressionExample").getOrCreate()

# 加载数据
data = spark.read.format("libsvm").load("hdfs://path/to/file")

2. 数据预处理

Spark MLlib 的线性回归算法要求输入的特征向量必须是一个 `DenseVector` 类型，因此需要对原始数据进行转换。可以使用 `VectorAssembler` 将多个特征列合并成一个特征向量列。

from pyspark.ml.feature import VectorAssembler

# 合并特征列
assembler = VectorAssembler(inputCols=["col1", "col2", "col3"], outputCol="features")
data = assembler.transform(data)

3. 划分数据集

为了评估模型的性能，通常将数据集划分为训练集和测试集。可以使用 `randomSplit` 方法将数据集随机划分为两个子集。

train_data, test_data = data.randomSplit([0.7, 0.3], seed=12345)

4. 训练模型

使用 `LinearRegression` 类创建一个线性回归模型，并使用训练数据对其进行训练。

from pyspark.ml.regression import LinearRegression

# 创建线性回归模型
lr = LinearRegression(featuresCol="features", labelCol="label")

# 训练模型
model = lr.fit(train_data)

5. 模型评估

使用测试数据对模型进行评估，可以使用 `RegressionEvaluator` 计算模型的均方误差和 R2 值。

from pyspark.ml.evaluation import RegressionEvaluator

# 预测结果
predictions = model.transform(test_data)

# 评估模型
evaluator = RegressionEvaluator(predictionCol="prediction", labelCol="label",
                                metricName="mse")
mse = evaluator.evaluate(predictions)
print("测试集均方误差 = %.3f" % mse)

evaluator = RegressionEvaluator(predictionCol="prediction", labelCol="label",
                                metricName="r2")
r2 = evaluator.evaluate(predictions)
print("测试集 R2 值 = %.3f" % r2)

完整的代码如下：

from pyspark.sql import SparkSession
from pyspark.ml.feature import VectorAssembler
from pyspark.ml.regression import LinearRegression
from pyspark.ml.evaluation import RegressionEvaluator

spark = SparkSession.builder.appName("LinearRegressionExample").getOrCreate()

# 加载数据
data = spark.read.format("libsvm").load("hdfs://path/to/file")

# 合并特征列
assembler = VectorAssembler(inputCols=["col1", "col2", "col3"], outputCol="features")
data = assembler.transform(data)

# 划分数据集
train_data, test_data = data.randomSplit([0.7, 0.3], seed=12345)

# 创建线性回归模型
lr = LinearRegression(featuresCol="features", labelCol="label")

# 训练模型
model = lr.fit(train_data)

# 预测结果
predictions = model.transform(test_data)

# 评估模型
evaluator = RegressionEvaluator(predictionCol="prediction", labelCol="label",
                                metricName="mse")
mse = evaluator.evaluate(predictions)
print("测试集均方误差 = %.3f" % mse)

evaluator = RegressionEvaluator(predictionCol="prediction", labelCol="label",
                                metricName="r2")
r2 = evaluator.evaluate(predictions)
print("测试集 R2 值 = %.3f" % r2)

参考资料：

- [Spark MLlib 官方文档](https://spark.apache.org/docs/latest/ml-guide.html)