Spark长运行作业优化：API调用型UDF性能提升方案咨询

原代码中每条数据通过UDF单独发起API请求，存在大量重复请求、网络IO开销大、并发度低的问题，直接导致作业耗时极长。以下是针对性的优化方案：

1. 批量处理唯一的(city, state)组合，避免重复API调用 #

如果数据中存在大量相同的city+state组合，每个组合单独请求API完全是资源浪费。先提取所有唯一的组合，一次性请求完所有数据后，再通过关联操作把邮编映射回原数据集。

from pyspark.sql import functions as F

# 提取唯一的城市-州组合
unique_city_state = (trxDF.join(broadcast(citiesDF), "city_id")
                     .filter(F.col("state_abv").isNotNull())
                     .select("city", "state_abv")
                     .dropDuplicates()
                     .collect())

# 批量请求API，构建映射字典
zip_map = {}
for row in unique_city_state:
    city = row["city"]
    state = row["state_abv"]
    url = f"http://api.zippopotam.us/us/{state}/{city.replace(' ', '%20')}"
    response = requests.get(url)
    if response.status_code == 200:
        data = response.json()
        if 'places' in data and len(data['places']) > 0:
            zip_map[(city, state)] = data['places'][0]['post code']
        else:
            zip_map[(city, state)] = None
    else:
        zip_map[(city, state)] = None

# 将映射转为Spark DataFrame并关联原数据集
zip_df = spark.createDataFrame(zip_map.items(), ["city_state", "zip_code"])
zip_df = zip_df.withColumn("city", F.split(F.col("city_state"), ",")[0])\
               .withColumn("state_abv", F.split(F.col("city_state"), ",")[1])\
               .drop("city_state")

finalDF = (trxDF.join(broadcast(citiesDF), "city_id")
           .filter(F.col("state_abv").isNotNull())
           .join(broadcast(zip_df), ["city", "state_abv"]))

2. 使用异步请求提升API调用并发度 #

批量同步请求还是串行执行，效率极低。改用异步HTTP库（如aiohttp）同时发起多个请求，大幅缩短API调用的总耗时。

import aiohttp
import asyncio

async def fetch_async(session, city, state):
    url = f"http://api.zippopotam.us/us/{state}/{city.replace(' ', '%20')}"
    async with session.get(url) as response:
        if response.status == 200:
            data = await response.json()
            return (city, state, data['places'][0]['post code'] if 'places' in data and data['places'] else None)
        return (city, state, None)

async def batch_fetch(city_state_list):
    async with aiohttp.ClientSession() as session:
        tasks = [fetch_async(session, city, state) for city, state in city_state_list]
        return await asyncio.gather(*tasks)

# 转换为列表格式
city_state_list = [(row["city"], row["state_abv"]) for row in unique_city_state]

# 执行异步批量请求
loop = asyncio.get_event_loop()
results = loop.run_until_complete(batch_fetch(city_state_list))

# 构建映射DataFrame并关联
zip_df = spark.createDataFrame(results, ["city", "state_abv", "zip_code"])

finalDF = (trxDF.join(broadcast(citiesDF), "city_id")
           .filter(F.col("state_abv").isNotNull())
           .join(broadcast(zip_df), ["city", "state_abv"]))

3. 加入缓存机制，复用已有结果 #

如果作业需要重复运行，或者存在大量重复的city+state组合，可以用缓存（如本地LRU缓存、Redis）存储已请求过的结果，彻底避免重复API调用。

示例（本地LRU缓存）：

from functools import lru_cache
import requests

# LRU缓存，maxsize设置最大缓存条目数
@lru_cache(maxsize=1000)
def fetch_cached(city, state):
    url = f"http://api.zippopotam.us/us/{state}/{city.replace(' ', '%20')}"
    response = requests.get(url)
    if response.status_code == 200:
        data = response.json()
        return data['places'][0]['post code'] if 'places' in data and data['places'] else None
    return None

# 批量请求时复用缓存结果
zip_map = {}
for city, state in city_state_list:
    zip_map[(city, state)] = fetch_cached(city, state)

4. 调整Spark执行参数，优化作业并行度 #

• 增加Executor资源：提交作业时通过--executor-cores 4 --executor-memory 8g调整Executor的核数和内存，提升并行处理能力。
• 调整分区数：在join或filter操作后，使用repartition(n)调整分区数，避免分区过多（小文件多）或过少（并行度不足）的问题。
• 优化写入操作：如果zip_code基数极大，使用partitionBy会产生大量小文件，可先通过coalesce(n)合并分区，或者改用bucketBy替代partitionBy：

# 合并分区后写入，减少小文件
finalDF.coalesce(50)\
       .write.mode("overwrite")\
       .partitionBy("zip_code")\
       .parquet(finalPath)

# 或者用bucketBy替代partitionBy（适合高基数字段）
finalDF.write.mode("overwrite")\
       .bucketBy(100, "zip_code")\
       .saveAsTable("zip_trx_table")

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