如何用Lambda分割S3中的大CSV文件并保留表头行？

核心思路 #

要实现带表头的CSV分割，关键是先提取原文件的表头行，再将数据行分割为多个片段，最后给每个片段文件添加表头后上传到目标S3桶。以下提供Python和bash两种实现方式，适配Lambda运行环境。

方式一：Python实现（更易维护） #

1. 配置Lambda权限 #

确保Lambda执行角色拥有以下权限：

• 源S3桶的s3:GetObject权限
• 目标S3桶的s3:PutObject权限

2. 完整Lambda代码 #

import boto3
import os
from pathlib import Path

s3 = boto3.client('s3')
# 可自定义分割行数，比如每10000行一个文件
SPLIT_ROW_COUNT = 10000

def lambda_handler(event, context):
    # 从S3事件中获取源文件信息
    source_bucket = event['Records'][0]['s3']['bucket']['name']
    source_key = event['Records'][0]['s3']['object']['key']
    # 解析文件名（去掉后缀）
    file_name = Path(source_key).stem
    file_ext = Path(source_key).suffix

    # 下载源文件到Lambda临时目录
    local_file_path = f"/tmp/{source_key.split('/')[-1]}"
    s3.download_file(source_bucket, source_key, local_file_path)

    # 提取表头行
    with open(local_file_path, 'r') as f:
        header = f.readline()
        # 读取剩余数据行
        data_lines = f.readlines()

    # 分割数据行并生成文件
    total_splits = (len(data_lines) + SPLIT_ROW_COUNT - 1) // SPLIT_ROW_COUNT
    for i in range(total_splits):
        start_idx = i * SPLIT_ROW_COUNT
        end_idx = start_idx + SPLIT_ROW_COUNT
        chunk_lines = data_lines[start_idx:end_idx]

        # 生成输出文件名（如uniqueid_01.csv）
        output_file_name = f"{file_name}_{str(i+1).zfill(2)}{file_ext}"
        output_local_path = f"/tmp/{output_file_name}"

        # 写入表头+数据
        with open(output_local_path, 'w') as f:
            f.write(header)
            f.writelines(chunk_lines)

        # 上传到目标S3桶（替换为你的目标桶名）
        target_bucket = "your-target-bucket"
        s3.upload_file(output_local_path, target_bucket, output_file_name)

    # 清理临时文件（可选）
    os.remove(local_file_path)
    for i in range(total_splits):
        os.remove(f"/tmp/{file_name}_{str(i+1).zfill(2)}{file_ext}")

    return {
        'statusCode': 200,
        'body': f"分割完成，共生成{total_splits}个文件"
    }

注意事项 #

• 调整SPLIT_ROW_COUNT参数控制每个文件的行数
• Lambda的临时目录/tmp默认最大512MB，若文件过大，可在Lambda配置中扩展临时存储至10GB
• 若处理超大型文件（超过临时存储容量），建议结合S3 Select分批读取数据，避免内存溢出

方式二：Bash脚本实现（适配原参考逻辑） #

如果偏好bash脚本，可修改原逻辑，添加表头合并步骤：

1. Lambda配置 #

• 选择运行环境为Amazon Linux 2023
• 执行角色同样需要S3读写权限

2. 脚本代码 #

#!/bin/bash

# 从Lambda事件中解析参数（需在Lambda配置中启用环境变量或直接传入）
SOURCE_BUCKET=$1
SOURCE_KEY=$2
TARGET_BUCKET=$3
SPLIT_LINES=10000

# 下载源文件到临时目录
aws s3 cp s3://$SOURCE_BUCKET/$SOURCE_KEY /tmp/input.csv

# 提取表头
head -n 1 /tmp/input.csv > /tmp/header.csv

# 分割数据行（跳过表头）
tail -n +2 /tmp/input.csv | split -l $SPLIT_LINES - /tmp/split_

# 遍历分割文件，添加表头并上传
count=1
for file in /tmp/split_*; do
    # 生成带表头的文件
    output_file=$(basename $SOURCE_KEY .csv)_$(printf "%02d" $count).csv
    cat /tmp/header.csv $file > /tmp/$output_file
    # 上传到目标桶
    aws s3 cp /tmp/$output_file s3://$TARGET_BUCKET/$output_file
    # 清理临时文件
    rm $file /tmp/$output_file
    count=$((count+1))
done

# 清理剩余临时文件
rm /tmp/input.csv /tmp/header.csv

注意事项 #

• 需确保Lambda环境中安装了aws-cli（Amazon Linux环境默认已预装）
• 同样注意临时存储容量限制，避免处理超大型文件时出错

内容的提问来源于stack exchange，提问作者Ju Hae Lee