如何在PySpark中为列的数值计数分配区间并生成对应列

实现步骤 #

1. 计算目标列的数值总数（含重复值）

先统计列A、列B的非空值总数（含重复值，即该列所有值的数量）：

from pyspark.sql import functions as F

# 假设已有目标DataFrame df
count_a = df.select(F.count(F.col("A"))).collect()[0][0]
count_b = df.select(F.count(F.col("B"))).collect()[0][0]

2. 定义区间映射逻辑

可以通过自定义函数匹配计数对应的区间，也可以直接用when表达式实现：

方法一：自定义函数映射区间

先定义区间规则，再编写匹配函数：

# 定义指定区间池
interval_rules = [
    (1000, 3000, "(1000-3000)"),
    (3001, 7000, "(3001-7000)"),
    (7001, 20000, "(7001-20000)"),
    (20001, 500000, "(20001-500000)")
]

def match_interval(count_val):
    for lower, upper, interval_str in interval_rules:
        if lower <= count_val <= upper:
            return interval_str
    return None  # 处理计数不在区间内的情况

# 获取列A、B对应的区间字符串
interval_c = match_interval(count_a)
interval_d = match_interval(count_b)

方法二：直接用when条件表达式

如果区间数量较少，也可以直接用Spark的条件判断语法：

# 为列A匹配区间并生成列C
df = df.withColumn("C", 
    F.when((count_a >= 1000) & (count_a <= 3000), "(1000-3000)")
     .when((count_a >= 3001) & (count_a <= 7000), "(3001-7000)")
     .when((count_a >= 7001) & (count_a <= 20000), "(7001-20000)")
     .when((count_a >= 20001) & (count_a <= 500000), "(20001-500000)")
     .otherwise(None)
)

# 同理生成列D
df = df.withColumn("D", 
    F.when((count_b >= 1000) & (count_b <= 3000), "(1000-3000)")
     .when((count_b >= 3001) & (count_b <= 7000), "(3001-7000)")
     .when((count_b >= 7001) & (count_b <= 20000), "(7001-20000)")
     .when((count_b >= 20001) & (count_b <= 500000), "(20001-500000)")
     .otherwise(None)
)

3. 添加区间结果列

将匹配到的区间作为常量列添加到原DataFrame中（方法一的后续操作）：

df = df.withColumn("C", F.lit(interval_c))
df = df.withColumn("D", F.lit(interval_d))

补充说明 #

• 如果需要统计列中去重后的数值数量，只需将F.count(F.col("A"))替换为F.countDistinct(F.col("A"))即可。
• 若计数不在指定区间池内，代码默认返回None，可根据需求修改otherwise的返回值。

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