SparkCore实现离线IDMapping

最近在开发一个ID Mapping业务系统——识别数据上报中社交账号的关联关系，找到系统中哪些社交账号属于现实世界中的同一个人。简单来讲，如果同一条上报数据中出现了两个社交账号（比如一个手机号和一个QQ号），就认为这两个社交账号在现实世界属于同一个人。那么，如何计算这个关联关系呢？

一开始我们解决这个问题的思路很直接：现实世界的每个人在系统中用唯一的UUID标识，每次社交账号（Account）上报，在Redis中记录一个UUID->Account的正向索引，同时记录一个Account->UUID的反向索引。每个UUID下可能有多个社交账号，每个社交账号只会归属于一个UUID。每次收到2个社交账号相关的上报时，先通过Account->UUID的反向索引查到这两个社交账号对应的UUID，如果两个账号分别属于两个不同的UUID，就把这两个UUID合并为一个新的UUID，同时原来归属于这两个UUID的所有社交账号都归属于新的UUID；如果查到一个UUID，那么把这两个账号归属于这个UUID；如果未查到UUID，则新生成一个UUID，最后把新的正向索引和反向索引再写回Redis即可。

当然实际的ID Mapping可能有更复杂的关联关系，并不是两个社交账号在同一条上报数据中出现这么简单的逻辑，这就不在本文讨论的范围内了。

这种方式乍一看也没什么问题，但仔细想来却忽略了一个很基本的问题：这种方式只能不断的将账号关联在一起，而不能解除关联。举例来说，给定A-B、B-C两组社交账号的关联关系，A、B、C应该全部关联在一个UUID下，此时若B点被删除，或B-C关联关系解除，系统无法将A和C解除关联。而这种解除关联的场景在业务系统中也是很常见的，比如在企业的客户管理系统中，往往会记录一个客户的一些社交账号，有时企业会删除客户的一些社交账号，甚至删除某一个客户。

其实这个业务抽象成一个数据结构的问题，就是典型的不相交集问题。每一个社交账号看做一个点，两个社交账号的关联关系看做两点之间的无向边，上面的问题就变成了在一个图里划分不相交集的问题。既然很难处理动态变化的图，就每天批量计算一下某一时刻所有账号的关联关系吧。本文就是要介绍一下如何用Spark的RDD API实现静态图不相交集的计算。

一、问题定义

前言描述了问题的背景，这里再明确定义下本文要解决的问题。

算法输入是一张离线Hive表，每行有两个字段Pi、Pj，表示无向图中节点Pi和Pj之间存在一条边。输出同样得到一张Hive表，每行也是Pi、Pj两个字段，Pj均表示不相交集的根节点，即所有Pj相同的行中Pi的集合加上Pj就构成了与图中其他点不相关联的一个独立集合。如图一所示，对于左侧的输入，计算结果将得到右侧的输出，可以看出P1、P2、P3、P6、P8、P9构成一个独立集合，P4、P5、P7构成另一个独立集合。

图一算法输入与输出

二、计算过程

1. 使每行数据中Pi > Pj

为了保证迭代过程最终可以收敛，不妨将图中所有边都当做有向边处理，方向都是节点ID较大的节点指向节点ID较小的节点，这样最终计算得到的不相交集必是以集合中ID最小的点为根，即所有节点都指向所在集合中ID最小的点。因此，不妨将原始数据中的每一行当做由Pi指向Pj的有向边，若Pj>Pi，则交换Pi和Pj。如图二所示，这一步修改了第三行和第五行数据。

图二使Pi > Pj

edge_rdd = edge_rdd.map(row => if (row._1 > row._2) (row._1, row._2) else (row._2. row._1))

2. 保证Pi不重复

第一步处理完之后，原始数据中还会存在一个问题：多条边相交于一个ID较大的节点，这会导致ID较大的节点成为潜在的根节点。解决这个问题需要将局部相交于ID较大节点的边转化为相交于ID最小的节点。比如存在P6->P1和P6->P3两条边，这两条边交于P6，P1、P3、P6组成一个独立集合。需要将这个关联关系转换为P6->P1和P3->P1两条边，即以P1为根节点。前一步的处理已经保证数据中每一行都满足Pi>Pj，因此多条边交于ID较大的节点等价于多行的Pi相同。所以只需要在保证原有关联关系的条件下将表处理为Pi不重复即可。

这一步保证了Pi中的节点不会作为根节点，所有有向边都由叶子节点指向潜在的根节点。

图三保证Pi不重复

在算法实现上，首先将Pi相同的行聚合在一起，然后输出这个集合中每个点指向集合中ID最小点的有向边。

while(edge_rdd.keys.count() != edge_rdd.keys.distinct.count()) {
    edge_rdd = edge_rdd.groupbykey()
        .flatMap(row => {
            val vertex_list = row._2.toSeq.sorted.distinct
            val result = new ArrayBuffer[(String, String)]()
            result.append(row => (row._1, vertex_list.head))
            vertex_list.tail.foreach(vertex => result.append((vertex, vertex_list.head)))
            result
        })
}

