如何利用PySpark分布式计算优化小数据量适用的逐行处理逻辑

问题分析 #

原有代码通过collect()将全量用户-产品数据拉取到Driver内存，在本地循环生成每个用户的产品关联矩阵，这种做法在数据量较大时会导致Driver内存溢出，完全浪费了Spark的分布式计算优势。

核心改造思路 #

1. 避免本地拉取：所有计算逻辑基于Spark DataFrame API实现，全程不使用collect()，让计算任务在集群Executor节点分布式执行。
2. 分布式构建关联矩阵：通过Spark的join和groupBy算子替代本地循环，全局统计产品间的共现次数（关联强度）。
3. 分布式生成TopN推荐：利用Spark窗口函数完成分组排序，直接在集群中筛选每个产品的TopN关联产品。

改造后代码示例 #

假设原始数据为用户-产品交互表（字段：user_id, product_id），以下是完整分布式实现：

1. 构建分布式产品共现矩阵

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

# 读取并去重用户-产品交互数据（避免重复计数）
user_product_df = spark.read.table("user_product_interactions").select("user_id", "product_id").distinct()

# 自连接生成同一用户下的产品配对（排除产品自身配对）
product_pair_df = user_product_df.alias("a") \
    .join(user_product_df.alias("b"), on="user_id") \
    .filter(F.col("a.product_id") != F.col("b.product_id")) \
    .select(
        F.col("a.product_id").alias("product1"),
        F.col("b.product_id").alias("product2")
    )

# 统计每对产品的共现次数（关联强度）
cooccurrence_matrix = product_pair_df.groupBy("product1", "product2") \
    .agg(F.count("*").alias("cooccurrence_count"))

2. 分布式生成每个产品的TopN推荐

# 定义窗口规则：按product1分组，按共现次数降序排序
rank_window = Window.partitionBy("product1").orderBy(F.desc("cooccurrence_count"))

# 筛选每个产品的Top10关联产品（可调整N值）
product_topn_recs = cooccurrence_matrix \
    .withColumn("rank", F.row_number().over(rank_window)) \
    .filter(F.col("rank") <= 10) \
    .select("product1", "product2", "cooccurrence_count", "rank")

# 输出结果
product_topn_recs.show()

3. 扩展：个性化用户TopN推荐（可选）

如果需要为每个用户生成个性化推荐，同样用分布式逻辑实现，避免本地循环：

# 先聚合每个用户已交互的产品集合
user_interacted_df = user_product_df.groupBy("user_id") \
    .agg(F.collect_set("product_id").alias("interacted_products"))

# 关联用户已交互产品与共现矩阵，计算推荐候选的得分
user_candidate_df = user_interacted_df.alias("u") \
    .join(cooccurrence_matrix.alias("c"), F.array_contains(F.col("u.interacted_products"), F.col("c.product1"))) \
    .filter(~F.array_contains(F.col("u.interacted_products"), F.col("c.product2"))) \
    .groupBy("user_id", "c.product2") \
    .agg(F.sum("c.cooccurrence_count").alias("rec_score"))

# 每个用户取Top10推荐
user_rank_window = Window.partitionBy("user_id").orderBy(F.desc("rec_score"))
user_topn_recs = user_candidate_df \
    .withColumn("rank", F.row_number().over(user_rank_window)) \
    .filter(F.col("rank") <= 10) \
    .select("user_id", F.col("product2").alias("recommended_product"), "rec_score", "rank")

# 输出结果
user_topn_recs.show()

关键优势 #

• 无内存瓶颈：全程未将数据拉取到Driver，支持TB级别的大规模数据处理。
• 分布式算力利用：所有聚合、排序、筛选操作都在集群节点并行执行，性能远超本地循环。
• 代码简洁可维护：基于Spark内置API实现，逻辑清晰，便于后续扩展和优化。

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