如何在PySpark中按各变量分位数分组并计算均值/中位数

实现步骤 #

1. 预计算分组变量的分位数阈值
   确定需要处理的分组变量（如age、income）和目标分位数（如10%、25%、50%），通过PySpark的approxQuantile方法获取每个变量对应的分位数阈值，存储为字典便于后续调用。

2. 为数据添加分位数分组标签
   使用条件判断函数when，根据每个变量的分位数阈值，为每条数据打上对应的分组标签（如age_10th_percentile），并将所有分组标签整合到一个统一的列中，同时保留原始的credit score数据。

3. 按分组标签聚合计算指标
   以分组标签为键进行分组，分别计算每组credit score的均值和中位数（中位数使用percentile_approx实现近似计算，适合大数据场景），最后整理结果格式匹配预期输出。

完整代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, avg, percentile_approx, lit, concat

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

# 示例数据（替换为你的实际数据）
data = [
    (25, 30000, 700),
    (30, 40000, 710),
    (45, 60000, 750),
    (22, 25000, 710),
    (35, 50000, 730),
    (50, 70000, 740),
    (18, 20000, 550),
    (28, 35000, 560),
    (40, 55000, 580),
    (20, 22000, 590),
    (32, 45000, 610),
    (48, 65000, 640)
]
df = spark.createDataFrame(data, ["age", "income", "credit_score"])

# 定义参数
group_vars = ["age", "income"]
quantiles = [0.1, 0.25, 0.5]
quantile_labels = ["10th_percentile", "25th_percentile", "50th_percentile"]

# 步骤1：计算每个分组变量的分位数阈值
quantile_thresholds = {}
for var in group_vars:
    thresholds = df.approxQuantile(var, quantiles, 0.01)  # 0.01是相对误差
    quantile_thresholds[var] = dict(zip(quantile_labels, thresholds))

# 步骤2：生成所有分位数分组的标签列
group_cols = []
for var in group_vars:
    for label, threshold in quantile_thresholds[var].items():
        # 标记属于当前分位数分组的数据
        group_col = when(col(var) <= threshold, concat(lit(f"{var}_"), lit(label))).alias(f"{var}_{label}")
        group_cols.append(group_col)

# 将所有分组标签合并为一个列，过滤空值
df_with_groups = df.select("credit_score", *group_cols)
df_melted = df_with_groups.select(
    "credit_score",
    *[col(c).alias("group") for c in df_with_groups.columns if c != "credit_score"]
).na.drop(subset=["group"])

# 步骤3：按分组计算均值和中位数
result_df = df_melted.groupBy("group").agg(
    avg("credit_score").alias("mean(credit score)"),
    percentile_approx("credit_score", 0.5, 10000).alias("median(credit score)")
).orderBy("group")

# 展示结果
result_df.show(truncate=False)

说明 #

• approxQuantile的第三个参数是相对误差，设置为0.01在保证精度的同时提升计算效率，适合大数据集。
• 若需要更精确的中位数，可调整percentile_approx的第三个参数（样本量），但会增加计算开销。
• 示例数据仅作演示，实际使用时替换为你的业务数据表即可。

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