PySpark DataFrame去重时如何保留自定义排序后的顺序？

你的核心问题是dropDuplicates会触发数据洗牌，无法保证排序后的顺序。针对6000行的小数据集，用窗口函数+行号过滤的方案可以完美解决，无需转换为Pandas DataFrame。

具体实现步骤 #

1. 保留你原来的自定义排序规则，用create_map处理mode列的自定义排序权重
2. 用Window函数按去重字段分组，在分组内按你的排序规则生成行号
3. 过滤出每个分组内行号为1的记录，即为排序后需要保留的唯一记录

完整代码 #

from pyspark.sql.functions import col, create_map, lit, row_number
from pyspark.sql.window import Window
from itertools import chain

custom_dict = {'car': 0, 'airplane': 1}
map_col = create_map([lit(x) for x in chain(*custom_dict.items())])

# 定义窗口：按去重字段分组，按你的规则排序
window_spec = Window.partitionBy('index_id', 'text', 'file_id') \
                   .orderBy(
                       col("index_id").asc(),
                       col("sim_score").asc(),
                       col("cond").desc(),
                       col("cond1").desc(),
                       map_col.getItem(col("mode")).asc()
                   )

# 添加行号并过滤
df_result = df.withColumn("row_num", row_number().over(window_spec)) \
              .filter(col("row_num") == 1) \
              .drop("row_num")

原理说明 #

• partitionBy('index_id', 'text', 'file_id')：将相同去重字段的记录分到同一组
• orderBy(...)：在每个分组内严格按照你指定的规则排序
• row_number()：给每个分组内的记录按排序顺序编号，第一个符合要求的记录行号为1
• 过滤行号为1的记录，就保留了每个重复组中排序最靠前的那条，完全符合你的需求

这个方案对6000行的数据集非常高效，不会触发大规模洗牌，也不需要切换到Pandas。

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