圖解Spark Graphx基於connectedComponents函數實現連通圖底層原理

原創/朱季謙

第一次寫這麼長的graphx源碼解讀，還是比較晦澀，有較多不足之處，爭取改進。

一、連通圖說明

連通圖是指圖中的任意兩個頂點之間都存在路徑相連而組成的一個子圖。

用一個圖來說明，例如，下面這個叫graph的大圖裏，存在兩個連通圖。

左邊是一個連接圖，該子圖裏每個頂點都存在路徑相連，包括了頂點：{(5L, "Eve"), (7L, "Grace"), (1L, "Alice"), (2L, "Bob"), (3L, "Charlie")}。

右邊同樣是一個連接圖，該子圖裏每個頂點都存在路徑相連，包括了頂點：{(8L, "Henry")，(9L, "Ivy")，(6L, "Frank")}。

在現實生活裏，這兩個子圖就相當某個社區裏的關係網，在Spark Graphx裏，經常需要處理這類關係網的操作，那麼，在一個圖裏如何得到各個子圖的數據呢？

這時，就可以使用到Spark Graphx的connectedComponents函數，網上關於它的介紹，基本都是說它是Graphx三大圖算法之一的連通組件。

連通組件是指圖中的一組頂點，每個頂點之間都存在路徑互相關聯，也就是前面提到圖中的子圖概念。

通俗解釋，就是通過這個函數，可以將每個頂點都關聯到連通圖裏的最小頂點，例如，前面提到的子圖{(8L, "Henry")，(9L, "Ivy")，(6L, "Frank")}，在通過connectedComponents函數處理後，就可以得到每個頂點關聯到該子網的最小頂點ID。該子圖裏的最小頂點ID是6L，那麼，可以處理成以下數據{(8L，6L)，（9L，6L），(6L，6L)}。既然屬於同一個子圖的各個頂點都關聯到一個共同的最小頂點，不就意味着，通過該最小頂點，是可以按照分組的操作，將同一個最小頂點的數據都分組到一塊，這樣，就能提取出同一個子圖的頂點集合了。

二、案例說明

基於以上的圖頂點和邊數據，創建一個Graphx圖——

val conf = new SparkConf().setMaster("local[*]").setAppName("graphx")
val ss = SparkSession.builder().config(conf).getOrCreate()

// 創建頂點RDD
val vertices = ss.sparkContext.parallelize(Seq(
  (1L, "Alice"),
  (2L, "Bob"),
  (3L, "Charlie"),
  (5L, "Eve"),
  (6L, "Frank"),
  (7L, "Grace"),
  (8L, "Henry"),
  (9L, "Ivy")
))

// 創建邊RDD
val edges = ss.sparkContext.parallelize(Seq(
  Edge(5L, 7L, "friend"),
  Edge(5L, 1L, "friend"),
  Edge(1L, 2L, "friend"),
  Edge(2L, 3L, "friend"),
  Edge(6L, 9L, "friend"),
  Edge(9L, 8L, "friend")
))

//創建一個Graph圖
val graph = Graph(vertices, edges, null)

調用圖graph的connectedComponents函數，順便打印一下效果，可以看到，左邊子圖{(5L, "Eve"), (7L, "Grace"), (1L, "Alice"), (2L, "Bob"), (3L, "Charlie")}裏的各個頂點都關聯到了最小頂點1，右邊子圖{(8L, "Henry")，(9L, "Ivy")，(6L, "Frank")}裏的各個頂點都關聯到了最小頂點6。

val cc = graph.connectedComponents()
cc.vertices.foreach(println)

打印的結果——
(2,1)
(6,6)
(7,1)
(1,1)
(9,6)
(8,6)
(3,1)
(5,1)

注意一點，connectedComponents是可以傳參的，傳入的數字，是代表各個頂點最高可以連通迭代到多少步去尋找所在子圖裏的最小頂點。

舉個例子，可能就能明白了，假如，給connectedComponents傳參爲1，那麼代碼執行打印後，如下——

val cc = graph.connectedComponents(1)
cc.vertices.foreach(println)

