PySpark中pandas_udf的groupby分组数变化是否影响内存与任务分配

对Worker节点任务分配的影响 #

任务分配的核心逻辑不会发生改变，仅会影响负载分布：

• Spark groupby操作触发shuffle时，默认的shuffle分区数由配置项spark.sql.shuffle.partitions控制，默认值为200，你调整后分组数为100，远小于默认分区数，因此每个分组会被分配到独立的shuffle分区中，对应一个独立的task调度到Worker节点执行
• 原分组只有2个时，仅会有2个task有实际计算逻辑，大部分Worker节点处于空载状态；调整为100个分组后，会有100个task被分发到不同Worker节点执行，集群资源利用率会更高，负载更均匀

对内存消耗的影响 #

你的示例场景下完全不会引发内存问题，大数据量场景下反而会降低内存风险：

• GROUPED_MAP类型的pandas UDF要求将整个分组的全量数据加载到Executor的Python进程内存中再做处理，你将分组从2个拆为100个后，单个分组的数据量仅为原来的1/50，单个UDF执行时需要占用的内存大幅降低，反而减少了Python进程OOM的概率
• 仅有的额外开销是每个分组的pandas DataFrame初始化成本，100个分组的开销几乎可以忽略不计，只有当分组量级达到十万以上时才会出现明显的调度、初始化 overhead

额外说明 #

如果后续你进一步调大分组数量超过spark.sql.shuffle.partitions的配置值，才会出现单个task处理多个分组的情况，此时单个task的内存占用为其负责的所有分组内存总和，需要注意内存阈值评估。

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