如何用PySpark按Key分组执行函数且避免全组加载到内存（含报错排查）

给定数据生成代码：

data = [{"id": random.randint(1, 10), "content": random.choice(string.ascii_letters)} for _ in range(0, 1000)]

需求为按id对数据条目分组，并对每个组执行类似store_content(group)的函数（例如将所有id为1的条目通过该函数存储）。注意每组数据可能量很大，调用store_content时不希望将其物化/存储到内存中。

如何用PySpark实现该需求？

import string
import tempfile
from random import randint, choice
from functools import partial

import pyspark
from pyspark.sql import SparkSession
from more_itertools import peekable


def generate_samples(N):
    """生成数据样本"""
    return [{"id": randint(1, 10), 
            "content": choice(string.ascii_letters)} for _ in range(N)]


def save_content(partition, uri):
    """将每个分区保存到不同文件夹"""
    peeked_it = peekable(partition)
    key = peeked_it.peek(0)[0][0]
    values = next(peeked_it)
    
    with open(f"{uri}/{key}.txt", "w") as file:
        file.write("".join(values))


spark = SparkSession.builder.appName("test").getOrCreate()
rdd = spark.sparkContext.parallelize(generate_samples(1000)).map(lambda x: (x["id"], x))

with tempfile.TemporaryDirectory() as tempdir:
    rdd = rdd.partitionBy(numPartitions=10).glom()
    rdd.foreachPartition(partial(save_content, uri=tempdir))

Caused by: org.apache.spark.api.python.PythonException: Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/pyspark/python/lib/pyspark.zip/pyspark/worker.py", line 686, in main
    process()
  File "/usr/local/lib/python3.7/dist-packages/pyspark/python/lib/pyspark.zip/pyspark/worker.py", line 676, in process
    out_iter = func(split_index, iterator)
  File "/usr/local/lib/python3.7/dist-packages/pyspark/rdd.py", line 3472, in pipeline_func
    return func(split, prev_func(split, iterator))
  File "/usr/local/lib/python3.7/dist-packages/pyspark/rdd.py", line 3472, in pipeline_func
    return func(split, prev_func(split, iterator))
  File "/usr/local/lib/python3.7/dist-packages/pyspark/rdd.py", line 3472, in pipeline_func
    return func(split, prev_func(split, iterator))
  [Previous line repeated 1 more time]
  File "/usr/local/lib/python3.7/dist-packages/pyspark/rdd.py", line 540, in func
    return f(iterator)
  File "/usr/local/lib/python3.7/dist-packages/pyspark/rdd.py", line 1178, in func
    r = f(it)
  File "<ipython-input-9-7e18b1132f02>", line 25, in save_content
TypeError: sequence item 0: expected str instance, tuple found

期望结果：生成如1.txt的文件，其中包含所有id为1的字符。

问题分析 #

原代码存在以下核心问题：

• 对partitionBy+glom()后的分区结构理解错误：glom()会把分区内的所有(id, 字典)元组打包成列表，而你在save_content中的取值逻辑完全错误，导致拿到的是元组而非字符串。
• 错误假设partitionBy(10)会让每个分区只包含一个id的数据：Spark默认哈希分区只能保证同id数据在同一分区，但一个分区可能包含多个id的数据。
• save_content中用"".join(values)时，values是元组而非字符串序列，触发类型错误。

修正后的实现方案 #

方案一：groupByKey+流式写入（适合中小数据量） #

直接通过groupByKey按id分组，然后对每个组的迭代器进行流式写入，避免将整个组加载到内存：

import string
import tempfile
from random import randint, choice
from functools import partial

import pyspark
from pyspark.sql import SparkSession


def generate_samples(N):
    """生成数据样本"""
    return [{"id": randint(1, 10), 
            "content": choice(string.ascii_letters)} for _ in range(N)]


def store_content(key, content_iter, uri):
    """流式写入每个id对应的文件，不缓存全量数据"""
    with open(f"{uri}/{key}.txt", "w") as f:
        for content in content_iter:
            f.write(content)


spark = SparkSession.builder.appName("test").getOrCreate()
# 直接提取id和content，减少数据传输量
rdd = spark.sparkContext.parallelize(generate_samples(1000)) \
    .map(lambda x: (x["id"], x["content"]))

with tempfile.TemporaryDirectory() as tempdir:
    # groupByKey后每个元素是(id, content迭代器)，直接传入函数流式写入
    rdd.groupByKey().foreach(lambda x: store_content(x[0], x[1], tempdir))

方案二：自定义分区器+分区内流式处理（适合大数据量） #

如果数据量极大，想避免groupByKey的shuffle开销，可以自定义分区器让每个id对应一个分区，然后在分区内直接处理：

import string
import tempfile
from random import randint, choice
from functools import partial

import pyspark
from pyspark.sql import SparkSession
from pyspark import Partitioner


class IdPartitioner(Partitioner):
    """自定义分区器，每个id对应一个独立分区"""
    def __init__(self, num_ids):
        self.num_ids = num_ids
    
    def numPartitions(self):
        return self.num_ids
    
    def getPartition(self, key):
        # 将id转换为分区索引（分区从0开始）
        return key - 1


def generate_samples(N):
    """生成数据样本"""
    return [{"id": randint(1, 10), 
            "content": choice(string.ascii_letters)} for _ in range(N)]


def save_partition(partition, uri):
    """处理单个分区，流式写入对应id的文件"""
    current_key = None
    file_handle = None
    try:
        for key, content in partition:
            if key != current_key:
                # 切换id时关闭之前的文件
                if file_handle:
                    file_handle.close()
                current_key = key
                file_handle = open(f"{uri}/{current_key}.txt", "w")
            file_handle.write(content)
    finally:
        # 确保最后一个文件被正确关闭
        if file_handle:
            file_handle.close()


spark = SparkSession.builder.appName("test").getOrCreate()
rdd = spark.sparkContext.parallelize(generate_samples(1000)) \
    .map(lambda x: (x["id"], x["content"]))

with tempfile.TemporaryDirectory() as tempdir:
    # 使用自定义分区器，每个id对应一个分区
    rdd = rdd.partitionBy(IdPartitioner(num_ids=10))
    rdd.foreachPartition(partial(save_partition, uri=tempdir))

关键说明 #

• 两种方案均通过迭代器流式写入文件，不会将整个组的数据加载到内存，符合大数据量场景需求。
• 方案二通过自定义分区器减少了shuffle操作，性能更优，适合超大规模数据处理。
• 修正了原代码中对数据结构的错误处理，直接提取content字符串，避免了类型错误。

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