PySpark DataFrame并行写入双数据存储的正确实现方案咨询

错误原因分析 #

你遇到的TypeError: can't pickle _thread.lock objects是因为ProcessPoolExecutor使用多进程时需要序列化所有任务参数，而你的存储客户端（client1/client2）内部包含不可序列化的线程锁对象，无法通过pickle传递给子进程。

改用Thread库后程序无限运行，大概率是两个问题：

1. Spark的rows是迭代器，只能遍历一次，直接传给两个线程会导致第二个线程拿到空数据，客户端上传操作卡住；
2. 原代码中误写了两次storage1.result()，没有等待storage2的任务完成，导致程序无法正常退出。

正确实现方案 #

方案1：用ThreadPoolExecutor（推荐，线程安全场景） #

如果你的存储客户端支持线程安全，直接用线程池代替进程池，同时注意先将rows转为列表（避免迭代器耗尽），并正确等待两个任务完成：

修改Manager类代码：

def partition_funtion(self, arg1, arg2, ...):
    def process_rows(rows):
        rows_list = list(rows)  # 转换为列表，确保两个任务都能拿到完整数据
        # 使用线程池，避免多进程的序列化问题
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            future1 = executor.submit(upload_data, client1, rows_list)
            future2 = executor.submit(upload_data, client2, rows_list)
            # 等待两个任务完成，捕获可能的异常
            try:
                future1.result()
                future2.result()
            except Exception as e:
                # 处理上传失败的情况，比如日志记录
                print(f"Upload failed: {str(e)}")
                raise
    return process_rows

def upload_data(client, data):
    client.upload(data)

方案2：线程池+任务内初始化客户端（非线程安全/不可序列化场景） #

如果客户端不支持线程安全，或者无法被序列化，不要直接传递客户端对象，而是在每个任务中根据配置重新创建客户端：

def partition_funtion(self, config1, config2, ...):
    def process_rows(rows):
        rows_list = list(rows)
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            # 传递可序列化的配置，而非客户端实例
            future1 = executor.submit(upload_with_config, config1, rows_list)
            future2 = executor.submit(upload_with_config, config2, rows_list)
            future1.result()
            future2.result()
    return process_rows

def upload_with_config(storage_config, data):
    # 根据配置创建客户端
    client = create_client(storage_config)  # 替换为你的客户端初始化逻辑
    client.upload(data)

方案3：利用Spark原生并行能力（最简洁） #

Spark本身已经是分布式并行处理框架，完全可以避免在单个partition内开线程/进程，直接对DataFrame执行两次写入操作，让Spark集群并行处理两个任务：

# 方式1：直接调用Spark的写入API（如果支持你的存储类型）
df.write.format("your_storage_format").options(**storage1_options).save("storage1_target")
df.write.format("your_storage_format").options(**storage2_options).save("storage2_target")

# 方式2：如果是自定义客户端，分两次调用foreachPartition
def write_to_storage1(rows):
    client = create_client(storage1_config)
    client.upload(list(rows))

def write_to_storage2(rows):
    client = create_client(storage2_config)
    client.upload(list(rows))

df.foreachPartition(write_to_storage1)
df.foreachPartition(write_to_storage2)

这个方案不需要额外的线程/进程管理，完全利用Spark的分布式能力，代码更简洁，也避免了序列化问题。

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