如何确保含API调用的UDF在Azure Databricks多Worker节点并行运行

你的认知误区 #

普通UDF本身支持并行执行，但并行度完全由Spark DataFrame的分区数决定，而非UDF本身。你的200行DataFrame默认分区数可能只有1（或远小于Worker数量），导致所有计算都集中在单个Worker的单个任务里，无法利用多Worker资源。

核心原因 #

1. 分区数不足：Spark任务的并行度等于DataFrame的分区数，小数据集默认分区数极少，无法触发多Worker调度。
2. UDF低效初始化：你的UDF每行调用都重新创建ClientSecretCredential和ContainerClient，不仅浪费资源，还会增加延迟，进一步掩盖并行效果。

优化方案 #

1. 调整DataFrame分区数 #

通过repartition手动设置分区数，建议与Worker数量匹配（或略多），确保每个Worker都能分配到任务：

# 按Worker数量（6）设置分区数
containerInfo_df = containerInfo_df.repartition(6)

2. 优化UDF实现 #

改用广播变量复用Credential实例，结合**Pandas UDF（矢量化UDF）**提升执行效率：

from pyspark.sql import functions as F
from pyspark.sql.functions import pandas_udf
from azure.storage.blob import ContainerClient
from azure.identity import ClientSecretCredential
import pandas as pd

# 初始化Credential并广播到所有Executor（每个Executor只需一个实例）
tenant_id = "xxx"
appRegId = "xxx"
appRegSecret = "xxx"
credential = ClientSecretCredential(tenant_id=tenant_id, client_id=appRegId, client_secret=appRegSecret)
broadcast_credential = spark.sparkContext.broadcast(credential)

@pandas_udf("long")
def calculate_container_size(storage_accounts: pd.Series, containers: pd.Series) -> pd.Series:
    def get_total_size(storage_acc, cont_name):
        try:
            container_client = ContainerClient(
                account_url=f"https://{storage_acc}.blob.core.windows.net",
                container_name=cont_name,
                credential=broadcast_credential.value
            )
            # 用生成器表达式求和，减少内存占用
            return sum(blob.size for blob in container_client.list_blobs())
        except Exception as e:
            print(f"处理容器 {storage_acc}/{cont_name} 出错: {str(e)}")
            return 0
    
    # 批量处理每行数据
    return pd.Series([get_total_size(sa, cnt) for sa, cnt in zip(storage_accounts, containers)])

# 先分区再计算
containerInfo_df = containerInfo_df.repartition(6) \
    .withColumn("Size", calculate_container_size(F.col("Storage_Account_Name"), F.col("Container_Name")))

注：Subscription_Id未在函数逻辑中使用，可以从参数中移除以简化代码。

3. 验证并行效果 #

运行前可以先查看分区数确认设置生效：

print(f"当前DataFrame分区数: {containerInfo_df.rdd.getNumPartitions()}")

执行后查看Spark UI的Jobs页面，确认任务数与分区数一致，且多个Worker被调度。

额外建议 #

• 如果使用ADLS Gen2，也可以考虑通过Spark直接读取容器路径，利用fs.ls和聚合计算大小，避免调用SDK：

  def get_container_size_spark(storage_acc, cont_name):
      path = f"abfss://{cont_name}@{storage_acc}.dfs.core.windows.net/"
      return spark.read.format("binaryFile").load(path).agg(F.sum("length")).first()[0] or 0
  

• 确保集群自动缩放设置正确，Worker节点能根据任务负载自动扩容。

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