如何在Airflow中用Pandas处理S3文件？解决跨任务传递难题

针对你遇到的S3下载文件、Pandas DataFrame无法通过XCom序列化传递的问题，除了主机文件系统传递路径的方式，还有以下几种更优的实践方案：

方案一：用分布式存储替代本地文件系统 #

直接把中间文件存储到S3或MinIO这类分布式存储，而非主机本地路径，规避多Worker集群环境下本地文件无法跨节点访问的问题：

• 下载S3文件时，可直接在内存处理后转存为Parquet格式（比HDF5兼容性更好、压缩率更高）上传回S3
• 任务间传递S3上的文件路径（如s3://bucket/path/to/data.parquet），后续任务通过该路径直接读取

方案二：扩展XCom序列化能力（仅适用于小体量数据） #

Airflow默认XCom仅支持JSON序列化，但可以自定义序列化器支持Pandas DataFrame或二进制数据：

• 对于几MB级别的小DataFrame，可注册序列化器将其转为JSON或MsgPack格式传递
• 注意：1GB级别的DataFrame绝对不能用XCom传递，会严重占用元数据库资源，拖垮Airflow性能

方案三：合并轻量任务减少中间存储开销 #

如果部分任务逻辑简单（比如read_file_into_dataframe+prepare_dataframe_for_working），可以合并为单个任务，减少中间文件的读写次数，只在必要节点（如验证、入库前）持久化数据

优化后的代码示例（基于方案一） #

以下是结合S3存储中间Parquet文件的修改版本：

import pandas as pd
import boto3
from airflow import DAG
from airflow.decorators import task, task_group
from airflow.utils.dates import days_ago
from io import BytesIO

# 初始化S3客户端
s3_client = boto3.client('s3')

def save_df_to_s3(df, bucket_name, s3_path):
    """将DataFrame保存为Parquet格式到S3"""
    buffer = BytesIO()
    df.to_parquet(buffer, engine='pyarrow')
    buffer.seek(0)
    s3_client.put_object(Bucket=bucket_name, Key=s3_path, Body=buffer)
    return f"s3://{bucket_name}/{s3_path}"

def read_df_from_s3(s3_path):
    """从S3读取Parquet格式的DataFrame"""
    bucket_name, key = s3_path.replace("s3://", "").split("/", 1)
    response = s3_client.get_object(Bucket=bucket_name, Key=key)
    return pd.read_parquet(BytesIO(response['Body'].read()), engine='pyarrow')

@task
def download_from_s3(filename: str, bucket_name: str):
    """直接从S3读取Excel到内存，转存为Parquet后返回S3路径"""
    response = s3_client.get_object(Bucket=bucket_name, Key=filename)
    df = pd.read_excel(BytesIO(response['Body'].read()))
    # 转存为Parquet到S3临时路径
    parquet_key = f"temp/{filename.replace('.xlsx', '.parquet')}"
    return save_df_to_s3(df, bucket_name, parquet_key)

@task
def prepare_dataframe_for_working(s3_path: str):
    """读取S3上的Parquet，处理后再存回S3"""
    df = read_df_from_s3(s3_path)
    # 执行数据处理：设置索引、移除空行空列
    df = df.set_index('your_index_column')
    df = df.dropna(axis=0, how='all').dropna(axis=1, how='all')
    # 覆盖原文件或保存到新路径
    bucket_name = s3_path.split("/")[2]
    key = s3_path.split("/", 3)[3]
    return save_df_to_s3(df, bucket_name, key)

def _validate(df: pd.DataFrame):
    """自定义验证逻辑"""
    # 示例：检查是否存在缺失值
    return df.isnull().any().any()

@task
def validate(s3_path: str):
    """验证数据，返回是否有错误"""
    df = read_df_from_s3(s3_path)
    validation_errors = _validate(df)
    if validation_errors:
        # 保存错误信息到S3
        error_key = f"errors/{s3_path.split('/')[-1].replace('.parquet', '_errors.txt')}"
        s3_client.put_object(
            Bucket=s3_path.split("/")[2],
            Key=error_key,
            Body="Validation failed: missing values detected"
        )
        return True
    return False

@task
def save(s3_path: str):
    """将DataFrame写入Postgres"""
    df = read_df_from_s3(s3_path)
    # 执行Postgres写入逻辑，例如使用SQLAlchemy
    # df.to_sql('your_table', engine, if_exists='append', index=False)

@task_group
def process_file(filename: str, bucket_name: str):
    s3_parquet_path = download_from_s3(filename, bucket_name)
    processed_path = prepare_dataframe_for_working(s3_parquet_path)
    validation_result = validate(processed_path)
    # 仅验证通过时执行保存任务
    processed_path >> save(processed_path)

@task
def final():
    """最终收尾任务"""
    pass

def dummy_data():
    return [
        {'filename': 'file_1.xlsx', 'bucket_name': 'bucket1'},
        {'filename': 'file_2.xlsx', 'bucket_name': 'bucket2'},
    ]

with DAG('SO_process_file', start_date=days_ago(0), schedule=None, default_args={}, catchup=False):
    tasks = []
    for item in dummy_data():
        tsk = process_file(**item)
        tasks.append(tsk)
    tasks >> final()

关键注意事项 #

• 大文件处理：1GB级别的DataFrame必须用分布式存储传递，绝对禁止用XCom，避免压垮Airflow元数据库
• 格式选择：优先用Parquet替代HDF5，兼容性、压缩率和查询性能更优
• 多Worker环境：绝对不能依赖主机本地文件系统，所有中间数据必须放在共享存储中，否则跨Worker任务会找不到文件
• 依赖优化：可利用验证任务返回的布尔值控制后续保存任务的执行（如仅验证通过才触发入库）

内容的提问来源于stack exchange，提问作者Альберт Александров