从图三种可以看出，每次迭代之后，可能会产生新的Pi相同的数据，因此需要用迭代的方式，多次执行以上逻辑，迭代的终止条件就是Pi中的点不重复。

3. 将每一行中的Pj替换为集合中最小的节点ID

最后一步就是算法的核心，通过自关联，将所有叶子节点关联到根节点上。算法原理很好理解：若同时存在Pz->Py的有向边和Py->Px的有向边，就将Pz->Py替换为Pz->Px。

如图四所示，第一次迭代由P8->P2和P2->P1两条边得到P8->P1，由P9->P8和P8->P2两条边得到P9->P2；第二次迭代由P9->P2和P2->P1两条得到P9->P1。

如果我们定义P2与根节点P1的距离是1，P8和根节点P1的距离是2，P9和根节点P1的距离是3，依次类推。该算法每n次迭代可以将距离为2n之内的节点全部关联到根节点上。

图四将Pj替换为根节点

for (i <- 1 to iterateNum) {
    val inverse_edge = edge_rdd.map(row => (row._2, row._1))
    val edge_rdd = inverse_edge.leftOuterJoin(edge_rdd).map(row => (row._2._1, if (row._2._2.isDefined) row._2._2.get else row._1))
}

通过以上3个步骤的处理，可以看到原始数据集被划分成了2个不相交集，根节点分别为P1和P4。

三、执行优化

问题也并不是这样顺利的就解决了，将上述逻辑转化为工程代码时还遇到了一些其他问题，下面也分享下遇到的问题，以及采取的优化方案。

1. 推测执行

由于数据量比较大，计算时会使用较大的并发来执行，对于同一个job经常出现大部分task都执行很快，只是在等2-3个task的执行，而观察发现并没有明显的数据倾斜。查阅相关资料了解到这种情况可能是由于各个executor上的运行环境不同（有可能同时运行了其他任务、或者是硬件原因），导致计算时间差异较大，这个问题可以通过推测执行解决。

spark.speculation=true
spark.speculation.interval=100
spark.speculation.multiplier=1.5

2. checkpoint

由于该算法是一个迭代算法，执行流程较长，有时会出现某个executor丢失，使RDD中的一部分丢失，导致整个任务需要重新计算，甚至失败。查阅相关资料后，最终通过checkpoint的方式解决了这个问题。checkpoint算子将RDD写入到HDFS某个目录下，因此也一个Action，所以一般会先执行cache再执行checkpoint。

edge_rdd = edge_rdd.groupbykey().flatMap(row => {...}).cache
edge_rdd.checkpoint()

3. RDD cache释放

代码运行过程中还发现任务会占用很多内存，远比预期大的多，通过查看Spark任务的Storage页，发现其实是迭代的方式导致了“内存泄漏”。

在迭代的过程中，算法对每一次迭代得到的edge_rdd进行了cache，而事实上每次计算出新的edge_rdd后，前一次迭代的cache就没用了。但如果没有手动释放的话，这些RDD的cache在任务终止之前都不会被释放掉，会一直占用着集群内存，导致“内存泄漏”。从图五中可以看出，每次迭代都会生成一个RDD，并cache在内存中，如果迭代次数比较多，这部分内存浪费对集群资源的占用就很可观了。甚至如果新的RDD没有内存可以cache，会导致RDD的重复计算，这样会严重影响任务执行的时间。

图五没有手动释放RDD，导致“内存泄漏”

这个问题可以通过在每次计算生成新的RDD时手动unpersist上一个RDD来解决，在内存无效时立刻释放掉这部分内存。

val tmp = edge_rdd.groupbykey().flatMap(row => {...}).cache
tmp.checkpoint()
edge_rdd.unpersist()
edge_rdd = tmp

四、执行性能

用上述算法计算业务收集到的社交账号关联关系，数据量在5000万条左右，第二步计算需要7次迭代收敛，第三步计算需要3-4次迭代收敛。程序运行使用16核64G内存的分布式Spark运行环境，迭代过程中partition个数为64，整体运行时间在20分钟左右，基本达到了业务使用的要求。

五、附加完整实现代码

package com.yang.spark.idmapping

import com.yang.spark.utils.SparkUtils


/**
  * vertices edges
  */
object IncrementIDMapping {

