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 本质区别

// 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 (纯粹行号)

七、性能总结

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

作者：starzy | AI Data Engineer / 大数据技术实践者
博客：blog.starzy.cn | GitHub：starzy1990.github.io
专注 AI Agent · LangGraph · RAG · 大数据架构 · 数据工程实践