如何分批连接两张无法放入内存的超大型SQLAlchemy表格？

问题描述 #

我正在使用SQLAlchemy，需要连接两张表格。其中一张表格有5亿行数据，若一次性执行查询会超时，且无法将整张表载入内存，唯一的选择是分批处理。

我的问题在于，需要连接两张表，但不能直接同时分批获取两张表的数据，因为无法确保对应ID会出现在同一批次中。我需要完成这个连接操作，但不知道如何在不完整查询至少一张表的情况下实现。

以下是我的代码：

some_table = connector.connect_to_table(metadata=m, table="SOME_TABLE")
other_table = connector.connect_to_table(metadata=m, table="OTHER_TABLE")

query = select([some_table.c.id,
                some_table.c.type,
                other_table.c.some_other_type]
               ).join(vehicles_individual_query,
                      (some_table.c.id== other_table.c.id)
                      )

conn = self.engine.connect().execution_options(
        stream_results=True)

for chunk_dataframe in pd.read_sql(
        query, conn, chunksize=1000):

    print(f"Got dataframe w/{len(chunk_dataframe)} rows")
    # ... do something with dataframe ...

显然，由于连接条件的存在，即使按chunksize=1000获取数据，仍需处理整张other_table。请问如何实现分批连接？

方法1：按ID范围分片查询 #

核心是将数据量较大的表（比如SOME_TABLE）按ID区间拆分，每次仅查询指定ID范围内的行并关联另一张表，避免一次性加载全表。

实现步骤：

1. 先获取目标表的ID最小/最大值，确定分片范围
2. 按固定步长拆分ID区间，循环执行每批次的连接查询

代码示例：

from sqlalchemy import func

some_table = connector.connect_to_table(metadata=m, table="SOME_TABLE")
other_table = connector.connect_to_table(metadata=m, table="OTHER_TABLE")

conn = self.engine.connect()

# 获取ID的边界值
min_id = conn.execute(select(func.min(some_table.c.id))).scalar()
max_id = conn.execute(select(func.max(some_table.c.id))).scalar()

# 每批处理的ID步长，根据数据库性能调整
batch_id_step = 100000
current_start_id = min_id

while current_start_id <= max_id:
    current_end_id = current_start_id + batch_id_step
    # 构造当前批次的查询：仅查询ID在[start, end)区间的行并关联
    batch_query = (
        select([some_table.c.id, some_table.c.type, other_table.c.some_other_type])
        .join(other_table, some_table.c.id == other_table.c.id)
        .where(some_table.c.id >= current_start_id, some_table.c.id < current_end_id)
    )

    # 分批读取当前区间的数据
    for chunk_df in pd.read_sql(batch_query, conn, chunksize=1000):
        print(f"Got dataframe w/{len(chunk_df)} rows")
        # ... 处理数据逻辑 ...

    current_start_id = current_end_id

方法2：键集分页（Key Set Pagination） #

如果ID存在断号（比如有数据删除），按范围分片会导致每批数据量波动大，此时可以用键集分页：以上一批的最后一个ID作为下一批的起始标记，保证每批数据量稳定。

代码示例：

some_table = connector.connect_to_table(metadata=m, table="SOME_TABLE")
other_table = connector.connect_to_table(metadata=m, table="OTHER_TABLE")

conn = self.engine.connect()
last_processed_id = None
# 每批返回的行数，按需调整
batch_row_limit = 100000

while True:
    base_query = (
        select([some_table.c.id, some_table.c.type, other_table.c.some_other_type])
        .join(other_table, some_table.c.id == other_table.c.id)
        .order_by(some_table.c.id)
    )

    if last_processed_id is not None:
        # 从上次处理的最后一个ID之后开始查询
        base_query = base_query.where(some_table.c.id > last_processed_id)

    # 限制当前批次的行数
    batch_query = base_query.limit(batch_row_limit)
    batch_df = pd.read_sql(batch_query, conn)

    if batch_df.empty:
        break  # 数据处理完成

    print(f"Got dataframe w/{len(batch_df)} rows")
    # ... 处理数据逻辑 ...

    # 更新最后处理的ID
    last_processed_id = batch_df['id'].iloc[-1]

关键注意事项 #

• 确保some_table.c.id和other_table.c.id都建立了索引，否则每批次的连接查询会极慢
• 根据数据库的负载和性能调整批次大小，避免单次查询超时或查询次数过多
• 如果使用分布式数据库，可以结合数据库自身的分片机制进一步优化

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