打印的結果——
(2,1)
(5,1)
(8,8)
(7,5)
(1,1)
(9,6)
(6,6)
(3,2)

你會發現，各個頂點的連通組件即關聯所在子圖的最小頂點，大多都變了，這是因爲設置參數爲1 後，各個頂點沿着邊去迭代尋找連通組件時，只能迭代一步，相當本頂點只能走到一度鄰居頂點，然後將本頂點和鄰居頂點比較，誰最小，最小的當作連通組件。

以下圖說明，就是頂點(7L, "Grace")迭代一步去尋找最小頂點做連通組件，只能迭代到頂點(5L, "Eve")，沒法迭代到 (1L, "Alice")，這時頂點(7L, "Grace")就會拿自身與頂點(5L, "Eve")比較，發現5L更小，就會用5L當作自己的連通組件做關聯，即(7,5)。

當然，實際底層的源碼實現，並非是通過迭代多少步去尋找最小頂點，它的實現方式更精妙，站在原地就可以收集到所能迭代最大次數範圍內的最小頂點。

如果connectedComponents沒有設置參數，就會默認最大迭代次數是Int.MaxValue，2 的 31 次方 - 1即2147483647。

在實際業務當中，可以通過設置參數來避免在過大規模的子圖裏做耗時過長的迭代操作

接下來，就可以通過連通組件做分組，將具有共同連通組件的頂點分組到一塊，這樣就知道哪些頂點屬於同一子圖了。

val cc = graph.connectedComponents()
val group = cc.vertices.map{
  case (verticeId, minId) => (minId, verticeId)
}.groupByKey()

group.foreach(println)


打印結果——
(1,CompactBuffer(1, 2, 3, 5, 7))
(6,CompactBuffer(8, 9, 6))

基於這個函數，就可以得到哪些頂點在一張關係網裏了。

三、connectedComponents源碼解析

先來看一下connectedComponents函數源碼，在ConnectedComponents單例對象裏，可以看到，如果沒有傳參的話，默認迭代次數是Int.MaxValue，如果傳參的話，就使用參數的maxIterations做迭代次數——

/**
*無參數
*/
def connectedComponents(): Graph[VertexId, ED] = {
  ConnectedComponents.run(graph)
}


def run[VD: ClassTag, ED: ClassTag](graph: Graph[VD, ED]): Graph[VertexId, ED] = {
    run(graph, Int.MaxValue)
}



/**
*有參數
*/
def connectedComponents(maxIterations: Int): Graph[VertexId, ED] = {
    ConnectedComponents.run(graph, maxIterations)
}

在run方法裏，主要是做了一些函數和常量的準備工作，然後將這些函數和常量傳給單例對象Pregel的apply方法。apply是單例對象的特殊方法，就像Java類裏的構造方法一樣，創建對象時可以直接被調用。Pregel(ccGraph, initialMessage,maxIterations, EdgeDirection.Either)(......)最後調用的就是Pregel裏的apply方法。

def run[VD: ClassTag, ED: ClassTag](graph: Graph[VD, ED],
                                    maxIterations: Int): Graph[VertexId, ED] = {
  require(maxIterations > 0, s"Maximum of iterations must be greater than 0," +
    s" but got ${maxIterations}")
  //step1 初始化圖，將各頂點id設置爲頂點屬性，圖頂點結構（vid,vid）
  val ccGraph = graph.mapVertices { case (vid, _) => vid }
  
  //step2 處理圖裏的每一個三元組邊對象，該對象edge包含了源頂點(srcId，srcAttr)和目標頂點(dstId，dstAttr)的信息,及邊屬性attr，即（srcId，srcAttr，dstId，dstAttr，attr）
  def sendMessage(edge: EdgeTriplet[VertexId, ED]): Iterator[(VertexId, VertexId)] = {
    //如果源頂點屬性小於目標頂點屬性
    if (edge.srcAttr < edge.dstAttr) {
      //保存（目標頂點id,源頂點屬性），這裏的源頂點屬性等於源頂點id,其實保存的是（目標頂點id,源頂點id）
      Iterator((edge.dstId, edge.srcAttr))
       //如果源頂點屬性大於目標頂點屬性
    } else if (edge.srcAttr > edge.dstAttr) {
      //保存（源頂點id,目標頂點id）
      Iterator((edge.srcId, edge.dstAttr))
    } else {
      //如果兩個頂點屬性相同，說明已經在同一個子網裏，不需要處理
      Iterator.empty
    }
  }
  //step3 設置一個初始最大值，用於在初始化階段，比較每個頂點的屬性，這樣頂點屬性值在最初階段就相當是最小頂點
  val initialMessage = Long.MaxValue
  
  //step4 將上面設置的常量和函數當作參數傳給Pregel，其中EdgeDirection.Either表示處理包括出度和入度的頂點。
  val pregelGraph = Pregel(ccGraph, initialMessage,
    maxIterations, EdgeDirection.Either)(
    //將最初頂點的屬性attr與initialMessage比較，相當是子圖的0次迭代尋找最小頂點
    vprog = (id, attr, msg) => math.min(attr, msg),
    //上面定義的sendMessage方法
    sendMsg = sendMessage,
    //處理各個頂點收到的消息，然後將最小的頂點保存
    mergeMsg = (a, b) => math.min(a, b))
  ccGraph.unpersist()
  pregelGraph
}

