PySpark值计数性能优化求助：小数据集去重计数耗时超5分钟

太懂这种憋屈了——Spark读十亿级数据快得飞起，但偏偏碰到这种针对性的计数操作就掉链子，哪怕小数据集都能卡5分钟以上。我之前也踩过类似的坑，给你几个亲测有效的优化思路：

针对特定值计数（比如values='1'的行数） #

• 利用分区裁剪+谓词下推：如果你的数据集是按values分区的（比如Parquet/ORC分区表），直接读取对应分区能跳过99%的数据，速度会暴增：

  // 假设values是分区列，直接过滤分区
  val value1Count = spark.read.table("your_dataset")
    .where("values = '1'")
    .count()
  

  如果values不是分区列，确保你的数据源支持谓词下推（列式存储基本都支持），这样过滤逻辑会被推到数据源层面执行，不用把全表数据拉到Spark再过滤。

针对分组计数（groupBy("values").count()） #

这类操作慢大多是因为shuffle开销过大或者数据倾斜，可以从这两点入手：

• 调整shuffle分区数：Spark默认的spark.sql.shuffle.partitions是200，对于小型数据集来说这个数太多，会产生大量小文件拖慢速度；对于大数据集又可能不够。根据你的数据量调整，比如小数据集设为10-20，大数据集设为300-500：

  // 先设置合适的shuffle分区数
  spark.conf.set("spark.sql.shuffle.partitions", 30)
  df.groupBy("values")
    .count()
    .orderBy("count")
    .show()
  

• 解决数据倾斜：如果某个values值的行数特别多（比如占了总数据的30%以上），会导致单个shuffle任务卡壳。可以用加盐法拆分倾斜key，先局部聚合再全局聚合：

  import org.apache.spark.sql.functions.{rand, sum, cast}
  import org.apache.spark.sql.types.IntegerType
  
  df.withColumn("salt", (rand() * 10).cast(IntegerType)) // 给倾斜key加10个随机盐值
    .groupBy("values", "salt")
    .count() // 局部聚合
    .groupBy("values")
    .agg(sum("count").alias("count")) // 全局聚合
    .orderBy("count")
    .show()
  

通用优化技巧 #

• 合理缓存数据集：如果这个DataFrame后续还要多次执行计数操作，先把它缓存到内存+磁盘（避免OOM），触发缓存后再执行计数会快很多：

  val cachedDf = df.cache()
  // 触发缓存（执行一个action算子）
  cachedDf.count()
  // 之后的计数操作都会用缓存的数据
  val value1Count = cachedDf.where("values = '1'").count()
  cachedDf.groupBy("values").count().orderBy("count").show()
  // 用完记得释放缓存
  cachedDf.unpersist()
  

• 小型数据集可以尝试RDD的countByValue：这个方法会把结果直接拉到Driver端，适合数据量不大的场景，省去shuffle的开销：

  // 直接获取所有values的计数map
  val countMap = df.select("values")
    .rdd.map(row => row.getString(0))
    .countByValue()
  
  // 获取特定值的计数
  val value1Count = countMap.getOrElse("1", 0L)
  // 排序输出
  countMap.toList.sortBy(_._2).foreach(println)
  

  注意：如果数据量很大，这个方法会导致Driver端内存溢出，谨慎使用！

最后补充一句：你提到count是惰性算子，其实慢的根源不是惰性本身，而是执行时的shuffle、全表扫描或者数据倾斜这些环节。上面的优化都是针对这些瓶颈点来的，根据你的实际场景选对应的方法就行。

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