machine-learning - PySpark ALS 模型中的自定义 ColdStartStrategy

Question

我正在使用 movielens-100k 数据集。目前，PySpark ALS 模型中的 ColdStartStrategy 仅支持 NaN 和 Drop。但我想有一种自定义的方式来做到这一点。我想用该用户对该电影的平均评分填充 Nan 值。

movielens = sc.textFile('u.data').map(lambda x: tuple(x.split('\t'))).map(lambda x: tuple([float(x[0]), float(x[1]), float(x[2])]))
data = movielens.toDF(['userid','itemid','rating'])
data.take(10)
train, test = data.randomSplit([0.75,0.25],1111)
train.cache()
test.cache()
    # we use the cross validator to tune the hyperparameters
als = ALS(
         userCol="userid", 
         itemCol="itemid",
         ratingCol="rating", 
         coldStartStrategy="drop"
)

param_grid = ParamGridBuilder() \
            .addGrid(als.rank, [10, 100]) \
            .addGrid(als.regParam, [.1]) \
            .addGrid(als.maxIter, [10]) \
            .build()

evaluator = RegressionEvaluator(
           metricName="mse", 
           labelCol="rating", 
           predictionCol="prediction")

cv = CrossValidator(estimator=als, estimatorParamMaps=param_grid, evaluator=evaluator, numFolds=3, parallelism = 6)
model = cv.fit(train)

这是我drop用作冷启动策略的代码。如何使用平均值而不是删除行？