Spark 核心之二次排序、分组取 TopN 优化分析

摘要 ：为什么你的分组 TopN 任务跑了 30 分钟还 OOM？答案通常是四个字------groupByKey。本文将二次排序和分组 TopN 作为两条优化主线，从反模式剖析、三种方案对比（groupByKey / reduceByKey+mapPartitions / SQL 窗口函数）、groupByKey vs reduceByKey 本质区别、完整代码实战、Spark SQL 窗口函数最佳实践五个维度，配合 1 张原创深色架构图，助你彻底攻克 TopN 优化的每一个细节。

关键词：Spark 二次排序, TopN, groupByKey, reduceByKey, 窗口函数, row_number, mapPartitions, 分组优化

---

一、开篇：groupByKey 是 TopN 的超级陷阱

// ❌ 这是 95% 的新手会写的代码------也是 OOM 的根源
rdd.groupByKey()                          // 全量数据按 Key 分组
   .mapValues(_.toList.sortBy(-_._2).take(3))  // ⚠ 所有 Value 加载到内存！
// 问题: 如果某个 Key 有 1000 万条数据 → Executor OOM

groupByKey 的问题 ：它不做 Map 端预聚合，将全部 KV 对原封不动地 Shuffle 到 Reducer，然后全部加载到内存------数据倾斜时直接炸掉。

---

二、优化方案全景图

---

三、方案 1：二次排序（Secondary Sort）

3.1 自定义排序键

// 按 classId 升序，score 降序
case class SecondarySortKey(classId: String, score: Int) 
  extends Ordered[SecondarySortKey] {
  override def compare(that: SecondarySortKey): Int = {
    val cmp1 = this.classId.compareTo(that.classId)
    if (cmp1 != 0) cmp1 else -this.score.compareTo(that.score) // score 降序
  }
}

// 使用
rdd.map { case (classId, name, score) => 
    (SecondarySortKey(classId, score), (name, score))
  }
  .sortByKey()  // 按自定义键全局排序

3.2 逐组处理（避免 groupByKey）

// 利用 sortByKey 后的有序性，mapPartitions 内逐组处理
rdd.mapPartitions { iter =>
  var currentClass = ""
  var buffer = ListBuffer[(String, Int)]()
  iter.flatMap { case (key, (name, score)) =>
    if (key.classId != currentClass) {
      val result = buffer.toList
      currentClass = key.classId
      buffer = ListBuffer((name, score))
      if (result.nonEmpty) result.iterator else Iterator.empty
    } else {
      buffer += ((name, score))
      Iterator.empty
    }
  }
}

---

四、方案 2：分组 TopN

4.1 ❌ 反模式：groupByKey

rdd.groupByKey().mapValues(_.toList.sortBy(-_._2).take(3))
// 每个 Key 的 Value 全部加载到单节点内存 → OOM

4.2 ⚠ 中等方案：reduceByKey + 局部 TopN

rdd.map { case (key, v) => (key, List(v)) }
   .reduceByKey((a, b) => (a ++ b).sortBy(-_._2).take(3))
// Map 端先局部合并（Combiner） → Shuffle 数据量大幅减少
// 但 RDD API 写起来较繁琐

4.3 ✅ 推荐方案：Spark SQL 窗口函数

-- 每个班级取成绩 Top 3
SELECT class, name, score
FROM (
  SELECT *, 
    row_number() OVER (PARTITION BY class ORDER BY score DESC) AS rn
  FROM students
) t
WHERE rn <= 3

// DataFrame API
import org.apache.spark.sql.expressions.Window
val windowSpec = Window.partitionBy("class").orderBy($"score".desc)
df.withColumn("rn", row_number().over(windowSpec))
  .where($"rn" <= 3)
  .drop("rn")

---

五、groupByKey vs reduceByKey 本质区别

维度	groupByKey	reduceByKey
mapSideCombine	❌ false	✅ true
Shuffle 数据量	全部 KV 对	Map 端预聚合后少量数据
内存压力	高（全量加载）	低（逐批合并）
适用场景	需要对 Value 做非结合性操作	可结合、可交换的聚合
等价窗口函数	collect_list()	sum/count/max/min + 自定义 UDAF

// reduceByKey 的内部 Combiner 机制
rdd.reduceByKey(_ + _)
// 源码等价于:
// combineByKey(createCombiner, mergeValue, mergeCombiners)
// Map 端先 mergeValue → mergeCombiners 分批合并

---

六、实战：每个班级成绩 Top 3

// 方案 A: RDD reduceByKey（较复杂）
rdd.map { case (cls, name, score) => (cls, List((name, score))) }
   .reduceByKey((a, b) => (a ++ b).sortBy(-_._2).take(3))
   .flatMapValues(identity)

// 方案 B: Spark SQL 窗口函数（推荐）
val w = Window.partitionBy("class").orderBy($"score".desc)
df.withColumn("rn", row_number().over(w)).where($"rn" <= 3)

// 方案 C: rank/dense_rank（处理并列）
// rank():    1,1,3,4  (并列占位)
// dense_rank(): 1,1,2,3 (并列不占位)
// row_number(): 1,2,3,4 (纯粹行号)

---

七、性能总结

方案	Shuffle 量	内存安全	推荐度
groupByKey	100%	❌ 高危	⛔ 禁止
reduceByKey	10-30%	⚠ 可控	可用
窗口函数	10-30%	✅ 最优	⭐ 首选

金句：groupByKey 是 Spark 新手的第一大坑------它把"分组"当成目的，却忘了分组只是手段。真正的目的是在减少 Shuffle 的前提下拿到想要的结果。记住：能用 reduceByKey 绝不用 groupByKey，能用窗口函数就用窗口函数。

---

作者：starzy | AI Data Engineer / 大数据技术实践者

博客：blog.starzy.cn | GitHub：starzy1990.github.io

专注 AI Agent · LangGraph · RAG · 大数据架构 · 数据工程实践