如何获取由一个列分组并由另一个pandas DataFrame排序的记录的位置

Question

我有一个非常大的DataFrame，大约有100行，看起来像这样:

query     score1    score2   key
0  query0  97.149704  1.317513  key1
1  query1  86.344880  1.337784  key2
2  query2  85.192480  1.312714  key3
3  query1  86.240326  1.317513  key4
4  query2  85.192480  1.312714  key5
...

我想按"query"分组数据帧，然后获得"score1"和"score2"排序的每一行的位置(越高越好)，因此输出应该看起来像这样-

query     score1    score2   key  pos1  pos2
0  query0  97.149704  1.317513  key1     0     0
1  query1  86.344880  1.237784  key2     0     1
2  query2  85.192480  1.312714  key3     1     0
3  query1  86.240326  1.317513  key4     1     0
4  query2  85.492410  1.212714  key5     0     1

目前，我有一个函数看起来像这样:

def func(query, df, score1=True):
mini_df = df[df["query"] == query]
mini_df.reset_index(drop=True, inplace=True)
col_name = "pos_score2"
if score1:
col_name = "pos_score1"
mini_df[col_name] = mini_df.index
return mini_df

我从main()调用:

p = Pool(cpu_count())
df_list = list(p.starmap(func, zip(queries, repeat(df))))
df = pd.concat(df_list, ignore_index=True)

但是需要很长时间。我在96个cpu和512G内存的Intel Xeon机器上运行这个，它仍然需要超过24小时。实现这一目标的更快的方法是什么?

Answer 1

使用groupby和rank:

df[['pos1', 'pos2']] = (df.groupby('query')[['score1', 'score2']]
.rank(method='max', ascending=False)
.sub(1).astype(int))
print(df)
# Output
query     score1    score2   key  pos1  pos2
0  query0  97.149704  1.317513  key1     0     0
1  query1  86.344880  1.237784  key2     0     1
2  query2  85.192480  1.312714  key3     1     0
3  query1  86.240326  1.317513  key4     1     0
4  query2  85.492410  1.212714  key5     0     1