spark性能優化之DataSource表limit操作下推實現kudu limit查詢性能千倍提升

問題描述

我們的 spark基於DataSource V1版本，整合了kudu表，可以直接使用sql操作讀寫kudu表。目前我們的kudu-1.7.0版本，隨着kudu表的使用場景不斷增加，kudu的查詢的性能也暴露出來很多問題。此外，隨着kudu版本的升級，支持了許多新特性。比如，1.9版本的kudu支持了limit操作，且limit的性能非常高，基本不會隨着數據的增長而增長，查詢時間保持在1s以內。
目前，如果spark執行一個select from limit的操作，如果查詢的是kudu表，會進行一個全表掃描，將結果全部返回，在spark這邊在各分區進行locallimit，再將最終結果進行globalLimit， 返回一個limit的結果。 這樣的問題非常明顯：
1.無法利用kudu的limit特性加快查詢
2. 如果kudu源表的數據非常大，查詢時間會隨着數據的增大迅速上升，且非常容易造成OOM
因此，如果能夠在查詢kudu時，將limit操作下推到數據源，利用kudu本身的limit特性，就會非常地快。
我們經過測試，查詢一個10億條數據的kudu表，執行一個limit 操作大概會花掉30-40min的時間，這簡直讓人無法忍受。

解決方案

由於基於DataSource V1版本，因此我們在BaseRelation添加了一個變量，limit，用於記錄是否有limit操作

 var limit: Long = Long.MaxValue

然後我在PushDownOperatorsToDataSource規則中新增了一個方法，用於對DataSource表進行limit下推：

 /**
    * created by XXXX on 2019/07/22
    *
    * to push down limit to datasources.
    * Cases about to [[BaseRelation]] is to push down limit of datasources implemented by datasource API V1
    *  We do not push limit when order by kudu table , because kudu client does not support order by
    *
    * @param logicalPlan
    * @return
    */
  private def pushDownLimitToDataSources(logicalPlan: LogicalPlan): LogicalPlan = logicalPlan.transformDown {
    case l: GlobalLimit => {
      val localLimit = l.child.asInstanceOf[LocalLimit]
      val limitValue = localLimit.limitExpr match {
        case IntegerLiteral(limit) => limit
        case _ => -1
      }
      val newPlan: LogicalPlan = localLimit.child match {
        // A datasource table only with limit specified
        case r@LogicalRelation(baseRelation, _, _, _) =>
          //          logWarning(s"push down limit $limitValue to kudu datasource $r on select * without where")
          baseRelation.limit = limitValue
          r
        //  a data source table with Filter and Limit specified
        case f@Filter(condition, r@LogicalRelation(baseRelation, _, _, _)) =>
          if (supportsAllFiltersAndPredicates(condition, baseRelation)) {
            //            logWarning(s"push down limit $limitValue to kudu datasource $r for select * with where")
            baseRelation.limit = limitValue
          }
          f
        case p@Project(_, Filter(condition, r@LogicalRelation(baseRelation, _, _, _))) =>
          if (supportsAllFiltersAndPredicates(condition, baseRelation)) {
            //            logWarning(s"push down limit $limitValue to kudu datasource $r on select with where ")
            baseRelation.limit = limitValue
          }
          p
        case p@Project(_, r@LogicalRelation(baseRelation, _, _, _)) =>
          //          logWarning(s"push down limit $limitValue to kudu datasource $r on select some columns without where ")
          baseRelation.limit = limitValue
          p

        case p: LogicalPlan =>
          p
      }
      l.copy(l.limitExpr, child = localLimit.copy(child = newPlan))
    }
  }

這裏使用模式匹配對幾種可以進行limit下推的情況進行的匹配，並且，如果DataSource存在謂詞下推的情況，需要判斷其是否支持所有的謂詞下推，如果不支持，該情況下是不能進行limit下推的，否則會導致結果不正確。

我們把它加到規則的最後:

 override def apply(plan: LogicalPlan): LogicalPlan = {
    // Note that, we need to collect the target operator along with PROJECT node, as PROJECT may
    // appear in many places for column pruning.
    // TODO: Ideally column pruning should be implemented via a plan property that is propagated
    // top-down, then we can simplify the logic here and only collect target operators.
    val filterPushed = plan transformUp {
      case FilterAndProject(fields, condition, r@DataSourceV2Relation(_, reader)) =>
        val (candidates, nonDeterministic) =
          splitConjunctivePredicates(condition).partition(_.deterministic)

        val stayUpFilters: Seq[Expression] = reader match {
          case r: SupportsPushDownCatalystFilters =>
            r.pushCatalystFilters(candidates.toArray)

          case r: SupportsPushDownFilters =>
            // A map from original Catalyst expressions to corresponding translated data source
            // filters. If a predicate is not in this map, it means it cannot be pushed down.
            val translatedMap: Map[Expression, sources.Filter] = candidates.flatMap { p =>
              DataSourceStrategy.translateFilter(p).map(f => p -> f)
            }.toMap

            // Catalyst predicate expressions that cannot be converted to data source filters.
            val nonConvertiblePredicates = candidates.filterNot(translatedMap.contains)

            // Data source filters that cannot be pushed down. An unhandled filter means
            // the data source cannot guarantee the rows returned can pass the filter.
            // As a result we must return it so Spark can plan an extra filter operator.
            val unhandledFilters = r.pushFilters(translatedMap.values.toArray).toSet
            val unhandledPredicates = translatedMap.filter { case (_, f) =>
              unhandledFilters.contains(f)
            }.keys

            nonConvertiblePredicates ++ unhandledPredicates

          case _ => candidates
        }

        val filterCondition = (stayUpFilters ++ nonDeterministic).reduceLeftOption(And)
        val withFilter = filterCondition.map(Filter(_, r)).getOrElse(r)
        if (withFilter.output == fields) {
          withFilter
        } else {
          Project(fields, withFilter)
        }
    }

    // TODO: add more push down rules.

    val columnPruned = pushDownRequiredColumns(filterPushed, filterPushed.outputSet)
    //After push down filters, we may push down LIMIT too
    // 這裏調用下推規則
    val limitPushed = pushDownLimitToDataSources(columnPruned)
    // After column pruning, we may have redundant PROJECT nodes in the query plan, remove them.
    RemoveRedundantProject(limitPushed)
  }

上面代碼的倒數第二句，我們進行了limit下推，這樣我們就可以在讀取kudu的時候，拿到limit的設置了，最後，我們在讀取kudu生成kuduRDD時，將limit參數設置到scanner中即可。

測試結果

最終，我們使用下推的後limit操作查詢苦讀， 在查詢10億條數據的kudu表時，用時不會超過1s，而原來需要超過半個小時，這是3000+倍性能提升，如果數據持續增大，性能提升將更大。

總結

最近想寫博客總結下工作中遇到的這些問題，又去翻看了此前這部分的代碼，發現這個方案一個不太好的地方，當時由於工作不久，代碼風格上還沒有養成好的習慣，現在一看，我覺得在一個trait上加一個變量這種編碼方式非常不優雅，更爲優雅的方式，是新增一個抽象類並聲明一個變量，讓BaseRelation繼承這個抽象類，這樣會看起來舒服些。畢竟在接口上定義變量，不是很符合常規的編程習慣。
但是這個功能讓我對於spark的性能優化有了更深入的理解，是一個很好的練手機會，也讓我在工作中受到了正向的肯定，還是一次不出的經歷。