| /* |
| * 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.flink.api.common.functions.Function |
| import org.apache.flink.streaming.api.functions.ProcessFunction |
| import org.apache.flink.streaming.api.transformations.{OneInputTransformation, StreamTransformation} |
| import org.apache.flink.table.api.types.{DataType, DataTypes, InternalType, RowType, TypeConverters} |
| import org.apache.flink.table.api.{TableConfig, TableEnvironment, TableException, TableSchema} |
| import org.apache.flink.table.calcite.FlinkTypeFactory |
| import org.apache.flink.table.codegen.CodeGenUtils._ |
| import org.apache.flink.table.codegen.operator.OperatorCodeGenerator.generateOneInputStreamOperator |
| import org.apache.flink.table.dataformat.{BaseRow, GenericRow, JoinedRow} |
| import org.apache.flink.table.functions.utils.UserDefinedFunctionUtils.getEvalMethodSignature |
| import org.apache.flink.table.functions.utils.{TableSqlFunction, UserDefinedFunctionUtils} |
| import org.apache.flink.table.plan.nodes.logical.FlinkLogicalTableFunctionScan |
| import org.apache.flink.table.plan.schema.FlinkTableFunction |
| import org.apache.flink.table.plan.util.CorrelateUtil |
| import org.apache.flink.table.runtime.OneInputSubstituteStreamOperator |
| import org.apache.flink.table.runtime.conversion.DataStructureConverters._ |
| import org.apache.flink.table.runtime.util.StreamRecordCollector |
| import org.apache.flink.table.typeutils.{BaseRowTypeInfo, TypeUtils} |
| |
| import org.apache.calcite.rel.`type`.RelDataType |
| import org.apache.calcite.rex._ |
| import org.apache.calcite.sql.SemiJoinType |
| import org.apache.flink.table.util.TableSchemaUtil |
| |
| import scala.collection.JavaConversions._ |
| |
| |
| object CorrelateCodeGenerator { |
| |
| private[flink] def generateCorrelateTransformation( |
| tableEnv: TableEnvironment, |
| operatorCtx: CodeGeneratorContext, |
| inputTransformation: StreamTransformation[BaseRow], |
| inputRelType: RelDataType, |
| projectProgram: Option[RexProgram], |
| scan: FlinkLogicalTableFunctionScan, |
| condition: Option[RexNode], |
| outDataType: RelDataType, |
| joinType: SemiJoinType, |
| parallelism: Int, |
| retainHeader: Boolean, |
| expression: (RexNode, List[String], Option[List[RexNode]]) => String, |
| ruleDescription: String): StreamTransformation[BaseRow] = { |
| val config = tableEnv.getConfig |
| val funcRel = scan.asInstanceOf[FlinkLogicalTableFunctionScan] |
| val rexCall = funcRel.getCall.asInstanceOf[RexCall] |
| val sqlFunction = rexCall.getOperator.asInstanceOf[TableSqlFunction] |
| // we need result Type to do code generation |
| val arguments = UserDefinedFunctionUtils.transformRexNodes(rexCall.operands) |
| val argTypes = getEvalMethodSignature( |
| sqlFunction.getTableFunction, |
| rexCall.operands |
| .map(_.getType) |
| .map(FlinkTypeFactory.toInternalType).toArray) |
| val udtfExternalType = sqlFunction |
| .getFunction |
| .asInstanceOf[FlinkTableFunction] |
| .getExternalResultType(arguments, argTypes) |
| val pojoFieldMapping = Some(UserDefinedFunctionUtils.getFieldInfo(udtfExternalType)._2) |
| val inputType = FlinkTypeFactory.toInternalRowType(inputRelType) |
| val (returnType, swallowInputOnly ) = if (projectProgram.isDefined) { |
| val program = projectProgram.get |
| val selects = program.getProjectList.map(_.getIndex) |
| val inputFieldCnt = program.getInputRowType.getFieldCount |
| val swallowInputOnly = selects(0) > inputFieldCnt && |
| (inputFieldCnt - outDataType.getFieldCount == inputRelType.getFieldCount) |
| // partial output or output right only |
| (FlinkTypeFactory.toInternalRowType(outDataType), swallowInputOnly) |
| } else { |
| // completely output left input + right |
| (FlinkTypeFactory.toInternalRowType(outDataType), false) |
| } |
| // adjust indicies of InputRefs to adhere to schema expected by generator |
| val changeInputRefIndexShuttle = new RexShuttle { |
| override def visitInputRef(inputRef: RexInputRef): RexNode = { |
| new RexInputRef(inputRelType.getFieldCount + inputRef.getIndex, inputRef.getType) |
| } |
| } |
| |
| val collectorCtx = CodeGeneratorContext(config, true) |
| val collector = generateCollector( |
| collectorCtx, |
| config, |
| inputType, |
| projectProgram, |
| swallowInputOnly, |
| udtfExternalType, |
| returnType, |
| condition.map(_.accept(changeInputRefIndexShuttle)), |
| pojoFieldMapping, |
| retainHeader) |
| |
| val substituteStreamOperator = generateOperator( |
| operatorCtx, |
| collectorCtx, |
| config, |
| inputType, |
| projectProgram, |
| swallowInputOnly, |
| udtfExternalType.toInternalType, |
| returnType, |
| joinType, |
| rexCall, |
| pojoFieldMapping, |
| ruleDescription, |
| classOf[ProcessFunction[BaseRow, BaseRow]], |
| collector, |
| retainHeader) |
| |
| new OneInputTransformation( |
| inputTransformation, |
| CorrelateUtil.correlateOpName( |
| inputRelType, |
| rexCall, |
| sqlFunction, |
| outDataType, |
| expression), |
| substituteStreamOperator, |
| TypeConverters.toBaseRowTypeInfo(returnType), |
| parallelism) |
| } |
| |
| /** |
| * Generates the flat map operator to run the user-defined table function. |
| */ |
| private[flink] def generateOperator[T <: Function]( |
| ctx: CodeGeneratorContext, |
| collectorCtx: CodeGeneratorContext, |
| config: TableConfig, |
| inputType: RowType, |
| projectProgram: Option[RexProgram], |
| swallowInputOnly: Boolean = false, |
| udtfType: InternalType, |
| returnType: RowType, |
| joinType: SemiJoinType, |
| rexCall: RexCall, |
| pojoFieldMapping: Option[Array[Int]], |
| ruleDescription: String, |
| functionClass: Class[T], |
| udtfCollector: GeneratedCollector, |
| retainHeader: Boolean = true): OneInputSubstituteStreamOperator[BaseRow, BaseRow] = { |
| ctx.references ++= collectorCtx.references |
| val exprGenerator = new ExprCodeGenerator(ctx, false, config.getNullCheck) |
| .bindInput(inputType) |
| .bindSecondInput(udtfType, inputFieldMapping = pojoFieldMapping) |
| |
| // 1.compile and init udtf collector |
| val udtfCollectorTerm = newName("udtfCollectorTerm") |
| ctx.addReusableMember(s"private ${udtfCollector.name} $udtfCollectorTerm = null;") |
| ctx.addReusableInnerClass(udtfCollector.name, udtfCollector.code) // add a inner class. |
| |
| val call = exprGenerator.generateExpression(rexCall) |
| val openUDTFCollector = |
| s""" |
| |$udtfCollectorTerm = new ${udtfCollector.name}(); |
| |$udtfCollectorTerm.setCollector( |
| | new ${classOf[StreamRecordCollector[_]].getCanonicalName}( |
| | ${CodeGeneratorContext.DEFAULT_OPERATOR_COLLECTOR_TERM })); |
| |${call.resultTerm}.setCollector($udtfCollectorTerm); |
| |""".stripMargin |
| ctx.addReusableOpenStatement(openUDTFCollector) |
| |
| // 2. call udtf |
| var body = |
| s""" |
| |$udtfCollectorTerm.setInput(${exprGenerator.input1Term}); |
| |$udtfCollectorTerm.reset(); |
| |${call.code} |
| |""".stripMargin |
| |
| // 3. left join |
| if (joinType == SemiJoinType.LEFT) { |
| if (swallowInputOnly) { |
| // and the returned row table function is empty, collect a null |
| val nullRowTerm = CodeGenUtils.newName("nullRow") |
| ctx.addOutputRecord(toRowType(udtfType), classOf[GenericRow], nullRowTerm) |
| ctx.addReusableNullRow(nullRowTerm, TypeUtils.getArity(udtfType)) |
| val header = if (retainHeader) { |
| s"$nullRowTerm.setHeader(${exprGenerator.input1Term}.getHeader());" |
| } else { |
| "" |
| } |
| body += |
| s""" |
| |boolean hasOutput = $udtfCollectorTerm.isCollected(); |
| |if (!hasOutput) { |
| | $header |
| | $udtfCollectorTerm.getCollector().collect($nullRowTerm); |
| |} |
| |""".stripMargin |
| } else if (projectProgram.isDefined) { |
| // output partial fields of left and right |
| val outputTerm = CodeGenUtils.newName("projectOut") |
| ctx.addOutputRecord(returnType, classOf[GenericRow], outputTerm) |
| |
| val header = if (retainHeader) { |
| s"$outputTerm.setHeader(${CodeGeneratorContext.DEFAULT_INPUT1_TERM}.getHeader());" |
| } else { |
| "" |
| } |
| val projectionExpression = generateProjectResultExpr( |
| ctx, |
| config, |
| inputType, |
| udtfType, |
| pojoFieldMapping, |
| udtfAlwaysNull = true, |
| returnType, |
| outputTerm, |
| projectProgram.get) |
| |
| body += |
| s""" |
| |boolean hasOutput = $udtfCollectorTerm.isCollected(); |
| |if (!hasOutput) { |
| | ${projectionExpression.code} |
| | $header |
| | $udtfCollectorTerm.getCollector().collect($outputTerm); |
| |} |
| |""".stripMargin |
| |
| } else { |
| // output all fields of left and right |
| // in case of left outer join and the returned row of table function is empty, |
| // fill all fields of row with null |
| val joinedRowTerm = CodeGenUtils.newName("joinedRow") |
| val nullRowTerm = CodeGenUtils.newName("nullRow") |
| ctx.addOutputRecord(returnType, classOf[JoinedRow], joinedRowTerm) |
| ctx.addReusableNullRow(nullRowTerm, TypeUtils.getArity(udtfType)) |
| val header = if (retainHeader) { |
| s"$joinedRowTerm.setHeader(${exprGenerator.input1Term}.getHeader());" |
| } else { |
| "" |
| } |
| body += |
| s""" |
| |boolean hasOutput = $udtfCollectorTerm.isCollected(); |
| |if (!hasOutput) { |
| | $joinedRowTerm.replace(${exprGenerator.input1Term}, $nullRowTerm); |
| | $header |
| | $udtfCollectorTerm.getCollector().collect($joinedRowTerm); |
| |} |
| |""".stripMargin |
| |
| } |
| } else if (joinType != SemiJoinType.INNER) { |
| throw new TableException(s"Unsupported SemiJoinType: $joinType for correlate join.") |
| } |
| |
| val genOperator = generateOneInputStreamOperator[BaseRow, BaseRow]( |
| ctx, |
| ruleDescription, |
| body, |
| "", |
| inputType, |
| config) |
| new OneInputSubstituteStreamOperator[BaseRow, BaseRow]( |
| genOperator.name, |
| genOperator.code, |
| references = ctx.references) |
| } |
| |
| private def toRowType(fromType: InternalType): RowType = { |
| val tableSchema = TableSchemaUtil.fromDataType(fromType) |
| val fieldNames = tableSchema.getColumnNames |
| val fieldTypes = tableSchema.getTypes |
| new RowType(fieldTypes.toArray[DataType], fieldNames) |
| } |
| |
| private def generateProjectResultExpr( |
| ctx: CodeGeneratorContext, |
| config: TableConfig, |
| input1Type: RowType, |
| udtfType: InternalType, |
| udtfPojoFieldMapping: Option[Array[Int]], |
| udtfAlwaysNull: Boolean, |
| returnType: RowType, |
| outputTerm: String, |
| program: RexProgram): GeneratedExpression = { |
| val projectExprGenerator = new ExprCodeGenerator(ctx, udtfAlwaysNull, config.getNullCheck) |
| .bindInput(input1Type, CodeGeneratorContext.DEFAULT_INPUT1_TERM) |
| if (udtfAlwaysNull) { |
| val udtfNullRow = CodeGenUtils.newName("udtfNullRow") |
| ctx.addReusableNullRow(udtfNullRow, TypeUtils.getArity(udtfType)) |
| |
| projectExprGenerator.bindSecondInput( |
| toRowType(udtfType), |
| udtfNullRow, |
| inputFieldMapping = udtfPojoFieldMapping) |
| } else { |
| projectExprGenerator.bindSecondInput( |
| udtfType, |
| inputFieldMapping = udtfPojoFieldMapping) |
| } |
| val projection = program.getProjectList.map(program.expandLocalRef) |
| val projectionExprs = projection.map(projectExprGenerator.generateExpression) |
| projectExprGenerator.generateResultExpression( |
| projectionExprs, returnType, classOf[GenericRow], outputTerm) |
| } |
| |
| /** |
| * Generates table function collector. |
| */ |
| private[flink] def generateCollector( |
| ctx: CodeGeneratorContext, |
| config: TableConfig, |
| inputType: RowType, |
| projectProgram: Option[RexProgram], |
| swallowInputOnly: Boolean, |
| udtfExternalType: DataType, |
| resultType: RowType, |
| condition: Option[RexNode], |
| pojoFieldMapping: Option[Array[Int]], |
| retainHeader: Boolean = true): GeneratedCollector = { |
| val inputTerm = CodeGeneratorContext.DEFAULT_INPUT1_TERM |
| val udtfInputTerm = CodeGeneratorContext.DEFAULT_INPUT2_TERM |
| |
| val udtfType = udtfExternalType.toInternalType |
| val exprGenerator = new ExprCodeGenerator(ctx, false, config.getNullCheck).bindInput( |
| udtfType, inputTerm = udtfInputTerm, inputFieldMapping = pojoFieldMapping) |
| |
| val udtfBaseRowType = toRowType(udtfType) |
| val udtfResultExpr = exprGenerator.generateConverterResultExpression( |
| udtfBaseRowType, classOf[GenericRow]) |
| |
| val body = if (projectProgram.isDefined) { |
| // partial output |
| if (swallowInputOnly) { |
| // output right only |
| val header = if (retainHeader) { |
| s"${udtfResultExpr.resultTerm}.setHeader($inputTerm.getHeader());" |
| } else { |
| "" |
| } |
| s""" |
| |${udtfResultExpr.code} |
| |$header |
| |getCollector().collect(${udtfResultExpr.resultTerm}); |
| """.stripMargin |
| } else { |
| val outputTerm = CodeGenUtils.newName("projectOut") |
| ctx.addOutputRecord(resultType, classOf[GenericRow], outputTerm) |
| |
| val header = if (retainHeader) { |
| s"$outputTerm.setHeader($inputTerm.getHeader());" |
| } else { |
| "" |
| } |
| val projectionExpression = generateProjectResultExpr( |
| ctx, |
| config, |
| inputType, |
| udtfType, |
| pojoFieldMapping, |
| udtfAlwaysNull = false, |
| resultType, |
| outputTerm, |
| projectProgram.get) |
| |
| s""" |
| |$header |
| |${projectionExpression.code} |
| |getCollector().collect(${projectionExpression.resultTerm}); |
| """.stripMargin |
| } |
| } else { |
| // completely output left input + right |
| val joinedRowTerm = CodeGenUtils.newName("joinedRow") |
| ctx.addOutputRecord(resultType, classOf[JoinedRow], joinedRowTerm) |
| |
| val header = if (retainHeader) { |
| s"$joinedRowTerm.setHeader($inputTerm.getHeader());" |
| } else { |
| "" |
| } |
| s""" |
| |${udtfResultExpr.code} |
| |$joinedRowTerm.replace($inputTerm, ${udtfResultExpr.resultTerm}); |
| |$header |
| |getCollector().collect($joinedRowTerm); |
| """.stripMargin |
| } |
| |
| val collectorCode = if (condition.isEmpty) { |
| body |
| } else { |
| |
| val filterGenerator = new ExprCodeGenerator(ctx, false, config.getNullCheck) |
| .bindInput(inputType, inputTerm) |
| .bindSecondInput(udtfType, udtfInputTerm, pojoFieldMapping) |
| val filterCondition = filterGenerator.generateExpression(condition.get) |
| |
| s""" |
| |${filterCondition.code} |
| |if (${filterCondition.resultTerm}) { |
| | $body |
| |} |
| |""".stripMargin |
| } |
| |
| CollectorCodeGenerator.generateTableFunctionCollector( |
| ctx, |
| "TableFunctionCollector", |
| collectorCode, |
| inputType, |
| udtfType, |
| config, |
| inputTerm = inputTerm, |
| collectedTerm = udtfInputTerm, |
| converter = genToInternal(ctx, udtfExternalType)) |
| } |
| |
| } |
