2025/1/28

import org.apache.spark.ml.regression.LinearRegression
import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.sql.SparkSession

object LinearRegressionExample {
def main(args: Array[String]): Unit = {
val spark = SparkSession.builder()
.appName("LinearRegressionExample")
.master("local[*]")
.getOrCreate()

// 加载LibSVM格式的数据
val data = spark.read.format("libsvm").load("data/sample_libsvm_data.txt")

// 划分训练集和测试集
val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3), seed = 1234L)

// 创建线性回归模型
val lr = new LinearRegression()
.setMaxIter(10)
.setRegParam(0.3)
.setElasticNetParam(0.8)

// 训练模型
val model = lr.fit(trainingData)

// 预测
val predictions = model.transform(testData)
predictions.select("features", "label", "prediction").show(10)

// 评估模型
val evaluator = new RegressionEvaluator()
.setLabelCol("label")
.setPredictionCol("prediction")
.setMetricName("rmse")
val rmse = evaluator.evaluate(predictions)
println(s"Root Mean Squared Error (RMSE) on test data = $rmse")

spark.stop()
}
}