如何在Databricks中高效读取8亿条Delta表记录并按API限流要求批量写入第三方API？

看起来你在从Airflow迁移到Databricks Serverless时遇到了并行度不足、性能不达预期以及资源溢出的问题，这确实是大数据批量API写入中常见的挑战。我们可以从Spark作业的分区策略、Delta表读取优化、API调用的并行控制以及Serverless资源配置这几个维度来解决问题，帮你接近甚至达到Airflow时期的6M/分钟的写入速度。

一、先修复当前代码中的核心问题 #

1. 避免用count()触发全表扫描 #

你的代码中通过filtered_df.count()计算分区数，这会先触发一次全表Action操作——对于8亿条数据的表来说，这个count本身就会消耗大量资源，还会拖慢作业启动速度。更关键的是，这种方式计算出的分区数（8亿/1000=80万）会导致Spark调度压力剧增，小分区过多反而会降低并行效率，甚至引发内存/磁盘溢出。

优化方案：

• 放弃手动计算分区数，直接根据目标并行度设置分区数：Airflow用130线程实现130k条/秒的写入，对应Spark中我们可以设置650~1300个分区（留冗余量，避免单分区任务过重），每个分区对应一个独立的API处理进程。
• 利用Delta表的元数据快速获取近似行数（无需全表扫描）：

  # 通过Delta元数据获取近似行数
  from delta.tables import DeltaTable
  delta_table = DeltaTable.forName(spark, source_table)
  approx_count = delta_table.detail().select("numRecords").first()[0]
  

2. 不要强制按batch_size重分区 #

你当前的逻辑是让每个分区刚好对应1个1000条的batch，这会导致分区数爆炸。正确的做法是：让每个分区包含多个batch，比如每个分区处理100_{200个batch（即10}20万条数据），这样可以减少Spark的调度开销，同时让每个分区的API调用进程持续工作，避免频繁的资源启停。

二、优化Delta表的读取性能 #

针对8亿条数据的Delta表，读取阶段的优化能直接减少后续的资源消耗：

1. 只读取API需要的字段：不要全表读取，用select指定必要字段，减少数据传输和内存占用：

   df = spark.read.format("delta")\
       .table(source_table)\
       .select("day_type", "id", "api_required_field1", "api_required_field2")
   

2. 利用Delta的优化特性：提前对源表执行OPTIMIZE和ZORDER，让过滤day_type is not null时能跳过大量无关文件：

   OPTIMIZE tableName ZORDER BY (day_type);
   

3. 关闭不必要的Schema合并：读取时添加option("mergeSchema", "false")，减少Schema解析开销。

三、重写foreachPartition的API调用逻辑 #

你的process_partition函数是性能瓶颈的核心，必须从连接复用、限流控制、错误处理三个方面优化：

优化后的process_partition示例 #

import requests
from ratelimit import limits, sleep_and_retry
from tenacity import retry, stop_after_attempt, wait_exponential

def process_partition(partition):
    # 每个分区只创建1个API会话，复用HTTP连接（避免频繁创建连接的开销）
    session = requests.Session()
    # 给每个分区设置限流：比如每个分区每秒最多发送10次请求
    # 650个分区的话，总请求数为6500次/秒=6.5M条/分钟，接近目标
    CALLS_PER_SECOND = 10

    # 带重试的API发送函数
    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
    @sleep_and_retry
    @limits(calls=CALLS_PER_SECOND, period=1)
    def send_batch(batch):
        api_url = "https://your-3rd-party-api.com/batch-endpoint"
        try:
            response = session.post(
                api_url,
                json={"events": batch},
                timeout=10  # 设置超时，避免任务挂起
            )
            response.raise_for_status()
        except requests.exceptions.HTTPError as e:
            # 处理API错误，比如写入死信表
            print(f"Batch failed with HTTP error: {e}")
            # 可选：将失败批次写入Delta死信表
            # spark.createDataFrame(batch).write.mode("append").format("delta").table("dead_letter_table")
            raise

    batch = []
    for row in partition:
        # 只转换API需要的字段，减少数据体积
        event = {
            "id": row.id,
            "day_type": row.day_type,
            "field1": row.api_required_field1
        }
        batch.append(event)
        if len(batch) >= 1000:
            send_batch(batch)
            batch = []
    # 处理最后一个不满1000条的小批次
    if batch:
        send_batch(batch)
    session.close()

