如何控制Pyspark中PandasUDF在worker节点的最大并发执行数量

要控制每个worker节点上同时运行的Pandas UDF实例数，可通过以下几种组合方案实现，按落地成本从低到高排序：

方案1：资源配置硬限制（最稳妥，无代码侵入） #

Spark任务调度基于executor计算槽位分配，直接限制每个worker的可用task槽位到你设定的阈值n，就能从调度层控制UDF最大并发数：

1. 关闭动态资源分配，避免executor数量自动伸缩导致的调度波动：

   spark.dynamicAllocation.enabled false
   

2. 按你的集群规模配置executor参数，对应你举的5个worker、单worker最大并发4的场景，配置如下：

   # 总executor数和worker数一致，保证每个worker只跑1个executor
   spark.executor.instances 5
   # 每个executor的核数设为你要的单worker最大并发数n
   spark.executor.cores 4
   # 每个task占用1核，保证同一时间最多4个task并行
   spark.task.cpus 1
   

以上配置会让集群总可用task槽位固定为20，刚好匹配20个分组的处理需求，不会出现单worker并行UDF数超过4的情况。

方案2：自定义哈希分区保证分组分配均匀（配合资源配置效果最优） #

你之前使用repartitionByRange出现分配不均，核心原因是范围分区的边界计算可能受数据分布影响出现倾斜，可改为哈希分区严格按分组键均匀分配：

1. 先将数据重分区为worker数 * 单worker最大并发数个分区（你的场景为5*4=20），保证每个分区对应1个分组：

   df = df.repartition(20, "GROUP_IX")
   # 同步设置shuffle分区数，避免groupBy操作自动重分区
   spark.conf.set("spark.sql.shuffle.partitions", 20)
   
   df.groupBy('GROUP_IX').applyInPandas(my_pandas_udf, some_output_schema)
   

2. 配合方案1的资源配置，即可保证每个worker刚好分配到4个task，每个task处理1个分组，负载完全均匀。

方案3：任务与worker强制绑定（仅极端场景使用） #

如果需要完全控制特定分组的执行节点，可通过自定义RDD的位置偏好实现，强制Spark将任务调度到指定worker，不过代码侵入性较高，非必要不建议使用。

注意：如果单分组数据量较大，需要同步调高executor内存配置，避免UDF执行时出现OOM。

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