Keep

Author: cxzl25

paimon-spark/paimon-spark-4.1/pom.xml

@@ -153,4 +153,3 @@ under the License.
         </plugins>
     </build>
 </project>
-

paimon-spark/paimon-spark-4.1/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/MergePaimonScalarSubqueries.scala

@@ -90,4 +90,3 @@ object MergePaimonScalarSubqueries extends MergePaimonScalarSubqueriesBase {
     ScalarSubquery(plan, exprId = exprId)
   }
 }
-

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/KeepHelperBase.scala

@@ -0,0 +1,91 @@
+/*
+ * 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.commands
+
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.plans.logical.MergeRows.{Instruction, Keep}
+
+import scala.util.Try
+
+/**
+ * Helper object for creating MergeRows.Keep instructions compatible with different Spark versions.
+ *
+ * Different Spark versions have different Keep signatures:
+ *   - Spark 4.0: Keep(condition, output)
+ *   - Spark 4.1+: Keep(context, condition, output)
+ *
+ * This object uses reflection to determine which signature to use.
+ */
+object KeepHelper {
+
+  // Check if Keep has 3-parameter constructor (Spark 4.1+)
+  private lazy val keepHas3Params: Boolean = {
+    Try {
+      classOf[Keep].getConstructors.exists(_.getParameterCount == 3)
+    }.getOrElse(false)
+  }
+
+  // Context objects for Spark 4.1+
+  private lazy val updateContext: Option[AnyRef] = getContextObject("Update")
+  private lazy val copyContext: Option[AnyRef] = getContextObject("Copy")
+  private lazy val insertContext: Option[AnyRef] = getContextObject("Insert")
+
+  private def getContextObject(name: String): Option[AnyRef] = {
+    Try {
+      val contextClass =
+        Class.forName(s"org.apache.spark.sql.catalyst.plans.logical.MergeRows$$$name$$")
+      contextClass.getField("MODULE$").get(null).asInstanceOf[AnyRef]
+    }.toOption
+  }
+
+  private def createKeepWithContext(
+      context: AnyRef,
+      condition: Expression,
+      output: Seq[Expression]): Instruction = {
+    val constructor = classOf[Keep].getConstructors.find(_.getParameterCount == 3).get
+    constructor.newInstance(context, condition, output).asInstanceOf[Instruction]
+  }
+
+  /** Creates a Keep instruction for Update context (matched rows being updated). */
+  def createKeepForUpdate(condition: Expression, output: Seq[Expression]): Instruction = {
+    if (keepHas3Params && updateContext.isDefined) {
+      createKeepWithContext(updateContext.get, condition, output)
+    } else {
+      Keep(condition, output)
+    }
+  }
+
+  /** Creates a Keep instruction for Copy context (matched rows being copied as-is). */
+  def createKeepForCopy(condition: Expression, output: Seq[Expression]): Instruction = {
+    if (keepHas3Params && copyContext.isDefined) {
+      createKeepWithContext(copyContext.get, condition, output)
+    } else {
+      Keep(condition, output)
+    }
+  }
+
+  /** Creates a Keep instruction for Insert context (not matched rows being inserted). */
+  def createKeepForInsert(condition: Expression, output: Seq[Expression]): Instruction = {
+    if (keepHas3Params && insertContext.isDefined) {
+      createKeepWithContext(insertContext.get, condition, output)
+    } else {
+      Keep(condition, output)
+    }
+  }
+}

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/MergeIntoPaimonDataEvolutionTable.scala

@@ -36,7 +36,6 @@ import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, Equ
 import org.apache.spark.sql.catalyst.expressions.Literal.{FalseLiteral, TrueLiteral}
 import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftOuter}
 import org.apache.spark.sql.catalyst.plans.logical._
-import org.apache.spark.sql.catalyst.plans.logical.MergeRows.Keep
 import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
 import org.apache.spark.sql.functions.{col, udf}
 import org.apache.spark.sql.types.StructType
@@ -322,10 +321,12 @@ case class MergeIntoPaimonDataEvolutionTable(
         matchedInstructions = rewrittenUpdateActions
           .map(
             action => {
-              Keep(action.condition.getOrElse(TrueLiteral), action.assignments.map(a => a.value))
-            }) ++ Seq(Keep(TrueLiteral, output)),
+              KeepHelper.createKeepForUpdate(
+                action.condition.getOrElse(TrueLiteral),
+                action.assignments.map(a => a.value))
+            }) ++ Seq(KeepHelper.createKeepForCopy(TrueLiteral, output)),
         notMatchedInstructions = Nil,
-        notMatchedBySourceInstructions = Seq(Keep(TrueLiteral, output)),
+        notMatchedBySourceInstructions = Seq(KeepHelper.createKeepForCopy(TrueLiteral, output)),
         checkCardinality = false,
         output = output,
         child = readPlan
@@ -361,10 +362,12 @@ case class MergeIntoPaimonDataEvolutionTable(
         matchedInstructions = realUpdateActions
           .map(
             action => {
-              Keep(action.condition.getOrElse(TrueLiteral), action.assignments.map(a => a.value))
-            }) ++ Seq(Keep(TrueLiteral, output)),
+              KeepHelper.createKeepForUpdate(
+                action.condition.getOrElse(TrueLiteral),
+                action.assignments.map(a => a.value))
+            }) ++ Seq(KeepHelper.createKeepForCopy(TrueLiteral, output)),
         notMatchedInstructions = Nil,
-        notMatchedBySourceInstructions = Seq(Keep(TrueLiteral, output)).toSeq,
+        notMatchedBySourceInstructions = Seq(KeepHelper.createKeepForCopy(TrueLiteral, output)),
         checkCardinality = false,
         output = output,
         child = joinPlan
@@ -399,7 +402,7 @@ case class MergeIntoPaimonDataEvolutionTable(
       matchedInstructions = Nil,
       notMatchedInstructions = notMatchedActions.map {
         case insertAction: InsertAction =>
-          Keep(
+          KeepHelper.createKeepForInsert(
             insertAction.condition.getOrElse(TrueLiteral),
             insertAction.assignments.map(
               a =>