Spark按某列分组后统计另一列各值占比的实现方案

错误原因分析 #

你之前定义窗口时同时按prod_name和colour分区，计算得到的结果仅为单产品下单颜色的记录数，缺少单产品的总记录数作为分母，因此无法直接算出占比。

正确实现方案 #

这里提供两种PySpark环境下的实现方式，你可以根据数据量选择：

方案1：窗口函数实现（适合小数据量快速开发） #

from pyspark.sql import functions as F, Window

# 定义两个窗口：分别用于统计单产品单颜色数量、单产品总数量
w_color = Window.partitionBy("prod_name", "colour")
w_prod = Window.partitionBy("prod_name")

result_df = df.withColumn("color_cnt", F.count("*").over(w_color)) \
              .withColumn("prod_total_cnt", F.count("*").over(w_prod)) \
              .withColumn("percentage", F.round(F.col("color_cnt") / F.col("prod_total_cnt"), 2)) \
              .select("prod_name", "colour", "percentage") \
              .dropDuplicates()

方案2：分组聚合实现（适合大数据量，性能更优） #

from pyspark.sql import functions as F

# 统计每个产品下各颜色的记录数
color_count_df = df.groupBy("prod_name", "colour").agg(F.count("*").alias("color_cnt"))
# 统计每个产品的总记录数
prod_total_df = df.groupBy("prod_name").agg(F.count("*").alias("prod_total"))

# 关联两个统计结果计算占比
result_df = color_count_df.join(prod_total_df, on="prod_name", how="left") \
                          .withColumn("percentage", F.round(F.col("color_cnt") / F.col("prod_total"), 2)) \
                          .select("prod_name", "colour", "percentage")

两种方案输出结果都和你的预期一致，percentage字段默认保留两位小数，你可以根据需求调整round函数的参数修改精度。

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