如何在无Shuffle情况下为Spark DataFrame分区应用Pandas UDF

问题 #

在Spark 3.3.0中，尝试为DataFrame的每个分区单独应用pandas UDF以避免Shuffle，但运行以下代码时出现大量Shuffle，执行计划包含Sort阶段：

from pyspark.sql.functions import spark_partition_id
query = df.groupBy(spark_partition_id())\
  .applyInPandas(lambda x: pd.DataFrame([x.shape]), "n_rows long, n_cols long")
query.explain()

对应的物理计划：

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- FlatMapGroupsInPandas [SPARK_PARTITION_ID()#1562], <lambda>(id#0L, date#1L, feature#2, partition_id#926)#1561, [nr#1563L, nc#1564L]
   +- Sort [SPARK_PARTITION_ID()#1562 ASC NULLS FIRST], false, 0
      +- Exchange hashpartitioning(SPARK_PARTITION_ID()#1562, 200), ENSURE_REQUIREMENTS, [id=#748]
         +- Project [SPARK_PARTITION_ID() AS SPARK_PARTITION_ID()#1562, id#0L, date#1L, feature#2, partition_id#926]
            +- Scan ExistingRDD[id#0L,date#1L,feature#2,partition_id#926]

而运行类似的df.groupBy(spark_partition_id()).count().explain()时，执行计划无Sort阶段，且实际Shuffle数据量极小：

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[_nondeterministic#1532], functions=[count(1)])
   +- Exchange hashpartitioning(_nondeterministic#1532, 200), ENSURE_REQUIREMENTS, [id=#704]
      +- HashAggregate(keys=[_nondeterministic#1532], functions=[partial_count(1)])
         +- Project [SPARK_PARTITION_ID() AS _nondeterministic#1532]
            +- Scan ExistingRDD[id#0L,date#1L,feature#2,partition_id#926]

请问该现象的原因是什么？如何实现无Shuffle的分区级pandas UDF应用？

原因分析 #

1. 算子执行逻辑差异：

   • applyInPandas对应的物理算子是FlatMapGroupsInPandas，该算子要求同一分组的数据必须在同一个分区内且有序，因此Spark会先执行Exchange（Shuffle）将相同spark_partition_id()的数据聚到同一分区，再执行Sort保证分组内有序——即使理论上同一原始分区的spark_partition_id()值完全相同，Spark优化器也未针对该场景做特殊豁免，仍会触发Shuffle+Sort流程。
   • count()使用的是HashAggregate算子，支持本地部分聚合：先在每个原始分区计算partial_count，再Shuffle汇总结果。由于部分聚合后数据量极小，实际Shuffle开销可以忽略，且HashAggregate不需要分组数据有序，因此执行计划中无Sort阶段。

2. groupBy的固有行为：Spark的groupBy操作默认会触发Shuffle重分区，无论分组键是什么，这是其分布式分组的核心逻辑，而spark_partition_id()属于非确定性函数，优化器无法做无需Shuffle的特殊处理。

解决方案 #

要实现无Shuffle的分区级pandas UDF处理，应使用mapInPandas替代groupBy(spark_partition_id()).applyInPandas——mapInPandas直接对每个原始分区的迭代器应用函数，无需分组操作，完全规避Shuffle。

示例代码：

import pandas as pd

# 定义分区处理函数：接收分区迭代器，返回结果迭代器
def process_single_partition(partition_iter):
    for pd_df in partition_iter:
        # 对单个分区的pandas DataFrame执行处理逻辑
        yield pd.DataFrame([pd_df.shape], columns=["n_rows", "n_cols"])

# 应用mapInPandas
query = df.mapInPandas(process_single_partition, schema="n_rows long, n_cols long")
query.explain()

对应的物理计划将不会包含Exchange和Sort阶段，直接在原始分区上处理：

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- MapInPandas <function process_single_partition at 0x...>, [n_rows#L, n_cols#L]
   +- Scan ExistingRDD[id#0L,date#1L,feature#2,partition_id#926]

关键说明 #

• mapInPandas的输入是分区的迭代器，每个元素是对应分区转换后的pandas DataFrame；输出是结果的迭代器，每个元素是处理后的pandas DataFrame。
• 该方式严格保证每个分区的数据仅在本地处理，完全没有Shuffle开销，完美匹配“为每个分区单独应用pandas UDF”的需求。

内容的提问来源于stack exchange，提问作者B.M