Merge pull request #723 from tandav/master

more readable collect_set in _get_top_k_items
recommenders-team · Apr 9, 2019 · 8bc76fd · 8bc76fd
2 parents 1a3d846 + b8ea92c
commit 8bc76fd
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/reco_utils/evaluation/spark_evaluation.py b/reco_utils/evaluation/spark_evaluation.py
@@ -371,9 +371,8 @@ def _get_top_k_items(
             row_number().over(window_spec).alias("rank")
         )
         .where(col("rank") <= k)
-        .withColumn(col_prediction, F.collect_list(col_item).over(Window.partitionBy(col_user)))
-        .select(col_user, col_prediction)
-        .dropDuplicates([col_user, col_prediction])
+        .groupby(col_user)
+        .agg(F.collect_list(col_item).alias(col_prediction))
     )
 
     return items_for_user