将BigQuery数据追加至DataFrame的高效实现方法

架构老林

2026-7-2

问题：BigQuery数据转DataFrame的性能瓶颈优化

我在将BigQuery数据追加至DataFrame时遇到了严重的性能问题：查询本身仅需约20秒，但整个脚本运行耗时长达10分钟。

我的脚本代码：

groups = []
current_group = []
current_sum = 0

for index, row in df.iterrows():
    if current_sum + row['storeid'] <= 1000:
        current_group.append(row['tenantId'])
        current_sum += row['storeid']
    else:
        groups.append(current_group)
        current_group = [row['tenantId']]
        current_sum = row['storeid']

groups.append(current_group)
print(len(groups))

project1 = 'project'

timeline_data = []
for i in range(len(groups)):
    store_selection_t =  tenant_stores[tenant_stores.tenantId.isin(groups[i])]
    store_selection = list(set(store_selection_t.storeid))
    query = f"""
    WITH LatestEntries AS (
      SELECT
        *,
        ROW_NUMBER() OVER (PARTITION BY barcode ORDER BY Date(timestamp) DESC) AS row_num
      FROM
        `{project1}.dataset.table`
    )
    SELECT
      *
    FROM
      LatestEntries
    WHERE
      row_num = 1
    and storeid IN (@store_selection)
    """

    random_store_selection_str = ",".join(["'{}'".format(value) for value in store_selection])
    query = query.replace("@store_selection", random_store_selection_str)

    job_config = bigquery.QueryJobConfig()
    labels_shifted_aggregated = client.query(query, job_config=job_config).to_dataframe()
    timeline_data.append(labels_shifted_aggregated)

timeline_data = pd.concat(timeline_data)

查询本身耗时极短且返回3200万行数据，问题出在以下步骤：

labels_shifted_aggregated = client.query(query, job_config=job_config).to_dataframe()
timeline_data.append(labels_shifted_aggregated)

我曾尝试基于from google.cloud.bigquery_storage_v1 import BigQueryReadClient的解决方案，但完全无效。请问是否存在.to_dataframe()方法的替代方案？

解决方案

1. 核心优化：避免循环分批查询，改为一次性拉取

你的性能瓶颈本质是循环执行大量小查询+多次DataFrame内存拼接，即使单查询速度快，累加的网络开销、内存分配开销会被指数级放大。建议重构逻辑，一次性查询所有目标storeid：

# 先收集所有需要的storeid
all_storeids = set()
for i in range(len(groups)):
    store_selection_t = tenant_stores[tenant_stores.tenantId.isin(groups[i])]
    all_storeids.update(store_selection_t.storeid)
all_storeids = list(all_storeids)

# 一次性执行查询，使用参数化查询避免字符串拼接的安全风险和性能损耗
query = f"""
WITH LatestEntries AS (
  SELECT
    *,
    ROW_NUMBER() OVER (PARTITION BY barcode ORDER BY Date(timestamp) DESC) AS row_num
  FROM
    `{project1}.dataset.table`
)
SELECT
  *
FROM
  LatestEntries
WHERE
  row_num = 1
  and storeid IN UNNEST(@store_selection)
"""

job_config = bigquery.QueryJobConfig(
    query_parameters=[
        bigquery.ArrayQueryParameter("store_selection", "STRING", all_storeids)
    ]
)

# 一次性拉取全量数据到DataFrame
timeline_data = client.query(query, job_config=job_config).to_dataframe()

2. 正确使用BigQuery Storage API作为`to_dataframe()`替代方案

如果必须分批拉取，确保使用Storage API的正确姿势（之前的尝试可能存在会话配置错误）：

from google.cloud.bigquery_storage_v1 import BigQueryReadClient
from google.cloud.bigquery_storage_v1.types import ReadSession, TableReference

def query_to_dataframe_via_storage(client, query):
    # 执行查询并获取临时结果表
    query_job = client.query(query)
    query_job.result()  # 等待查询完成
    dest_table = query_job.destination

    # 初始化Storage客户端并创建读取会话
    storage_client = BigQueryReadClient()
    table_ref = TableReference(
        project_id=dest_table.project,
        dataset_id=dest_table.dataset_id,
        table_id=dest_table.table_id,
    )

    read_session = storage_client.create_read_session(
        parent=f"projects/{dest_table.project}",
        read_session=ReadSession(
            table=table_ref,
            data_format=ReadSession.DataFormat.ARROW,  # 使用Arrow格式加速传输
        ),
        max_stream_count=4,  # 根据机器性能调整并发流数量
    )

    # 读取所有流数据并合并为DataFrame
    dfs = []
    for stream in read_session.streams:
        reader = storage_client.read_rows(stream.name)
        dfs.append(reader.to_dataframe())
    
    # 删除临时表（可选，避免占用存储）
    client.delete_table(dest_table)
    return pd.concat(dfs, ignore_index=True)

# 在循环中替换原to_dataframe()调用
labels_shifted_aggregated = query_to_dataframe_via_storage(client, query)
timeline_data.append(labels_shifted_aggregated)

3. 优化DataFrame拼接逻辑

如果必须分批处理，避免每次append到列表再concat，可以用以下方式减少内存开销：

使用dask.dataframe并行处理分批数据，最后再合并为pandas DataFrame：

import dask.dataframe as dd

timeline_data = []
for i in range(len(groups)):
    # ... 执行查询获取单批次df
    timeline_data.append(dd.from_pandas(df, npartitions=1))

# 并行合并后转为pandas DataFrame
timeline_data = dd.concat(timeline_data).compute()

4. 其他辅助优化点

避免SELECT *，只查询业务需要的列，减少数据传输量
确保BigQuery表的storeid列存在索引，加速WHERE条件过滤
如果数据需要重复使用，将查询结果写入永久表而非临时表，避免重复计算

内容的提问来源于stack exchange，提问作者Serge de Gosson de Varennes