如何通过上游哈希分桶消除Spark groupby.applyInPandas排序

安全阿岳

2026-8-27

问题原因

分桶预排序后仍然残留Sort算子，本质是Spark Catalyst优化器的排序需求校验逻辑和分桶存储的元数据上报机制不匹配，核心原因有三点：

分桶扫描默认仅向优化器上报分区分布属性：执行计划中的BucketedScan仅代表优化器确认存储层的哈希分桶规则和groupBy要求的HashClusteredDistribution完全匹配，因此可以消除Exchange和Shuffle相关算子。但Foundry环境下默认的Parquet分桶写入，不会将写入时sort_by指定的排序规则注册到Catalyst的输出排序元数据中，优化器默认判定扫描输出的数据是无序的，因此会插入Sort算子，满足applyInPandas（对应物理算子FlatMapGroupsInPandas）对同组数据连续有序的硬性要求。
分桶sort_by的排序保证范围有限：即使元数据记录了排序信息，如果写入分桶时单个分桶生成了多个Parquet文件，sort_by仅能保证单个文件内部按指定键有序，多个文件之间的分组键范围可能重叠，优化器无法确认整个分区范围内数据全局有序，仍然会保留全分区Sort算子。
低版本Spark的硬编码限制：Foundry部分运行时基于Spark 3.1/3.2构建，这些版本中FlatMapGroupsInPandas的排序校验逻辑存在硬编码，不会信任数据源上报的排序属性，无论上游是否已经预排序，都会强制插入Sort算子。

可行优化方案

方案1：开启排序元数据传递，从根源消除Sort

首先在分桶写入作业中配置参数，保证每个分桶仅生成单个文件，避免跨文件排序范围重叠：

spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "false")
# 写入前按分桶键重分区，分区数与桶数一致
df = df.repartition(2, "id_1", "id_2")
# 设置单文件记录数上限远大于单桶预估数据量，避免文件切分
df.write.option("maxRecordsPerFile", "100000000") \
    .bucketBy(2, "id_1", "id_2") \
    .sortBy("id_1", "id_2") \
    .saveAsTable("example_df_bucketed")

在下游运行applyInPandas的作业中，开启分桶扫描的排序顺序传播配置：

spark.conf.set("spark.sql.optimizer.metadataOnly.bucketScans.allowSortOrderPropagation", "true")

配置完成后优化器可以识别到分桶扫描输出的有序性，直接满足applyInPandas的排序要求，自动移除Sort算子。

方案2：用原生Spark算子替代applyInPandas
示例中的组内减均值逻辑完全不需要Pandas UDF实现，直接用原生窗口函数即可：
```
from pyspark.sql import Window, functions as F
result_df = df.withColumn("v", F.col("v") - F.avg("v").over(Window.partitionBy("id_1", "id_2")))
```
原生窗口算子对分桶预排序的兼容性远好于Python UDF，只要分桶分区匹配，会自动消除Shuffle、Exchange、Sort三类算子，性能比applyInPandas高3~10倍。

方案3：自定义分区内迭代绕开内置Sort
如果业务逻辑必须用Pandas实现且无法调整Spark配置，可以直接绕开groupby.applyInPandas的内置逻辑，在mapPartitions中自行完成分组和Pandas处理：

import pandas as pd
def process_partition(iterator):
    buffer = []
    current_key = None
    for row in iterator:
        key = (row.id_1, row.id_2)
        if current_key is None:
            current_key = key
        if key != current_key:
            # 处理已缓存的上一分组
            pdf = pd.DataFrame(buffer)
            pdf["v"] = pdf["v"] - pdf["v"].mean()
            for rec in pdf.itertuples(index=False):
                yield rec
            buffer = []
            current_key = key
        buffer.append(row.asDict())
    # 处理最后一个分组
    if buffer:
        pdf = pd.DataFrame(buffer)
        pdf["v"] = pdf["v"] - pdf["v"].mean()
        for rec in pdf.itertuples(index=False):
            yield rec

# 读取分桶数据后直接做分区级处理，不需要groupby
result_df = spark.read.table("example_df_bucketed").mapPartitions(process_partition)

这种方式完全不会触发Shuffle和Sort算子，所有计算都在分区内完成，性能表现和消除Sort后的applyInPandas基本一致。

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