Spark优化两个等大DataFrame的Join操作：解决OOM问题

核心问题解析 #

你当前执行df_a.join(df_b, on=["category", "timestamp"])时，Spark默认会基于**category+timestamp的组合哈希值**做分区shuffle。由于timestamp是唯一值，同一category下的不同记录会生成不同的组合哈希，导致这些记录被分散到不同Worker节点，跨节点的数据传输和内存存储压力最终引发OOM。你的思路（让同一category的数据集中在同一个Worker，仅在节点内匹配timestamp）是正确的，但Spark不会自动实现这个逻辑，需要手动干预。

具体优化步骤 #

1. 预分区对齐两个DataFrame #

对df_a和df_b都按category重分区，确保同一category的所有数据落在同一个分区：

# 按category重分区，可根据集群规模指定分区数，比如200
df_a_repart = df_a.repartition(200, "category")
df_b_repart = df_b.repartition(200, "category")

这样做的核心优势是：Join操作时，同一category的数据已经在同一个Worker节点，无需跨节点shuffle，仅在节点内部完成timestamp的匹配，大幅降低内存和网络开销。

2. 执行预分区后的Join #

在重分区后的DataFrame上执行Join：

df_c = df_a_repart.join(df_b_repart, on=["category", "timestamp"])

3. 应对数据倾斜（可选） #

如果某些category的数据量极大（比如个别category占总数据的30%以上），单分区数据量过大仍会引发OOM，可以进一步拆分大分区：

from pyspark.sql import functions as F

# 结合category和timestamp的哈希分桶，拆分大分区
df_a_repart = df_a.repartition(200, "category", (F.hash("timestamp") % 10))
df_b_repart = df_b.repartition(200, "category", (F.hash("timestamp") % 10))

这样同一category下的数据会被拆分成10个小分区，分散到多个Worker处理，避免单节点内存过载。

4. 辅助内存参数调整 #

• 调大Executor内存：spark.executor.memory=16g（根据集群实际资源调整）
• 开启自适应执行：spark.sql.adaptive.enabled=true，让Spark自动调整分区大小和并行度
• 调整shuffle分区数：spark.sql.shuffle.partitions=200（默认200，可根据总数据量灵活调整）

补充说明 #

预分区优化的本质是将Join的shuffle粒度从「组合键哈希」降低到「category哈希」，减少跨节点数据传输量，同时让每个Worker仅处理自身分区内的category数据，从根源上缓解OOM问题。这种方案完全适配你的场景——两个DataFrame规模相当、category存在大量公共值。

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