如何将Airflow PostgresOperator输出保存为Pandas DataFrame

PostgresOperator本身仅用于执行SQL语句，默认不会返回查询结果集，所以无法直接将结果转为DataFrame，可通过以下两种常用方案实现需求：

方案1：使用PythonOperator + PostgresHook（最推荐） #

PostgresHook封装了直接将查询结果转为Pandas DataFrame的方法，操作更灵活：

1. 首先导入依赖包：

from airflow.operators.python import PythonOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
import pandas as pd

2. 定义查询转DataFrame的处理函数：

def fetch_sales_data(**context):
    # 初始化Postgres连接，使用你原有配置的conn_id
    pg_hook = PostgresHook(postgres_conn_id="db_conn_id")
    # 直接执行SQL并返回DataFrame，也可以传入sql参数为文件路径：sql="fetching_data.sql"
    sales_df = pg_hook.get_pandas_df(sql="select cast_id, prod_id, name from sales;")
    
    # 后续处理示例：
    # 1. 小数据集可以转JSON推送到XCom供后续任务使用
    context['ti'].xcom_push(key="sales_query_result", value=sales_df.to_json(orient="records"))
    # 2. 大数据量直接落地为本地CSV/对象存储文件
    sales_df.to_csv("/local/path/sales_result.csv", index=False)

3. 替换原有PostgresOperator为PythonOperator：

section_1 = PythonOperator(
    task_id='task_id',
    default_args=args,
    python_callable=fetch_sales_data,
    provide_context=True,
    dag=dag
)

方案2：自定义PostgresOperator子类 #

如果需要保留PostgresOperator的使用习惯，可以重写其execute方法处理返回结果：

from airflow.providers.postgres.operators.postgres import PostgresOperator
import pandas as pd

class PostgresToDFOperator(PostgresOperator):
    def execute(self, context):
        hook = self.get_db_hook()
        df = hook.get_pandas_df(self.sql)
        # 自定义处理逻辑，示例为保存到本地CSV
        df.to_csv("/local/path/sales_result.csv", index=False)
        # 可返回结果供XCom传递
        return df.to_json()

# 调用自定义Operator
section_1 = PostgresToDFOperator(
    task_id='task_id',
    default_args=args,
    postgres_conn_id="db_conn_id",
    sql="fetching_data.sql",
    dag=dag
)

注意事项 #

• 数据集较大时不要使用XCom传递DataFrame：Airflow默认XCom单条数据大小限制为48KB，超出会触发报错，大数据量建议直接落地到文件存储。
• 确保Airflow Worker节点已安装pandas依赖包，否则会出现运行时导入错误。
• Airflow 1.x版本的PostgresHook引入路径为from airflow.hooks.postgres_hook import PostgresHook，和2.x版本略有差异。

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