PySpark：基于多阈值批量生成聚合列（替代循环实现）

需求说明 #

现有PySpark DataFrame包含列：user、score、country、risk/safe、payment_id，给定阈值列表[10, 20, 30]，需按country分组，为每个阈值生成两列统计结果：

• 分数≥阈值的风险支付数占全部风险支付数的比例
• 至少有一次分数≥阈值的独立风险用户数占全部独立风险用户数的比例

当前通过外部循环实现，希望改为在单个DataFrame操作中完成所有计算。

原代码问题分析 #

你提供的原代码存在几个问题：

1. 列名转义错误：risk\\safe应直接写risk/safe
2. 仅计算了支付占比，未处理用户占比的统计
3. 循环会覆盖df变量，最终仅保留最后一个阈值的结果
4. 聚合时用count('country')统计支付数不准确，建议用payment_id（分组键country不会为null，无法准确统计支付数量）

解决方案代码 #

以下是无需外部循环、单DataFrame内完成所有统计的实现：

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

thresholds = [10, 20, 30]

# 1. 先过滤出风险支付数据，避免后续重复计算
risk_data = df.filter(F.col("risk/safe") == "risk")

# 2. 用窗口函数标记每个用户是否有至少一次达标支付（针对每个阈值）
window_spec = Window.partitionBy("country", "user")
for thresh in thresholds:
    # 标记单条支付是否达标
    risk_data = risk_data.withColumn(f"pay_pass_{thresh}", F.when(F.col("score") >= thresh, 1).otherwise(0))
    # 标记该用户在当前国家下是否至少有一次达标支付
    risk_data = risk_data.withColumn(f"user_pass_{thresh}", F.max(f"pay_pass_{thresh}").over(window_spec))

# 3. 按country分组，一次性计算所有阈值的两个比例
final_result = risk_data.groupBy("country").agg(
    # 生成所有阈值的支付占比列
    *[(F.sum(f"pay_pass_{thresh}") / F.count("payment_id")).alias(f"pay_ratio_over_{thresh}") for thresh in thresholds],
    # 生成所有阈值的用户占比列
    *[(F.countDistinct(F.when(F.col(f"user_pass_{thresh}") == 1, F.col("user"))) / F.countDistinct("user")).alias(f"user_ratio_over_{thresh}") for thresh in thresholds]
)

# 查看结果
final_result.show()

代码说明 #

• 窗口函数的作用：通过Window.partitionBy("country", "user")，对每个国家下的每个用户，标记其是否有至少一次支付满足阈值条件（用max取用户内的达标标记，只要有1次达标就记为1）
• 列表推导式批量生成聚合逻辑：避免多次循环操作DataFrame，一次性生成所有阈值的统计列，提升效率
• 准确统计维度：用count("payment_id")统计支付总数，countDistinct("user")统计独立用户数，确保统计维度准确

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