基于Polars的大规模数据集NDCG高效计算优化咨询

我现在需要在Polars里实现针对超大规模数据集的高效NDCG计算。

NDCG的核心思路是计算DCG和IDCG，这里我先跳过增益部分，只聚焦于折扣项的计算——这部分依赖于理想排序和推荐排序中对应项的排名位置。

我目前遇到的难点在于，如何准确且高效地计算理想列表和推荐列表中交集项的位置。举个例子：

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

两者的交集是{b, d, e}，对应的理想位置是idx_ideal=[2,4,5]（从1开始计数），推荐位置是idx_proposed=[2,1,4]。

我的目标是给包含(user, ideal, proposed)列的DataFrame，添加idx_ideal和idx_proposed列，最终得到的DataFrame结构是(user, ideal, proposed, idx_ideal, idx_proposed)。

我目前的实现思路是先处理理想列表的位置，再用同样的逻辑处理推荐列表，最后将结果合并：

# 先计算理想列表的位置，之后用同样逻辑计算推荐列表的位置再做关联
(
    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'))
)

但这种方法会生成一个只包含子集行的额外DataFrame，之后还需要关联回原数据，感觉效率不高。而且接下来我还需要对(idx_ideal, idx_proposed)做展开操作，才能计算每个用户的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] │
└──────┴──────────────────────┴──────────────────────┘

备注：内容来源于stack exchange，提问作者Sindbag