Оптимизированный расчет Polars NDCG ⇐ Python

[Anonymous]

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

Код: Выделить всё

ideal:   a b c d e f
propsed: d b g e h

поэтому у нас есть пересечение элементов {b, d, e с idx_ideal=[2,4,5] (начиная с 1) и idx_propsed= [2,1,4]
Итак, я хочу вычислить эти idx_propose и idx_ideal для фрейма данных со столбцами (пользовательский, идеальный, предложенный),чтобы полученный DF имел столбцы: (пользователь, идеал, предложенный, idx_ideal, idx_propsed)

Код: Выделить всё

# so its only ideal idx, then I create extra for proposed idx and join them
(
df
.explode('ideal')
.with_columns(idx=pl.int_range(pl.len()).over('user'))
.filter(pl.col('ideal').is_in(pl.col('proposed')))
.group_by('user', maintain_order=True)
.agg(pl.col('idx'))
)

Я разбиваю Ideal по пользователю и нахожу позицию, добавляя дополнительный столбец idx и оставляя только Ideal=proposition строки, но это приводит к дополнительному DF с подмножеством строк, к которым мне придется присоединиться назад, что, вероятно, не очень оптимально. Затем мне нужно еще раз вычислить его для предложенной стороны.
Более того, на следующем шаге мне придется развернуть (idx_ideal, idx_propose), чтобы вычислить IDCG пользователя, DCG и NDCG.
Не могли бы вы помочь мне оптимизировать эти вычисления?
Я думаю, мне следует использовать так, чтобы пользователи не взаимодействовали друг с другом и отдельные строки могли обрабатываться отдельно.
Вот генератор случайных данных

Код: Выделить всё

import polars as pl
import random

num_users = 100_000
min_len = 10
max_len = 200
item_range = 10_000

def generate_user_data():
length = random.randint(min_len, max_len)
ideal = random.sample(range(item_range), length)

length = random.randint(min_len, max_len)
predicted = random.sample(range(item_range), length)

return ideal, predicted

data = []
for user_id in range(num_users):
ideal, predicted = generate_user_data()
data.append({
'user': user_id,
'ideal': ideal,
'proposed': predicted
})

df = pl.DataFrame(data)
print(df.head())

Код: Выделить всё

shape: (5, 3)
┌──────┬──────────────────────┬──────────────────────┐
│ user ┆ ideal                ┆ proposed             │
│ ---  ┆ ---                  ┆ ---                  │
│ i64  ┆ list[i64]            ┆ list[i64]            │
╞══════╪══════════════════════╪══════════════════════╡
│ 0    ┆ [9973, 313, … 5733]  ┆ [8153, 3461, … 4602] │
│ 1    ┆ [3756, 9053, … 1014] ┆ [435, 9407, … 6159]  │
│ 2    ┆ [8152, 1615, … 2873] ┆ [5078, 9006, … 8157] │
│ 3    ┆ [6104, 2929, … 2606] ┆ [5110, 790, … 363]   │
│ 4    ┆ [1863, 6801, … 271]  ┆ [5571, 5555, … 5591] │
└──────┴──────────────────────┴──────────────────────┘

Подробнее здесь: https://stackoverflow.com/questions/791 ... alculation