step1 初始化圖，將各頂點id設置爲頂點屬性，圖頂點結構（vid,vid）——

 val ccGraph = graph.mapVertices { case (vid, _) => vid }

寫一個簡單的代碼驗證一下即可知道得到的ccGraph處理後頂點是否爲（vid，vid）結構了。

// 創建頂點RDD
val vertices = ss.sparkContext.parallelize(Seq(
  (1L, "Alice"),
  (2L, "Bob"),
  (3L, "Charlie"),
  (5L, "Eve"),
  (6L, "Frank"),
  (7L, "Grace"),
  (8L, "Henry"),
  (9L, "Ivy")
))

// 創建邊RDD
val edges = ss.sparkContext.parallelize(Seq(
  Edge(5L, 7L, "friend"),
  Edge(5L, 1L, "friend"),
  Edge(1L, 2L, "friend"),
  Edge(2L, 3L, "friend"),
  Edge(6L, 9L, "friend"),
  Edge(9L, 8L, "friend")
))

//創建一個Graph圖
val graph = Graph(vertices, edges, null)
graph.mapVertices{case  (vid,_) => vid}.vertices.foreach(println)

打印結果——
(2,2)
(5,5)
(3,3)
(6,6)
(7,7)
(8,8)
(1,1)
(9,9)

可見，ccGraph的圖頂點已經被處理成（vid,vid），即（頂點id, 頂點屬性），方便用於在sendMessage方法做屬性判斷處理。

step2 sendMessage處理圖裏的每一個三元組邊對象

前面處理的ccGraph頂點數據變成（頂點id, 頂點屬性）就是爲了放在這裏做處理，這裏的if (edge.srcAttr < edge.dstAttr) 相當是if (edge.srcId < edge.dstId)。

這個方法是基於邊的三元組做處理，將同一邊的源頂點和目標頂點比較，篩選出兩個頂點最小的頂點，然後針對最大的頂點，保留（最大頂點，最小頂點屬性）這樣的數據。

  def sendMessage(edge: EdgeTriplet[VertexId, ED]): Iterator[(VertexId, VertexId)] = {
    //如果源頂點屬性小於目標頂點屬性
    if (edge.srcAttr < edge.dstAttr) {
      //保存（目標頂點id,源頂點屬性），這裏的源頂點屬性等於源頂點id,其實保存的是（目標頂點id,源頂點id）
      Iterator((edge.dstId, edge.srcAttr))
       //如果源頂點屬性大於目標頂點屬性
    } else if (edge.srcAttr > edge.dstAttr) {
      //保存（源頂點id,目標頂點id）
      Iterator((edge.srcId, edge.dstAttr))
    } else {
      //如果兩個頂點屬性相同，說明已經在同一個子網裏，不需要處理
      Iterator.empty
    }
  }

