如何在PySpark中分布式生成无间隙全局行号？

你当前代码出现行号间隙、甚至分区2行号min大于max的异常，核心原因是spark_partition_id()返回的是物理分区ID，它和repartitionByRange按Data排序的逻辑分区顺序并不一定匹配——物理分区ID只是集群存储的分区编号，和数据的逻辑排序无关，导致你按物理分区ID计算的前缀偏移量完全错误。

下面给你几个可行的解决方案：

方法一：修正原方案，按数据逻辑顺序计算偏移量 #

先获取每个物理分区对应的Data范围，按Data的最小值确定分区的逻辑顺序，再计算正确的前缀偏移量（即当前分区之前所有分区的总记录数），这样就能生成连续无间隙的全局行号。

from pyspark.sql import Window
from pyspark.sql.functions import col, row_number, sum, spark_partition_id, min

empDF = spark.read.csv("dbfs:/FileStore/large_data.csv", header=True)

# 按Data范围分区，同时记录物理分区ID和每个分区的Data最小值（用于确定逻辑顺序）
empDFPartId = empDF.repartitionByRange(4, col("Data")) \
    .withColumn("partition_id", spark_partition_id()) \
    .withColumn("data_min", min("Data").over(Window.partitionBy("partition_id"))) \
    .cache()

# 计算每个分区的记录数，按Data最小值排序后生成前缀偏移量
partitionsDF = empDFPartId \
    .select("partition_id", "data_min") \
    .distinct() \
    .join(empDFPartId.groupBy("partition_id").count(), on="partition_id") \
    .orderBy("data_min") \
    .withColumn(
        "prefix_offset",
        sum("count").over(Window.orderBy("data_min").rowsBetween(Window.unboundedPreceding, Window.currentRow)) - col("count")
    )

# 关联偏移量，计算全局连续行号
empDFRowNumber = empDFPartId \
    .join(partitionsDF, on="partition_id") \
    .withColumn(
        "row_num_in_part",
        row_number().over(Window.partitionBy("partition_id").orderBy("Data"))
    ) \
    .withColumn("global_row_number", col("prefix_offset") + col("row_num_in_part"))

# 输出结果
empDFRowNumber \
    .write \
    .mode("overwrite") \
    .option("header", "true") \
    .csv("dbfs:/FileStore/partition_fixed")

方法二：用RDD的zipWithIndex生成连续行号 #

如果你的数据已经通过repartitionByRange保证了分区间Data范围有序、分区内Data有序，直接转成RDD使用zipWithIndex是性能最优的方案——它避免了多次Shuffle，直接在每个分区内生成局部索引，再拼接成全局连续索引。

from pyspark.sql import Row

empDF = spark.read.csv("dbfs:/FileStore/large_data.csv", header=True)

# 按Data范围分区，转成RDD
sorted_rdd = empDF.repartitionByRange(4, col("Data")).rdd

# zipWithIndex从0开始计数，加1转为从1开始的行号
indexed_rdd = sorted_rdd.zipWithIndex().map(lambda x: Row(**x[0].asDict(), global_row_number=x[1]+1))

# 转回DataFrame
empDFRowNumber = spark.createDataFrame(indexed_rdd)

# 输出结果
empDFRowNumber \
    .write \
    .mode("overwrite") \
    .option("header", "true") \
    .csv("dbfs:/FileStore/partition_rdd")

注意：repartitionByRange本身保证了分区间的Data是按顺序划分的，只要你的数据在分区内是按Data排序的（如果原数据没排序，在转RDD前可以加.orderBy("Data")），生成的全局行号就是连续且有序的。

方法三：全局Window的row_number()（仅小数据用） #

如果数据量很小，直接用全局排序的Window函数最省事，但这种方法会触发全量Shuffle，把所有数据集中到单个节点排序，大数据场景下会直接拖垮集群，谨慎使用。

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

empDF = spark.read.csv("dbfs:/FileStore/large_data.csv", header=True)

empDFRowNumber = empDF \
    .withColumn("global_row_number", row_number().over(Window.orderBy("Data")))

empDFRowNumber \
    .write \
    .mode("overwrite") \
    .option("header", "true") \
    .csv("dbfs:/FileStore/partition_global")

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