Spark Mlib ALS 交替最小二乘算法（学习笔记）

import org.apache.log4j.{Level, Logger}
import org.apache.spark.ml.recommendation.{ALS, ALSModel}
import org.apache.spark.sql.SparkSession

object als1 {

  def main(args: Array[String]): Unit = {
    // 0. 构建Spark对象
//    val spark = SparkSession.builder()
//      .appName("ALS")
//      .getOrCreate()
    val spark = SparkSession.builder()
      .appName("ALSRecommendation")
      .master("local[*]")  // 使用本地所有可用核心
      .getOrCreate()


    Logger.getRootLogger.setLevel(Level.WARN)

    // 1. 读取样本数据
    val data = spark.read.option("header", "false").csv("H:\\zcy\\ALSrecommend\\src\\main\\scala\\luoke\\test.data")
    val ratings = data.toDF("userId", "movieId", "rating")
      .selectExpr("cast(userId as int) as userId", "cast(movieId as int) as movieId", "cast(rating as float) as rating")

    // 2. 建立模型
    val rank = 10
    val numIterations = 20
    val als = new ALS()
      .setUserCol("userId")
      .setItemCol("movieId")
      .setRatingCol("rating")
      .setRank(rank)
      .setMaxIter(numIterations)
      .setRegParam(0.01)

    val model = als.fit(ratings)

    // 3. 预测结果
    val usersProducts = ratings.select("userId", "movieId")

    val predictions = model.transform(usersProducts)

    // 计算均方误差 (MSE)
    val ratesAndPreds = ratings.join(predictions, Seq("userId", "movieId"))
    val MSE = ratesAndPreds
      .selectExpr("power(rating - prediction, 2) as error")
      .agg(org.apache.spark.sql.functions.mean("error"))
      .first()
      .getDouble(0)

    println("Mean Squared Error = " + MSE)

    // 4. 保存/加载模型
    model.write.overwrite().save("myModelPath")
    val sameModel = ALSModel.load("myModelPath")

    spark.stop() // 结束Spark会话
  }
}

运行结果：
Mean Squared Error = 1.904915633588189E-5