如何在PySpark中维护DataFrame行顺序？现有函数致行重排

我希望在PySpark中维护DataFrame的行顺序。现有一个用于生成行ID列的函数（如下所示），但该函数会导致DataFrame中的行发生重排，请问是否有办法实现行顺序维护的需求？

def createRowIdColumn( new_column, position, start_value):
    row_count = df.count()
    row_ids = spark.range(int(start_value), int(start_value) + row_count, 1).toDF(new_column)
    window = Window.orderBy(lit(1))
    df_row_ids = row_ids.withColumn("row_num", row_number().over(window) - 1)
    df_with_row_num = df.withColumn("row_num", row_number().over(window) - 1)

    if position == "Last Column":
        result = df_with_row_num.join(df_row_ids, on="row_num").drop("row_num")
    else:
        result = df_row_ids.join(df_with_row_num, on="row_num").drop("row_num")
       
    return result.orderBy(new_column)

问题分析 #

原函数导致行重排的核心原因：

• Window.orderBy(lit(1))没有实际排序依据，Spark分布式计算环境下，分区内的行顺序不确定，生成的row_num无法对应原始行顺序。
• 最终返回时调用orderBy(new_column)，强制按行ID排序，直接打乱了原DataFrame的行顺序。

解决方案 #

要维护原始行顺序，必须依赖一个能唯一标识原始行顺序的临时列。以下是改进后的实现：

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

def createRowIdColumn(df, new_column, position, start_value):
    # 添加临时全局递增ID，用于保留原始行顺序
    df_with_temp_order = df.withColumn("_temp_order", F.monotonically_increasing_id())
    
    # 基于临时顺序列生成连续的目标行ID
    global_window = Window.orderBy("_temp_order")
    df_with_row_id = df_with_temp_order.withColumn(
        new_column,
        F.row_number().over(global_window) + int(start_value) - 1
    ).drop("_temp_order")
    
    # 调整新列的位置
    if position == "Last Column":
        column_order = df.columns + [new_column]
    else:
        column_order = [new_column] + df.columns
    
    return df_with_row_id.select(column_order)

关键说明 #

1. monotonically_increasing_id()的作用：
   这个函数会为每个分区分配不重叠的ID范围，同一分区内的行ID连续递增，能保证全局范围内的行顺序和原始DataFrame的顺序一致（前提是原DataFrame未经过shuffle操作，若已shuffle则原始顺序已丢失）。
2. 生成连续行ID：
   通过row_number().over(Window.orderBy("_temp_order"))基于临时顺序列生成连续序号，加上start_value-1得到从指定值开始的连续行ID。
3. 列位置调整：
   通过select直接调整列顺序，避免了原函数中join操作带来的性能开销和顺序不确定性。

注意事项 #

• 如果原DataFrame已经执行过shuffle类操作（如groupBy、orderBy、join），原始行顺序已无法恢复，此时生成的ID仅基于当前DataFrame的分区顺序。
• monotonically_increasing_id()生成的是64位非连续ID，必须结合row_number才能得到连续的行ID序列。

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