PySpark极小DataFrame执行explode后无法存Parquet，repartition无效

原问题 #

我的初始DataFrame df体量极小，共包含10 000 000行数据，仅占用单个60Mb的分区即可存储。repeat列存储的是数组，数组大小范围为1到10^9。我需要对repeat列执行explode操作，示例代码为df = df.select("*", F.explode("repeat"))，但执行后会生成数十亿条记录，无法将结果保存为Parquet文件。我已尝试在explode操作执行前后分别调用repartition，但问题仍未得到解决，请问该如何处理？

问题原因 #

• 原始DataFrame仅存在1个分区。explode属于窄依赖操作，Spark会先在当前分区完成全量计算，再触发后续的shuffle重分区流程。如果先直接做explode再repartition，单分区计算时就会生成数十亿条记录，直接撑爆Executor内存，根本无法走到shuffle环节。
• 即便在explode前提前调用repartition打散原始行，也极易出现数据倾斜：少数携带超大数组的行被分配到同一分区，单个分区的explode输出量远超其他分区的承载能力，触发OOM或任务长时间卡死。

解决方案 #

分步拆分数组，避免单分区生成超量数据 #

不要对完整数组直接执行explode，先把原数组按固定块大小拆分成多个子数组，把子数组打散到不同分区后再做二次explode，从根源上避免单分区输出量过大：

from pyspark.sql import functions as F

# 单块元素大小根据Executor内存调整，建议控制在10w-100w区间
BLOCK_SIZE = 100000

# 将原repeat数组按固定块切片，生成子数组列表
df_blocked = df.withColumn(
    "repeat_blocks",
    F.expr(f"""
        transform(
            sequence(0, int(size(repeat)/{BLOCK_SIZE})),
            block_idx -> slice(repeat, block_idx * {BLOCK_SIZE} + 1, {BLOCK_SIZE})
        )
    """)
)

# 先explode拆分出的子块，按总输出量设置足够的分区数，再explode块内的实际元素
# 分区数估算逻辑：预计explode后总行数 / 单分区合理输出量（建议100w-1000w/分区）
df_exploded = df_blocked\
    .select("*", F.explode("repeat_blocks").alias("current_block"))\
    .repartition(200) \
    .select("*", F.explode("current_block").alias("repeat_value"))\
    .drop("repeat_blocks", "current_block")

哪怕遇到长度达到1e9的极端大数组，也会被拆分为10000个大小10w的子块，这些子块会被随机打散到不同分区，每个分区仅需处理10w级别的元素，完全不会触发OOM。

Parquet写出优化 #

写出时对齐默认存储块大小，避免生成过大文件或大量小文件：

df_exploded.write\
    .mode("overwrite")\
    .option("parquet.block.size", 128 * 1024 * 1024)\
    .parquet("你的存储路径")

倾斜兜底优化 #

如果存在极个别长度超过1e6的超大数组行，可以给这类行的子块增加随机打散key，进一步避免超大块集中到同一分区：

df_prepped = df_blocked.withColumn(
    "shuffle_key",
    F.when(F.size("repeat") > 1000000, F.rand(5).cast("int"))
     .otherwise(F.lit(1))
)
df_exploded = df_prepped.repartition(400, "shuffle_key")\
    .select("*", F.explode("repeat_blocks").alias("current_block"))\
    .select("*", F.explode("current_block").alias("repeat_value"))\
    .drop("repeat_blocks", "current_block", "shuffle_key")

避坑注意点 #

• 不要在explode流程完成前执行count()、collect()这类全量触发计算的操作，否则会直接触发单分区的超大计算量，导致Driver或Executor OOM。
• 分区数不要设置过大，否则会生成大量小Parquet文件，降低后续查询效率，按单分区100w-1000w输出量调整即可。

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