Airflow与S3交互[编辑 | 编辑源代码]

介绍[编辑 | 编辑源代码]

Apache Airflow 是一个开源的工作流自动化工具，用于编排复杂的数据管道。Amazon S3 (Simple Storage Service) 是一种广泛使用的对象存储服务。Airflow 与 S3 的交互允许用户从 S3 存储桶读取数据、写入数据，或触发基于 S3 事件的工作流。这种集成在数据工程和机器学习场景中非常常见，例如：

从 S3 加载原始数据到数据处理管道
将处理后的结果保存回 S3
监控 S3 文件变化并触发后续任务

核心概念[编辑 | 编辑源代码]

Airflow 的 S3 钩子 (Hook)[编辑 | 编辑源代码]

Airflow 提供了 S3Hook，这是一个高级接口，封装了与 S3 交互的常用操作。它基于 boto3 库实现。

连接配置[编辑 | 编辑源代码]

要使用 S3Hook，需先在 Airflow 中配置 S3 连接： 1. 在 Airflow Web UI 的 Admin → Connections 中添加新连接 2. 选择连接类型为 Amazon Web Services 3. 提供 AWS 访问密钥 (Access Key ID 和 Secret Access Key) 或使用 IAM 角色

代码示例[编辑 | 编辑源代码]

基本文件操作[编辑 | 编辑源代码]

from airflow.providers.amazon.aws.hooks.s3 import S3Hook

def upload_to_s3():
    hook = S3Hook(aws_conn_id='my_aws_conn')
    hook.load_file(
        filename='/local/path/to/file.csv',
        key='s3_key/path/file.csv',
        bucket_name='my-bucket',
        replace=True
    )

def download_from_s3():
    hook = S3Hook(aws_conn_id='my_aws_conn')
    hook.download_file(
        key='s3_key/path/file.csv',
        bucket_name='my-bucket',
        local_path='/local/download/path/'
    )

使用 S3 Sensor[编辑 | 编辑源代码]

S3 Sensor 可以监测 S3 中文件的存在或变化：

from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor

s3_sensor = S3KeySensor(
    task_id='check_s3_file',
    bucket_key='data/*.csv',
    bucket_name='my-bucket',
    aws_conn_id='my_aws_conn',
    timeout=18*60*60,
    poke_interval=120
)

高级用法[编辑 | 编辑源代码]

使用 S3 前缀和通配符[编辑 | 编辑源代码]

Airflow 支持 S3 通配符查询：

hook = S3Hook(aws_conn_id='my_aws_conn')
files = hook.list_keys(
    bucket_name='my-bucket',
    prefix='data/year=2023/month=*/',
    delimiter='/'
)

大文件分块上传[编辑 | 编辑源代码]

对于大文件，可以使用分块上传：

hook = S3Hook(aws_conn_id='my_aws_conn')
hook.load_file(
    filename='/path/to/large_file.bin',
    key='large_files/data.bin',
    bucket_name='my-bucket',
    replace=True,
    encrypt=True,
    acl_policy='private',
    multipart=True,
    num_parallel_uploads=4
)

实际案例[编辑 | 编辑源代码]

数据管道示例[编辑 | 编辑源代码]

以下 DAG 展示了一个完整的数据处理流程：

1. 监测 S3 中新上传的 CSV 文件 2. 下载文件到本地 3. 处理数据 4. 上传结果回 S3

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
from datetime import datetime

def process_data(**context):
    # 获取触发文件的信息
    s3_key = context['ti'].xcom_pull(task_ids='check_new_file')
    # 处理逻辑...

with DAG('s3_data_pipeline', start_date=datetime(2023,1,1)) as dag:
    check_file = S3KeySensor(
        task_id='check_new_file',
        bucket_key='input/*.csv',
        bucket_name='data-bucket',
        aws_conn_id='aws_default',
        mode='poke',
        timeout=300
    )
    
    download_task = PythonOperator(
        task_id='download_from_s3',
        python_callable=download_from_s3
    )
    
    process_task = PythonOperator(
        task_id='process_data',
        python_callable=process_data,
        provide_context=True
    )
    
    upload_task = PythonOperator(
        task_id='upload_to_s3',
        python_callable=upload_to_s3
    )
    
    check_file >> download_task >> process_task >> upload_task

性能优化[编辑 | 编辑源代码]

并行操作[编辑 | 编辑源代码]

对于大量小文件，可以使用多线程并行处理：

from concurrent.futures import ThreadPoolExecutor

def process_multiple_files():
    hook = S3Hook(aws_conn_id='my_aws_conn')
    keys = hook.list_keys(bucket_name='my-bucket', prefix='data/')
    
    def process_key(key):
        local_path = f'/tmp/{key.split("/")[-1]}'
        hook.download_file(key, 'my-bucket', local_path)
        # 处理文件...
    
    with ThreadPoolExecutor(max_workers=8) as executor:
        executor.map(process_key, keys)

内存管理[编辑 | 编辑源代码]

处理大文件时，可以使用流式传输避免内存溢出：

hook = S3Hook(aws_conn_id='my_aws_conn')
s3_obj = hook.get_key('large_file.bin', 'my-bucket')

with open('/local/path.bin', 'wb') as f:
    for chunk in s3_obj.iter_chunks(chunk_size=1024*1024):  # 1MB chunks
        f.write(chunk)