這個方法的作用，就是找出同一條邊上哪個頂點最小，例如下圖中，2L比3L小，那麼2L是這條邊上最小的頂點，將以最大點關聯最小點的方式(edge.dstId, edge.srcAttr)即（3L，2L）保存下來。最後會將（3L，2L）中的_.2也就是2L發送給頂點（3L，3L），而頂點（3L，3L）後續需要做的事情是，是將這一輪收到的消息即最小頂點2L與現在的屬性3L值通過math.min(a, b)做比較，保留最小頂點當作屬性值，即變成了（3L，2L）。

可見，在子圖裏，每一輪迭代後，各個頂點的屬性值都可能會被更新接收到的最小頂點值，這就是連通組件迭代的精妙。

這個方法會在後面的Pregel對象裏用到。

step3 設置一個初始最大值，用於比較後初始化每個頂點最初的屬性值

val initialMessage = Long.MaxValue需要與vprog = (id, attr, msg) => math.min(attr, msg)結合來看，相當在0次迭代時，將頂點（id，attr）的屬性值與initialMessage做比較，理論上，肯定是attr比較小，就意味着初始化時，頂點關聯的最小頂點就是attr，在這裏，就相當關聯的最小頂點是它本身，相當於子圖做了0次迭代處理。

step4 執行Pregel的構造函數apply方法

可以看到，前面創建的ccGraph，initialMessage，maxIterations（最大迭代次數），EdgeDirection.Either都當作參數傳給了Pregel。

val pregelGraph = Pregel(ccGraph, initialMessage,
    maxIterations, EdgeDirection.Either)(
    //將最初頂點的屬性attr與initialMessage比較，相當是子圖的0次迭代尋找最小頂點
    vprog = (id, attr, msg) => math.min(attr, msg),
    //上面定義的sendMessage方法
    sendMsg = sendMessage,
    //處理各個頂點收到的消息，然後將最小的頂點保存
    mergeMsg = (a, b) => math.min(a, b))

該Pregel對象底層主要就是對一系列的三元組邊的源頂點和目標頂點做比較，將兩頂點最小的頂點值發送給該條邊最大的頂點，最大的頂點收到消息後，會比較當前屬性與收到的最小頂點值比較，然後保留最小值。這樣，每一輪迭代，可能關聯的屬性值都會一直變化，不斷保留歷史最小頂點值，直到迭代完成。最後，就可以實現通過connectedComponents得到每個頂點都關聯到最小頂點的數據。

三、Pregel源碼解析

Pregel是一個圖處理模型和計算框架，核心思想是將一系列頂點之間的消息做傳遞和狀態更新操作，並以迭代的方式進行計算。讓我們繼續深入看一下它的底層實現。

以下是保留主要核心代碼的函數——

def apply[VD: ClassTag, ED: ClassTag, A: ClassTag]
   (graph: Graph[VD, ED],
    initialMsg: A,
    maxIterations: Int = Int.MaxValue,
    activeDirection: EdgeDirection = EdgeDirection.Either)
   (vprog: (VertexId, VD, A) => VD,
    sendMsg: EdgeTriplet[VD, ED] => Iterator[(VertexId, A)],
    mergeMsg: (A, A) => A)
  : Graph[VD, ED] =
{
  ......
  //step1
  var g = graph.mapVertices((vid, vdata) => vprog(vid, vdata, initialMsg))
  ......
  //step2
  var messages = GraphXUtils.mapReduceTriplets(g, sendMsg, mergeMsg)
  ......
  //step3
  var activeMessages = messages.count()
  var prevG: Graph[VD, ED] = null
  var i = 0
  //step4
  while (activeMessages > 0 && i < maxIterations) {
    prevG = g
    g = g.joinVertices(messages)(vprog)
    val oldMessages = messages
    messages = GraphXUtils.mapReduceTriplets(
      g, sendMsg, mergeMsg, Some((oldMessages, activeDirection)))
    activeMessages = messages.count()
    i += 1
  }

  g
}

這行 var g = graph.mapVertices((vid, vdata) => vprog(vid, vdata, initialMsg))代碼，需要聯繫到前面傳過來的參數，它的真實面目其實是這樣的——

var g = graph.mapVertices((vid, vdata) => {
  	(id, attr, initialMsg) => math.min(attr, initialMsg)
})

