使用PySpark DataFrame的foreach写入S3时遇PicklingError问题求助

我有一个PySpark DataFrame（base_df），希望将其每一行数据以JSON格式存储到AWS S3桶中。为此实现了如下write_final_json函数：

def write_final_json(self, row)->None:
        print("write_final_json")
        print(row)
        output_dir = self.args['DESTINATION_S3_PATH']
        """
        Args:
        output_dir: S3 path or linux dir based on env
        row: record from datafrome

        Purpose: Write final json output to S3
        """
        if not os.path.exists(f"{output_dir}/{row.batch_id}"):
            os.makedirs(f"{output_dir}/{row.batch_id}")
        print(f"Writing {output_dir}/{row.batch_id}/{row.member_id}.json")
        with open(f"{output_dir}/{row.batch_id}/{row.member_id}.json", "w") as f:
            f.write(row)

并通过以下方式为每行调用该方法：

base_df.foreach(lambda row: self.write_final_json(row))

但出现如下错误：

PicklingError: Could not serialize object: RuntimeError: It appears that you are attempting to reference SparkContext from a broadcast variable, action, or transformation. SparkContext can only be used on the driver, not in code that it run on workers. For more information, see SPARK-5063.

• 序列化冲突：foreach属于分布式执行的算子，会把传入的代码序列化后发送到Worker节点执行。你调用的是类方法self.write_final_json，这会触发整个类实例（self）的序列化，但类实例中极可能隐含引用了SparkContext/SparkSession——而SparkContext是无法被序列化的，它只能在Driver端使用，因此抛出序列化错误。
• S3操作方式错误：os.path.exists、os.makedirs、open都是本地文件系统操作API，无法直接识别S3的s3://路径，这类API只能处理本地磁盘路径。

方案一：使用Spark原生API（推荐） #

Spark提供了直接写入JSON到S3的原生能力，无需手动逐行处理，性能和兼容性更优：

# 按batch_id分区写入，每个分区下生成JSON文件
base_df.write \
    .mode("overwrite") \
    .partitionBy("batch_id") \
    .json(self.args['DESTINATION_S3_PATH'])

如果需要每个member_id对应单独文件，可以通过 repartition 控制：

# 先按batch_id和member_id重分区，确保每个member_id对应一个输出文件
base_df.repartition("batch_id", "member_id") \
    .write \
    .mode("overwrite") \
    .partitionBy("batch_id", "member_id") \
    .json(self.args['DESTINATION_S3_PATH'])

注意：过度生成小文件会影响S3性能，需根据业务场景权衡。

方案二：修复自定义函数（仅适用于特殊业务场景） #

如果必须保留自定义逐行逻辑，需要做两点修改：

1. 避免序列化类实例，改用独立函数
2. 使用AWS SDK操作S3，替换本地文件API

import boto3
import json

def write_final_json(row, output_dir):
    s3 = boto3.client('s3')
    # 拆分S3路径为桶名和文件前缀
    s3_path_parts = output_dir.replace("s3://", "").split("/", 1)
    bucket = s3_path_parts[0]
    prefix = s3_path_parts[1] if len(s3_path_parts) > 1 else ""
    # 构造S3对象键
    object_key = f"{prefix}/{row.batch_id}/{row.member_id}.json"
    # 将Spark Row对象转为JSON字符串
    json_content = json.dumps(row.asDict())
    # 写入S3
    s3.put_object(Bucket=bucket, Key=object_key, Body=json_content)

# 调用时直接传递参数，避免引用self
target_s3_path = self.args['DESTINATION_S3_PATH']
base_df.foreach(lambda row: write_final_json(row, target_s3_path))

注意：需确保Worker节点安装了boto3依赖，且配置了有效的AWS访问凭证。

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