Spark如何通过UDF分发少量计算密集型任务到多Worker节点并行执行

• 初始数据分区数过少：小体量数据集直接调用createDataFrame生成的DataFrame默认分区数为1，没有主动打散的情况下所有数据都会在同一个分区处理。
• 重分区数据倾斜：仅调用repartition(5)不指定分区键时，Spark默认哈希分区规则可能出现极端倾斜，所有60行数据落到同一个分区。
• 自适应查询执行（AQE）自动合并小分区：Spark 3.x默认开启AQE功能，会自动合并数据量过小的分区，你拆分的5个分区每个仅12行，会被AQE判定为小分区合并为单个分区，最终仅生成1个任务在单个Worker上执行。
• 并行度配置不匹配：默认spark.task.cpus配置为1，若你的函数需要占用整个Worker的所有核心，可能出现资源抢占或者任务调度异常。

1. 强制打散数据到指定数量分区 #

使用随机数作为分区键调用重分区，保证数据完全均匀分配到5个分区，同时添加校验逻辑确认分区分布：

from pyspark.sql.functions import rand, spark_partition_id, count

# 按随机数重分区，保证数据均匀分布
sdf = sdf.repartition(num_executors, rand())

# 校验分区分布，正常输出应该为5个分区，每个分区行数在12左右
sdf.groupBy(spark_partition_id()).agg(count("*").alias("row_count")).show()

2. 调整Spark并行度配置 #

关闭自适应执行避免小分区被合并，同时设置单任务占用的CPU数和单Executor核心数匹配，防止同一个Worker同时跑多个任务抢占资源：

# 关闭自适应执行，禁止合并小分区
spark.conf.set("spark.sql.adaptive.enabled", "false")
# 将<单Executor核心数>替换为你的Worker节点实际分配的CPU核心数，比如8核就填8
spark.conf.set("spark.task.cpus", <单Executor核心数>)

3. 替换行级UDF为分区级处理（可选但更稳定） #

行级UDF每次处理一行数据，改为按分区处理可以更明确的让Spark将每个分区作为独立任务调度到不同Executor：

def process_partition(partition_rows):
    for row in partition_rows:
        input_path = row["data"]
        debug_msg = my_complex_function(input_path)
        yield (input_path, debug_msg)

# 按分区处理后转回DataFrame
result_rdd = sdf.rdd.mapPartitions(process_partition)
sdf_new = result_rdd.toDF(["data", "output"])
display(sdf_new)

注意：如果你的函数内部已经实现了并行化逻辑，绝对不要让单个Executor同时运行多个任务，否则会出现CPU资源抢占，导致运行时间大幅增加甚至任务失败

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