如何为PySpark DataFrame中每个值分配间隔5%的分位数区间

实现步骤 #

• 计算分位数阈值：用PySpark的approxQuantile方法计算目标列的5%、10%...95%分位数，适配分布式数据集场景。
• 构建分箱规则：基于分位数阈值设置边界，覆盖所有数据的区间范围，同时定义每个区间对应的上限标签（如55-60%区间对应标签60）。
• 添加分箱列：通过Bucketizer或条件语句，将每个数据值映射到对应的分位数区间上限。

完整代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when
from pyspark.ml.feature import Bucketizer

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

# 创建示例DataFrame
data = [5,6,9,87,2,3,5,7,2,6,5,2,3,4,69,4]
df = spark.createDataFrame([(val,) for val in data], ["data"])

# 计算5%间隔的分位数（5到95，步长5）
percentiles = list(range(5, 100, 5))
# 转换为0-1的比例值，适配approxQuantile参数要求
quantiles = df.approxQuantile("data", [p/100 for p in percentiles], 0.0)

# 构建分箱边界：包含最小/最大边界确保所有数据被覆盖
bins = [-float("inf")] + quantiles + [float("inf")]
# 每个区间对应的上限标签（5,10,...,95,100）
labels = percentiles + [100]

# 使用Bucketizer完成分箱
bucketizer = Bucketizer(splits=bins, inputCol="data", outputCol="bucket_idx")
df_bucketed = bucketizer.transform(df)

# 将桶索引映射为区间上限值
for idx, label in enumerate(labels):
    df_bucketed = df_bucketed.withColumn(
        "percentile_upper",
        when(col("bucket_idx") == idx, label).otherwise(col("percentile_upper"))
    )

# 输出预期格式的结果
df_bucketed.select("data", "percentile_upper").show()

关键说明 #

• approxQuantile的第三个参数为相对误差，设为0.0时计算精确分位数（小数据集适用，大数据集可适当调大以提升性能）。
• 分箱边界加入-inf和inf，确保极端值也能被正确分箱：小于5%分位数的归为5%区间，大于95%分位数的归为100%区间。
• 最终percentile_upper列即为预期的区间上限，比如值5对应60时，代表该值处于55%-60%的分位数区间。

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