如何对PySpark DataFrame排序分块并在连续块上执行处理任务？

一、验证DataFrame是否按指定列排序 #

原表通过Iceberg的ALTER TABLE ... WRITE ORDERED BY按SORT_COLUMNS写入，理论上数据全局有序，可通过以下两种方式验证：

高效分区极值验证法（适合超大规模数据） #

通过分布式计算每个分区的排序列极值，再检查分区间的连续性：

from pyspark.sql import functions as F

# 计算每个分区的排序列最大/最小值
partition_sort_stats = df.withColumn("partition_id", F.spark_partition_id()) \
    .groupBy("partition_id") \
    .agg(
        F.min("SORT_COLUMNS").alias("min_sort_val"),
        F.max("SORT_COLUMNS").alias("max_sort_val")
    ) \
    .orderBy("partition_id")

# 检查分区间排序连续性
sorted_stats_list = partition_sort_stats.collect()
is_globally_sorted = True
for i in range(1, len(sorted_stats_list)):
    prev_max = sorted_stats_list[i-1]["max_sort_val"]
    curr_min = sorted_stats_list[i]["min_sort_val"]
    if curr_min < prev_max:
        is_globally_sorted = False
        print(f"分区{i}与前序分区{i-1}的排序列不连续")
        break

print(f"DataFrame是否按{SORT_COLUMNS}全局排序: {is_globally_sorted}")

抽样行验证法（轻量快速） #

随机抽取部分行，检查排序顺序：

from pyspark.sql.window import Window

# 抽样1%数据检查排序
sampled_df = df.sample(fraction=0.01, seed=42)
window_spec = Window.orderBy("SORT_COLUMNS")
sampled_df_with_prev = sampled_df.withColumn(
    "prev_sort_val",
    F.lag("SORT_COLUMNS").over(window_spec)
)

# 统计无序行数
unsorted_count = sampled_df_with_prev.filter(F.col("prev_sort_val") > F.col("SORT_COLUMNS")).count()
is_sorted = unsorted_count == 0

二、分割为连续排序块 #

针对mapPartitions无法保证分区连续的问题，通过重新分区实现连续块划分：

优先策略：按约100万行分块 #

利用全局递增行号实现等行数量分块，再重新分区保证块的连续性：

# 添加全局递增行号（无重复但可能有间隙，不影响分块连续性）
df_with_row_id = df.withColumn("row_id", F.monotonically_increasing_id())

# 计算块ID：每100万行一个块
chunk_size = 1000000
df_with_chunk_id = df_with_row_id.withColumn(
    "chunk_id",
    F.floor(F.col("row_id") / chunk_size)
)

# 按chunk_id重新分区，确保每个分区对应一个连续排序块
chunked_df = df_with_chunk_id.repartition(F.col("chunk_id"))

备选策略：按GROUP_COLUMNS分组分块 #

若GROUP_COLUMNS是SORT_COLUMNS的前缀，直接按分组列重新分区即可保证同组数据连续：

# 按GROUP_COLUMNS重新分区，同组数据集中在一个分区
chunked_df = df.repartition(*GROUP_COLUMNS)

若GROUP_COLUMNS不是前缀，需先按SORT_COLUMNS, GROUP_COLUMNS排序再分区（注意：超大规模数据排序性能开销大，仅在必选时使用）。

三、对每个连续块执行process_chunk #

重新分区后，每个分区对应一个连续排序块，可安全使用mapPartitions或foreachPartition处理：

用mapPartitions处理并返回结果 #

def process_chunk(iterator):
    # 将迭代器转为列表，执行自定义处理逻辑
    chunk_data = list(iterator)
    # 示例：对每行数据做转换
    processed_rows = [row.asDict() for row in chunk_data]
    return iter(processed_rows)

# 应用处理函数并转回DataFrame
processed_df = chunked_df.rdd.mapPartitions(process_chunk).toDF(chunked_df.schema)

用foreachPartition处理（无需返回结果） #

适合数据落地、批量计算等场景：

def process_chunk(df_iterator):
    for chunk in df_iterator:
        # 示例：将块数据写入Iceberg表
        chunk.write.mode("append").format("iceberg").saveAsTable("database.processed_table")

chunked_df.foreachPartition(process_chunk)

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