SPARK中的wholeStageCodegen全代码生成--GenerateUnsafeProjection.createCode说明

Posted 2022-09-21 鸿乃江边鸟

背景

对于在在RangeExec中出现的GenerateUnsafeProjection.createCode的方法进行说明

分析

对应的代码为：

  def createCode(
      ctx: CodegenContext,
      expressions: Seq[Expression],
      useSubexprElimination: Boolean = false): ExprCode = 
    val exprEvals = ctx.generateExpressions(expressions, useSubexprElimination)
    val exprSchemas = expressions.map(e => Schema(e.dataType, e.nullable))

    val numVarLenFields = exprSchemas.count 
      case Schema(dt, _) => !UnsafeRow.isFixedLength(dt)
      // TODO: consider large decimal and interval type
    

    val rowWriterClass = classOf[UnsafeRowWriter].getName
    val rowWriter = ctx.addMutableState(rowWriterClass, "rowWriter",
      v => s"$v = new $rowWriterClass($expressions.length, $numVarLenFields * 32);")

    // Evaluate all the subexpression.
    val evalSubexpr = ctx.subexprFunctionsCode

    val writeExpressions = writeExpressionsToBuffer(
      ctx, ctx.INPUT_ROW, exprEvals, exprSchemas, rowWriter, isTopLevel = true)
//   println(s"writeExpressions: $writeExpressions")
    val code =
      code"""
         |$rowWriter.reset();
         |$evalSubexpr
         |$writeExpressions
       """.stripMargin
    // `rowWriter` is declared as a class field, so we can access it directly in methods.
//    println(s"code: $code")
    ExprCode(code, FalseLiteral, JavaCode.expression(s"$rowWriter.getRow()", classOf[UnsafeRow]))
  

其中 expressions的值为Seq(BoundReference(0, long, false))
useSubexpreElimination为false

• val exprEvals = ctx.generateExpressions(expressions, useSubexprElimination)
  这里只是代码生成，exprEvals的值就是range_value_0
  因为useSubexprElimination 是false,所以不会进行公共代码的消除
• val exprSchemas = expressions.map(e => Schema(e.dataType, e.nullable))
  得到对应的表达式的schema
• val numVarLenFields =
  计算出非固定长度字段的个数，用于初始化UnsafeRowWriter
• val rowWriter =
  定义并初始化rowWriter，该rowWriter是全局范围的，生成的代码如下：

   private org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[] range_mutableStateArray_0 = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[3];
   public void init(int index, scala.collection.Iterator[] inputs) 
   ...
   range_mutableStateArray_0[0] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);
   
  

• val evalSubexpr = ctx.subexprFunctionsCode
  这里为空字符串
• val writeExpressions = writeExpressionsToBuffer

  private def writeExpressionsToBuffer(
      ctx: CodegenContext,
      row: String,
      inputs: Seq[ExprCode],
      schemas: Seq[Schema],
      rowWriter: String,
      isTopLevel: Boolean = false): String = 
    val resetWriter = if (isTopLevel) 
      // For top level row writer, it always writes to the beginning of the global buffer holder,
      // which means its fixed-size region always in the same position, so we don't need to call
      // `reset` to set up its fixed-size region every time.
      if (inputs.map(_.isNull).forall(_ == FalseLiteral)) 
        // If all fields are not nullable, which means the null bits never changes, then we don't
        // need to clear it out every time.
        ""
       else 
        s"$rowWriter.zeroOutNullBytes();"
      
     else 
      s"$rowWriter.resetRowWriter();"
    
  
    val writeFields = inputs.zip(schemas).zipWithIndex.map 
      case ((input, Schema(dataType, nullable)), index) =>
        val dt = UserDefinedType.sqlType(dataType)
  
        val setNull = dt match 
          case t: DecimalType if t.precision > Decimal.MAX_LONG_DIGITS =>
            // Can't call setNullAt() for DecimalType with precision larger than 18.
            s"$rowWriter.write($index, (Decimal) null, $t.precision, $t.scale);"
          case CalendarIntervalType => s"$rowWriter.write($index, (CalendarInterval) null);"
          case _ => s"$rowWriter.setNullAt($index);"
        
  
        val writeField = writeElement(ctx, input.value, index.toString, dt, rowWriter)
        if (!nullable) 
          s"""
             |$input.code
             |$writeField.trim
           """.stripMargin
         else 
          s"""
             |$input.code
             |if ($input.isNull) 
             |  $setNull.trim
             | else 
             |  $writeField.trim
             |
           """.stripMargin
        
    
  
    val writeFieldsCode = if (isTopLevel && (row == null || ctx.currentVars != null)) 
      // TODO: support whole stage codegen
      writeFields.mkString("\\n")
     else 
      assert(row != null, "the input row name cannot be null when generating code to write it.")
      ctx.splitExpressions(
        expressions = writeFields,
        funcName = "writeFields",
        arguments = Seq("InternalRow" -> row))
    
    s"""
       |$resetWriter
       |$writeFieldsCode
     """.stripMargin
  
  

  • val resetWriter =
    因为inputs为null为false，所以resetWriter的值为空字符串

  • val writeFields =
    因为inputs的类型是LONG类型，所以对应到val writeField = writeElement(ctx, input.value, index.toString, dt, rowWriter)代码为：
    case _ => s"$writer.write($index, $input);",所以生成的代码为：

     range_mutableStateArray_0[0].write(0, range_value_0)
    

  • val writeFieldsCode =以及后面的代码组装
    对每一个变量的赋值按照换行符进行分隔。

• val code =
  组装成ExprCode的code部分，生成的代码如下：

  range_mutableStateArray_0[0].reset();
  
  range_mutableStateArray_0[0].write(0, range_value_0);
  

最后ExprCode的完整部分如下：

 ExprCode(range_mutableStateArray_0[0].reset();
range_mutableStateArray_0[0].write(0, range_value_0);,false,(range_mutableStateArray_0[0].getRow()))