如何为Airflow任务设置BigQuery临时表以传递Dataframe

前置条件 #

确认Airflow环境已安装apache-airflow-providers-google依赖包（版本≥8.0.0避免接口兼容问题），且已配置好有权限读写BigQuery的GCP连接。 #

核心逻辑说明 #

原单进程内DataFrame直接传参的逻辑，拆分为Airflow多Task后，通过「上游Task处理完数据写入BigQuery临时表→把临时表名通过Airflow XCom传递给下游→下游Task读取临时表生成DataFrame继续处理」的流程实现数据传递。
注意不要使用BigQuery的会话级临时表，这类表仅在当前连接会话内可见，Airflow不同Task的BigQuery连接会话相互独立，无法跨Task访问，必须将临时表创建在指定的普通数据集下，通过过期规则实现自动清理。 #

具体实现步骤 #

• 提前在BigQuery中创建专门用于存放临时表的数据集，设置默认表过期时间为24小时，避免冗余数据占用存储。临时表统一按{项目ID}.{临时数据集名}.tmp_{DAGID}_{执行时间}_{步骤标识}规则命名，避免不同DAG、不同执行批次的表名冲突。
• 第一步：CSV加载任务实现

from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
import pandas as pd

def load_csv_to_bq(**context):
    # 读取CSV生成DataFrame，支持从本地/GCS/S3等路径读取
    df = pd.read_csv("你的CSV文件路径")
    # 生成当前步骤临时表名
    dag_id = context["dag"].dag_id
    exec_date = context["execution_date"].strftime("%Y%m%d%H%M%S")
    tmp_table = f"your_project_id.tmp_dataset.tmp_{dag_id}_{exec_date}_load_csv"
    # 写入BigQuery
    bq_hook = BigQueryHook(gcp_conn_id="你配置的GCP连接ID", use_legacy_sql=False)
    df.to_gbq(
        destination_table=tmp_table,
        project_id=bq_hook.project_id,
        if_exists="replace",
        credentials=bq_hook.get_credentials(),
        progress_bar=False # 生产环境关闭进度条避免日志冗余
    )
    # 将临时表名推入XCom供下游任务取用
    context["ti"].xcom_push(key="load_csv_tmp_table", value=tmp_table)

• 第二步：中间转换任务实现

def transform_step(**context):
    ti = context["ti"]
    # 拉取上游任务的临时表名
    upstream_tmp_table = ti.xcom_pull(key="load_csv_tmp_table", task_ids="load_csv_task_id")
    bq_hook = BigQueryHook(gcp_conn_id="你配置的GCP连接ID", use_legacy_sql=False)
    # 从BigQuery读取数据生成DataFrame
    df = bq_hook.get_pandas_df(sql=f"SELECT * FROM `{upstream_tmp_table}`")
    # 执行原有转换逻辑
    transformed_df = your_original_transform_function(df)
    # 生成当前步骤临时表名
    dag_id = context["dag"].dag_id
    exec_date = context["execution_date"].strftime("%Y%m%d%H%M%S")
    current_tmp_table = f"your_project_id.tmp_dataset.tmp_{dag_id}_{exec_date}_transform_step1"
    # 写入新的临时表
    transformed_df.to_gbq(
        destination_table=current_tmp_table,
        project_id=bq_hook.project_id,
        if_exists="replace",
        credentials=bq_hook.get_credentials(),
        progress_bar=False
    )
    # 推入XCom供下游任务取用
    ti.xcom_push(key="transform_step1_tmp_table", value=current_tmp_table)

• 第三步：最终输出任务
  读取最后一步转换生成的临时表，根据需求写入正式BigQuery表、导出到GCS等目标位置即可，所有中间临时表会按数据集的默认过期规则自动删除，无需手动清理。

优化建议 #

• 单表数据量超过100万行时，不建议使用pandas读写BigQuery，可改用BigQuery原生LOAD JOB直接从GCS加载CSV文件，转换逻辑也改为SQL直接操作BigQuery表，避免本地内存溢出同时大幅提升处理性能。
• 可给每个临时表单独设置更短的过期时间，写入后执行SQL：ALTER TABLE 你的临时表名 SET OPTIONS(expiration_timestamp = TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL 12 HOUR))即可。

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