如何使用PySpark重命名AWS Glue脚本输出到S3位置的文件？

实现逻辑 #

Glue作业基于Spark引擎运行，默认输出文件为part-xxx前缀的分片文件，附带_SUCCESS等元数据文件。以下代码直接集成在Glue作业末尾，作业写完数据后自动执行S3文件重命名，无需额外触发服务。

前置要求 #

• Glue作业关联的IAM角色已配置目标S3路径的s3:GetObject、s3:PutObject、s3:DeleteObject权限
• 无需额外安装依赖，Glue运行环境默认自带boto3、PySpark依赖

参考代码 #

import sys
import boto3
from awsglue.context import GlueContext
from pyspark.context import SparkContext

# 初始化Glue上下文
sc = SparkContext()
glueContext = GlueContext(sc)
spark = glueContext.spark_session

# ---------------------- 你的原有ETL逻辑 ----------------------
# 示例：假设有处理后的DataFrame需要输出到S3
# df = 你的处理后的DataFrame
output_s3_path = "s3://你的桶名/输出路径/"
# df.coalesce(1).write.mode("overwrite").csv(output_s3_path, header=True)
# -----------------------------------------------------------

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

# 解析S3路径的桶名和前缀
def parse_s3_path(s3_path):
    path_parts = s3_path.replace("s3://", "").split("/")
    bucket = path_parts.pop(0)
    prefix = "/".join(path_parts)
    return bucket, prefix

bucket_name, prefix = parse_s3_path(output_s3_path)

# 列出输出路径下的所有文件
response = s3.list_objects_v2(Bucket=bucket_name, Prefix=prefix)

# 过滤Spark生成的元数据文件，只保留数据分片文件
data_files = [obj for obj in response.get('Contents', []) 
              if not obj['Key'].endswith(('_SUCCESS', '.crc')) 
              and '/_committed_' not in obj['Key'] 
              and '/_started_' not in obj['Key']]

# 自定义重命名规则，示例：按业务标识+序号命名，可自行修改
target_file_prefix = "业务标识_20240520"
file_suffix = ".csv" # 对应你输出的文件格式，parquet的话改后缀即可

for index, file_obj in enumerate(data_files):
    old_key = file_obj['Key']
    # 生成新文件名，单文件的话可以去掉序号
    new_key = f"{prefix}{target_file_prefix}{f'_{index+1}' if len(data_files) >1 else ''}{file_suffix}"
    
    # S3没有原生重命名接口，执行复制+删除操作
    s3.copy_object(
        Bucket=bucket_name,
        CopySource={'Bucket': bucket_name, 'Key': old_key},
        Key=new_key
    )
    # 复制成功后删除原分片文件
    s3.delete_object(Bucket=bucket_name, Key=old_key)

# 可选：删除多余的_SUCCESS文件
s3.delete_object(Bucket=bucket_name, Key=f"{prefix}_SUCCESS")

注意事项 #

• 如果你需要输出单个合并文件，可以在写数据前调用df.coalesce(1)，仅适用于数据量小于单个Executor内存上限的场景，大文件不建议合并避免OOM
• 跨桶重命名会产生数据复制流量，大文件场景会延长作业运行时间，建议同桶路径下操作
• 可以自行修改target_file_prefix的命名规则，比如加入作业运行ID、日期参数等

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