PySpark长表转宽表时Databricks与Colab结果不一致且输出不稳定问题解决

场景说明 #

• 核心目标：将单值组成的长表PySpark DataFrame，按固定边长转换为正方形结构的宽表
• 环境表现差异：
  • Google Colab 执行相同逻辑可得到稳定、符合预期的结果
  • Databricks 生产集群（DBR 7.0、Spark 3.0.0，64核224GB内存配置）执行时，每次运行的输出结果随机变化
• 原始实现代码：

from pyspark.sql.types import DoubleType
from pyspark.sql import SparkSession, functions as F, Window
from math import sqrt

spark = SparkSession.builder.appName('abc').getOrCreate()
df = spark.createDataFrame([1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0], DoubleType()).toDF("Number")
c = int(sqrt(df.count())) # 计算得到边长c=3

rnum = F.row_number().over(Window.orderBy(F.lit(1)))
out = (df.withColumn("Rnum",((rnum-1)%c).cast("Integer"))
 .withColumn("idx",F.row_number().over(Window.partitionBy("Rnum").orderBy("Rnum")))
.groupby("Rnum").pivot("idx").agg(F.first("Number")))

• 符合预期的稳定输出：

+----+---+---+---+
|Rnum|  1|  2|  3|
+----+---+---+---+
|   0|1.0|4.0|7.0|
|   1|2.0|5.0|8.0|
|   2|3.0|6.0|9.0|
+----+---+---+---+

• 集群上的随机异常输出示例：

+----+---+---+---+
|Rnum|  1|  2|  3|
+----+---+---+---+
|   0|3.0|7.0|1.0|
|   1|4.0|2.0|5.0|
|   2|8.0|6.0|9.0|
+----+---+---+---+

问题根源定位 #

• 非确定性窗口排序失效：根据Spark窗口函数的设计规范，当排序键存在重复值时，相同排序值的行的先后顺序不受保障，结果会随执行环境动态变化。原代码中的窗口使用常量作为排序键（如orderBy(F.lit(1))、orderBy("Rnum")），所有行的排序值完全一致，直接导致row_number()计算结果不可复现。
  • Colab 环境因数据量极小，默认单分区存储数据，行顺序恰好和原始写入顺序一致，因此能跑出正确结果；Databricks 生产集群会根据配置将数据拆分到多个并行分区，常量排序下无固定规则，自然出现随机输出。
• 非确定性聚合取值：原代码使用F.first("Number")做聚合，当分组内的行顺序没有稳定规则时，每次运行取到的第一行值都可能不同，进一步放大结果的随机性。

可稳定复现的修正方案 #

核心修改原则：所有窗口排序必须使用唯一、确定、和业务顺序强绑定的列，彻底消除排序过程的不确定性。
修正后的实现代码：

from pyspark.sql.types import DoubleType
from pyspark.sql import SparkSession, functions as F, Window
from math import sqrt

spark = SparkSession.builder.appName('abc').getOrCreate()
df = spark.createDataFrame([1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0], DoubleType()).toDF("Number")
c = int(sqrt(df.count()))

# 基于唯一值Number生成全局连续行号，固定全行顺序
global_row_window = Window.orderBy("Number")
df_with_global_row = df.withColumn("global_row", F.row_number().over(global_row_window))

# 基于固定行号计算分组键Rnum
df_with_rnum = df_with_global_row.withColumn(
    "Rnum", ((F.col("global_row") - 1) % c).cast("Integer")
)

# 分组内排序使用全局固定行号，保证组内顺序稳定
group_row_window = Window.partitionBy("Rnum").orderBy("global_row")
out = (df_with_rnum
       .withColumn("idx", F.row_number().over(group_row_window))
       .groupby("Rnum").pivot("idx").agg(F.first("Number"))
       .orderBy("Rnum")
      )

修正方案的稳定性保障：

• 全局排序使用唯一的Number列，每次执行排序结果完全一致，不受分区数、调度策略影响
• 分组内排序使用全局生成的唯一行号，不存在同分区内排序值重复的情况，row_number()计算结果100%可复现
• 聚合阶段因分组内行顺序固定，F.first("Number")取值永远一致，最终输出和Colab上的预期结果完全匹配

注意：如果原始数据没有天然的唯一有序列，可以在数据读取后先通过F.monotonically_increasing_id()生成唯一标识列，后续所有排序逻辑都基于这个标识列完成，依然可以获得稳定可复现的结果。

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