| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| package org.apache.flink.table.codegen |
| |
| import org.apache.calcite.rex._ |
| import org.apache.calcite.sql.`type`.{ReturnTypes, SqlTypeName} |
| import org.apache.flink.streaming.api.functions.ProcessFunction |
| import org.apache.flink.table.api.TableException |
| import org.apache.flink.table.api.types._ |
| import org.apache.flink.table.calcite.{FlinkTypeFactory, RexAggBufferVariable, RexAggLocalVariable, RexDistinctKeyVariable} |
| import org.apache.flink.table.codegen.CodeGenUtils._ |
| import org.apache.flink.table.codegen.GeneratedExpression.{NEVER_NULL, NO_CODE} |
| import org.apache.flink.table.dataformat._ |
| import org.apache.flink.table.functions.sql.{ProctimeSqlFunction, StreamRecordTimestampSqlFunction} |
| import org.apache.flink.table.typeutils.TypeUtils |
| import org.apache.flink.util.Preconditions.checkArgument |
| |
| import scala.collection.JavaConversions._ |
| |
| /** |
| * This code generator is mainly responsible for generating codes for a given calcite [[RexNode]]. |
| * It can also generate type conversion codes for the result converter. |
| */ |
| class ExprCodeGenerator(ctx: CodeGeneratorContext, nullableInput: Boolean, nullCheck: Boolean) |
| extends RexVisitor[GeneratedExpression] { |
| |
| // check if nullCheck is enabled when inputs can be null |
| if (nullableInput && !nullCheck) { |
| throw new CodeGenException("Null check must be enabled if entire rows can be null.") |
| } |
| |
| /** |
| * term of the [[ProcessFunction]]'s context, can be changed when needed |
| */ |
| var contextTerm = "ctx" |
| |
| /** |
| * information of the first input |
| */ |
| var input1Type: InternalType = _ |
| var input1Term: String = _ |
| var input1FieldMapping: Option[Array[Int]] = None |
| |
| /** |
| * information of the optional second input |
| */ |
| var input2Type: Option[InternalType] = None |
| var input2Term: Option[String] = None |
| var input2FieldMapping: Option[Array[Int]] = None |
| |
| /** |
| * Bind the input information, should be called before generating expression. |
| */ |
| def bindInput( |
| inputType: InternalType, |
| inputTerm: String = CodeGeneratorContext.DEFAULT_INPUT1_TERM, |
| inputFieldMapping: Option[Array[Int]] = None): ExprCodeGenerator = { |
| input1Type = inputType |
| input1Term = inputTerm |
| input1FieldMapping = inputFieldMapping |
| this |
| } |
| |
| /** |
| * In some cases, the expression will have two inputs (e.g. join condition and udtf). We should |
| * bind second input information before use. |
| */ |
| def bindSecondInput( |
| inputType: InternalType, |
| inputTerm: String = CodeGeneratorContext.DEFAULT_INPUT2_TERM, |
| inputFieldMapping: Option[Array[Int]] = None): ExprCodeGenerator = { |
| input2Type = Some(inputType) |
| input2Term = Some(inputTerm) |
| input2FieldMapping = inputFieldMapping |
| this |
| } |
| |
| protected lazy val input1Mapping: Array[Int] = input1FieldMapping match { |
| case Some(mapping) => mapping |
| case _ => (0 until TypeUtils.getArity(input1Type)).toArray |
| } |
| |
| protected lazy val input2Mapping: Array[Int] = input2FieldMapping match { |
| case Some(mapping) => mapping |
| case _ => input2Type match { |
| case Some(input) => (0 until TypeUtils.getArity(input)).toArray |
| case _ => Array[Int]() |
| } |
| } |
| |
| /** |
| * Generates an expression from a RexNode. If objects or variables can be reused, they will be |
| * added to reusable code sections internally. |
| * |
| * @param rex Calcite row expression |
| * @return instance of GeneratedExpression |
| */ |
| def generateExpression(rex: RexNode): GeneratedExpression = { |
| rex.accept(this) |
| } |
| |
| /** |
| * Generates an expression that converts the first input (and second input) into the given type. |
| * If two inputs are converted, the second input is appended. If objects or variables can |
| * be reused, they will be added to reusable code sections internally. The evaluation result |
| * will be stored in the variable outRecordTerm. |
| * |
| * @param returnType conversion target type. Inputs and output must have the same arity. |
| * @param outRecordTerm the result term |
| * @param outRecordWriterTerm the result writer term |
| * @param reusedOutRow If objects or variables can be reused, they will be added to reusable |
| * code sections internally. |
| * @return instance of GeneratedExpression |
| */ |
| def generateConverterResultExpression( |
| returnType: RowType, |
| returnTypeClazz: Class[_ <: BaseRow], |
| outRecordTerm: String = CodeGeneratorContext.DEFAULT_OUT_RECORD_TERM, |
| outRecordWriterTerm: String = CodeGeneratorContext.DEFAULT_OUT_RECORD_WRITER_TERM, |
| reusedOutRow: Boolean = true, |
| fieldCopy: Boolean = false, |
| rowtimeExpression: Option[RexNode] = None) |
| : GeneratedExpression = { |
| val input1AccessExprs = input1Mapping.map { |
| case DataTypes.ROWTIME_STREAM_MARKER | |
| DataTypes.ROWTIME_BATCH_MARKER if rowtimeExpression.isDefined => |
| // generate rowtime attribute from expression |
| generateExpression(rowtimeExpression.get) |
| case DataTypes.ROWTIME_STREAM_MARKER | |
| DataTypes.ROWTIME_BATCH_MARKER => |
| throw new TableException("Rowtime extraction expression missing. Please report a bug.") |
| case DataTypes.PROCTIME_STREAM_MARKER => |
| // attribute is proctime indicator. |
| // we use a null literal and generate a timestamp when we need it. |
| generateNullLiteral(DataTypes.PROCTIME_INDICATOR, nullCheck) |
| case DataTypes.PROCTIME_BATCH_MARKER => |
| // attribute is proctime field in a batch query. |
| // it is initialized with the current time. |
| generateCurrentTimestamp(ctx) |
| case idx => |
| // get type of result field |
| generateInputAccess( |
| ctx, |
| input1Type, |
| input1Term, |
| idx, |
| nullableInput, |
| nullCheck, |
| fieldCopy) |
| } |
| |
| val input2AccessExprs = input2Type match { |
| case Some(ti) => |
| input2Mapping.map(idx => generateInputAccess( |
| ctx, |
| ti, |
| input2Term.get, |
| idx, |
| nullableInput, |
| nullCheck) |
| ).toSeq |
| case None => Seq() // add nothing |
| } |
| |
| generateResultExpression( |
| input1AccessExprs ++ input2AccessExprs, |
| returnType, |
| returnTypeClazz, |
| outRow = outRecordTerm, |
| outRowWriter = Some(outRecordWriterTerm), |
| reusedOutRow = reusedOutRow) |
| } |
| |
| /** |
| * Generates an expression from a sequence of other expressions. The evaluation result |
| * may be stored in the variable outRecordTerm. |
| * |
| * @param fieldExprs field expressions to be converted |
| * @param returnType conversion target type. Type must have the same arity than fieldExprs. |
| * @param outRow the result term |
| * @param outRowWriter the result writer term for BinaryRow. |
| * @param reusedOutRow If objects or variables can be reused, they will be added to reusable |
| * code sections internally. |
| * @param outRowAlreadyExists Don't need addReusableRecord if out row already exists. |
| * @return instance of GeneratedExpression |
| */ |
| def generateResultExpression( |
| fieldExprs: Seq[GeneratedExpression], |
| returnType: RowType, |
| returnTypeClazz: Class[_ <: BaseRow], |
| outRow: String = CodeGeneratorContext.DEFAULT_OUT_RECORD_TERM, |
| outRowWriter: Option[String] = Some(CodeGeneratorContext.DEFAULT_OUT_RECORD_WRITER_TERM), |
| reusedOutRow: Boolean = true, |
| outRowAlreadyExists: Boolean = false): GeneratedExpression = { |
| val fieldExprIdxToOutputRowPosMap = fieldExprs.indices.map(i => i -> i).toMap |
| generateResultExpression(fieldExprs, fieldExprIdxToOutputRowPosMap, returnType, |
| returnTypeClazz, outRow, outRowWriter, reusedOutRow, outRowAlreadyExists) |
| } |
| |
| /** |
| * Generates an expression from a sequence of other expressions. The evaluation result |
| * may be stored in the variable outRecordTerm. |
| * |
| * @param fieldExprs field expressions to be converted |
| * @param fieldExprIdxToOutputRowPosMap Mapping index of fieldExpr in `fieldExprs` |
| * to position of output row. |
| * @param returnType conversion target type. Type must have the same arity than fieldExprs. |
| * @param outRow the result term |
| * @param outRowWriter the result writer term for BinaryRow. |
| * @param reusedOutRow If objects or variables can be reused, they will be added to reusable |
| * code sections internally. |
| * @param outRowAlreadyExists Don't need addReusableRecord if out row already exists. |
| * @return instance of GeneratedExpression |
| */ |
| def generateResultExpression( |
| fieldExprs: Seq[GeneratedExpression], |
| fieldExprIdxToOutputRowPosMap: Map[Int, Int], |
| returnType: RowType, |
| returnTypeClazz: Class[_ <: BaseRow], |
| outRow: String, |
| outRowWriter: Option[String], |
| reusedOutRow: Boolean, |
| outRowAlreadyExists: Boolean) |
| : GeneratedExpression = { |
| // initial type check |
| if (returnType.getArity != fieldExprs.length) { |
| throw new CodeGenException( |
| s"Arity [${returnType.getArity}] of result type [$returnType] does not match " + |
| s"number [${fieldExprs.length}] of expressions [$fieldExprs].") |
| } |
| if (fieldExprIdxToOutputRowPosMap.size != fieldExprs.length) { |
| throw new CodeGenException( |
| s"Size [${returnType.getArity}] of fieldExprIdxToOutputRowPosMap does not match " + |
| s"number [${fieldExprs.length}] of expressions [$fieldExprs].") |
| } |
| // type check |
| fieldExprs.zipWithIndex foreach { |
| // timestamp type(Include TimeIndicator) and generic type can compatible with each other. |
| case (fieldExpr, i) |
| if fieldExpr.resultType.isInstanceOf[GenericType[_]] || |
| fieldExpr.resultType.isInstanceOf[TimestampType] => |
| if (returnType.getInternalTypeAt(i).getClass != fieldExpr.resultType.getClass |
| && !returnType.getInternalTypeAt(i).isInstanceOf[GenericType[_]]) { |
| throw new CodeGenException( |
| s"Incompatible types of expression and result type, Expression[$fieldExpr] type is " + |
| s"[${fieldExpr.resultType}], result type is [${returnType.getInternalTypeAt(i)}]") |
| } |
| case (fieldExpr, i) if fieldExpr.resultType != returnType.getInternalTypeAt(i) => |
| throw new CodeGenException( |
| s"Incompatible types of expression and result type. Expression[$fieldExpr] type is " + |
| s"[${fieldExpr.resultType}], result type is [${returnType.getInternalTypeAt(i)}]") |
| case _ => // ok |
| } |
| |
| def getOutputRowPos(fieldExprIdx: Int): Int = |
| fieldExprIdxToOutputRowPosMap.getOrElse(fieldExprIdx, |
| throw new CodeGenException(s"Illegal field expr index: $fieldExprIdx")) |
| |
| def objectArrayRowWrite( |
| genUpdate: (Int, GeneratedExpression) => String): GeneratedExpression = { |
| val initReturnRecord = if (outRowAlreadyExists) { |
| "" |
| } else { |
| ctx.addOutputRecord(returnType, returnTypeClazz, outRow, reused = reusedOutRow) |
| } |
| val resultBuffer: Seq[String] = fieldExprs.zipWithIndex map { |
| case (fieldExpr, i) => |
| val idx = getOutputRowPos(i) |
| if (nullCheck) { |
| s""" |
| |${fieldExpr.code} |
| |if (${fieldExpr.nullTerm}) { |
| | $outRow.setNullAt($idx); |
| |} else { |
| | ${genUpdate(idx, fieldExpr)}; |
| |} |
| |""".stripMargin |
| } |
| else { |
| s""" |
| |${fieldExpr.code} |
| |${genUpdate(idx, fieldExpr)}; |
| |""".stripMargin |
| } |
| } |
| |
| val statement = |
| s""" |
| |$initReturnRecord |
| |${resultBuffer.mkString("")} |
| """.stripMargin.trim |
| |
| GeneratedExpression(outRow, "false", statement, returnType, |
| codeBuffer = resultBuffer, preceding = s"$initReturnRecord") |
| } |
| |
| returnTypeClazz match { |
| case cls if cls == classOf[BinaryRow] => |
| outRowWriter match { |
| case Some(writer) => // binary row writer. |
| val initReturnRecord = if (outRowAlreadyExists) { |
| "" |
| } else { |
| ctx.addOutputRecord(returnType, returnTypeClazz, outRow, outRowWriter, reusedOutRow) |
| } |
| val resetWriter = if (nullCheck) s"$writer.reset();" else s"$writer.resetCursor();" |
| val completeWriter: String = s"$writer.complete();" |
| val resultBuffer = fieldExprs.zipWithIndex map { |
| case (fieldExpr, i) => |
| val t = returnType.getInternalTypeAt(i) |
| val idx = getOutputRowPos(i) |
| val writeCode = binaryWriterWriteField(ctx, idx, fieldExpr.resultTerm, writer, t) |
| if (nullCheck) { |
| s""" |
| |${fieldExpr.code} |
| |if (${fieldExpr.nullTerm}) { |
| | ${binaryWriterWriteNull(idx, writer, t)}; |
| |} else { |
| | $writeCode; |
| |} |
| |""".stripMargin.trim |
| } else { |
| s""" |
| |${fieldExpr.code} |
| |$writeCode; |
| |""".stripMargin.trim |
| } |
| } |
| |
| val statement = |
| s""" |
| |$initReturnRecord |
| |$resetWriter |
| |${resultBuffer.mkString("\n")} |
| |$completeWriter |
| |""".stripMargin.trim |
| GeneratedExpression(outRow, "false", statement, returnType, |
| codeBuffer = resultBuffer, |
| preceding = s"$initReturnRecord\n$resetWriter", |
| flowing = s"$completeWriter") |
| |
| case None => // update to binary row (setXXX). |
| checkArgument(outRowAlreadyExists) |
| val resultBuffer = fieldExprs.zipWithIndex map { |
| case (fieldExpr, i) => |
| val t = returnType.getInternalTypeAt(i) |
| val idx = getOutputRowPos(i) |
| val writeCode = binaryRowFieldSetAccess(idx, outRow, t, fieldExpr.resultTerm) |
| if (nullCheck) { |
| s""" |
| |${fieldExpr.code} |
| |if (${fieldExpr.nullTerm}) { |
| | ${binaryRowSetNull(idx, outRow, t)}; |
| |} else { |
| | $writeCode; |
| |} |
| |""".stripMargin.trim |
| } else { |
| s""" |
| |${fieldExpr.code} |
| |$writeCode; |
| |""".stripMargin.trim |
| } |
| } |
| GeneratedExpression( |
| outRow, "false", resultBuffer.mkString(""), returnType) |
| } |
| |
| case cls if cls == classOf[GenericRow] => |
| objectArrayRowWrite((i: Int, expr: GeneratedExpression) => |
| s"$outRow.update($i, ${expr.resultTerm})") |
| |
| case cls if cls == classOf[BoxedWrapperRow] => |
| objectArrayRowWrite((i: Int, expr: GeneratedExpression) => |
| boxedWrapperRowFieldUpdateAccess(i, expr.resultTerm, outRow, expr.resultType)) |
| } |
| } |
| |
| override def visitInputRef(inputRef: RexInputRef): GeneratedExpression = { |
| // if inputRef index is within size of input1 we work with input1, input2 otherwise |
| val input = if (inputRef.getIndex < TypeUtils.getArity(input1Type)) { |
| (input1Type, input1Term) |
| } else { |
| (input2Type.getOrElse(throw new CodeGenException("Invalid input access.")), |
| input2Term.getOrElse(throw new CodeGenException("Invalid input access."))) |
| } |
| |
| val index = if (input._2 == input1Term) { |
| inputRef.getIndex |
| } else { |
| inputRef.getIndex - TypeUtils.getArity(input1Type) |
| } |
| |
| generateInputAccess(ctx, input._1, input._2, index, nullableInput, nullCheck) |
| } |
| |
| override def visitTableInputRef(rexTableInputRef: RexTableInputRef): GeneratedExpression = |
| visitInputRef(rexTableInputRef) |
| |
| override def visitFieldAccess(rexFieldAccess: RexFieldAccess): GeneratedExpression = { |
| val refExpr = rexFieldAccess.getReferenceExpr.accept(this) |
| val index = rexFieldAccess.getField.getIndex |
| val fieldAccessExpr = generateFieldAccess( |
| ctx, |
| refExpr.resultType, |
| refExpr.resultTerm, |
| index, |
| nullCheck) |
| |
| val resultTypeTerm = primitiveTypeTermForType(fieldAccessExpr.resultType) |
| val defaultValue = primitiveDefaultValue(fieldAccessExpr.resultType) |
| val Seq(resultTerm, nullTerm) = ctx.newReusableFields( |
| Seq("result", "isNull"), |
| Seq(resultTypeTerm, "boolean")) |
| |
| val resultCode = if (nullCheck) { |
| s""" |
| |${refExpr.code} |
| |if (${refExpr.nullTerm}) { |
| | $resultTerm = $defaultValue; |
| | $nullTerm = true; |
| |} |
| |else { |
| | ${fieldAccessExpr.code} |
| | $resultTerm = ${fieldAccessExpr.resultTerm}; |
| | $nullTerm = ${fieldAccessExpr.nullTerm}; |
| |} |
| |""".stripMargin |
| } else { |
| s""" |
| |${refExpr.code} |
| |${fieldAccessExpr.code} |
| |$resultTerm = ${fieldAccessExpr.resultTerm}; |
| |""".stripMargin |
| } |
| |
| GeneratedExpression(resultTerm, nullTerm, resultCode, fieldAccessExpr.resultType) |
| } |
| |
| override def visitLiteral(literal: RexLiteral): GeneratedExpression = { |
| val resultType = FlinkTypeFactory.toInternalType(literal.getType) |
| val value = literal.getValue3 |
| generateLiteral(ctx, literal.getType, resultType, value, nullCheck) |
| } |
| |
| override def visitCorrelVariable(correlVariable: RexCorrelVariable): GeneratedExpression = { |
| GeneratedExpression(input1Term, NEVER_NULL, NO_CODE, input1Type) |
| } |
| |
| override def visitLocalRef(localRef: RexLocalRef): GeneratedExpression = localRef match { |
| case localVar: RexAggBufferVariable => |
| val resultTerm = localVar.getName |
| val nullTerm = resultTerm + "IsNull" |
| val pType = primitiveTypeTermForType(localVar.internalType) |
| ctx.addReusableMember(s"$pType $resultTerm;") |
| ctx.addReusableMember(s"boolean $nullTerm;") |
| GeneratedExpression(resultTerm, nullTerm, "", localVar.internalType) |
| case local: RexAggLocalVariable => |
| GeneratedExpression(local.fieldTerm, local.nullTerm, NO_CODE, local.internalType) |
| case value: RexDistinctKeyVariable => |
| val inputExpr = ctx.getReusableInputUnboxingExprs(input1Term, 0) match { |
| case Some(expr) => expr |
| case None => |
| val pType = primitiveTypeTermForType(value.internalType) |
| val defaultValue = primitiveDefaultValue(value.internalType) |
| val resultTerm = newName("field") |
| val nullTerm = newName("isNull") |
| val code = |
| s""" |
| |$pType $resultTerm = $defaultValue; |
| |boolean $nullTerm = true; |
| |if ($input1Term != null) { |
| | $nullTerm = false; |
| | $resultTerm = ($pType) $input1Term; |
| |} |
| """.stripMargin |
| val expr = GeneratedExpression(resultTerm, nullTerm, code, value.internalType) |
| ctx.addReusableInputUnboxingExprs(input1Term, 0, expr) |
| expr |
| } |
| // hide the generated code as it will be executed only once |
| GeneratedExpression(inputExpr.resultTerm, inputExpr.nullTerm, NO_CODE, inputExpr.resultType) |
| case _ => throw new CodeGenException("Local variables are not supported yet.") |
| } |
| |
| override def visitRangeRef(rangeRef: RexRangeRef): GeneratedExpression = |
| throw new CodeGenException("Range references are not supported yet.") |
| |
| override def visitDynamicParam(dynamicParam: RexDynamicParam): GeneratedExpression = |
| throw new CodeGenException("Dynamic parameter references are not supported yet.") |
| |
| override def visitCall(call: RexCall): GeneratedExpression = { |
| |
| // special case: time materialization |
| if (call.getOperator == ProctimeSqlFunction) { |
| return generateProctimeTimestamp(contextTerm, ctx) |
| } |
| |
| if (call.getOperator == StreamRecordTimestampSqlFunction) { |
| return generateRowtimeAccess(contextTerm, ctx) |
| } |
| |
| val resultType = FlinkTypeFactory.toInternalType(call.getType) |
| |
| // convert operands and help giving untyped NULL literals a type |
| val operands = call.getOperands.zipWithIndex.map { |
| |
| // this helps e.g. for AS(null) |
| // we might need to extend this logic in case some rules do not create typed NULLs |
| case (operandLiteral: RexLiteral, 0) if |
| operandLiteral.getType.getSqlTypeName == SqlTypeName.NULL && |
| call.getOperator.getReturnTypeInference == ReturnTypes.ARG0 => |
| generateNullLiteral(resultType, nullCheck) |
| |
| case (o@_, idx) => o.accept(this) |
| } |
| |
| generateCallExpression( |
| ctx, call.getOperator, operands, resultType, nullCheck) |
| } |
| |
| override def visitOver(over: RexOver): GeneratedExpression = |
| throw new CodeGenException("Aggregate functions over windows are not supported yet.") |
| |
| override def visitSubQuery(subQuery: RexSubQuery): GeneratedExpression = |
| throw new CodeGenException("Subqueries are not supported yet.") |
| |
| override def visitPatternFieldRef(fieldRef: RexPatternFieldRef): GeneratedExpression = |
| throw new CodeGenException("Pattern field references are not supported yet.") |
| } |
