圖解Spark排序算子sortBy的核心源碼

原創/朱季謙

一、案例說明

以前剛開始學習Spark的時候，在練習排序算子sortBy的時候，曾發現一個有趣的現象是，在使用排序算子sortBy後直接打印的話，發現打印的結果是亂序的，並沒有出現完整排序。

例如，有一個包含多個（姓名，金額）結構的List數據，將這些數據按照金額降序排序時，代碼及打印效果如下：

val money = ss.sparkContext.parallelize(
  List(("Alice", 9973),
    ("Bob", 6084),
    ("Charlie", 3160),
    ("David", 8588),
    ("Emma", 8241),
    ("Frank", 117),
    ("Grace", 5217),
    ("Hannah", 5811),
    ("Ivy", 4355),
    ("Jack", 2106))
)
money.sortBy(x =>x._2, false).foreach(println)


打印結果——
(Ivy,4355)
(Grace,5217)
(Jack,2106)
(Frank,117)
(Emma,8241)
(Alice,9973)
(Charlie,3160)
(Bob,6084)
(Hannah,5811)
(David,8588)

可見，這樣的執行結果並沒有按照金額進行降序排序。但是，如果使用collect或者重新將分區設置爲1以及直接將結果進行save保存時，發現結果都是能夠按照金額進行降序排序。（注意一點，按照save保存結果，雖然可能生成很多part-00000 ～part-00005的文件，但從part-00000到part-00005，內部數據其實也按照金額進行了降序排序）。

money.sortBy(x =>x._2, false).collect().foreach(println)
或者
money.repartition(1).sortBy(x =>x._2, false).foreach(println)
或者
money.sortBy(x =>x._2, false).saveAsTextFile("result")

最後結果——
(Alice,9973)
(David,8588)
(Emma,8241)
(Bob,6084)
(Hannah,5811)
(Grace,5217)
(Ivy,4355)
(Charlie,3160)
(Jack,2106)
(Frank,117)

二、sortBy源碼分析

爲何單獨通過sortBy後對數據打印，是亂序的，而在sortBy之後通過collect、save或者重分區爲1個分區repartition(1)，數據就是有序的呢？

帶着這個疑問，去看一下sortBy底層源碼——

def sortBy[K](
    f: (T) => K,
    ascending: Boolean = true,
    numPartitions: Int = this.partitions.length)
    (implicit ord: Ordering[K], ctag: ClassTag[K]): RDD[T] = withScope {
  this.keyBy[K](f)
      .sortByKey(ascending, numPartitions)
      .values
}

可以看到，核心源碼是 this.keyBy[K](f).sortByKey(ascending, numPartitions).values，我會將該源碼分成this.keyBy[K](f) ， sortByKey(ascending, numPartitions)以及values三部分講解——

2.1、逐節分析sortBy源碼之一：this.keyByK

this.keyBy[K](f)這行代碼是基於_.sortBy(x =>x._2, false)傳進來的x =>x._2重新生成一個新RDD數據，可以進入到其底層源碼看一下——

def keyBy[K](f: T => K): RDD[(K, T)] = withScope {
  val cleanedF = sc.clean(f)
  map(x => (cleanedF(x), x))
}

若執行的是_.sortBy(x =>x._2, false)，那麼f: T => K匿名函數就是x =>x._2，因此，keyBy函數的裏面代碼真面目是這樣——

map(x => (sc.clean(x =>x._2), x))

sc.clean(x =>x._2)這個clean相當是對傳入的函數做序列化，因爲最後會將這個函數得到結果當作排序key分發到不同分區節點做排序，故而涉及到網絡傳輸，因此做序列化後就方便在分佈式計算中在不同節點之間傳遞和執行函數，clean最終底層實現是這行代碼SparkEnv.get.closureSerializer.newInstance().serialize(func)，感興趣可以深入研究。

keyBy最終會生成一個新的RDD，至於這個結構是怎樣的，通過原先的測試數據調用keyBy打印一下就一目瞭然——

val money = ss.sparkContext.parallelize(
  List(("Alice", 9973),
    ("Bob", 6084),
    ("Charlie", 3160),
    ("David", 8588),
    ("Emma", 8241),
    ("Frank", 117),
    ("Grace", 5217),
    ("Hannah", 5811),
    ("Ivy", 4355),
    ("Jack", 2106))
)
money.keyBy(x =>x._2).foreach(println)

打印結果——
(5217,(Grace,5217))
(5811,(Hannah,5811))
(8588,(David,8588))
(8241,(Emma,8241))
(9973,(Alice,9973))
(3160,(Charlie,3160))
(4355,(Ivy,4355))
(2106,(Jack,2106))
(117,(Frank,117))
(6084,(Bob,6084))

由此可知，原先這樣("Alice", 9973)結構的RDD，通過keyBy源碼裏的map(x => (sc.clean(x =>x._2), x))代碼，最終會生成這樣結構的數據（x._2，x）也就是，(9973,(Alice,9973)), 就是重新將需要排序的字段金額當作了新RDD的key。

