在Databricks中用PySpark迭代追加Spark DataFrame效率低求优化

原代码的核心性能瓶颈有两个：

1. 串行API调用：循环里逐个发起请求，等待前一个请求返回才处理下一个，大量时间浪费在等待响应上
2. 循环Union操作：Spark DataFrame是不可变对象，每次Union都会生成新的执行计划，多次Union会导致计划膨胀，触发不必要的Shuffle和计算开销

优化思路 #

1. 并行调用API：利用多线程同时发起多个API请求，大幅减少总等待时间
2. 批量创建DataFrame：收集所有API返回的数据后，一次性创建最终DataFrame，彻底避免循环Union的开销

优化代码示例 #

1. 定义并行API请求函数 #

首先封装单个API请求的逻辑，方便并行调用：

from concurrent.futures import ThreadPoolExecutor

def fetch_single_api(header_data):
    # 生成API请求参数
    input_data = get_api_function(header_data)
    # 发起请求
    response = postrequest(input_data)
    result = response.json()["result"]
    # 返回列名和数据行
    return result["Headers"], result["Data"]

2. 并行获取所有API数据 #

# 原表头数据列表（假设是包含header键的字典列表）
list1 = [{"header": ...}, ...]

# 并行请求API，max_workers根据API并发限制调整
with ThreadPoolExecutor(max_workers=10) as executor:
    # 提交所有请求任务
    futures = [executor.submit(fetch_single_api, item["header"]) for item in list1]
    # 收集所有成功返回的结果
    all_api_results = [future.result() for future in futures]

3. 批量生成最终DataFrame #

根据API返回列的一致性，分两种场景处理：

场景A：所有API返回的列完全一致

# 初始化列名和数据列表
target_columns = None
all_data_rows = []

for cols, rows in all_api_results:
    if target_columns is None:
        target_columns = cols
    # 可选：校验列一致性，避免数据错误
    if cols != target_columns:
        raise ValueError(f"API返回列不一致：预期{target_columns}，实际{cols}")
    all_data_rows.extend(rows)

# 一次性创建最终DataFrame
df_final = spark.createDataFrame(all_data_rows, target_columns)

场景B：API返回的列可能不一致（自动对齐列）

如果不同API返回的列有差异，可以将每行数据转为字典，Spark会自动对齐所有列，缺失值填充Null：

all_data_dicts = []

for cols, rows in all_api_results:
    # 将每行数据转为字典，键为列名，值为对应数据
    for row in rows:
        row_dict = dict(zip(cols, row))
        all_data_dicts.append(row_dict)

# Spark自动识别所有列，缺失列填充Null
df_final = spark.createDataFrame(all_data_dicts)

额外优化建议 #

• 控制并发数：max_workers不要设置过大，避免触发API的限流机制，建议参考API文档的并发限制
• 错误处理：在fetch_single_api中添加异常捕获，处理请求超时、失败的情况，比如添加重试逻辑或记录错误日志
• 减少JSON解析次数：原代码中两次调用response.json()，优化为解析一次后复用结果，节省CPU资源
• 大数据量分批处理：如果API返回数据量极大，可将请求分成多批，每批收集数据后合并一次DataFrame，避免内存溢出

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