  def main(args: Array[String]): Unit = {

    //最大迭代次数
    val step3MaxIterateNum = 100

    val spark = SparkUtils.initSession(isLocal = false, this.getClass.getSimpleName)

    spark.sqlContext.setConf("spark.sql.adaptive.maxNumPostShufflePartitions", "1000")

    /*val inputRDD: RDD[(String, String)] = spark.sparkContext.makeRDD(Seq(
      ("100", "101"),
      ("100", "105"),
      ("101", "102"),
      ("102", "103"),
      ("104", "105"),
      ("105", "106"),
      ("106", "103"),
      ("107", "108"),
      ("108", "109"),
      ("110", "107"),
      ("111", "112"),
      ("113", "114"),
      ("114", "112"),
      ("115", "116"),
      ("117", "118"),
      ("117", "116")
    ))*/

    val inputRDD = spark
      .sql(
        s"""
           |select id, ccid as id2
           |from hdp_jinrong_tech_dw.dw_wb_ajk_idmapping_output_idrevtable
           |where dt = '20200618'
           |union all
           |select id1, id2
           |from hdp_jinrong_tech_ods.temp_ajk_phone_idmapping_input_data_20200619
        """.stripMargin
      )
      .rdd

    /**
      * 1. 使每行数据中(id1, id2) 满足 id1 > id2
      * 为了保证迭代过程最终可以收敛，不妨将图中所有边都当做有向边处理，方向都是节点ID较大的节点指向节点ID较小的节点，
      * 这样最终计算得到的不相交集必是以集合中ID最小的点为根，即所有节点都指向所在集合中ID最小的点。
      * 因此，不妨将原始数据中的每一行当做由id1指向id2的有向边，若id2 > id1，则交换id1和id2。
      */
    var edgeRDD = inputRDD.repartition(1000)
      .map(row => (row.getAs[String]("id1"), row.getAs[String]("id2")))
      .map { case (id1, id2) => if (id1 > id2) (id1, id2) else (id2, id1) }
    /*println("=================step2: edgeRDD=====================")
    edgeRDD.take(20).foreach(println)*/

    /**
      * 2. 保证id1不重复
      * 第一步处理完之后，原始数据中还会存在一个问题：多条边相交于一个ID较大的节点，
      * 这会导致ID较大的节点成为潜在的根节点。解决这个问题需要将局部相交于ID较大节点的边转化为相交于ID最小的节点。
      * 比如存在P6->P1和P6->P3两条边，这两条边交于P6，P1、P3、P6组成一个独立集合。
      * 需要将这个关联关系转换为P6->P1和P3->P1两条边，即以P1为根节点。
      * 前一步的处理已经保证数据中每一行都满足id1 > id2，因此多条边交于ID较大的节点等价于多行的id1相同。
      * 所以只需要在保证原有关联关系的条件下将表处理为id1不重复即可。
      * 这一步保证了id1中的节点不会作为根节点，所有有向边都由叶子节点指向潜在的根节点。
      */
    var step2IterateNum = 0
    while (edgeRDD.keys.count() != edgeRDD.keys.distinct().count()) {
      step2IterateNum = step2IterateNum + 1
      println(s"================step2: iterateNum: $step2IterateNum======================")
      edgeRDD = edgeRDD.groupByKey(1000)
        .flatMap {
          case (id, ids) =>
            val vertexList = ids.toSeq.sorted
            val head = vertexList.head
            val result = new scala.collection.mutable.HashSet[(String, String)]
            result.add((id, head))
            vertexList.tail.foreach(vertex => result.add((vertex, head)))
            result
        }
    }

    /*println("================step2: edgeRDD======================")
    edgeRDD.take(20).foreach(println)*/

    /**
      * 3. 将每一行中的id2替换为集合中最小的节点ID
      * 第三步就是算法的核心，通过自关联，将所有叶子节点关联到根节点上。
      * 算法原理很好理解：若同时存在id3 -> id2的有向边和id2 -> id1的有向边，就将id3 -> id2替换为id3 -> id1。
      * 第一次迭代由P8->P2和P2->P1两条边得到P8->P1，由P9->P8和P8->P2两条边得到P9->P2；
      * 第二次迭代由P9->P2和P2->P1两条得到P9->P1。
      * 如果我们定义P2与根节点P1的距离是1，P8和根节点P1的距离是2，P9和根节点P1的距离是3，依次类推。
      * 该算法每n次迭代可以将距离为2^^n之内的节点全部关联到根节点上。
      */
    import scala.util.control.Breaks._
    breakable {
      for (step3IterateNum <- 1 to step3MaxIterateNum) {
        println(s"================step3: iterateNum: $step3IterateNum======================")
        val reverseEdgeRDD = edgeRDD.map { case (id1, id2) => (id2, id1) }
        /*println("================step3: reverseEdgeRDD======================")
        reverseEdgeRDD.take(20).foreach(println)*/
        val joinRDD = reverseEdgeRDD.leftOuterJoin(edgeRDD, 1000)
        /*println("================step3: joinRDD======================")
        joinRDD.take(20).foreach(println)*/
        val innerCount = joinRDD.filter{case (k, (_, v2)) => v2.isDefined && k != v2.get}.count()
        if(0 != innerCount) {
          edgeRDD = joinRDD.map { case (k, (v1, v2)) => (v1, if (v2.isDefined) v2.get else k) }
        } else {
          break
        }
      }
    }

    /*println("================step3: edgeRDD======================")
    edgeRDD.take(20).foreach(println)*/

    import spark.implicits._
    val resultDF = edgeRDD.toDF("id1", "id2")
    resultDF.createOrReplaceTempView("idmapping_output_temp_view")

    spark.sql(
      s"""
         |insert overwrite table hdp_jinrong_tech_dw.dw_wb_ajk_idmapping_output_idrevtable
         |partition (dt = '20200619')
         |select *
         |from idmapping_output_temp_view
       """.stripMargin
    )

  }

}