简介

Spark是目前最流行的分布式大数据批处理框架，使用Spark可以轻易地实现上百G甚至T级别数据的SQL运算，例如单行特征计算或者多表的Join拼接。

OpenMLDB是针对AI场景优化的开源数据库项目，实现了数据与计算一致性的离线MPP场景和在线OLTP场景计算引擎。其实MPP引擎可基于Spark实现，并通过拓展Spark源码实现数倍性能提升。

Spark本身实现也非常高效，基于Antlr实现的了标准ANSI SQL的词法解析、语法分析，还有在Catalyst模块中实现大量SQL静态优化，然后转成分布式RDD计算，底层数据结构是使用了Java Unsafe API来自定义内存分布的UnsafeRow，还依赖Janino JIT编译器为计算方法动态生成优化后的JVM bytecode。但在拓展性上仍有改进空间，尤其针对机器学习计算场景的需求虽能满足但不高效，本文以LastJoin为例介绍OpenMLDB如何通过拓展Spark源码来实现数倍甚至数十倍性能提升。

机器学习场景LastJoin

LastJoin是一种AI场景引入的特殊拼表类型，是LeftJoin的变种，在满足Join条件的前提下，左表的每一行只拼取右表符合一提交的最后一行。LastJoin的语义特性，可以保证拼表后输出结果的行数与输入的左表一致。在机器学习场景中就是维持了输入的样本表数量一致，不会因为拼表等数据操作导致最终的样本数量增加或者减少，这种方式对在线服务支持比较友好也更符合科学家建模需求。

OpenMLDB: 拓展Spark源码实现高性能Join_OpenMLDB

以技术保护的角度考虑，LastJoin的设计和实现均为第四范式（北京）技术有限公司的专利，公开号为111611245A，公开日为2020-09-01。包含LastJoin功能的OpenMLDB项目代码以Apache 2.0协议在Github中开源，所有用户都可放心使用。

OpenMLDB: 拓展Spark源码实现高性能Join_第四范式_02

基于Spark的LastJoin实现

由于LastJoin类型并非ANSI SQL中的标准，因此在SparkSQL等主流计算平台中都没有实现，为了实现类似功能用户只能通过更底层的DataFrame或RDD等算子来实现。基于Spark算子实现LastJoin的思路是首先对左表添加索引列，然后使用标准LeftOuterJoin，最后对拼接结果进行reduce和去掉索引行，虽然可以实现LastJoin语义但性能还是有很大瓶颈。

相比于兼容SQL功能和语法，Spark的另一个特点是用户可以通过map、reduce、groupby等接口和自定义UDF的方式来实现标准SQL所不支持的数值计算逻辑。但Join功能用户却无法通过DataFrame或者RDD API来拓展实现，因为拼表的实现是在Spark Catalyst物理节点中实现的，涉及了shuffle后多个internal row的拼接，以及生成Java源码字符串进行JIT的过程，而且根据不同的输入表数据量，Spark内部会适时选择BrocastHashJoin、SortMergeJoin或ShuffleHashJoin来实现，普通用户无法用RDD API来拓展这些拼表实现算法。

在OpenMLDB项目中可以查看完整的Spark LastJoin实现，代码地址为 https://github.com/4paradigm/OpenMLDB/blob/main/java/openmldb-batch/src/main/scala/com/_4paradigm/openmldb/batch/nodes/JoinPlan.scala 。

第一步是对输入的左表进行索引列扩充，扩充方式有多种实现，只要添加的索引列每一行有unique id即可，下面是第一步的实现代码。

  // Add the index column for Spark DataFrame
  def addIndexColumn(spark: SparkSession, df: DataFrame, indexColName: String, method: String): DataFrame = {
    logger.info("Add the indexColName(%s) to Spark DataFrame(%s)".format(indexColName, df.toString()))

    method.toLowerCase() match {
      case "zipwithuniqueid" | "zip_withunique_id" => addColumnByZipWithUniqueId(spark, df, indexColName)
      case "zipwithindex" | "zip_with_index" => addColumnByZipWithIndex(spark, df, indexColName)
      case "monotonicallyincreasingid" | "monotonically_increasing_id" =>
        addColumnByMonotonicallyIncreasingId(spark, df, indexColName)
      case _ => throw new HybridSeException("Unsupported add index column method: " + method)
    }

  }

  def addColumnByZipWithUniqueId(spark: SparkSession, df: DataFrame, indexColName: String = null): DataFrame = {
    logger.info("Use zipWithUniqueId to generate index column")
    val indexedRDD = df.rdd.zipWithUniqueId().map {
      case (row, id) => Row.fromSeq(row.toSeq :+ id)
    }
    spark.createDataFrame(indexedRDD, df.schema.add(indexColName, LongType))
  }

  def addColumnByZipWithIndex(spark: SparkSession, df: DataFrame, indexColName: String = null): DataFrame = {
    logger.info("Use zipWithIndex to generate index column")
    val indexedRDD = df.rdd.zipWithIndex().map {
      case (row, id) => Row.fromSeq(row.toSeq :+ id)
    }
    spark.createDataFrame(indexedRDD, df.schema.add(indexColName, LongType))
  }

  def addColumnByMonotonicallyIncreasingId(spark: SparkSession,
                                           df: DataFrame, indexColName: String = null): DataFrame = {
    logger.info("Use monotonicallyIncreasingId to generate index column")
    df.withColumn(indexColName, monotonically_increasing_id())
  }

更多文章请关注《万象专栏》