2.2、逐節分析sortBy源碼之二：sortByKey

通過 this.keyBy[K](f)得到結構爲（x._2，x）的RDD後，可以看到，雖然我們前面調用money.sortBy(x =>x._2, false)來排序，但底層本質還是調用了另一個排序算子sortByKey，它有兩個參數，一個是布爾值的ascending，true表示按升序排序，false表示按降序排序，我們這裏傳進來的是false。另一個參數numPartitions，表示分區數，可以通過定義的rdd.partitions.size知道所在環境分區數。

進入到sortByKey源碼裏，通過以下函數註釋，就可以知道sortByKey函數都做了什麼事情——

/**
 * Sort the RDD by key, so that each partition contains a sorted range of the elements. Calling
 * `collect` or `save` on the resulting RDD will return or output an ordered list of records
 * (in the `save` case, they will be written to multiple `part-X` files in the filesystem, in
 * order of the keys).
 *
 *按鍵對RDD進行排序，以便每個分區包含一個已排序的元素範圍。
 在結果RDD上調用collect或save將返回或輸出一個有序的記錄列表
 (在save情況下，它們將按照鍵的順序寫入文件系統中的多個part-X文件)。
 */
// TODO: this currently doesn't work on P other than Tuple2!
def sortByKey(ascending: Boolean = true, numPartitions: Int = self.partitions.length)
    : RDD[(K, V)] = self.withScope
{
  val part = new RangePartitioner(numPartitions, self, ascending)
  new ShuffledRDD[K, V, V](self, part)
    .setKeyOrdering(if (ascending) ordering else ordering.reverse)
}

到這裏，基於註解就可以知道sortByKey做了什麼事情了——

第一步，每個分區按鍵對RDD進行shuffle洗牌後將相同Key劃分到同一個分區，進行排序。

第二步，在調用collect或save後，會對各個已經排序好的各個分區進行合併，最終得到一個完整的排序結果。

這就意味着，若沒有調用collect或save將各個分區結果進行彙總返回給master驅動進程話，雖然分區內的數據是排序的，但分區間就不一定是有序的。這時若直接foreach打印，因爲打印是並行執行的，即使分區內有序，但並行一塊打印就亂七八糟了。

可以寫段代碼驗證一下，各個分區內是不是有序的——

money.sortBy(x => x._2, false).foreachPartition(x => {
  val partitionId = TaskContext.get.partitionId
  //val index = UUID.randomUUID()
  x.foreach(x => {
    println("分區號" + partitionId + "：   " + x)
  })
})

打印結果——
分區號2：   (Ivy,4355)
分區號2：   (Charlie,3160)
分區號2：   (Jack,2106)
分區號2：   (Frank,117)

分區號1：   (Bob,6084)
分區號1：   (Hannah,5811)
分區號1：   (Grace,5217)

分區號0：   (Alice,9973)
分區號0：   (David,8588)
分區號0：   (Emma,8241)

設置環境爲3個分區，可見每個分區裏的數據已經是降序排序了。

若是隻有一個分區的話，該分區裏的數據也會變成降序排序，這就是爲何money.repartition(1).sortBy(x =>x._2, false).foreach(println)得到的數據也是排序結果。

sortBy主要流程如下，假設運行環境有3個分區，讀取的數據去創建一個RDD的時候，會按照默認Hash分區器將數據分到3個分區裏。

在調用sortBy後，RDD會通過 this.keyBy[K](f)重新生成一個新的RDD，例如結構如(8588, (David,8588))，接着進行shuffle操作，把RDD數據洗牌打散，將相應範圍的key重新分到同一個分區裏，意味着，同key值的數據就會分發到了同一個分區，例如下圖的(2106, (Jack,2106))，(999, (Alice,999))，(999, (Frank,999))，(999, (Hannah,999))含同一個Key都在一起了。

shuffleRDD中，使用mapPartitions會對每個分區的數據按照key進行相應的升序或者降序排序，得到分區內有序的結果集。

2.3、逐節分析sortBy源碼之三：.values

sortBy底層源碼裏 this.keyBy[K](f).sortByKey(ascending, numPartitions).values，在sortByKey之後，最後調用了.values。源碼.values裏面是def values: RDD[V] = self.map(_._2)，就意味着，排序完成後，只返回x._2的數據，用於排序生成的RDD。類似排序過程中RDD是(5217,(Grace,5217))這樣結構，排序後，若只返回x._2，就只返回(Grace,5217)這樣結構的RDD即可。

可以看到，shuffleRDD將相應範圍的key重新分到同一個分區裏，例如，0～100劃到分區0，101～200劃分到分區1，201～300劃分到分區2，這樣還有一個好處——當0，1，2分區內部的數據已經有序時，這時從整體按照0，1，2分區全局來看，其實就已經是全局有序了，當然，若要實現全局有序，還需要將其合併返回給驅動程序。

三、合併各個分區的排序，返回全局排序

調用collect或save就是把各個分區結果進行彙總，相當做了一個歸併排序操作——

以上，就是關於sortBy核心源碼的講解。