Python迭代器-大数据量的处理

@半良人2024-12-05 8:07

一生成器的实际使用（大量数据的导出）

python 复制代码

#分批导出数据然后分批写入excel
import pandas as pd
import openpyxl
from openpyxl.utils.dataframe import dataframe_to_rows


def execute_query(query):
    # 假设这是执行 SQL 查询的函数
    # 返回查询结果
    pass


def fetch_products_in_batches(batch_size=1000):
    offset = 0
    while True:
        query = f"SELECT * FROM products LIMIT {batch_size} OFFSET {offset}"
        results = execute_query(query)
        if not results:
            break
        yield results
        offset += batch_size


def export_to_excel(generator, filename):
    # 创建一个新的 Excel 工作簿
    wb = openpyxl.Workbook()
    ws = wb.active

    # 写入表头
    first_batch = next(generator)
    df = pd.DataFrame(first_batch)
    for r in dataframe_to_rows(df, index=False, header=True):
        ws.append(r)

    # 写入后续批次的数据
    for batch in generator:
        df = pd.DataFrame(batch)
        for r in dataframe_to_rows(df, index=False, header=False):
            ws.append(r)

    # 保存工作簿
    wb.save(filename)


# 主逻辑
generator = fetch_products_in_batches()
export_to_excel(generator, 'products.xlsx')