圖解Spark排序運算元sortBy的核心原始碼

朱季謙發表於2023-09-18
圖解Spark排序原始碼

image

原創/朱季謙

一、案例說明

以前剛開始學習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。

image


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進行相應的升序或者降序排序，得到分割槽內有序的結果集。
image


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即可。

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


三、合併各個分割槽的排序，返回全域性排序

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

image

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

相關文章