Pandas按列分组后按其他列排序获取行位置的高效方法

你当前的实现慢的核心原因有三个，和硬件配置无关：

• 多进程调用时传入全量DataFrame，会产生极大的序列化/反序列化、跨进程内存拷贝开销，96核场景下内存带宽会被完全占满，绝大多数时间都在复制数据而非实际计算
• 逐query执行df[df["query"] == query]是全表线性扫描，假设存在k个不同query，总时间复杂度为O(k*n)，1亿行规模下会退化成接近O(n²)的计算量
• 所有筛选、排序、索引赋值逻辑都在Python层执行，没有用到数据处理库内置的C/Rust层向量化优化，单步操作的开销极高。

不需要手写多进程，直接用框架内置的分组排名算子即可，单进程/内置多线程就能在几分钟内跑完1亿行数据，比你当前方案快100倍以上。

方案1：纯Pandas实现（零额外依赖） #

首先先压缩数据类型降低内存占用、加速计算，再直接调用内置groupby.rank完成分组排名，结果和你的需求完全对齐：

# 第一步：优化数据类型，内存占用可降低60%以上
df["query"] = df["query"].astype("category")
# 如果业务允许分数精度到小数点后6位，用float32代替默认float64
df["score1"] = df["score1"].astype("float32")
df["score2"] = df["score2"].astype("float32")

# 第二步：直接计算分组内排名，-1转为0起始位置
# method="first"表示分数相同时按数据原始出现顺序分配连续位置
df["pos1"] = df.groupby("query")["score1"].rank(ascending=False, method="first").astype("int32") - 1
df["pos2"] = df.groupby("query")["score2"].rank(ascending=False, method="first").astype("int32") - 1

注：如果需要相同分数获得相同并列排名，把method参数改成"min"即可，可根据业务需求调整。

方案2：Polars实现（极致性能，比Pandas快3~10倍） #

Polars是Rust编写的多线程数据处理库，自带惰性计算、内存优化，处理1亿行规模数据的速度远快于Pandas，不需要你手写任何并行逻辑：

import polars as pl

# 转成惰性表达式自动做计算优化，内置多线程并行
ldf = pl.from_pandas(df).lazy()
ldf = ldf.with_columns(
    pl.col("score1").rank(method="ordinal", descending=True).over("query").sub(1).cast(pl.Int32).alias("pos1"),
    pl.col("score2").rank(method="ordinal", descending=True).over("query").sub(1).cast(pl.Int32).alias("pos2")
)
# 执行计算，如需转回Pandas直接调用to_pandas()即可
df = ldf.collect().to_pandas()

• 绝对不要在Python层手写逐组循环、逐组筛选的逻辑，所有分组计算优先用库内置的算子，这些算子都是在底层语言层面实现的向量化逻辑，性能比手写Python代码高几个数量级
• 处理亿行数据前一定要先做数据类型压缩，字符串列转分类类型、数值列用满足精度的最小字长，能大幅降低内存占用、提升计算速度
• 多进程适合Python层的重型计算，不适合Pandas这类本身已经在C层做了优化的操作，错误的多进程用法反而会因为拷贝开销让性能下降几十上百倍。

内容的提问来源于stack exchange，提问作者ghostryder