使用spark为大数据集生成推荐模型

pip install pyspark

userId,itemId,rating
1,101,5.0
1,102,3.0
2,101,4.0
...

from pyspark.sql import SparkSession

# 初始化 SparkSession
spark = SparkSession.builder \
    .appName("RecommendationSystem") \
    .getOrCreate()

# 读取数据
data = spark.read.csv("ratings.csv", header=True, inferSchema=True)

# 显示前几行数据
data.show()

from pyspark.ml.recommendation import ALS
from pyspark.sql.types import StructType, StructField, IntegerType, FloatType

# 定义 Schema（如果未在读取时推断）
schema = StructType([
    StructField("userId", IntegerType(), True),
    StructField("itemId", IntegerType(), True),
    StructField("rating", FloatType(), True)
])

data = spark.read.csv("ratings.csv", schema=schema, header=True)

# 将数据拆分为训练集和测试集
(training, test) = data.randomSplit([0.8, 0.2], seed=42)

# 定义 ALS 参数
als = ALS(userCol="userId", itemCol="itemId", ratingCol="rating",
          coldStartStrategy="drop", nonnegative=True, implicitPrefs=False)

# 训练模型
model = als.fit(training)

from pyspark.ml.evaluation import RegressionEvaluator

# 生成预测
predictions = model.transform(test)

# 初始化评估器
evaluator = RegressionEvaluator(metricName="rmse", labelCol="rating", predictionCol="prediction")

# 计算 RMSE
rmse = evaluator.evaluate(predictions)
print(f"Root-mean-square error = {rmse}")

# 为用户 ID 为 1 的用户生成前 10 个推荐
user_recs = model.recommendForAllUsers(10)
user_recs.filter(user_recs.userId == 1).show(truncate=False)

# 为物品 ID 为 101 的物品生成推荐
item_recs = model.recommendForAllItems(10)
item_recs.filter(item_recs.itemId == 101).show(truncate=False)

from pyspark.ml.tuning import CrossValidator, ParamGridBuilder

# 定义参数网格
param_grid = ParamGridBuilder() \
    .addGrid(als.rank, [10, 20, 30]) \
    .addGrid(als.maxIter, [10, 20]) \
    .addGrid(als.regParam, [0.01, 0.1]) \
    .build()

# 定义交叉验证器
cross_validator = CrossValidator(estimator=als,
                                 estimatorParamMaps=param_grid,
                                 evaluator=evaluator,
                                 numFolds=3)

# 训练模型
cv_model = cross_validator.fit(training)

# 获取最佳模型
best_model = cv_model.bestModel

# 保存模型
best_model.save("hdfs:///path/to/save/model")

# 加载模型
from pyspark.ml.recommendation import ALSModel

loaded_model = ALSModel.load("hdfs:///path/to/save/model")