基于PySpark高效查找列表列表中最频繁组合的方案咨询

问题背景 #

在Databricks PySpark Notebook中需处理以下场景：

• 待匹配组合：380万种三元组组合
• 目标子列表：17万个
• 当前Python算法效率极低：处理10个子列表需1分钟，全量预计耗时280小时，需将处理时间压缩至30分钟内（理想10分钟）

当前低效代码：

from collections import defaultdict

combinations_of_n = [ ('a','b','c'), ('e','f','g'), ('h','i','j') ] # 3.8 million combinations
df_list = [ ['b','v','e','a','b','c'], ['g','g','a','b','c','f','b'], ['i','k','l','a','i','k'] ] # 170000 lists

# Count occurrences of each combination in lst_of_lsts
combination_count = defaultdict(int)

for sublist in df_list:
    for comb in combinations_of_n:
        if all(elem in sublist for elem in comb):
            combination_count[comb] += 1

# Find the top 5 most frequent combinations
top_combinations = sorted(combination_count.items(), key=lambda x: x[1], reverse=True)[:5]

# Print the results
print("Top 5 most frequent combinations:")
for comb, count in top_combinations:
    print(f"{comb}: {count} occurrences")

优化思路与实现 #

一、Python本地代码优化（适合小数据量过渡） #

• 预转换子列表为集合：将每个子列表转成set，把elem in sublist的O(n)查询降为O(1)
• 批量过滤组合：对每个子列表的集合，直接筛选出所有子集组合，而非逐个遍历380万组合

from collections import defaultdict

combinations_of_n = [('a','b','c'), ('e','f','g'), ('h','i','j')]
# 先把组合转成frozenset，方便后续子集判断
comb_set_map = {frozenset(comb): comb for comb in combinations_of_n}
df_list = [['b','v','e','a','b','c'], ['g','g','a','b','c','f','b'], ['i','k','l','a','i','k']]

combination_count = defaultdict(int)

for sublist in df_list:
    sub_set = set(sublist)
    # 筛选出所有是sub_set子集的组合
    for comb_set, original_comb in comb_set_map.items():
        if comb_set.issubset(sub_set):
            combination_count[original_comb] += 1

top_combinations = sorted(combination_count.items(), key=lambda x: x[1], reverse=True)[:5]

print("Top 5 most frequent combinations:")
for comb, count in top_combinations:
    print(f"{comb}: {count} occurrences")

二、PySpark分布式优化方案（推荐全量数据使用） #

利用Spark的分布式计算能力，将数据拆解为并行任务，核心思路是将子列表与组合进行关联匹配，而非嵌套遍历：

from pyspark.sql import SparkSession
from pyspark.sql.functions import collect_set, countDistinct, desc, size

# 初始化SparkSession
spark = SparkSession.builder.appName("CombinationFrequency").getOrCreate()

# 1. 处理子列表数据：转为DataFrame并生成唯一ID
df_list = [['b','v','e','a','b','c'], ['g','g','a','b','c','f','b'], ['i','k','l','a','i','k']]
sublist_df = spark.createDataFrame([(i, arr) for i, arr in enumerate(df_list)], ["sublist_id", "elements"])
# 展开子列表为单元素行，保留sublist_id
sublist_elements_df = sublist_df.selectExpr("sublist_id", "explode(elements) as element")
# 对每个子列表去重，得到唯一元素集合
sublist_unique_elements = sublist_elements_df.groupBy("sublist_id").agg(collect_set("element").alias("unique_elements"))

# 2. 处理待匹配组合数据：转为DataFrame
combinations_of_n = [('a','b','c'), ('e','f','g'), ('h','i','j')]
comb_df = spark.createDataFrame([(i, list(comb)) for i, comb in enumerate(combinations_of_n)], ["comb_id", "comb_elements"])
# 展开组合为单元素行，保留comb_id
comb_elements_df = comb_df.selectExpr("comb_id", "explode(comb_elements) as element")
# 计算每个组合的元素数量
comb_length_df = comb_df.select("comb_id", size("comb_elements").alias("comb_length"))

# 3. 关联匹配并统计次数
# 关联元素，得到组合与子列表的匹配对
matched_pairs = comb_elements_df.join(sublist_elements_df, on="element")
# 统计每个组合在子列表中匹配到的元素数量
matched_element_count = matched_pairs.groupBy("comb_id", "sublist_id").count()
# 筛选出所有元素都匹配的情况（count等于组合长度）
full_matches = matched_element_count.join(comb_length_df, on="comb_id").filter("count = comb_length")
# 统计每个组合的总匹配次数
comb_frequency = full_matches.groupBy("comb_id").agg(countDistinct("sublist_id").alias("occurrences"))

# 4. 关联原始组合，获取Top5结果
result_df = comb_frequency.join(comb_df, on="comb_id").orderBy(desc("occurrences")).limit(5)

# 打印结果
result_df.show(truncate=False)

# 关闭SparkSession
spark.stop()

三、额外优化点 #

• 调整数据分区：根据集群资源，设置spark.sql.shuffle.partitions为集群核心数的2-3倍，减少shuffle开销
• 广播小数据集：将组合相关DataFrame（comb_df、comb_elements_df）广播，避免重复传输
• 缓存中间结果：对sublist_unique_elements等重复使用的中间DataFrame调用cache()，减少重复计算

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