如何按指定列表顺序对PySpark DataFrame行进行排序？

要实现按指定列表col_A_itm_order = ['a2', 'a3', 'a1']的顺序对PySpark DataFrame的行排序，核心思路是给每个col_A的元素映射一个对应顺序的数值索引，然后按这个索引排序即可。以下是具体实现方案：

方法一：使用create_map生成自定义排序映射（推荐，适合长列表） #

这种方式通过创建元素与排序位置的映射关系，代码简洁且扩展性强，适合列表元素较多的场景。

完整代码示例 #

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

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

# 构建原始DataFrame
data = [("a1", "b1", "c1"), ("a2", "b2", "c2"), ("a3", "b3", "b3")]
df = spark.createDataFrame(data, schema=["col_A", "col_B", "col_c"])

# 定义自定义排序顺序
col_A_itm_order = ['a2', 'a3', 'a1']

# 创建元素到排序索引的映射：将列表元素对应到其在列表中的位置（0,1,2...）
sort_mapping = F.create_map(
    *[F.lit(item) for pair in enumerate(col_A_itm_order) for item in reversed(pair)]
)

# 添加临时排序索引列，按索引排序后删除临时列
sorted_df = df.withColumn("sort_index", sort_mapping[F.col("col_A")]) \
              .orderBy("sort_index") \
              .drop("sort_index")

# 查看结果
sorted_df.show()

代码说明 #

• enumerate(col_A_itm_order)生成元素与其位置索引的元组（如(0, 'a2')），通过reversed(pair)将元组反转成('a2', 0)的键值对形式。
• F.create_map接收一系列键值对字面量，生成Spark的映射类型列，通过sort_mapping[F.col("col_A")]可取出每个col_A元素对应的排序索引。
• 最后按临时的sort_index列排序，再删除该列即可得到目标顺序的DataFrame。

方法二：使用when链式判断（适合短列表） #

如果自定义排序的元素较少，可以直接用when函数逐个指定每个元素的排序值：

sorted_df = df.withColumn("sort_index", 
                          F.when(F.col("col_A") == "a2", 0)
                           .when(F.col("col_A") == "a3", 1)
                           .when(F.col("col_A") == "a1", 2)) \
              .orderBy("sort_index") \
              .drop("sort_index")

sorted_df.show()

这种方式逻辑直观，但当列表元素较多时代码会变得冗长，不如create_map高效。

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