Блог

Итак, у меня есть фрейм данных в Spark со следующими данными:

 user_id item   category score
-------------------------------
user_1  item1  categoryA  8
user_1  item2  categoryA  7
user_1  item3  categoryA  6
user_1  item4  categoryD  5
user_1  item5  categoryD  4
user_2  item6  categoryB  7
user_2  item7  categoryB  7
user_2  item8  categoryB  7
user_2  item9  categoryA  4
user_2  item10 categoryE  2
user_2  item11 categoryE  1
 

И я хочу отфильтровать его, поэтому я сохраняю только 2 элемента одной и той же категории для каждого пользователя на основе оценки, поэтому это должно выглядеть так:

 user_id item   category score
-------------------------------
user_1  item1  categoryA  8
user_1  item2  categoryA  7
user_1  item4  categoryD  5
user_1  item5  categoryD  4
user_2  item7  categoryB  7
user_2  item8  categoryB  7
user_2  item9  categoryA  4
user_2  item10 categoryE  2
user_2  item11 categoryE  1
 

Где элементы 3 и 6 были удалены из-за того, что они были 3-м элементом той же категории для пользователя и имели более низкий балл (или, если оценка одинакова, отбрасываются случайным образом, пока не останется только два элемента этой категории).

Я попытался сделать это с разделением над окном следующим образом:

 df2 = test.withColumn(
    'rank', 
    F.rank().over(Window.partitionBy('user_id','category').orderBy(F.desc('score')))
).filter('rank <= 2')
 

Это вернуло меня:

 user_id item   category score
-------------------------------
user_1  item1  categoryA  8
user_1  item2  categoryA  7
user_1  item4  categoryD  5
user_1  item5  categoryD  4
user_2  item6  categoryB  7
user_2  item7  categoryB  7
user_2  item8  categoryB  7
user_2  item9  categoryA  4
user_2  item10 categoryE  2
user_2  item11 categoryE  1
 

Это работает с пользователем 1, но не будет работать с пользователем 2, поскольку у него есть элементы той же категории с равным счетом, поэтому для элементов 6,7,8, которые принадлежат к CategoryB и имеют оценку 7, они не будут отфильтровываться. В этом крайнем случае я хотел бы отфильтровать одно случайным образом, чтобы иметь только 2.

Есть идеи о том, как выполнить этот тип фильтрации?

Вы можете использовать row_number поверх окна, разделенного по идентификатору пользователя и категории:

 df2 = df.withColumn(
    'rank',
    F.row_number().over(Window.partitionBy('user_id', 'category').orderBy(F.desc('score')))
).filter('rank <= 2').drop('rank').orderBy('user_id', 'item', 'category', F.desc('score'))

df2.show()
 ------- ------ --------- ----- 
|user_id|  item| category|score|
 ------- ------ --------- ----- 
| user_1| item1|categoryA|    8|
| user_1| item2|categoryA|    7|
| user_1| item4|categoryD|    5|
| user_1| item5|categoryD|    4|
| user_2|item10|categoryE|    2|
| user_2|item11|categoryE|    1|
| user_2| item7|categoryB|    7|
| user_2| item8|categoryB|    5|
| user_2| item9|categoryA|    4|
 ------- ------ --------- -----