SPARK中的wholeStageCodegen全代码生成--以aggregate代码生成为例说起（5）

  test("SortAggregate should be included in WholeStageCodegen") {
    val df = spark.range(10).agg(max(col("id")), avg(col("id")))
    withSQLConf("spark.sql.test.forceApplySortAggregate" -> "true") {
      val plan = df.queryExecution.executedPlan
      assert(plan.exists(p =>
        p.isInstanceOf[WholeStageCodegenExec] &&
          p.asInstanceOf[WholeStageCodegenExec].child.isInstanceOf[SortAggregateExec]))
      assert(df.collect() === Array(Row(9, 4.5)))
    }
  }
该sql形成的执行计划第一部分的全代码生成部分如下：

 WholeStageCodegenExec      SortAggregateExec(partial)     RangeExec        
  =========================================================================
  -> execute()
      |
   doExecute() --------->   inputRDDs() -----------------> inputRDDs() 
      |
   doCodeGen()
      |
      +----------------->   produce()
                              |
                           doProduce() 
                              |
                           doProduceWithoutKeys() -------> produce()
                                                              |
                                                          doProduce()
                                                              |
                           doConsume()<------------------- consume()
                              |
                           doConsumeWithoutKeys()
                              |并不是doConsumeWithoutKeys调用consume,而是由doProduceWithoutKeys调用
   doConsume()  <--------  consume()

override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = {
    if (groupingExpressions.isEmpty) {
      doConsumeWithoutKeys(ctx, input)
    } else {
      doConsumeWithKeys(ctx, input)
    }
  }

  private def doConsumeWithoutKeys(ctx: CodegenContext, input: Seq[ExprCode]): String = {
    // only have DeclarativeAggregate
    val functions = aggregateExpressions.map(_.aggregateFunction.asInstanceOf[DeclarativeAggregate])
    val inputAttrs = functions.flatMap(_.aggBufferAttributes) ++ inputAttributes
    // To individually generate code for each aggregate function, an element in `updateExprs` holds
    // all the expressions for the buffer of an aggregation function.
    val updateExprs = aggregateExpressions.map { e =>
      e.mode match {
        case Partial | Complete =>
          e.aggregateFunction.asInstanceOf[DeclarativeAggregate].updateExpressions
        case PartialMerge | Final =>
          e.aggregateFunction.asInstanceOf[DeclarativeAggregate].mergeExpressions
      }
    }
    ctx.currentVars = bufVars.flatten ++ input
    println(s"updateExprs: $updateExprs")
    val boundUpdateExprs = updateExprs.map { updateExprsForOneFunc =>
      bindReferences(updateExprsForOneFunc, inputAttrs)
    }
    val subExprs = ctx.subexpressionEliminationForWholeStageCodegen(boundUpdateExprs.flatten)
    val effectiveCodes = ctx.evaluateSubExprEliminationState(subExprs.states.values)
    val bufferEvals = boundUpdateExprs.map { boundUpdateExprsForOneFunc =>
      ctx.withSubExprEliminationExprs(subExprs.states) {
        boundUpdateExprsForOneFunc.map(_.genCode(ctx))
      }
    }
    val aggNames = functions.map(_.prettyName)
    val aggCodeBlocks = bufferEvals.zipWithIndex.map { case (bufferEvalsForOneFunc, i) =>
      val bufVarsForOneFunc = bufVars(i)
      // All the update code for aggregation buffers should be placed in the end
      // of each aggregation function code.
      println(s"bufVarsForOneFunc: $bufVarsForOneFunc")
      val updates = bufferEvalsForOneFunc.zip(bufVarsForOneFunc).map { case (ev, bufVar) =>
        s"""
           |${bufVar.isNull} = ${ev.isNull};
           |${bufVar.value} = ${ev.value};
         """.stripMargin
      }
      code"""
            |${ctx.registerComment(s"do aggregate for ${aggNames(i)}")}
            |${ctx.registerComment("evaluate aggregate function")}
            |${evaluateVariables(bufferEvalsForOneFunc)}
            |${ctx.registerComment("update aggregation buffers")}
            |${updates.mkString("\n").trim}
       """.stripMargin
    }
    val codeToEvalAggFuncs = generateEvalCodeForAggFuncs(
      ctx, input, inputAttrs, boundUpdateExprs, aggNames, aggCodeBlocks, subExprs)
    s"""
       |// do aggregate
       |// common sub-expressions
       |$effectiveCodes
       |// evaluate aggregate functions and update aggregation buffers
       |$codeToEvalAggFuncs
     """.stripMargin
  }

    Add(
    sum,
    coalesce(child.cast(sumDataType), Literal.default(sumDataType)),
    failOnError = useAnsiAdd),
  /* count = */ If(child.isNull, count, count + 1L)

List(ExprCode(boolean sortAgg_isNull_7 = true;
     double sortAgg_value_7 = -1.0;
     if (!sortAgg_bufIsNull_1) {
       sortAgg_sortAgg_isNull_9_0 = true;
    double sortAgg_value_9 = -1.0;
    do {
    boolean sortAgg_isNull_10 = false;
    double sortAgg_value_10 = -1.0;
    if (!false) {
     sortAgg_value_10 = (double) sortAgg_expr_0_0;
    }
   if (!sortAgg_isNull_10) {
     sortAgg_sortAgg_isNull_9_0 = false;
     sortAgg_value_9 = sortAgg_value_10;
     continue;
   }
   if (!false) {
     sortAgg_sortAgg_isNull_9_0 = false;
     sortAgg_value_9 = 0.0D;
     continue;
   }
   } while (false);
   sortAgg_isNull_7 = false; // resultCode could change nullability.
   sortAgg_value_7 = sortAgg_bufValue_1 + sortAgg_value_9;
           },sortAgg_isNull_7,sortAgg_value_7), 
   ExprCode(boolean sortAgg_isNull_13 = false;
   long sortAgg_value_13 = -1L;
   if (!false && false) {
     sortAgg_isNull_13 = sortAgg_bufIsNull_2;
     sortAgg_value_13 = sortAgg_bufValue_2;
   } else {
     boolean sortAgg_isNull_17 = true;
     long sortAgg_value_17 = -1L;
     if (!sortAgg_bufIsNull_2) {
   sortAgg_isNull_17 = false; // resultCode could change nullability.
   sortAgg_value_17 = sortAgg_bufValue_2 + 1L;
           }
     sortAgg_isNull_13 = sortAgg_isNull_17;
     sortAgg_value_13 = sortAgg_value_17;
   },sortAgg_isNull_13,sortAgg_value_13))

  sortAgg_bufIsNull_1 = sortAgg_isNull_7;
  sortAgg_bufValue_1 = sortAgg_value_7;
  sortAgg_bufIsNull_2 = sortAgg_isNull_13;
  sortAgg_bufValue_2 = sortAgg_value_13;

     sortAgg_doAggregate_max_0(sortAgg_expr_0_0);
     sortAgg_doAggregate_avg_0(sortAgg_expr_0_0);