PySpark按ID分组聚合各状态计数：PartitionBy是否适用？

问题描述 #

现有一张包含id和status字段的大数据表，需在PySpark中实现按id统计各状态的数量，将结果分别存入start-count、in-progress-count、end-count列。示例如下：

原表示例 #

目标结果表示例 #

能否用PartitionBy实现？ #

单独使用PartitionBy无法直接实现需求。PartitionBy主要用于窗口函数中划定计算的分组范围，它本身不具备聚合生成多列统计结果的能力。如果强行结合窗口函数使用，会额外产生不必要的计算步骤（如重复数据遍历、去重），效率远低于直接聚合方案。

最优实现方案：分组聚合+条件计数 #

直接通过groupBy按id分组，结合count(when(...))的条件计数逻辑，一次性生成目标统计列。这种方案仅需一次Shuffle操作，适合大数据场景，效率最高。

代码示例 #

# 导入必要依赖
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, count, when

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

# 加载/创建原表数据（此处为示例数据）
raw_data = [
    (1, "In Progress"), (1, "In Progress"),
    (2, "In Progress"), (2, "Start"),
    (3, "Start"), (3, "In Progress"), (3, "End")
]
raw_df = spark.createDataFrame(raw_data, schema=["id", "status"])

# 执行统计计算
result_df = raw_df.groupBy("id") \
    .agg(
        # 统计Start状态数量
        count(when(col("status") == "Start", 1)).alias("start-count"),
        # 统计In Progress状态数量
        count(when(col("status") == "In Progress", 1)).alias("in-progress-count"),
        # 统计End状态数量
        count(when(col("status") == "End", 1)).alias("end-count")
    )

# 查看结果
result_df.show()

补充说明 #

• 如果status字段存在大小写不一致（如start/Start），可先统一格式：

  from pyspark.sql.functions import lower
  raw_df = raw_df.withColumn("status", lower(col("status")))
  

  之后将条件判断改为col("status") == "start"即可。
• 该方案自动为未出现对应状态的id填充0值，无需额外处理空值。

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