如何按指定列表顺序在PySpark中筛选并排序数据？

因为Spark的分布式特性，DataFrame本身是无序的，isin方法仅负责筛选符合条件的行，不会保证结果顺序和输入列表一致。要实现按list_codes的顺序排序，有两种高效的方法：

方法1：利用映射字典+临时排序列 #

适合所有Spark版本，尤其适合列表元素较多的场景：

1. 先构建code到其在列表中索引的映射字典
2. 通过create_map生成Spark可识别的映射关系
3. 添加临时排序列，按该列排序后删除临时列

from pyspark.sql import functions as f

list_codes = ["code_123", "code_456"]
# 假设原DataFrame为df
code_order = {code: idx for idx, code in enumerate(list_codes)}
# 生成Spark映射对象
code_order_map = f.create_map([f.lit(item) for pair in code_order.items() for item in pair])

# 筛选+排序
result_df = df.filter(f.col("code").isin(list_codes)) \
              .withColumn("temp_sort_idx", code_order_map[f.col("code")]) \
              .orderBy("temp_sort_idx") \
              .drop("temp_sort_idx")

result_df.show()

方法2：使用array_position函数（Spark 2.4+） #

代码更简洁，直接利用Spark内置函数获取元素在数组中的位置进行排序：

from pyspark.sql import functions as f

list_codes = ["code_123", "code_456"]

# 筛选+按列表顺序排序
result_df = df.filter(f.col("code").isin(list_codes)) \
              .orderBy(f.array_position(f.array([f.lit(code) for code in list_codes]), f.col("code")))

result_df.show()

原理说明 #

array_position函数会返回目标值在指定数组中的索引位置（从1开始），按这个位置排序就能严格遵循list_codes的顺序。

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