diff --git a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala index 9e66661b1c1f..3c93b17e63ce 100644 --- a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala +++ b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala @@ -18,15 +18,21 @@ package org.apache.spark.sql.execution.datasources.v2 +import org.apache.paimon.spark.util.OptionUtils + import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, BasePredicate, Expression, Projection, UnsafeProjection} -import org.apache.spark.sql.catalyst.expressions.codegen.GeneratePredicate +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, BasePredicate, BindReferences, Expression, Projection, UnsafeProjection} +import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode, FalseLiteral, GeneratePredicate, JavaCode} +import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper import org.apache.spark.sql.catalyst.plans.logical.MergeRows._ import org.apache.spark.sql.catalyst.util.truncatedString -import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.{CodegenSupport, SparkPlan, UnaryExecNode} +import org.apache.spark.sql.types.BooleanType import org.roaringbitmap.longlong.Roaring64Bitmap +import scala.collection.mutable + case class MergeRowsExec( isSourceRowPresent: Expression, isTargetRowPresent: Expression, @@ -36,7 +42,8 @@ case class MergeRowsExec( checkCardinality: Boolean, output: Seq[Attribute], child: SparkPlan) - extends UnaryExecNode { + extends UnaryExecNode + with CodegenSupport { @transient override lazy val producedAttributes: AttributeSet = { AttributeSet(output.filterNot(attr => inputSet.contains(attr))) @@ -66,6 +73,222 @@ case class MergeRowsExec( child.execute().mapPartitions(processPartition) } + override def inputRDDs(): Seq[RDD[InternalRow]] = { + child.asInstanceOf[CodegenSupport].inputRDDs() + } + + override def needCopyResult: Boolean = { + val hasSplitInstruction = (matchedInstructions ++ notMatchedInstructions ++ + notMatchedBySourceInstructions).exists(_.isInstanceOf[Split]) + hasSplitInstruction || child.asInstanceOf[CodegenSupport].needCopyResult + } + + override def supportCodegen: Boolean = { + OptionUtils.mergeCodegenEnabled() && + conf.wholeStageEnabled && + CodeGenerator.isValidParamLength(CodeGenerator.calculateParamLength(child.output)) + } + + override protected def doProduce(ctx: CodegenContext): String = { + child.asInstanceOf[CodegenSupport].produce(ctx, this) + } + + override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = { + val funcName = ctx.freshName("mergeProcessRow") + val (args, params, paramExprs) = constructConsumeParameters(ctx, child.output, input) + val body = generateInstructionExecutionCode(ctx, paramExprs) + val addedFuncName = ctx.addNewFunction( + funcName, + s""" + |private void $funcName(${params.mkString(", ")}) throws java.io.IOException { + | $body + |} + """.stripMargin + ) + + s"$addedFuncName(${args.mkString(", ")});" + } + + private def generateCardinalityValidationCode( + ctx: CodegenContext, + rowIdOrdinal: Int, + input: Seq[ExprCode]): String = { + val bitmapClass = classOf[Roaring64Bitmap] + val rowIdBitmap = ctx.addMutableState( + bitmapClass.getName, + "matchedRowIds", + variable => s"$variable = new ${bitmapClass.getName}();") + val currentRowId = input(rowIdOrdinal) + + code""" + |${currentRowId.code} + |if ($rowIdBitmap.contains(${currentRowId.value})) { + | throw new RuntimeException("Should not happens"); + |} + |$rowIdBitmap.add(${currentRowId.value}); + """.stripMargin.toString + } + + private def generateInstructionExecutionCode( + ctx: CodegenContext, + inputExprs: Seq[ExprCode]): String = { + val sourcePresentExpr = generatePredicateCode(ctx, isSourceRowPresent, child.output, inputExprs) + val targetPresentExpr = generatePredicateCode(ctx, isTargetRowPresent, child.output, inputExprs) + val matchedInstructionsCode = generateInstructionsCode(ctx, matchedInstructions, inputExprs) + val notMatchedInstructionsCode = + generateInstructionsCode(ctx, notMatchedInstructions, inputExprs) + val notMatchedBySourceInstructionsCode = + generateInstructionsCode(ctx, notMatchedBySourceInstructions, inputExprs) + val cardinalityValidationCode = if (checkCardinality) { + val rowIdOrdinal = child.output.indexWhere(attr => conf.resolver(attr.name, ROW_ID)) + assert(rowIdOrdinal != -1, "Cannot find row ID attr") + generateCardinalityValidationCode(ctx, rowIdOrdinal, inputExprs) + } else { + "" + } + + s""" + |${sourcePresentExpr.code} + |${targetPresentExpr.code} + |if (${targetPresentExpr.value} && ${sourcePresentExpr.value}) { + | $cardinalityValidationCode + | $matchedInstructionsCode + |} else if (${sourcePresentExpr.value}) { + | $notMatchedInstructionsCode + |} else if (${targetPresentExpr.value}) { + | $notMatchedBySourceInstructionsCode + |} + """.stripMargin + } + + private def generateInstructionsCode( + ctx: CodegenContext, + instructions: Seq[Instruction], + inputExprs: Seq[ExprCode]): String = { + if (instructions.isEmpty) { + "" + } else { + val instructionCodes = + instructions.map(instruction => generateSingleInstructionCode(ctx, instruction, inputExprs)) + s""" + |${instructionCodes.mkString("\n")} + |return; + """.stripMargin + } + } + + private def generateSingleInstructionCode( + ctx: CodegenContext, + instruction: Instruction, + inputExprs: Seq[ExprCode]): String = { + instruction match { + case Keep(condition, outputExprs) => + val projectionExpr = generateProjectionCode(ctx, outputExprs, inputExprs) + val predicateExpr = generatePredicateCode(ctx, condition, child.output, inputExprs) + s""" + |${predicateExpr.code} + |if (${predicateExpr.value}) { + | ${consume(ctx, projectionExpr)} + | return; + |} + """.stripMargin + + case Discard(condition) => + val predicateExpr = generatePredicateCode(ctx, condition, child.output, inputExprs) + s""" + |${predicateExpr.code} + |if (${predicateExpr.value}) { + | return; + |} + """.stripMargin + + case Split(condition, outputExprs, otherOutputExprs) => + val projectionExpr = generateProjectionCode(ctx, outputExprs, inputExprs) + val otherProjectionExpr = generateProjectionCode(ctx, otherOutputExprs, inputExprs) + val predicateExpr = generatePredicateCode(ctx, condition, child.output, inputExprs) + s""" + |${predicateExpr.code} + |if (${predicateExpr.value}) { + | ${consume(ctx, projectionExpr)} + | ${consume(ctx, otherProjectionExpr)} + | return; + |} + """.stripMargin + + case other => + throw new RuntimeException("Unsupported instruction type: " + other.getClass.getSimpleName) + } + } + + private def withCodegenContext[T](ctx: CodegenContext, inputCurrentVars: Seq[ExprCode])( + block: => T): T = { + val originalCurrentVars = ctx.currentVars + val originalInputRow = ctx.INPUT_ROW + try { + ctx.currentVars = inputCurrentVars + block + } finally { + ctx.currentVars = originalCurrentVars + ctx.INPUT_ROW = originalInputRow + } + } + + private def generatePredicateCode( + ctx: CodegenContext, + predicate: Expression, + inputAttrs: Seq[Attribute], + inputCurrentVars: Seq[ExprCode]): ExprCode = { + withCodegenContext(ctx, inputCurrentVars) { + val boundPredicate = BindReferences.bindReference(predicate, inputAttrs) + val evaluatedPredicate = boundPredicate.genCode(ctx) + val predicateVar = ctx.freshName("predicateResult") + val code = code""" + |${evaluatedPredicate.code} + |boolean $predicateVar = !${evaluatedPredicate.isNull} && + | ${evaluatedPredicate.value}; + """.stripMargin + ExprCode(code, FalseLiteral, JavaCode.variable(predicateVar, BooleanType)) + } + } + + private def generateProjectionCode( + ctx: CodegenContext, + outputExprs: Seq[Expression], + inputCurrentVars: Seq[ExprCode]): Seq[ExprCode] = { + withCodegenContext(ctx, inputCurrentVars) { + val boundExprs = outputExprs.map(BindReferences.bindReference(_, child.output)) + boundExprs.map(_.genCode(ctx)) + } + } + + private def constructConsumeParameters( + ctx: CodegenContext, + attributes: Seq[Attribute], + variables: Seq[ExprCode]): (Seq[String], Seq[String], Seq[ExprCode]) = { + val arguments = mutable.ArrayBuffer[String]() + val parameters = mutable.ArrayBuffer[String]() + val paramVars = mutable.ArrayBuffer[ExprCode]() + + variables.zipWithIndex.foreach { + case (evaluatedVariable, index) => + val paramName = ctx.freshName(s"expr_$index") + val paramType = CodeGenerator.javaType(attributes(index).dataType) + arguments += evaluatedVariable.value.toString + parameters += s"$paramType $paramName" + val paramIsNull = if (!attributes(index).nullable) { + FalseLiteral + } else { + val isNull = ctx.freshName(s"exprIsNull_$index") + arguments += evaluatedVariable.isNull.toString + parameters += s"boolean $isNull" + JavaCode.isNullVariable(isNull) + } + paramVars += ExprCode(paramIsNull, JavaCode.variable(paramName, attributes(index).dataType)) + } + + (arguments.toSeq, parameters.toSeq, paramVars.toSeq) + } + private def processPartition(rowIterator: Iterator[InternalRow]): Iterator[InternalRow] = { val isSourceRowPresentPred = createPredicate(isSourceRowPresent) val isTargetRowPresentPred = createPredicate(isTargetRowPresent) diff --git a/paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTest.scala b/paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTest.scala new file mode 100644 index 000000000000..c91e114c273f --- /dev/null +++ b/paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTest.scala @@ -0,0 +1,30 @@ +/* + * 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.paimon.spark.sql + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.plans.logical.MergeRows.{Instruction, Keep} + +class MergeRowsCodegenTest extends MergeRowsCodegenTestBase { + + override protected def keepInstruction( + condition: Expression, + output: Seq[Expression]): Instruction = + Keep(condition, output) +} diff --git a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala index 9e66661b1c1f..3c93b17e63ce 100644 --- a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala +++ b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala @@ -18,15 +18,21 @@ package org.apache.spark.sql.execution.datasources.v2 +import org.apache.paimon.spark.util.OptionUtils + import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, BasePredicate, Expression, Projection, UnsafeProjection} -import org.apache.spark.sql.catalyst.expressions.codegen.GeneratePredicate +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, BasePredicate, BindReferences, Expression, Projection, UnsafeProjection} +import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode, FalseLiteral, GeneratePredicate, JavaCode} +import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper import org.apache.spark.sql.catalyst.plans.logical.MergeRows._ import org.apache.spark.sql.catalyst.util.truncatedString -import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.{CodegenSupport, SparkPlan, UnaryExecNode} +import org.apache.spark.sql.types.BooleanType import org.roaringbitmap.longlong.Roaring64Bitmap +import scala.collection.mutable + case class MergeRowsExec( isSourceRowPresent: Expression, isTargetRowPresent: Expression, @@ -36,7 +42,8 @@ case class MergeRowsExec( checkCardinality: Boolean, output: Seq[Attribute], child: SparkPlan) - extends UnaryExecNode { + extends UnaryExecNode + with CodegenSupport { @transient override lazy val producedAttributes: AttributeSet = { AttributeSet(output.filterNot(attr => inputSet.contains(attr))) @@ -66,6 +73,222 @@ case class MergeRowsExec( child.execute().mapPartitions(processPartition) } + override def inputRDDs(): Seq[RDD[InternalRow]] = { + child.asInstanceOf[CodegenSupport].inputRDDs() + } + + override def needCopyResult: Boolean = { + val hasSplitInstruction = (matchedInstructions ++ notMatchedInstructions ++ + notMatchedBySourceInstructions).exists(_.isInstanceOf[Split]) + hasSplitInstruction || child.asInstanceOf[CodegenSupport].needCopyResult + } + + override def supportCodegen: Boolean = { + OptionUtils.mergeCodegenEnabled() && + conf.wholeStageEnabled && + CodeGenerator.isValidParamLength(CodeGenerator.calculateParamLength(child.output)) + } + + override protected def doProduce(ctx: CodegenContext): String = { + child.asInstanceOf[CodegenSupport].produce(ctx, this) + } + + override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = { + val funcName = ctx.freshName("mergeProcessRow") + val (args, params, paramExprs) = constructConsumeParameters(ctx, child.output, input) + val body = generateInstructionExecutionCode(ctx, paramExprs) + val addedFuncName = ctx.addNewFunction( + funcName, + s""" + |private void $funcName(${params.mkString(", ")}) throws java.io.IOException { + | $body + |} + """.stripMargin + ) + + s"$addedFuncName(${args.mkString(", ")});" + } + + private def generateCardinalityValidationCode( + ctx: CodegenContext, + rowIdOrdinal: Int, + input: Seq[ExprCode]): String = { + val bitmapClass = classOf[Roaring64Bitmap] + val rowIdBitmap = ctx.addMutableState( + bitmapClass.getName, + "matchedRowIds", + variable => s"$variable = new ${bitmapClass.getName}();") + val currentRowId = input(rowIdOrdinal) + + code""" + |${currentRowId.code} + |if ($rowIdBitmap.contains(${currentRowId.value})) { + | throw new RuntimeException("Should not happens"); + |} + |$rowIdBitmap.add(${currentRowId.value}); + """.stripMargin.toString + } + + private def generateInstructionExecutionCode( + ctx: CodegenContext, + inputExprs: Seq[ExprCode]): String = { + val sourcePresentExpr = generatePredicateCode(ctx, isSourceRowPresent, child.output, inputExprs) + val targetPresentExpr = generatePredicateCode(ctx, isTargetRowPresent, child.output, inputExprs) + val matchedInstructionsCode = generateInstructionsCode(ctx, matchedInstructions, inputExprs) + val notMatchedInstructionsCode = + generateInstructionsCode(ctx, notMatchedInstructions, inputExprs) + val notMatchedBySourceInstructionsCode = + generateInstructionsCode(ctx, notMatchedBySourceInstructions, inputExprs) + val cardinalityValidationCode = if (checkCardinality) { + val rowIdOrdinal = child.output.indexWhere(attr => conf.resolver(attr.name, ROW_ID)) + assert(rowIdOrdinal != -1, "Cannot find row ID attr") + generateCardinalityValidationCode(ctx, rowIdOrdinal, inputExprs) + } else { + "" + } + + s""" + |${sourcePresentExpr.code} + |${targetPresentExpr.code} + |if (${targetPresentExpr.value} && ${sourcePresentExpr.value}) { + | $cardinalityValidationCode + | $matchedInstructionsCode + |} else if (${sourcePresentExpr.value}) { + | $notMatchedInstructionsCode + |} else if (${targetPresentExpr.value}) { + | $notMatchedBySourceInstructionsCode + |} + """.stripMargin + } + + private def generateInstructionsCode( + ctx: CodegenContext, + instructions: Seq[Instruction], + inputExprs: Seq[ExprCode]): String = { + if (instructions.isEmpty) { + "" + } else { + val instructionCodes = + instructions.map(instruction => generateSingleInstructionCode(ctx, instruction, inputExprs)) + s""" + |${instructionCodes.mkString("\n")} + |return; + """.stripMargin + } + } + + private def generateSingleInstructionCode( + ctx: CodegenContext, + instruction: Instruction, + inputExprs: Seq[ExprCode]): String = { + instruction match { + case Keep(condition, outputExprs) => + val projectionExpr = generateProjectionCode(ctx, outputExprs, inputExprs) + val predicateExpr = generatePredicateCode(ctx, condition, child.output, inputExprs) + s""" + |${predicateExpr.code} + |if (${predicateExpr.value}) { + | ${consume(ctx, projectionExpr)} + | return; + |} + """.stripMargin + + case Discard(condition) => + val predicateExpr = generatePredicateCode(ctx, condition, child.output, inputExprs) + s""" + |${predicateExpr.code} + |if (${predicateExpr.value}) { + | return; + |} + """.stripMargin + + case Split(condition, outputExprs, otherOutputExprs) => + val projectionExpr = generateProjectionCode(ctx, outputExprs, inputExprs) + val otherProjectionExpr = generateProjectionCode(ctx, otherOutputExprs, inputExprs) + val predicateExpr = generatePredicateCode(ctx, condition, child.output, inputExprs) + s""" + |${predicateExpr.code} + |if (${predicateExpr.value}) { + | ${consume(ctx, projectionExpr)} + | ${consume(ctx, otherProjectionExpr)} + | return; + |} + """.stripMargin + + case other => + throw new RuntimeException("Unsupported instruction type: " + other.getClass.getSimpleName) + } + } + + private def withCodegenContext[T](ctx: CodegenContext, inputCurrentVars: Seq[ExprCode])( + block: => T): T = { + val originalCurrentVars = ctx.currentVars + val originalInputRow = ctx.INPUT_ROW + try { + ctx.currentVars = inputCurrentVars + block + } finally { + ctx.currentVars = originalCurrentVars + ctx.INPUT_ROW = originalInputRow + } + } + + private def generatePredicateCode( + ctx: CodegenContext, + predicate: Expression, + inputAttrs: Seq[Attribute], + inputCurrentVars: Seq[ExprCode]): ExprCode = { + withCodegenContext(ctx, inputCurrentVars) { + val boundPredicate = BindReferences.bindReference(predicate, inputAttrs) + val evaluatedPredicate = boundPredicate.genCode(ctx) + val predicateVar = ctx.freshName("predicateResult") + val code = code""" + |${evaluatedPredicate.code} + |boolean $predicateVar = !${evaluatedPredicate.isNull} && + | ${evaluatedPredicate.value}; + """.stripMargin + ExprCode(code, FalseLiteral, JavaCode.variable(predicateVar, BooleanType)) + } + } + + private def generateProjectionCode( + ctx: CodegenContext, + outputExprs: Seq[Expression], + inputCurrentVars: Seq[ExprCode]): Seq[ExprCode] = { + withCodegenContext(ctx, inputCurrentVars) { + val boundExprs = outputExprs.map(BindReferences.bindReference(_, child.output)) + boundExprs.map(_.genCode(ctx)) + } + } + + private def constructConsumeParameters( + ctx: CodegenContext, + attributes: Seq[Attribute], + variables: Seq[ExprCode]): (Seq[String], Seq[String], Seq[ExprCode]) = { + val arguments = mutable.ArrayBuffer[String]() + val parameters = mutable.ArrayBuffer[String]() + val paramVars = mutable.ArrayBuffer[ExprCode]() + + variables.zipWithIndex.foreach { + case (evaluatedVariable, index) => + val paramName = ctx.freshName(s"expr_$index") + val paramType = CodeGenerator.javaType(attributes(index).dataType) + arguments += evaluatedVariable.value.toString + parameters += s"$paramType $paramName" + val paramIsNull = if (!attributes(index).nullable) { + FalseLiteral + } else { + val isNull = ctx.freshName(s"exprIsNull_$index") + arguments += evaluatedVariable.isNull.toString + parameters += s"boolean $isNull" + JavaCode.isNullVariable(isNull) + } + paramVars += ExprCode(paramIsNull, JavaCode.variable(paramName, attributes(index).dataType)) + } + + (arguments.toSeq, parameters.toSeq, paramVars.toSeq) + } + private def processPartition(rowIterator: Iterator[InternalRow]): Iterator[InternalRow] = { val isSourceRowPresentPred = createPredicate(isSourceRowPresent) val isTargetRowPresentPred = createPredicate(isTargetRowPresent) diff --git a/paimon-spark/paimon-spark-3.3/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTest.scala b/paimon-spark/paimon-spark-3.3/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTest.scala new file mode 100644 index 000000000000..c91e114c273f --- /dev/null +++ b/paimon-spark/paimon-spark-3.3/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTest.scala @@ -0,0 +1,30 @@ +/* + * 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.paimon.spark.sql + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.plans.logical.MergeRows.{Instruction, Keep} + +class MergeRowsCodegenTest extends MergeRowsCodegenTestBase { + + override protected def keepInstruction( + condition: Expression, + output: Seq[Expression]): Instruction = + Keep(condition, output) +} diff --git a/paimon-spark/paimon-spark-3.4/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala b/paimon-spark/paimon-spark-3.4/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala index 9e66661b1c1f..3c93b17e63ce 100644 --- a/paimon-spark/paimon-spark-3.4/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala +++ b/paimon-spark/paimon-spark-3.4/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MergeRowsExec.scala @@ -18,15 +18,21 @@ package org.apache.spark.sql.execution.datasources.v2 +import org.apache.paimon.spark.util.OptionUtils + import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, BasePredicate, Expression, Projection, UnsafeProjection} -import org.apache.spark.sql.catalyst.expressions.codegen.GeneratePredicate +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, BasePredicate, BindReferences, Expression, Projection, UnsafeProjection} +import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode, FalseLiteral, GeneratePredicate, JavaCode} +import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper import org.apache.spark.sql.catalyst.plans.logical.MergeRows._ import org.apache.spark.sql.catalyst.util.truncatedString -import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.{CodegenSupport, SparkPlan, UnaryExecNode} +import org.apache.spark.sql.types.BooleanType import org.roaringbitmap.longlong.Roaring64Bitmap +import scala.collection.mutable + case class MergeRowsExec( isSourceRowPresent: Expression, isTargetRowPresent: Expression, @@ -36,7 +42,8 @@ case class MergeRowsExec( checkCardinality: Boolean, output: Seq[Attribute], child: SparkPlan) - extends UnaryExecNode { + extends UnaryExecNode + with CodegenSupport { @transient override lazy val producedAttributes: AttributeSet = { AttributeSet(output.filterNot(attr => inputSet.contains(attr))) @@ -66,6 +73,222 @@ case class MergeRowsExec( child.execute().mapPartitions(processPartition) } + override def inputRDDs(): Seq[RDD[InternalRow]] = { + child.asInstanceOf[CodegenSupport].inputRDDs() + } + + override def needCopyResult: Boolean = { + val hasSplitInstruction = (matchedInstructions ++ notMatchedInstructions ++ + notMatchedBySourceInstructions).exists(_.isInstanceOf[Split]) + hasSplitInstruction || child.asInstanceOf[CodegenSupport].needCopyResult + } + + override def supportCodegen: Boolean = { + OptionUtils.mergeCodegenEnabled() && + conf.wholeStageEnabled && + CodeGenerator.isValidParamLength(CodeGenerator.calculateParamLength(child.output)) + } + + override protected def doProduce(ctx: CodegenContext): String = { + child.asInstanceOf[CodegenSupport].produce(ctx, this) + } + + override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = { + val funcName = ctx.freshName("mergeProcessRow") + val (args, params, paramExprs) = constructConsumeParameters(ctx, child.output, input) + val body = generateInstructionExecutionCode(ctx, paramExprs) + val addedFuncName = ctx.addNewFunction( + funcName, + s""" + |private void $funcName(${params.mkString(", ")}) throws java.io.IOException { + | $body + |} + """.stripMargin + ) + + s"$addedFuncName(${args.mkString(", ")});" + } + + private def generateCardinalityValidationCode( + ctx: CodegenContext, + rowIdOrdinal: Int, + input: Seq[ExprCode]): String = { + val bitmapClass = classOf[Roaring64Bitmap] + val rowIdBitmap = ctx.addMutableState( + bitmapClass.getName, + "matchedRowIds", + variable => s"$variable = new ${bitmapClass.getName}();") + val currentRowId = input(rowIdOrdinal) + + code""" + |${currentRowId.code} + |if ($rowIdBitmap.contains(${currentRowId.value})) { + | throw new RuntimeException("Should not happens"); + |} + |$rowIdBitmap.add(${currentRowId.value}); + """.stripMargin.toString + } + + private def generateInstructionExecutionCode( + ctx: CodegenContext, + inputExprs: Seq[ExprCode]): String = { + val sourcePresentExpr = generatePredicateCode(ctx, isSourceRowPresent, child.output, inputExprs) + val targetPresentExpr = generatePredicateCode(ctx, isTargetRowPresent, child.output, inputExprs) + val matchedInstructionsCode = generateInstructionsCode(ctx, matchedInstructions, inputExprs) + val notMatchedInstructionsCode = + generateInstructionsCode(ctx, notMatchedInstructions, inputExprs) + val notMatchedBySourceInstructionsCode = + generateInstructionsCode(ctx, notMatchedBySourceInstructions, inputExprs) + val cardinalityValidationCode = if (checkCardinality) { + val rowIdOrdinal = child.output.indexWhere(attr => conf.resolver(attr.name, ROW_ID)) + assert(rowIdOrdinal != -1, "Cannot find row ID attr") + generateCardinalityValidationCode(ctx, rowIdOrdinal, inputExprs) + } else { + "" + } + + s""" + |${sourcePresentExpr.code} + |${targetPresentExpr.code} + |if (${targetPresentExpr.value} && ${sourcePresentExpr.value}) { + | $cardinalityValidationCode + | $matchedInstructionsCode + |} else if (${sourcePresentExpr.value}) { + | $notMatchedInstructionsCode + |} else if (${targetPresentExpr.value}) { + | $notMatchedBySourceInstructionsCode + |} + """.stripMargin + } + + private def generateInstructionsCode( + ctx: CodegenContext, + instructions: Seq[Instruction], + inputExprs: Seq[ExprCode]): String = { + if (instructions.isEmpty) { + "" + } else { + val instructionCodes = + instructions.map(instruction => generateSingleInstructionCode(ctx, instruction, inputExprs)) + s""" + |${instructionCodes.mkString("\n")} + |return; + """.stripMargin + } + } + + private def generateSingleInstructionCode( + ctx: CodegenContext, + instruction: Instruction, + inputExprs: Seq[ExprCode]): String = { + instruction match { + case Keep(condition, outputExprs) => + val projectionExpr = generateProjectionCode(ctx, outputExprs, inputExprs) + val predicateExpr = generatePredicateCode(ctx, condition, child.output, inputExprs) + s""" + |${predicateExpr.code} + |if (${predicateExpr.value}) { + | ${consume(ctx, projectionExpr)} + | return; + |} + """.stripMargin + + case Discard(condition) => + val predicateExpr = generatePredicateCode(ctx, condition, child.output, inputExprs) + s""" + |${predicateExpr.code} + |if (${predicateExpr.value}) { + | return; + |} + """.stripMargin + + case Split(condition, outputExprs, otherOutputExprs) => + val projectionExpr = generateProjectionCode(ctx, outputExprs, inputExprs) + val otherProjectionExpr = generateProjectionCode(ctx, otherOutputExprs, inputExprs) + val predicateExpr = generatePredicateCode(ctx, condition, child.output, inputExprs) + s""" + |${predicateExpr.code} + |if (${predicateExpr.value}) { + | ${consume(ctx, projectionExpr)} + | ${consume(ctx, otherProjectionExpr)} + | return; + |} + """.stripMargin + + case other => + throw new RuntimeException("Unsupported instruction type: " + other.getClass.getSimpleName) + } + } + + private def withCodegenContext[T](ctx: CodegenContext, inputCurrentVars: Seq[ExprCode])( + block: => T): T = { + val originalCurrentVars = ctx.currentVars + val originalInputRow = ctx.INPUT_ROW + try { + ctx.currentVars = inputCurrentVars + block + } finally { + ctx.currentVars = originalCurrentVars + ctx.INPUT_ROW = originalInputRow + } + } + + private def generatePredicateCode( + ctx: CodegenContext, + predicate: Expression, + inputAttrs: Seq[Attribute], + inputCurrentVars: Seq[ExprCode]): ExprCode = { + withCodegenContext(ctx, inputCurrentVars) { + val boundPredicate = BindReferences.bindReference(predicate, inputAttrs) + val evaluatedPredicate = boundPredicate.genCode(ctx) + val predicateVar = ctx.freshName("predicateResult") + val code = code""" + |${evaluatedPredicate.code} + |boolean $predicateVar = !${evaluatedPredicate.isNull} && + | ${evaluatedPredicate.value}; + """.stripMargin + ExprCode(code, FalseLiteral, JavaCode.variable(predicateVar, BooleanType)) + } + } + + private def generateProjectionCode( + ctx: CodegenContext, + outputExprs: Seq[Expression], + inputCurrentVars: Seq[ExprCode]): Seq[ExprCode] = { + withCodegenContext(ctx, inputCurrentVars) { + val boundExprs = outputExprs.map(BindReferences.bindReference(_, child.output)) + boundExprs.map(_.genCode(ctx)) + } + } + + private def constructConsumeParameters( + ctx: CodegenContext, + attributes: Seq[Attribute], + variables: Seq[ExprCode]): (Seq[String], Seq[String], Seq[ExprCode]) = { + val arguments = mutable.ArrayBuffer[String]() + val parameters = mutable.ArrayBuffer[String]() + val paramVars = mutable.ArrayBuffer[ExprCode]() + + variables.zipWithIndex.foreach { + case (evaluatedVariable, index) => + val paramName = ctx.freshName(s"expr_$index") + val paramType = CodeGenerator.javaType(attributes(index).dataType) + arguments += evaluatedVariable.value.toString + parameters += s"$paramType $paramName" + val paramIsNull = if (!attributes(index).nullable) { + FalseLiteral + } else { + val isNull = ctx.freshName(s"exprIsNull_$index") + arguments += evaluatedVariable.isNull.toString + parameters += s"boolean $isNull" + JavaCode.isNullVariable(isNull) + } + paramVars += ExprCode(paramIsNull, JavaCode.variable(paramName, attributes(index).dataType)) + } + + (arguments.toSeq, parameters.toSeq, paramVars.toSeq) + } + private def processPartition(rowIterator: Iterator[InternalRow]): Iterator[InternalRow] = { val isSourceRowPresentPred = createPredicate(isSourceRowPresent) val isTargetRowPresentPred = createPredicate(isTargetRowPresent) diff --git a/paimon-spark/paimon-spark-3.4/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTest.scala b/paimon-spark/paimon-spark-3.4/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTest.scala new file mode 100644 index 000000000000..c91e114c273f --- /dev/null +++ b/paimon-spark/paimon-spark-3.4/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTest.scala @@ -0,0 +1,30 @@ +/* + * 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.paimon.spark.sql + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.plans.logical.MergeRows.{Instruction, Keep} + +class MergeRowsCodegenTest extends MergeRowsCodegenTestBase { + + override protected def keepInstruction( + condition: Expression, + output: Seq[Expression]): Instruction = + Keep(condition, output) +} diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java index 4dd9329d1c4c..682e09f82624 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java @@ -92,6 +92,13 @@ public class SparkConnectorOptions { + "dynamic partition columns. If false, the query output follows " + "the table schema order."); + public static final ConfigOption MERGE_CODEGEN_ENABLED = + key("write.merge.codegen.enabled") + .booleanType() + .defaultValue(false) + .withDescription( + "Whether to enable whole-stage code generation for merge row processing."); + public static final ConfigOption DATA_EVOLUTION_UPDATE_CONFLICT_RETRY_MAX_ATTEMPTS = key("write.data-evolution.update-conflict-retry.max-attempts") .intType() diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/util/OptionUtils.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/util/OptionUtils.scala index 1649a57eadb6..1cd2df01a1fc 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/util/OptionUtils.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/util/OptionUtils.scala @@ -114,6 +114,10 @@ object OptionUtils extends SQLConfHelper with Logging { getOptionString(SparkConnectorOptions.HIVE_STYLE_DYNAMIC_PARTITION_ENABLED).toBoolean } + def mergeCodegenEnabled(): Boolean = { + getOptionString(SparkConnectorOptions.MERGE_CODEGEN_ENABLED).toBoolean + } + def writeMergeSchemaExplicitCastEnabled(): Boolean = { getOptionString(SparkConnectorOptions.EXPLICIT_CAST).toBoolean } diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTestBase.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTestBase.scala new file mode 100644 index 000000000000..a45e91502b9b --- /dev/null +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/MergeRowsCodegenTestBase.scala @@ -0,0 +1,95 @@ +/* + * 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.paimon.spark.sql + +import org.apache.paimon.spark.{PaimonSparkTestBase, SparkConnectorOptions} + +import org.apache.spark.sql.{PaimonUtils, Row} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression, LessThan, Literal} +import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral +import org.apache.spark.sql.catalyst.plans.logical.MergeRows +import org.apache.spark.sql.catalyst.plans.logical.MergeRows.{Discard, Instruction, Split} +import org.apache.spark.sql.execution.WholeStageCodegenExec +import org.apache.spark.sql.execution.datasources.v2.MergeRowsExec +import org.apache.spark.sql.types.IntegerType + +abstract class MergeRowsCodegenTestBase extends PaimonSparkTestBase { + + import testImplicits._ + + protected def keepInstruction(condition: Expression, output: Seq[Expression]): Instruction + + test("merge row codegen requires Spark and Paimon flags") { + assert(!SparkConnectorOptions.MERGE_CODEGEN_ENABLED.defaultValue()) + + val paimonCodegenKey = + s"spark.paimon.${SparkConnectorOptions.MERGE_CODEGEN_ENABLED.key()}" + + Seq( + (false, true, false), + (true, false, false), + (true, true, true) + ).foreach { + case (paimonCodegenEnabled, sparkCodegenEnabled, expectedCodegen) => + withSparkSQLConf( + paimonCodegenKey -> paimonCodegenEnabled.toString, + "spark.sql.codegen.wholeStage" -> sparkCodegenEnabled.toString) { + val input = Seq( + (1, 10, true, true, 101L), + (2, 20, true, false, 102L), + (3, 30, false, true, 103L), + (4, 40, false, false, 104L), + (5, 50, true, true, 105L) + ).toDF("target_id", "source_value", "source_present", "target_present", MergeRows.ROW_ID) + val inputPlan = input.queryExecution.analyzed + val targetId = inputPlan.output(0) + val sourceValue = inputPlan.output(1) + val sourcePresent = inputPlan.output(2) + val targetPresent = inputPlan.output(3) + val output = Seq( + AttributeReference("id", IntegerType, nullable = false)(), + AttributeReference("value", IntegerType, nullable = false)()) + val mergeRows = MergeRows( + isSourceRowPresent = sourcePresent, + isTargetRowPresent = targetPresent, + matchedInstructions = Seq( + keepInstruction(LessThan(sourceValue, Literal(15)), Seq(targetId, sourceValue)), + Discard(TrueLiteral)), + notMatchedInstructions = + Seq(keepInstruction(TrueLiteral, Seq(sourceValue, sourceValue))), + notMatchedBySourceInstructions = + Seq(Split(TrueLiteral, Seq(targetId, Literal(-1)), Seq(targetId, Literal(-2)))), + checkCardinality = true, + output = output, + child = inputPlan + ) + val result = PaimonUtils.createDataset(spark, mergeRows) + val executedPlan = result.queryExecution.executedPlan + val mergeRowsIsCodegen = executedPlan.collectFirst { + case stage: WholeStageCodegenExec + if stage.collectFirst { case _: MergeRowsExec => true }.isDefined => + true + }.isDefined + + assert(mergeRowsIsCodegen == expectedCodegen, executedPlan) + checkAnswer(result, Seq(Row(1, 10), Row(20, 20), Row(3, -1), Row(3, -2))) + } + } + } +}