基于PySpark实现大规模检索任务的Reciprocal Rank Fusion重排序

核心逻辑 #

Reciprocal Rank Fusion（RRF）的核心计算逻辑是：对每个文档，其最终得分等于所有模型中该文档排名的倒数之和（加平滑系数避免除以0），公式为：
RRF(doc) = Σ(1 / (rank(doc, model_i) + k))
其中k为平滑系数，常规取值为60。

PySpark分步实现 #

1. 数据格式转换 #

将每个embedding模型输出的top k结果，转换为PySpark DataFrame结构，字段需包含：

• query_id：查询ID
• doc_id：文档ID
• rank：该模型下文档的排名（按相似度降序，第1名对应rank=1）

如果原始数据只有相似度得分，需先对每个查询的结果按得分降序排序，生成对应排名。

2. 分布式计算RRF得分 #

• 合并所有模型的结果DataFrame
• 按query_id和doc_id分组，累加每个文档的RRF得分
• 按query_id分区，对文档按RRF得分降序排序，提取top k结果

3. 示例代码 #

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 初始化Spark会话
spark = SparkSession.builder.appName("LargeScaleRRF").getOrCreate()

# 模拟模型输出的top k数据（替换为实际数据加载逻辑）
embed1_data = [("query_1", "doc_1", 1), ("query_1", "doc_94", 2), ("query_1", "doc_11", 3), ("query_1", "doc_10", 4)]
embed2_data = [("query_1", "doc_8", 1), ("query_1", "doc_10", 2), ("query_1", "doc_96", 3), ("query_1", "doc_5", 4)]
embed3_data = [("query_1", "doc_4", 1), ("query_1", "doc_37", 2), ("query_1", "doc_2", 3), ("query_1", "doc_11", 4)]

df_embed1 = spark.createDataFrame(embed1_data, ["query_id", "doc_id", "rank"])
df_embed2 = spark.createDataFrame(embed2_data, ["query_id", "doc_id", "rank"])
df_embed3 = spark.createDataFrame(embed3_data, ["query_id", "doc_id", "rank"])

# 合并所有模型数据
combined_df = df_embed1.union(df_embed2).union(df_embed3)

# 定义平滑系数
k = 60

# 计算每个文档的RRF得分
rrf_scores = combined_df.groupBy("query_id", "doc_id") \
    .agg(F.sum(1 / (F.col("rank") + k)).alias("rrf_score"))

# 按查询分组，取top k文档
window = Window.partitionBy("query_id").orderBy(F.desc("rrf_score"))
top_k_docs = rrf_scores.withColumn("rn", F.row_number().over(window)) \
    .filter(F.col("rn") <= 4)  # 替换为实际需要的k值

# 查看query_1的重排序结果
top_k_docs.filter(F.col("query_id") == "query_1").select("doc_id", "rrf_score").show()

# 关闭Spark会话
spark.stop()

4. 性能优化要点 #

• 分区优化：根据query_id对DataFrame分区，保证各分区查询数量均衡，减少shuffle开销
• 数据合并：加载原始数据时合并小文件，用repartition或coalesce调整分区数，避免任务过多
• 内存配置：根据集群资源调整Spark内存参数（如spark.executor.memory），减少磁盘溢出

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