如何高效为PySpark DataFrame的每100行批次添加唯一ID

要给PySpark DataFrame按每100行一批分配唯一ID，核心是先明确行的顺序，再通过行号计算批次ID，具体实现如下：

核心思路 #

1. 确定行顺序：PySpark是分布式计算框架，数据默认无固定顺序，必须先通过业务排序键（或临时ID）对DataFrame排序，保证批次划分的一致性。
2. 生成连续行号：用窗口函数row_number()生成全局连续的行号。
3. 计算批次ID：通过整数除法将每100行映射到同一个批次ID。

代码实现 #

假设你的DataFrame为df，有业务排序键order_col（若没有可替换为monotonically_increasing_id()作为临时排序依据）：

from pyspark.sql import Window
from pyspark.sql.functions import row_number, floor, col

# 定义全局排序窗口：按业务列排序（替换为你的排序字段）
window_spec = Window.orderBy("order_col")

# 添加行号并计算批次ID
result_df = df.withColumn("row_num", row_number().over(window_spec)) \
              .withColumn("batch_id", floor((col("row_num") - 1) / 100))

代码说明 #

• row_number().over(window_spec)：生成从1开始的连续行号，完全基于排序后的顺序。
• floor((row_num -1)/100)：将前100行（行号1-100）分配批次ID0，接下来100行分配1，以此类推。若需要批次ID从1开始，改为floor((col("row_num") - 1)/100) + 1即可。

注意事项 #

• 排序键的选择：如果业务有明确的顺序要求（比如时间、ID），必须用对应的字段排序；若无需严格顺序，可用monotonically_increasing_id()临时排序，但该ID基于数据分区生成，不会保证全局严格连续，仅用于批次划分足够。
• 大数据量性能优化：全局排序会触发shuffle操作，超大规模数据下可能影响性能。如果业务允许按分区内的100行划分批次，可将窗口改为Window.partitionBy("partition_col").orderBy("order_col")，这样只会在分区内排序，避免全局shuffle，但批次ID是分区内唯一的，需根据需求选择。
• 重复排序键的处理：若排序键存在重复，row_number()仍会为重复键的行分配不同的行号，确保每100行严格划分为一个批次。

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