PySpark UDF调用REST API写入Cosmos DB性能问题及优化咨询

问题根源 #

你的核心问题是Python UDF逐行发起请求+Spark高并行度导致API被瞬时请求量压垮：Spark会为DataFrame的每个分区启动多个任务，每个任务逐行执行UDF，3-4万条数据会瞬间触发数万次并发请求，直接打满API资源，导致作业挂起、API无响应。

最优方案：批量请求+分布式并行控制 #

你提到的两个选项都不是最优解，正确的做法是按分区批量打包数据，控制并发请求数，复用网络连接，具体实现如下：

1. 核心优化思路 #

• 批量提交：将多条数据打包成一个请求发送，大幅减少请求总数（比如每100条发一次，请求数从3万降到300）
• 分区级资源复用：每个分区只创建一次HTTP Session，避免重复建立连接的开销
• 控制并行度：通过重分区限制同时发起的批量请求数，匹配API的承载能力
• 精细化重试：只在服务器错误时重试，增加退避时间，避免加重API负载

2. 代码实现 #

使用mapPartitions替代UDF，按分区批量处理数据：

from pyspark.sql import functions as F
from pyspark.sql.types import ArrayType, IntegerType
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import json

# 定义分区级批量处理函数
def batch_save_partition(partition):
    # 每个分区初始化一次Session，复用连接
    retry_strategy = Retry(
        total=3,  # 减少重试次数，避免过度重试
        status_forcelist=[500, 502, 503, 504],  # 仅在服务器错误时重试
        method_whitelist=['POST'],
        backoff_factor=0.5  # 指数退避，缓解短时间内的请求压力
    )
    adapter = HTTPAdapter(max_retries=retry_strategy)
    session = requests.Session()
    session.mount('https://', adapter)
    session.mount('http://', adapter)
    session.keep_alive = True  # 开启连接复用，降低TCP握手开销

    batch_size = 100  # 根据API承载能力调整批量大小
    batch = []
    results = []

    for row in partition:
        batch.append({
            'A': row.A,
            'B': row.B,
            'C': row.C
        })
        # 达到批量大小就发送请求
        if len(batch) >= batch_size:
            try:
                response = session.post(
                    url=rest_api_url,
                    headers={'Authorization': 'Bearer ' + api_token, 'Content-Type': "application/json"},
                    data=json.dumps(batch)
                )
                # 给批量内每条数据返回对应状态码
                results.extend([response.status_code] * len(batch))
                batch = []
            except Exception as e:
                # 异常标记为-1，可根据需求调整
                results.extend([-1] * len(batch))
                batch = []
                print(f"Batch failed: {str(e)}")
    # 处理剩余不足批量大小的数据
    if batch:
        try:
            response = session.post(
                url=rest_api_url,
                headers={'Authorization': 'Bearer ' + api_token, 'Content-Type': "application/json"},
                data=json.dumps(batch)
            )
            results.extend([response.status_code] * len(batch))
        except Exception as e:
            results.extend([-1] * len(batch))
            print(f"Final batch failed: {str(e)}")
    session.close()
    return results

# 重分区控制并行度，比如设为10个分区（根据API承载能力调整）
df_repartitioned = df.repartition(10)
# 应用批量处理函数，展开结果与原数据关联
final_df = df_repartitioned.rdd.mapPartitions(batch_save_partition).zip(df_repartitioned.rdd).toDF(["status", "data"])\
    .select("status", "data.A", "data.B", "data.C")

3. 对比你的选项 #

• 转为Pandas DataFrame逐行发送：完全不推荐，单进程单线程处理速度极慢，3万条数据会耗时很久，且没有利用Spark的分布式能力。
• 拆分DataFrame分批发送：如果是逐行发送，问题依然存在；如果是批量发送，思路和上述方案一致，但mapPartitions更贴合Spark分布式模型，无需手动拆分，效率更高。

额外优化建议 #

1. 匹配API批量限制：查看Cosmos DB REST API的批量操作上限，调整batch_size到最大值，进一步减少请求数。
2. 监控API负载：调整分区数和批量大小时，观察API的CPU/内存利用率，维持在70%左右的合理区间。
3. 使用官方Cosmos DB Spark连接器：如果业务允许，直接使用官方连接器替代自定义REST请求，它内置了批量处理、重试、并发控制等优化，性能和稳定性更优：

   df.write.format("cosmos.oltp")\
       .option("spark.cosmos.endpoint", cosmos_endpoint)\
       .option("spark.cosmos.masterKey", cosmos_master_key)\
       .option("spark.cosmos.database", database_name)\
       .option("spark.cosmos.container", container_name)\
       .mode("append")\
       .save()
   

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