PySpark中monotonically_increasing_id本地与AWS EMR结果不一致问题

背景与现象 #

我编写了一个函数，通过给每行分配复合ID将数据按指定子集大小分组。本地运行完全正常，但部署到AWS EMR的PySpark环境后，结果出现巨大差异：

• 测试用例：77700行的DataFrame，子集长度设为50000
• 本地结果：2个子集，分别为50000行和27700行
• EMR结果：约26个子集，每个子集行数不超过3200行

原子集划分逻辑 #

partition_column = "partition"
partitioned_df = dataframe.withColumn(
    partition_column, floor(monotonically_increasing_id() / subset_length)
)
partitioned_df_ids = (
    partitioned_df.select(partition_column)
    .distinct()
    .rdd.flatMap(lambda x: x)
    .collect()
)
for partition_id in partitioned_df_ids:
    temp_df = partitioned_df.filter(col(partition_column) == partition_id)
    dataframe_refs.append(temp_df)

问题根源 #

原代码依赖monotonically_increasing_id()生成分组ID，但这个函数的特性是分布式环境下按分区递增，而非全局连续：

• 本地环境数据分区数少，ID近似连续，floor(id/subset_length)能生成正确的分组
• EMR环境中数据分区数多（默认按文件块或文件数分区），每个分区的ID起始值跨度极大，导致每个分组仅包含单个分区内的少量数据，最终生成大量子集

待审核解决方案代码 #

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

dataframe_refs = []
partition_window = Window.orderBy(F.lit(1))
ranges_to_subset_by = []
num_of_rows = dataframe.count()
num_of_splits = math.ceil(num_of_rows / subset_length)

remainder = num_of_rows
start = 0
for _ in range(num_of_splits):
    print(_)
    # 修正原代码的判断错误：最后一次循环用剩余行数
    end = start + subset_length if _ != num_of_splits - 1 else start + remainder
    ranges_to_subset_by.append(
        (start + 1, end)
    )
    remainder -= subset_length
    start = end

print(ranges_to_subset_by)

df = dataframe.withColumn("row_number", F.row_number().over(partition_window))
df.show()

for start, stop in ranges_to_subset_by:
    dataframe_refs.append(df.filter(F.col("row_number").between(start, stop)))

方案评价与优化点 #

1. 核心思路正确：通过全局行号+范围过滤实现分组，解决了monotonically_increasing_id()的分布式缺陷
2. 修复bug：原代码中_ != num_of_splits的判断错误，循环索引从0开始，最后一次循环应为_ == num_of_splits -1
3. 性能与稳定性提示：
   • Window.orderBy(F.lit(1))会触发全局排序，大数据量下性能较差
   • 若需要固定的分组结果，建议替换为实际业务字段排序（如主键），避免行号顺序随机变化

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