EMR上PySpark批量下载文件时如何确保遵守文件大小限制

问题描述 #

在Amazon EMR上运行PySpark应用，目标是从DataFrame的URL下载文件，累计达到指定大小后打包成tar，但当前实现每个文件单独创建tar，未实现批量打包。预期流程是：下载文件→跟踪累计大小→达标后打包→重置跟踪继续处理。

代码中的关键问题排查 #

1. DataFrame字段引用错误 #

你的DataFrame schema定义为["file"]，但循环中使用row.filename获取文件名，这会触发AttributeError。正确的字段引用应为row.file，否则循环会因无法获取文件名中断，可能触发异常后的错误逻辑（比如每次仅处理一个条目就执行打包）。

2. Repartition用法错误 #

df.repartition(spark_partition_id())是错误用法：

• spark_partition_id()是Spark SQL函数，不能直接作为repartition的参数。若需基于原分区处理，无需额外执行repartition，foreachPartition本身会对每个原分区执行函数。
• 错误的重分区操作可能导致数据被打散，每个分区仅包含单个文件，自然会出现每个文件单独打包的情况。

3. 文件大小计算或返回异常 #

如果download_file函数返回的file_size不符合预期，会直接触发单次打包：

• 若file_size始终大于等于FILE_SIZE_LIMIT，每次下载后都会立即触发tar创建。
• 若download_file未正确计算实际下载文件的大小（比如误用URL响应头的Content-Length而非本地文件真实大小），会导致累计大小判断失效。建议在download_file中用os.path.getsize(file_path)获取本地文件的真实大小。

4. 单文件大小超限的边界处理 #

当前逻辑中，若单个文件大小超过FILE_SIZE_LIMIT，会直接将其单独打包，这属于合理逻辑，但如果这不是你期望的行为，需添加单独处理分支（比如跳过或标记记录）。

修正后的核心代码示例 #

from pyspark.sql import SparkSession
import os
import uuid

# Initialize Spark session
spark = SparkSession.builder.appName("DownloadandTar").getOrCreate()

# Sample DataFrame of files
files = [("file1",), ("file2",), ("file3",)]
schema = ["file"]
df = spark.createDataFrame(files, schema=schema)

# Define the file size limit (e.g., 100MB)
FILE_SIZE_LIMIT = 100 * 1024 * 1024  # 100MB

def download_file(url):
    # 替换为真实下载逻辑
    file_path = f"/tmp/{uuid.uuid4()}"
    with open(file_path, "w") as f:
        f.write("test content")
    # 获取实际文件大小
    file_size = os.path.getsize(file_path)
    return file_path, file_size

def create_tar_file(file_list, tar_name):
    # 替换为真实tar打包逻辑
    print(f"Creating tar {tar_name} with files: {file_list}")
    # 可选：打包后删除源文件
    for f in file_list:
        os.remove(f)

def download_and_tar_files(partition):
    accumulated_size = 0
    files_to_tar = []

    for row in partition:
        file_url = row.file
        file_path, file_size = download_file(file_url)
        files_to_tar.append(file_path)
        accumulated_size += file_size

        if accumulated_size >= FILE_SIZE_LIMIT:
            tar_file_name = f"/tmp/{uuid.uuid4()}.tar"
            create_tar_file(files_to_tar, tar_file_name)
            files_to_tar = []
            accumulated_size = 0

    # 处理剩余文件
    if files_to_tar:
        tar_file_name = f"/tmp/{uuid.uuid4()}.tar"
        create_tar_file(files_to_tar, tar_file_name)

# 直接使用foreachPartition处理原分区
df.foreachPartition(download_and_tar_files)

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