嘿，你的场景在Spark大数据处理里真的太常见了！当前逐个计算KPI再union聚合的方式虽然能得到结果，但性能问题会很突出——每计算一个KPI就要单独扫描一遍原始DataFrame，50个KPI等于重复扫50次数据，这对大数据量来说完全是没必要的资源浪费，而且多次union和groupBy还会带来额外的shuffle开销。

其实Spark原生就有更优的实现模式：单遍扫描+条件聚合，也就是把所有KPI的计算逻辑整合到同一个groupBy操作中，用when/otherwise结合聚合函数完成所有统计，只需要扫描一次原始数据，性能提升非常明显。

核心思路 #

针对每个KPI，我们用when函数指定过滤条件，满足条件时取对应的值（或1用于计数），不满足时取0（或null，根据聚合类型调整），然后直接在groupBy后对这些条件化的字段执行sum/avg等聚合操作。这样所有KPI的计算都在同一个执行计划里完成，只需要一次数据扫描和shuffle。

代码示例 #

假设你的原始DataFrame是raw_df，分组列是col1, col2, col3, col4，我们来实现几个典型KPI的统计：

Scala版本

import org.apache.spark.sql.functions.{col, when, sum, avg, count}

// 可以把复杂条件抽成变量，提升代码可读性
val kpi1_condition = col("status") === "success" && col("value") > 100
val kpi3_condition = col("category") === "premium"

val grouped_df = raw_df.groupBy("col1", "col2", "col3", "col4")
  .agg(
    // KPI1: 满足条件A的数值总和
    sum(when(kpi1_condition, col("value")).otherwise(0)).alias("kpi1"),
    // KPI2: 满足条件B的记录数
    sum(when(col("type") === "payment", 1).otherwise(0)).alias("kpi2"),
    // KPI3: 满足条件C的数值平均值（处理分母为0的情况）
    when(sum(when(kpi3_condition, 1).otherwise(0)) === 0, 0)
      .otherwise(sum(when(kpi3_condition, col("score")).otherwise(0)) / sum(when(kpi3_condition, 1).otherwise(0)))
      .alias("kpi3"),
    // KPI4: 满足条件D的非空记录数
    count(when(col("is_valid") === true, col("id"))).alias("kpi4"),
    // 继续添加剩下的46个KPI...
  )

Python版本

from pyspark.sql import functions as F

# 抽离复杂条件
kpi1_condition = (F.col("status") == "success") & (F.col("value") > 100)
kpi3_condition = F.col("category") == "premium"

grouped_df = raw_df.groupBy("col1", "col2", "col3", "col4") \
    .agg(
        # KPI1: 满足条件A的数值总和
        F.sum(F.when(kpi1_condition, F.col("value")).otherwise(0)).alias("kpi1"),
        # KPI2: 满足条件B的记录数
        F.sum(F.when(F.col("type") == "payment", 1).otherwise(0)).alias("kpi2"),
        # KPI3: 满足条件C的数值平均值（处理分母为0的情况）
        F.when(F.sum(F.when(kpi3_condition, 1).otherwise(0)) == 0, 0)
          .otherwise(F.sum(F.when(kpi3_condition, F.col("score")).otherwise(0)) / F.sum(F.when(kpi3_condition, 1).otherwise(0)))
          .alias("kpi3"),
        # 其他KPI依次添加...
    )

为什么这个方案更优 #

• 减少IO开销：只需要扫描原始DataFrame一次，避免了50次重复扫描带来的磁盘/网络IO浪费。
• 降低shuffle次数：所有聚合操作在同一个groupBy里完成，只需要一次shuffle，而原来的方案每个KPI计算都可能产生shuffle，union后还要再shuffle一次。
• 更好的执行计划优化：Spark Catalyst优化器可以对单遍聚合的执行计划做更多优化（比如谓词下推、列裁剪），进一步提升性能。

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