基于数组交集匹配DataFrame列：Pandas内存溢出求PySpark方案

现有两个DataFrame：df1包含标识列A和字符串列表列B，df2包含标识列C和字符串列表列D。需求是根据D与B的最大元素交集，将df1的A列映射到df2；如果没有交集则设为Null。

当前采用Pandas结合numpy的方案在小数据量（约5K行）下正常运行，但处理87K行df2与567行df1时，在Databricks中频繁出现内存溢出错误（Spark驱动意外终止重启），需要高效的PySpark实现方案。

示例数据 #

df1数据 #

{'A': ['AA', 'BB', 'CC'], 'B': [['a', 'b', 'c', 'd'], ['a', 'f', 'g', 'c'], ['a', 'b', 'l', 'm']]}

df2数据 #

{'C': ['XX', 'YY', 'UU'], 'D': [['a', 'b', 'c', 'n'], ['a', 'm', 'r', 's'], ['e', 'h', 'l', 'j']]}

现有Pandas实现代码 #

b = df1['B'].apply(set).to_numpy()
d = df2['D'].apply(set).to_numpy()
# 计算两两交集
common = d[:,None] & b
# 获取每行的最大交集
vlen = np.vectorize(len)
idx = np.argmax(vlen(common), axis=1)
# 分配交集结果和原始ID
df2['common_items'] = common[np.arange(len(d)), idx]
df2['A'] = np.where(df2['common_items'].str.len()>0,
                    df1['A'].to_numpy()[idx], None)

PySpark高效实现方案 #

核心思路 #

避免全量笛卡尔积带来的内存压力，通过列表展开+分组统计交集大小，再筛选每个df2行对应的最大交集df1记录：

1. 将df1和df2的列表列展开为单元素行，降低关联复杂度；
2. 关联元素匹配的行，统计每个(df2行, df1行)组合的交集元素数量；
3. 用窗口函数筛选每个df2行中交集数量最大的df1记录，无匹配则设为Null。

代码实现 #

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 1. 展开df1的列表列，添加唯一行标识
df1_expanded = df1.withColumn("b_item", F.explode(F.col("B"))) \
                  .withColumn("df1_id", F.monotonically_increasing_id())

# 2. 展开df2的列表列，添加唯一行标识
df2_expanded = df2.withColumn("d_item", F.explode(F.col("D"))) \
                  .withColumn("df2_id", F.monotonically_increasing_id())

# 3. 关联匹配元素，统计每个df2行与df1行的交集元素数
# 广播df1（仅567行）大幅提升关联效率
match_counts = df2_expanded.join(F.broadcast(df1_expanded), F.col("d_item") == F.col("b_item"), "left") \
                           .groupBy("df2_id", "C", "D", "A", "B") \
                           .agg(F.count("d_item").alias("intersection_size"))

# 4. 窗口函数筛选每个df2行的最大交集记录
window_spec = Window.partitionBy("df2_id").orderBy(F.desc("intersection_size"), F.desc("A"))
ranked_matches = match_counts.withColumn("rank", F.row_number().over(window_spec)) \
                             .filter(F.col("rank") == 1) \
                             .select("C", "D",
                                     F.when(F.col("intersection_size") > 0, F.col("A")).alias("A"),
                                     F.array_intersect(F.col("D"), F.col("B")).alias("common_items"))

# 5. 补充无交集的df2行，保证结果完整性
final_result = df2.join(ranked_matches, on=["C", "D"], how="left") \
                  .select("C", "D",
                          F.coalesce(F.col("A"), F.lit(None)).alias("A"),
                          F.coalesce(F.col("common_items"), F.array()).alias("common_items"))

final_result.show(truncate=False)

关键优化点 #

• 广播小表：df1仅567行，用broadcast将其分发到所有executor，避免shuffle开销；
• 避免笛卡尔积：通过元素展开关联，替代全量集合运算，降低内存占用；
• 窗口函数去重：确保每个df2行仅保留最优匹配，若多个df1行交集大小相同，按A降序取第一个。

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