Vertex AI中BigQuery to_dataframe()性能问题排查与优化咨询

问题描述 #

在Vertex AI流水线中执行BigQuery查询时，我遇到了严重的性能问题。示例代码如下：

query = "SELECT * FROM `project.dataset.table`"  # 返回5,000,000行数据
query_job = client.query(query)
df = query_job.result().to_dataframe()

其中df = query_job.result().to_dataframe()操作在Vertex AI上耗时约50分钟，在本地机器仅需不到5分钟。当前资源配置为：

.set_cpu_limit("60")
.set_memory_limit("128G")

我使用to_dataframe()将约500万行数据加载到Pandas DataFrame中。

咨询问题 #

1. 为何该操作在Vertex AI上速度远慢于本地机器？
2. 可对Vertex AI配置或代码做哪些调整提升性能？

问题解答 #

1. 性能差异的核心原因 #

• 网络链路差异：本地机器大概率与BigQuery处于同一区域/可用区，网络链路更优；而Vertex AI流水线的执行环境可能与BigQuery跨区域部署，或是受到平台内部网络带宽限制，导致数据传输耗时剧增。
• 客户端配置与版本差异：本地使用的BigQuery Python客户端库可能是最新版本，默认并行下载线程数等配置更高效；Vertex AI环境可能搭载旧版库，或是默认配置未充分利用资源。
• 资源利用率不足：虽然配置了高CPU和内存，但to_dataframe()默认的低并发处理逻辑无法利用60核的CPU资源，导致CPU闲置，网络或数据序列化成为性能瓶颈；同时Vertex AI环境可能存在资源争用，实际分配的资源未达到配置值。
• 序列化损耗：Vertex AI的Python运行时可能因依赖库版本不匹配，在BigQuery结果转Pandas DataFrame的序列化/反序列化过程中产生额外性能开销。

2. 性能优化调整方案 #

代码层面优化

• 启用BigQuery Storage API：该API比传统REST API快数倍，需先在GCP控制台启用API，再安装依赖并修改代码：

pip install google-cloud-bigquery-storage

from google.cloud.bigquery import QueryJobConfig

job_config = QueryJobConfig(use_query_cache=True, use_bigquery_storage_api=True)
query_job = client.query(query, job_config=job_config)
df = query_job.result().to_dataframe()

• 先导出到GCS再加载：跳过直接查询转DataFrame的步骤，先将数据导出到Cloud Storage，再加载到Pandas，大幅降低传输耗时：

# 导出BigQuery表到GCS Parquet文件
destination_uri = "gs://your-bucket/temp_data.parquet"
extract_job = client.extract_table(
    "project.dataset.table",
    destination_uri,
    location="US"  # 需与BigQuery表区域一致
)
extract_job.result()

# 从GCS加载Parquet到DataFrame
import pandas as pd
df = pd.read_parquet(destination_uri)

• 优化查询与下载参数：增大page_size提升批量下载效率，同时避免SELECT *只查询所需列，减少数据传输量：

query = "SELECT col1, col2, col3 FROM `project.dataset.table`"  # 仅保留需要的列
query_job = client.query(query)
df = query_job.result(page_size=100000).to_dataframe(progress_bar_type=None)

Vertex AI配置优化

• 区域对齐：将Vertex AI流水线的执行区域设置为与BigQuery表所在区域完全一致，消除跨区域数据传输的延迟和带宽损耗。
• 资源配额与配比调整：检查Vertex AI的资源配额，确认60核CPU和128G内存是否实际分配到位；当前60核配128G内存的配比存在内存过剩情况，可调整为32核+64G这类更均衡的配置，避免资源浪费。
• 升级运行时环境：选择Vertex AI中最新版本的Python运行时，确保依赖库为最新稳定版，减少运行时层面的性能损耗。

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