Найти максимальный рейтинг от PySpark Dataframe

Question

У меня есть датафрейм (pyspark.sql.dataframe.DataFrame) user_recs типа DataFrame[user: int, recommendations: array<struct<item:int,rating:float>>].

Я хочу выбрать для каждого пользователя лучшую rating и соответствующую ей item.

Например, для номера пользователя 471 у меня есть следующий ряд

user_recs.where(user_recs.user == 471).select("recommendations.item", "recommendations.rating").collect()
>>>[Row(item=[0, 23, 4], rating=[0.005226806737482548, 0.0044402251951396465, 0.004139747936278582])]

Я хочу, чтобы результатом был кадр данных, аналогичный user_recs, но с максимальным рейтингом. Я хочу, чтобы число item было равно 0, поскольку его лучший рейтинг равен 0,005 *. (Рейтинг также должен быть включен в фрейм данных.

Answer 1

Один из подходов заключается в использовании метода Spark map и сортировке, чтобы получить максимальный рейтинг для пользователя.Это будет выглядеть примерно так:

from pyspark.sql.functions import udf

def top_rating(s):
    return sorted(s, lambda x: x.rating)[0]

top_rating_udf = udf(top_rating)
result = user_recs.withColumn('top_rated_item', top_rating_udf(col('recommendations'))

Я не проверял, но это должно быть близко.Обратите внимание, это фактически дает вам новый столбец с результатом, поэтому вам не нужно удалять другие столбцы, если вы не хотите.Вы также можете извлечь верхний рейтинг и верхний элемент в свои столбцы:

result.select('user', 'top_rated_item.item', 'top_rated_item.rating')