也就是前面step3裏提到的，這裏相當做了0次迭代，將attr當作頂點id關聯的最小頂點，初始化後，attr其實是頂點id本身。

var messages = GraphXUtils.mapReduceTriplets(g, sendMsg, mergeMsg)這行代碼中，主要定義了一個函數sendMsg和調用了aggregateMessagesWithActiveSet方法。

private[graphx] def mapReduceTriplets[VD: ClassTag, ED: ClassTag, A: ClassTag](
    g: Graph[VD, ED],
    mapFunc: EdgeTriplet[VD, ED] => Iterator[(VertexId, A)],
    reduceFunc: (A, A) => A,
    activeSetOpt: Option[(VertexRDD[_], EdgeDirection)] = None): VertexRDD[A] = {
  def sendMsg(ctx: EdgeContext[VD, ED, A]) {
    mapFunc(ctx.toEdgeTriplet).foreach { kv =>
      val id = kv._1
      val msg = kv._2
      if (id == ctx.srcId) {
        ctx.sendToSrc(msg)
      } else {
        assert(id == ctx.dstId)
        ctx.sendToDst(msg)
      }
    }
  }
  g.aggregateMessagesWithActiveSet(
    sendMsg, reduceFunc, TripletFields.All, activeSetOpt)
}

函數sendMsg裏需要看懂一點是，這裏的mapFunc(ctx.toEdgeTriplet)正是調用了前面定義的ConnectedComponents裏的sendMessage方法，因此，這個方法恢復原樣，是這樣的——

    def sendMsg(ctx: EdgeContext[VD, ED, A]) {
      (ctx.toEdgeTriplet => {
        case edge =>
        if (edge.srcAttr < edge.dstAttr) {
          Iterator((edge.dstId, edge.srcAttr))
        } else if (edge.srcAttr > edge.dstAttr) {
          Iterator((edge.srcId, edge.dstAttr))
        } else {
          Iterator.empty
        }
      }).foreach { kv =>
        val id = kv._1
        val msg = kv._2
        if (id == ctx.srcId) {
          ctx.sendToSrc(msg)
        } else {
          assert(id == ctx.dstId)
          ctx.sendToDst(msg)
        }
      }
    }

這個方法的作用，就是找出同一條邊上哪個頂點最小，例如下圖中，2L比3L小，那麼2L是這條邊上最小的頂點，將以最大點關聯最小點的方式(edge.dstId, edge.srcAttr)即（3L，2L）保存下來。最後會將（3L，2L）中的_.2也就是2L發送給頂點（3L，3L），而頂點（3L，3L）後續需要做的事情是，是將這一輪收到的消息即最小頂點2L與現在的屬性3L值通過math.min(a, b)做比較，保留最小頂點當作屬性值，即變成了（3L，2L）。

剩下aggregateMessagesWithActiveSet就是做聚合了，sendMsg就是上面的獲取最小頂點後發送給頂點的操作，reduceFunc對應的是 mergeMsg = (a, b) => math.min(a, b))，保留歷史最小頂點當作該頂點屬性。

g.aggregateMessagesWithActiveSet(
  sendMsg, reduceFunc, TripletFields.All, activeSetOpt)

最後這個while遍歷，如果設置了迭代次數，迭代次數就會傳至給maxIterations，activeMessages表示還有多少頂點需要處理。

  while (activeMessages > 0 && i < maxIterations) {
    prevG = g
    g = g.joinVertices(messages)(vprog)
    val oldMessages = messages
    messages = GraphXUtils.mapReduceTriplets(
      g, sendMsg, mergeMsg, Some((oldMessages, activeDirection)))
    activeMessages = messages.count()
    i += 1
  }

這個方法，就是不斷做迭代，不斷更新各個頂點屬性對應的最小頂點，直到迭代出子圖裏的最小頂點。

很精妙的一點設計是，每個頂點只需要不斷迭代，以三元組邊爲維度，互相將最小頂點發送給屬性值（頂點保留的上一輪最小頂點所做的屬性）較大的頂點，頂點只需要保留收到的消息裏最小的頂點更新爲屬性值即可。