Spark中GroupBy与聚合后分区数的决定因素及固定1分区疑问

一、决定GroupBy与聚合后分区数量的核心因素 #

• shuffle分区数配置：spark.sql.shuffle.partitions是控制shuffle阶段分区数的核心参数，Spark 3.x默认值为200。不过当聚合后结果数据量极小（比如仅少数分组），Spark会自动触发小分区合并优化，最终分区数会远低于该配置值。
• 聚合结果的数据规模：如果聚合后的分组数量少、单分组数据量小，Spark会合并空分区或极小分区，减少最终分区数；反之，若分组多、数据量大，分区数会接近spark.sql.shuffle.partitions的设置。
• 执行环境与资源：本地模式下，Spark会根据本地资源和数据量灵活调整分区；集群模式下则会结合集群节点数、核心数等资源配置，平衡并行度与资源利用率，不会轻易将分区合并到1个。
• 手动分区干预：若在聚合后调用repartition()或coalesce()，会直接覆盖自动调整的分区数；聚合前的分区数不直接决定聚合后的分区数，因为聚合操作必然触发shuffle，shuffle后的分区数由上述规则决定。

二、示例中聚合后生成1个分区的情况并非始终发生 #

你的示例中聚合后仅1个分区，本质是因为聚合结果只有2个分组，数据量极小，Spark自动合并了所有空分区和小分区。这种情况不会一直出现，具体场景差异如下：

1. 分组数量较多时：如果数据源包含大量不同分组（比如上千个），即使在本地模式，聚合后的分区数会接近spark.sql.shuffle.partitions的配置值，不会合并为1个。
2. 修改shuffle配置参数：若显式设置spark.sql.shuffle.partitions为更大的值（比如50），且分组数足够多，聚合后的分区数会对应增加；即使分组少，若分区内有实际数据，也可能保留多个分区（比如2个分组对应2个分区）。
3. 集群环境运行：在分布式集群中，Spark会优先保证并行度，不会轻易将所有数据合并到1个分区，除非结果数据量确实极小且资源允许。

验证示例（修改后代码） #

若增加分组数量并调整shuffle配置：

from pyspark.sql import SparkSession
from pyspark.sql.functions import sum

spark = (
    SparkSession.builder.master("local").appName("localapp")
    .config("spark.sql.shuffle.partitions", 10)
    .getOrCreate()
)

columns = ["language", "user_count"]
# 生成20个不同分组的测试数据
data = [(f"Lang_{i}", i) for i in range(20)] * 10

df = spark.createDataFrame(data).toDF(*columns)
df = df.repartition(200)
print(df.rdd.getNumPartitions()) # 200

df = df.groupby("language").agg(sum("user_count"))
print(df.rdd.getNumPartitions()) # 输出会接近10（或等于实际分组数20，取决于数据分布）

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