Python拆分SQL WHERE IN大列表 分批查询规避AWS Data API上限

AWS Aurora Serverless 的Data API对单次查询请求的payload大小、IN子句的参数数量均存在硬限制，当ID列表长度超过1000时就会触发阈值报错。

核心逻辑是将全量ID列表拆分为多个小批次，逐批查询后合并结果，批次大小建议设置为800-900，预留足够余量避免卡阈值。

修改后完整代码 #

import awswrangler as wr
import pandas as pd

# 数据库连接参数配置
engine = wr.data_api.rds.connect(
    resource_arn = resource_arn, 
    database=database_name, 
    secret_arn=secret_arn
)

# 读取S3 CSV获取ID列表
read_df = wr.s3.read_csv(path=s3_path_in)
requested_ids = read_df["ids"].tolist()

# 列表分片函数，将全量ID拆分为指定大小的子列表
def split_list(lst, chunk_size):
    for i in range(0, len(lst), chunk_size):
        yield lst[i:i + chunk_size]

# 每批查询800个ID，可根据实际ID的长度调整大小
chunk_size = 800
id_chunks = list(split_list(requested_ids, chunk_size))

# 逐批查询并存储结果
result_dfs = []
query_id = """
select c.*
from table1 b
INNER JOIN table2 c on b.id = c.id
where b.id in %s
"""
for chunk in id_chunks:
    chunk_tuple = tuple(chunk)
    chunk_df = wr.data_api.rds.read_sql_query(query_id % chunk_tuple, engine)
    result_dfs.append(chunk_df)

# 合并所有批次的查询结果
out_ids = pd.concat(result_dfs, ignore_index=True)

优化建议 #

• 可替换字符串拼接SQL的写法为参数化查询，避免SQL注入风险，awswrangler的read_sql_query支持params参数传入查询参数
• 若ID总量极大（超过10万），可结合Lambda异步调用或Step Functions做分批处理，避免超出Lambda最长执行时间限制
• 可针对Data API的限流报错添加指数退避重试逻辑，提升查询稳定性

内容的提问来源于stack exchange，提问作者Prof. Falken