PySpark中替代嵌套for循环的Join方案：预算与产品价格统计

数据结构 #

• 客户预算表：包含Audit ID、Customer ID、Budget、TimeStamp字段
• 产品价格表：包含Audit ID、Product ID、Price、TimeStamp字段
• 客户产品关联表：包含Customer ID、Product ID字段

业务需求 #

1. 统计客户预算大于对应产品价格的次数，需选取该客户预算时间戳之前的最新产品价格
2. 计算每个客户的预算与对应产品价格的最大差值

原实现问题 #

之前用Pandas嵌套for循环实现了逻辑，小数据集可以正常运行，但无法处理大数据量，需要改用PySpark实现（PySpark不支持嵌套循环，需用Join操作替代）

原Pandas代码（中文注释版） #

count_intances_budget_more_than_price = 0  # 统计预算大于价格的总次数
map_customer_id_max_delta = {}  # 存储每个客户的预算与价格最大差值
processed_product_for_customer = set()  # 标记当前客户已处理过的产品

for customer_row in customer_dataset:
    max_delta = 0
    # 如果客户已有记录，读取当前的最大差值
    if customer_row['Customer ID'] in map_customer_id_max_delta:
        max_delta = map_customer_id_max_delta.get(customer_row['Customer ID'])
    
    for product_row in product_dataset:
        # 仅处理与当前客户关联的产品
        if product_row['Product ID'] in map_customer_id_product_id[customer_row['Customer ID']]:
            # 每个产品仅处理一次
            if product_row['Product ID'] not in processed_product_for_customer:
                processed_product_for_customer.add(product_row['Product ID'])
                # 筛选预算时间戳之前的产品价格，且预算大于价格的情况
                if product_row['TimeStamp'] < customer_row['TimeStamp'] and product_row['Price'] < customer_row['Budget']:
                    count_intances_budget_more_than_price += 1
                    # 更新当前客户的最大差值
                    max_delta = max(max_delta, customer_row['Budget'] - product_row['Price'])
    
    # 保存当前客户的最大差值
    map_customer_id_max_delta[customer_row['Customer ID']] = max_delta
    # 清空已处理产品集合，准备处理下一个客户
    processed_product_for_customer.clear()

PySpark实现方案 #

第一步：关联客户预算与对应产品 #

先把客户预算表和客户产品关联表做内连接，得到每个客户预算记录对应的关联产品列表：

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

# 假设三张表已通过SparkSession加载，表名分别为customer_budget、product_price、customer_product_rel
customer_budget = spark.table("customer_budget")
product_price = spark.table("product_price")
customer_product_rel = spark.table("customer_product_rel")

# 关联客户预算和客户-产品关系
customer_product_budget = customer_budget.join(
    customer_product_rel,
    on="Customer ID",
    how="inner"
)

第二步：提取每个预算时间点前的最新产品价格 #

用窗口函数给每个客户-产品组合的价格记录按时间戳降序排名，只保留预算时间之前的第一条（最新）价格：

# 定义窗口：按客户ID、产品ID分组，产品时间戳倒序排列
window_spec = Window.partitionBy("Customer ID", "Product ID").orderBy(F.desc("product_price.TimeStamp"))

# 关联产品价格表，筛选出预算时间戳之前的记录，再取最新的那条
filtered_latest_prices = customer_product_budget.join(
    product_price,
    on="Product ID",
    how="inner"
).filter(
    product_price.TimeStamp < customer_product_budget.TimeStamp
).withColumn(
    "price_rank", F.row_number().over(window_spec)
).filter(
    F.col("price_rank") == 1
).drop("price_rank")

第三步：统计次数&计算客户最大差值 #

# 计算预算与价格的差值，只保留预算大于价格的有效记录
valid_records = filtered_latest_prices.withColumn(
    "delta", F.col("Budget") - F.col("Price")
).filter(
    F.col("delta") > 0
)

# 统计总次数
total_more_count = valid_records.count()

# 按客户分组，计算每个客户的最大差值
customer_max_delta = valid_records.groupBy("Customer ID").agg(
    F.max("delta").alias("max_budget_price_delta")
)

# 输出结果
print(f"客户预算大于对应产品价格的总次数：{total_more_count}")
customer_max_delta.show()

方案优势 #

• 完全基于Spark分布式计算框架，避免了嵌套循环的单点瓶颈，轻松处理大数据量
• 窗口函数精准筛选出每个预算时间点前的最新产品价格，和原逻辑完全对齐
• 全程用Join、分组聚合替代嵌套遍历，执行效率远高于Pandas的循环实现

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