应用Schema时PySpark的row_number列为何出现混乱？

我需要为Spark DataFrame的特定非技术列应用Schema。之前我通过Window和row_number添加了一个人工ID，目的是后续能把初始DataFrame中的其他技术列关联到新生成的DataFrame上。但应用Schema之后，生成的ID出现了混乱，有时正常有时异常，循环测试能看到不同的结果。

代码示例 #

from pyspark.sql.functions import row_number, lit, col, monotonically_increasing_id, sum
from pyspark.sql.window import Window
from pyspark.sql.types import StructType, StructField, IntegerType, StringType

# 示例DataFrame
data = [(1, "Alice"), (2, "Bob"), (3, "Charlie")]
df = spark.createDataFrame(data, ["id", "name"])

# 要应用的Schema
schema = StructType([
    StructField("id", IntegerType(), False),
    StructField("name", StringType(), False),
])

# 创建ID列
w = Window().orderBy(lit('A'))
df = df.withColumn('_special_surrogate_id', row_number().over(w))

# 扩展Schema，加入代理ID字段
surrogate_key_field = StructField("_special_surrogate_id", StringType(), False)
schema_with_surrogate = StructType(schema.fields + [surrogate_key_field])

# 循环测试，结果时好时坏
for i in range(11):
    
    df_filtered = df.select("id", "name", "_special_surrogate_id")   
    df_filtered = spark.createDataFrame(df_filtered.rdd, schema_with_surrogate)

    combined_df = df.withColumnRenamed("id", "id1").join(df_filtered.withColumnRenamed("id", "id2"), on="_special_surrogate_id")

    print(f"第{i}次迭代的差异数:")
    print(combined_df.withColumn("diff", (col("id1") != col("id2")).cast("integer")).agg(sum("diff")).collect()[0][0])

原因分析 #

问题的核心在于Spark RDD是无顺序保证的：

• 当你把DataFrame转换成RDD时，Spark不会保留原DataFrame的行顺序，每次转换时RDD内的元素顺序可能随机变化。
• 后续用spark.createDataFrame(df_filtered.rdd, schema_with_surrogate)重新构建DataFrame时，行的顺序已经和原DataFrame不一致，导致_special_surrogate_id和对应的业务列（id、name）错位，最终join后出现id不匹配的情况。
• 另外，你用Window().orderBy(lit('A'))生成row_number时，因为排序键是固定值，Spark的排序是不稳定的，分区内的行顺序可能随机，但这不是主要问题，核心还是RDD转换时的无序性。

解决方法 #

方法1：避免转换RDD，直接用DataFrame API应用Schema #

不需要转成RDD，直接通过列类型转换来匹配目标Schema，这样能完全保留原DataFrame的行顺序：

# 直接对列进行类型转换，匹配目标Schema
df_filtered = df.select(
    col("id").cast(IntegerType()).alias("id"),
    col("name").cast(StringType()).alias("name"),
    col("_special_surrogate_id").cast(StringType()).alias("_special_surrogate_id")
)

# 后续的join逻辑不变
combined_df = df.withColumnRenamed("id", "id1").join(df_filtered.withColumnRenamed("id", "id2"), on="_special_surrogate_id")

方法2：如果必须使用RDD转换，强制保留顺序 #

如果因为某些原因必须转RDD，可以通过zipWithIndex给RDD加上索引，再按索引排序来保证顺序：

df_filtered = df.select("id", "name", "_special_surrogate_id")

# 给RDD加上索引，排序后再去掉索引
rdd_with_index = df_filtered.rdd.zipWithIndex().map(lambda x: (x[1],) + x[0])
sorted_rdd = rdd_with_index.sortByKey().map(lambda x: x[1:])

# 用排序后的RDD构建DataFrame
df_filtered = spark.createDataFrame(sorted_rdd, schema_with_surrogate)

方法3：优化代理ID的生成（可选） #

如果你不需要连续的row_number，可以用monotonically_increasing_id()生成全局唯一的代理ID，这个ID不会因为顺序变化而错位，不过它是不连续的：

# 替换row_number为monotonically_increasing_id
df = df.withColumn('_special_surrogate_id', monotonically_increasing_id())

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