Python循环转PySpark构建DataFrames及Delta表写入报错求助

核心问题分析 #

原代码存在两个关键问题：

1. 串行API调用：单循环遍历5K个账户ID，逐个调用API，完全没有利用Spark的分布式并行能力，导致耗时过长。
2. 大量小DataFrame合并：每次创建仅含1条数据的DataFrame再通过unionAll合并，最终生成数千个小任务，触发spark.driver.maxResultSize超限报错。

优化实现代码 #

1. 并行化API调用与数据收集 #

利用Spark RDD的mapPartitions实现分布式API调用，在每个Executor分区内批量处理账户ID，避免单线程串行等待：

from pyspark.sql import Row

# 1. 将账户ID列表转换为Spark DataFrame，为分布式计算做准备
account_ids_df = spark.createDataFrame(account_header['providerAccounts'], schema="id string")

# 2. 定义分区处理函数：每个分区内批量调用API，收集账户和联系人数据
def process_account_partition(partition):
    account_records = []
    provider_records = []
    
    for row in partition:
        account_id = row.id
        # 调用API获取详情（可添加重试机制，见后续优化点）
        account_detail = get_accounts(account_id)
        
        # 处理账户数据，转为Row对象
        account_data = account_detail['providerAccount']
        account_records.append(Row(**account_data))
        
        # 处理联系人数据，添加账户ID关联字段（可选，方便后续关联查询）
        for provider in account_data['providers']:
            provider['account_id'] = account_id
            provider_records.append(Row(**provider))
    
    # 返回当前分区的账户和联系人结果
    return (account_records, provider_records)

# 3. 分布式处理分区，提取结果RDD
result_rdd = account_ids_df.rdd.mapPartitions(lambda p: [process_account_partition(p)])

# 拆分账户和联系人的RDD
account_rdd = result_rdd.flatMap(lambda x: x[0])
provider_rdd = result_rdd.flatMap(lambda x: x[1])

# 4. 转换为DataFrame（复用原有schema）
df_account = spark.createDataFrame(account_rdd, schema=account_schema)
df_provider = spark.createDataFrame(provider_rdd, schema=provider_schema)

2. 解决Delta写入的内存报错 #

针对spark.driver.maxResultSize超限问题，可通过以下两种方式处理：

• 调大Driver内存限制（根据集群资源调整）：

  spark.conf.set("spark.driver.maxResultSize", "8g")
  

• 提前重分区，减少单任务数据量：

  # 根据数据量设置分区数：5K账户设10-20分区，10K联系人设20-40分区
  df_account.repartition(15).write.format("delta").mode('overwrite').save(acct_fulldeltapath)
  df_provider.repartition(30).write.format("delta").mode('overwrite').save(px_fulldeltapath)
  

额外优化点 #

1. 批量API调用：如果目标API支持批量查询（比如一次传入多个account_id），直接修改process_account_partition实现批量调用，可大幅减少HTTP请求次数，提升效率。
2. API重试机制：给get_accounts添加重试逻辑，避免网络波动导致任务失败：

   from tenacity import retry, stop_after_attempt, wait_exponential
   
   @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
   def get_accounts(account_id):
       # 原API调用逻辑
       response = requests.get(f"your_api_endpoint/{account_id}")
       response.raise_for_status()
       return response.json()
   

3. Executor依赖分发：如果get_accounts依赖第三方库（如requests），需要通过spark.sparkContext.addPyFile或集群配置将依赖包分发到所有Executor节点。

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