Azure Databricks(PySpark)分布式API调用及Parquet存储方案咨询

核心思路 #

抛弃驱动端串行调用的模式，直接基于Spark分布式DataFrame操作：让每个Executor的Task处理分片数据，独立调用REST API，将响应转为结构化数据后，通过Spark原生的分布式写入能力存储到ADLS2，全程无需驱动端集中处理。

具体实现方案 #

方案1：Pandas UDF批量调用API（高效推荐） #

批量调用能减少HTTP连接建立开销，比单条调用效率更高，适合API支持批量请求或高并发场景。

1. 读取Azure SQL数据（保持分布式）

from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf, col
import pandas as pd
import requests
from typing import Iterator

# Databricks中无需手动初始化SparkSession，默认已配置
spark = SparkSession.builder.appName("DistributedAPIFetch").getOrCreate()

# 配置Azure SQL JDBC连接
jdbc_url = "jdbc:sqlserver://<your-sql-server>.database.windows.net:1433;database=<your-db-name>"
conn_props = {
    "user": "<sql-username>",
    "password": "<sql-password>",
    "driver": "com.microsoft.sqlserver.jdbc.SQLServerDriver"
}

# 仅读取目标id列，生成分布式DataFrame
df_ids = spark.read.jdbc(url=jdbc_url, table="<target-table>", columns=["id"], properties=conn_props)

2. 定义批量处理Pandas UDF

# 预定义API响应的Schema，替换为实际返回的JSON结构
response_schema = "struct<response_id:string, data_field1:string, data_field2:int>"

@pandas_udf(response_schema)
def batch_call_api(ids: pd.Series) -> pd.DataFrame:
    # 初始化请求Session，复用连接池
    session = requests.Session()
    responses = []
    
    for id_val in ids:
        try:
            api_url = f"https://<your-api-endpoint>/?id={id_val}"
            resp = session.get(api_url, timeout=15)
            resp.raise_for_status()
            json_data = resp.json()
            
            # 提取JSON中需要的字段，转为DataFrame行
            responses.append({
                "response_id": json_data.get("id"),
                "data_field1": json_data.get("field1"),
                "data_field2": json_data.get("field2")
            })
        except Exception as e:
            # 异常降级处理，记录错误标识
            responses.append({
                "response_id": str(id_val),
                "data_field1": "api_call_failed",
                "data_field2": -1
            })
    
    session.close()
    return pd.DataFrame(responses)

# 调用UDF，生成包含API响应的分布式DataFrame
df_api_responses = df_ids.withColumn("api_data", batch_call_api(col("id"))).select("api_data.*")

3. 写入ADLS2为Parquet

# ADLS2路径使用abfss协议（Databricks已集成ADLS2认证）
adls_output_path = "abfss://<container-name>@<storage-account>.dfs.core.windows.net/api-responses-parquet"

# 分布式写入Parquet，支持overwrite/append等模式
df_api_responses.write.mode("overwrite").parquet(adls_output_path)

方案2：普通PySpark UDF单条调用API #

如果API不支持批量请求，可使用普通UDF单条处理数据：

from pyspark.sql.functions import udf
from pyspark.sql.types import StructType, StructField, StringType, IntegerType

# 定义API响应的结构化Schema
response_schema = StructType([
    StructField("response_id", StringType(), nullable=True),
    StructField("data_field1", StringType(), nullable=True),
    StructField("data_field2", IntegerType(), nullable=True)
])

@udf(response_schema)
def single_call_api(id_val):
    try:
        api_url = f"https://<your-api-endpoint>/?id={id_val}"
        resp = requests.get(api_url, timeout=15)
        resp.raise_for_status()
        json_data = resp.json()
        return (json_data.get("id"), json_data.get("field1"), json_data.get("field2"))
    except Exception as e:
        return (str(id_val), "api_call_failed", -1)

# 生成响应DataFrame
df_api_responses = df_ids.withColumn("api_data", single_call_api(col("id"))).select("api_data.*")

# 写入ADLS2
df_api_responses.write.mode("overwrite").parquet(adls_output_path)

方案3：foreachPartition自定义分布式处理 #

如果需要更灵活的存储逻辑（比如直接操作ADLS2文件），可使用foreachPartition让每个分区在Executor上独立处理：

import pandas as pd
from azure.storage.filedatalake import DataLakeServiceClient

def process_partition(partition_rows):
    # 初始化ADLS2客户端（每个分区初始化一次，复用连接）
    adls_service_client = DataLakeServiceClient(
        account_url="https://<storage-account>.dfs.core.windows.net",
        credential="<storage-credential>"  # 可使用SAS token或服务主体认证
    )
    file_system_client = adls_service_client.get_file_system_client(file_system="<container-name>")
    
    # 收集分区内所有API响应
    partition_data = []
    session = requests.Session()
    
    for row in partition_rows:
        id_val = row.id
        try:
            api_url = f"https://<your-api-endpoint>/?id={id_val}"
            resp = session.get(api_url, timeout=15)
            resp.raise_for_status()
            partition_data.append(resp.json())
        except Exception as e:
            partition_data.append({"id": str(id_val), "error": str(e)})
    
    session.close()
    
    # 将分区数据转为DataFrame，批量写入Parquet
    if partition_data:
        df_part = pd.DataFrame(partition_data)
        # 生成唯一的分区文件名，避免冲突
        file_name = f"partition-{hash(tuple(partition_data[0].values()))}.parquet"
        file_client = file_system_client.create_file(file_name)
        
        # 将Pandas DF转为字节流写入ADLS2
        parquet_bytes = df_part.to_parquet()
        file_client.upload_data(parquet_bytes, overwrite=True)

# 执行分布式分区处理
df_ids.foreachPartition(process_partition)

关键优化建议 #

• 连接池复用：在UDF或foreachPartition中使用requests.Session()，避免重复建立HTTP连接，提升调用效率。
• 重试与限流：使用tenacity库实现API调用重试逻辑，同时根据API的并发限制调整Spark executor数量和核心数，避免触发限流。
• Schema预定义：提前明确API响应的Schema，避免Spark自动推断Schema带来的额外开销。
• 小文件合并：写入Parquet后，可通过repartition或coalesce合并小文件，提升后续查询效率。

关于foreach的存储说明 #

如果使用foreach（单Row处理），不建议直接单条写入Parquet——Parquet是列式存储，单条写入效率极低。正确做法是在foreachPartition中收集整个分区的响应数据，转为Pandas DataFrame后批量写入，既保证分布式处理，又能提升存储效率。

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