四、Databricks Serverless的资源配置调优 #

Serverless的性能和你选择的Compute规格直接相关，针对你的场景建议：

1. 选择大规格的Worker：比如用Large或XLarge的Serverless Compute，这类Worker有更多的CPU核心和内存，能支撑更高的并行度。
2. 调整Spark核心配置：

   spark = SparkSession.builder\
       .appName("OptimizeBatchProcessing")\
       .config("spark.sql.shuffle.partitions", "650")  # 和分区数保持一致
       .config("spark.default.parallelism", "650")\
       .config("spark.sql.adaptive.enabled", "true")  # 开启自适应执行计划
       .config("spark.sql.execution.arrow.pyspark.enabled", "true")  # 加速Python与Spark的数据传输
       .getOrCreate()
   

3. 避免过度分区：Serverless的自动扩缩容有上限，分区数不要超过1000，否则会导致部分分区等待资源，反而降低整体速度。

五、最终优化后的完整代码 #

from pyspark.sql import SparkSession
import requests
from ratelimit import limits, sleep_and_retry
from tenacity import retry, stop_after_attempt, wait_exponential

def create_spark_session():
    return SparkSession.builder\
        .appName("OptimizeBatchProcessing")\
        .config("spark.sql.shuffle.partitions", "650")\
        .config("spark.default.parallelism", "650")\
        .config("spark.sql.adaptive.enabled", "true")\
        .config("spark.sql.execution.arrow.pyspark.enabled", "true")\
        .getOrCreate()

def process_partition(partition):
    session = requests.Session()
    CALLS_PER_SECOND = 10

    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
    @sleep_and_retry
    @limits(calls=CALLS_PER_SECOND, period=1)
    def send_batch(batch):
        api_url = "https://your-3rd-party-api.com/batch-endpoint"
        try:
            response = session.post(api_url, json={"events": batch}, timeout=10)
            response.raise_for_status()
        except requests.exceptions.HTTPError as e:
            print(f"Batch failed: {e}")
            raise

    batch = []
    for row in partition:
        event = {
            "id": row.id,
            "day_type": row.day_type,
            "required_field": row.required_field
        }
        batch.append(event)
        if len(batch) >= 1000:
            send_batch(batch)
            batch = []
    if batch:
        send_batch(batch)
    session.close()

if __name__ == "__main__":
    spark = create_spark_session()
    source_table = "tableName"
    # 只读取必要字段，利用Delta的过滤优化
    df = spark.read.format("delta")\
        .option("mergeSchema", "false")\
        .table(source_table)\
        .select("id", "day_type", "required_field")
    filtered_df = df.filter(df["day_type"].isNotNull())
    # 设置目标分区数，匹配预期并行度
    partitioned_df = filtered_df.repartition(650)
    partitioned_df.foreachPartition(process_partition)

六、额外的监控与调建议 #

1. 通过Spark UI定位瓶颈：在Databricks作业中打开Spark UI，观察任务的执行时间、分区处理速度、内存/磁盘使用情况，判断是读取慢还是API写入慢。
2. 处理数据倾斜：如果发现部分分区处理时间远长于其他分区，说明存在数据倾斜，可以通过salting（加盐）的方式重新分区。
3. 测试小数据量验证：先用100万条数据测试优化后的代码，确认性能符合预期后再放大到全量数据。

备注：内容来源于stack exchange，提问作者pcincad