Spark多count distinct单表达式高效精确计数优化方案咨询

嘿，这个问题我太有共鸣了——之前处理PB级数据集的时候，也被Spark的Expand操作坑过，十亿级数据直接膨胀5倍，shuffle慢到怀疑人生。给你几个亲测有效的高效方案，都是能拿到精确值的：

方案1：用collect_set + size替代countDistinct(when(...)) #

这是最简单直接的方案，完全避免Expand操作。原理是利用collect_set自动去重的特性，先收集满足条件的字段值（忽略null），再用size统计去重后的数量，全程只需要一次聚合。

Scala API实现： #

import org.apache.spark.sql.functions.{collect_set, size, when, lit}

sourceData.groupBy($"number")
  .agg(
    size(collect_set(when($"status" === "A", $"user_id"))).alias("cnt_a"),
    size(collect_set(when($"status" === "B", $"user_id"))).alias("cnt_b"),
    size(collect_set(when($"status" === "C", $"user_id"))).alias("cnt_c"),
    size(collect_set(when($"status" === "D", $"user_id"))).alias("cnt_d"),
    size(collect_set(when($"status" === "E", $"user_id"))).alias("cnt_e")
  )

SQL API实现： #

SELECT 
  number,
  SIZE(COLLECT_SET(CASE WHEN status = 'A' THEN user_id END)) AS cnt_a,
  SIZE(COLLECT_SET(CASE WHEN status = 'B' THEN user_id END)) AS cnt_b,
  SIZE(COLLECT_SET(CASE WHEN status = 'C' THEN user_id END)) AS cnt_c,
  SIZE(COLLECT_SET(CASE WHEN status = 'D' THEN user_id END)) AS cnt_d,
  SIZE(COLLECT_SET(CASE WHEN status = 'E' THEN user_id END)) AS cnt_e
FROM sourceData
GROUP BY number;

优点：代码简洁，无需额外预处理，完全消除Expand操作，shuffle数据量和原数据集一致。
注意：如果单个number分组对应的去重字段（比如user_id）数量极大，可能会占用较多内存，此时可以调整Spark的内存参数（如spark.sql.shuffle.partitions、spark.executor.memory）优化。

方案2：拆分条件→聚合→Pivot（适合大分组场景） #

如果方案1的内存压力太大，可以先把满足条件的记录拆分成多行，再做去重聚合，最后转成列。这种方式的shuffle数据量是满足条件的记录数（而非原数据的5倍），适合大部分记录不满足所有条件的场景。

Scala API实现： #

import org.apache.spark.sql.functions.{array, explode, filter, lit, when}

// 第一步：生成所有满足的条件列表
val withConditions = sourceData.select(
  $"number",
  $"user_id",
  filter(
    array(
      when($"status" === "A", lit("A")),
      when($"status" === "B", lit("B")),
      when($"status" === "C", lit("C")),
      when($"status" === "D", lit("D")),
      when($"status" === "E", lit("E"))
    ),
    _.isNotNull
  ).alias("valid_conditions")
)

// 第二步：拆分条件为多行，过滤无效值
val exploded = withConditions.select(
  $"number",
  $"user_id",
  explode($"valid_conditions").alias("condition")
)

// 第三步：分组去重统计，再转成列
exploded.groupBy($"number", $"condition")
  .agg(countDistinct($"user_id").alias("cnt"))
  .groupBy($"number")
  .pivot($"condition", Seq("A", "B", "C", "D", "E")) // 指定顺序避免乱序
  .agg(first($"cnt"))
  .withColumnRenamed("A", "cnt_a")
  .withColumnRenamed("B", "cnt_b")
  .withColumnRenamed("C", "cnt_c")
  .withColumnRenamed("D", "cnt_d")
  .withColumnRenamed("E", "cnt_e")

优点：内存压力小，shuffle数据量更低（仅包含满足条件的记录）；
缺点：多了一次拆分和聚合操作，代码稍复杂。

为什么原代码会产生Expand？ #

原代码中每个countDistinct(when(...))都是独立的聚合任务，Spark为了并行处理这些任务，会用Expand算子把每条记录复制N份（N=countDistinct的数量），每份对应一个聚合任务的分组键，导致数据量直接膨胀N倍——对于十亿级数据来说，这个开销是致命的。

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