2025/1/26

import org.apache.spark.ml.classification.DecisionTreeClassifier
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.sql.SparkSession

object DecisionTreeExample {
def main(args: Array[String]): Unit = {
val spark = SparkSession.builder()
.appName("DecisionTreeExample")
.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 dt = new DecisionTreeClassifier()
.setLabelCol("label")
.setFeaturesCol("features")

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

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

// 评估模型
val evaluator = new MulticlassClassificationEvaluator()
.setLabelCol("label")
.setPredictionCol("prediction")
.setMetricName("accuracy")
val accuracy = evaluator.evaluate(predictions)
println(s"Test Error = ${(1 - accuracy) * 100}%")
println(s"Test Accuracy = $accuracy")

spark.stop()
}
}