PySpark中使用QuantileDiscretizer按group分组分桶且不使用for循环的实现方法

核心思路 #

• 利用Spark的applyInPandas分组处理能力，直接在每个SEQ_ID分组内根据该组的指定分桶数执行分桶逻辑，全程不需要Python侧的for循环，所有计算分布式执行

具体实现步骤 #

步骤1：关联原始数据和分桶配置表 #

首先把你的原始数据集和存储各SEQ_ID分桶数的配置表做关联，让每一行数据都带上对应分组的分桶数：

import pandas as pd
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, IntegerType, DoubleType

# df_raw 为存储SEQ_ID、RESULT的原始表
# df_bucket_cfg 为存储SEQ_ID、num_buckets的分桶配置表
df_joined = df_raw.join(df_bucket_cfg, on="SEQ_ID", how="inner")

步骤2：定义输出Schema与分组处理函数 #

定义输出的表结构，以及每个分组内的分桶处理逻辑：

# 自定义输出Schema，可根据你实际的字段调整，新增buckets列存储分桶结果
output_schema = StructType([
    StructField("SEQ_ID", IntegerType(), nullable=False),
    StructField("RESULT", DoubleType(), nullable=False),
    StructField("num_buckets", IntegerType(), nullable=False),
    StructField("buckets", IntegerType(), nullable=True)
])

def bucket_process(pdf: pd.DataFrame) -> pd.DataFrame:
    # 取当前分组的分桶数，同SEQ_ID的num_buckets一致，取第一个值即可
    bucket_cnt = pdf["num_buckets"].iloc[0]
    # 用qcut实现分位数分桶，和QuantileDiscretizer逻辑对齐
    # duplicates="drop"处理边界值重复导致的分桶数不足问题
    pdf["buckets"] = pd.qcut(
        pdf["RESULT"], 
        q=bucket_cnt, 
        labels=False, 
        duplicates="drop"
    )
    return pdf

步骤3：执行分组分桶 #

df_binned = df_joined.groupBy("SEQ_ID").applyInPandas(bucket_process, schema=output_schema)

注意事项 #

• 该方案依赖Spark 3.0及以上版本的applyInPandas能力，性能远高于Python侧循环处理后union的方案，适合大规模数据集
• 如果你需要完全对齐Spark原生QuantileDiscretizer的近似分桶逻辑，可以把分桶逻辑替换为在分组内调用approxQuantile计算分桶边界再映射，精度和性能可通过分位数近似参数调整

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