[SPARK-53097] Make WriteOperationV2 in SparkConnectPlanner side effect free

Author: heyihong

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -3146,16 +3146,12 @@ class SparkConnectPlanner(
     executeHolder.eventsManager.postFinished()
   }
 
-  /**
-   * Transforms the write operation and executes it.
-   *
-   * The input write operation contains a reference to the input plan and transforms it to the
-   * corresponding logical plan. Afterwards, creates the DataFrameWriter and translates the
-   * parameters of the WriteOperation into the corresponding methods calls.
-   *
-   * @param writeOperation
-   */
-  private def handleWriteOperationV2(writeOperation: proto.WriteOperationV2): Unit = {
+  private def runCommand(command: LogicalPlan, tracker: QueryPlanningTracker): Unit = {
+    val qe = new QueryExecution(session, command, tracker)
+    qe.assertCommandExecuted()
+  }
+
+  private def transformWriteOperationV2(writeOperation: proto.WriteOperationV2): LogicalPlan = {
     // Transform the input plan into the logical plan.
     val plan = transformRelation(writeOperation.getInput)
     // And create a Dataset from the plan.
@@ -3190,31 +3186,42 @@ class SparkConnectPlanner(
     writeOperation.getMode match {
       case proto.WriteOperationV2.Mode.MODE_CREATE =>
         if (writeOperation.hasProvider) {
-          w.using(writeOperation.getProvider).create()
-        } else {
-          w.create()
+          w.using(writeOperation.getProvider)
         }
+        w.createCommand()
       case proto.WriteOperationV2.Mode.MODE_OVERWRITE =>
-        w.overwrite(Column(transformExpression(writeOperation.getOverwriteCondition)))
+        w.overwriteCommand(Column(transformExpression(writeOperation.getOverwriteCondition)))
       case proto.WriteOperationV2.Mode.MODE_OVERWRITE_PARTITIONS =>
-        w.overwritePartitions()
+        w.overwritePartitionsCommand()
       case proto.WriteOperationV2.Mode.MODE_APPEND =>
-        w.append()
+        w.appendCommand()
       case proto.WriteOperationV2.Mode.MODE_REPLACE =>
         if (writeOperation.hasProvider) {
-          w.using(writeOperation.getProvider).replace()
-        } else {
-          w.replace()
+          w.using(writeOperation.getProvider)
         }
+        w.replaceCommand(orCreate = false)
       case proto.WriteOperationV2.Mode.MODE_CREATE_OR_REPLACE =>
         if (writeOperation.hasProvider) {
-          w.using(writeOperation.getProvider).createOrReplace()
-        } else {
-          w.createOrReplace()
+          w.using(writeOperation.getProvider)
         }
+        w.replaceCommand(orCreate = true)
       case other =>
         throw InvalidInputErrors.invalidEnum(other)
     }
+  }
+
+  /**
+   * Transforms the write operation and executes it.
+   *
+   * The input write operation contains a reference to the input plan and transforms it to the
+   * corresponding logical plan. Afterwards, creates the DataFrameWriter and translates the
+   * parameters of the WriteOperation into the corresponding methods calls.
+   *
+   * @param writeOperation
+   */
+  private def handleWriteOperationV2(writeOperation: proto.WriteOperationV2): Unit = {
+    val tracker = executeHolder.eventsManager.createQueryPlanningTracker()
+    runCommand(transformWriteOperationV2(writeOperation), tracker)
     executeHolder.eventsManager.postFinished()
   }
 

sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriterV2.scala

@@ -148,14 +148,17 @@ final class DataFrameWriterV2[T] private[sql](table: String, ds: Dataset[T])
 
   /** @inheritdoc */
   override def create(): Unit = {
-    runCommand(
-      CreateTableAsSelect(
-        UnresolvedIdentifier(tableName),
-        partitioning.getOrElse(Seq.empty) ++ clustering,
-        logicalPlan,
-        buildTableSpec(),
-        options.toMap,
-        false))
+    runCommand(createCommand())
+  }
+
+  private[sql] def createCommand(): LogicalPlan = {
+    CreateTableAsSelect(
+      UnresolvedIdentifier(tableName),
+      partitioning.getOrElse(Seq.empty) ++ clustering,
+      logicalPlan,
+      buildTableSpec(),
+      options.toMap,
+      false)
   }
 
   private def buildTableSpec(): UnresolvedTableSpec = {
@@ -186,28 +189,37 @@ final class DataFrameWriterV2[T] private[sql](table: String, ds: Dataset[T])
   /** @inheritdoc */
   @throws(classOf[NoSuchTableException])
   def append(): Unit = {
-    val append = AppendData.byName(
+    runCommand(appendCommand())
+  }
+
+  private[sql] def appendCommand(): LogicalPlan = {
+    AppendData.byName(
       UnresolvedRelation(tableName).requireWritePrivileges(Seq(INSERT)),
       logicalPlan, options.toMap)
-    runCommand(append)
   }
 
   /** @inheritdoc */
   @throws(classOf[NoSuchTableException])
   def overwrite(condition: Column): Unit = {
-    val overwrite = OverwriteByExpression.byName(
+    runCommand(overwriteCommand(condition))
+  }
+
+  private[sql] def overwriteCommand(condition: Column): LogicalPlan = {
+    OverwriteByExpression.byName(
       UnresolvedRelation(tableName).requireWritePrivileges(Seq(INSERT, DELETE)),
       logicalPlan, expression(condition), options.toMap)
-    runCommand(overwrite)
   }
 
   /** @inheritdoc */
   @throws(classOf[NoSuchTableException])
   def overwritePartitions(): Unit = {
-    val dynamicOverwrite = OverwritePartitionsDynamic.byName(
+    runCommand(overwritePartitionsCommand())
+  }
+
+  private[sql] def overwritePartitionsCommand(): LogicalPlan = {
+    OverwritePartitionsDynamic.byName(
       UnresolvedRelation(tableName).requireWritePrivileges(Seq(INSERT, DELETE)),
       logicalPlan, options.toMap)
-    runCommand(dynamicOverwrite)
   }
 
   /**
@@ -220,13 +232,17 @@ final class DataFrameWriterV2[T] private[sql](table: String, ds: Dataset[T])
   }
 
   private def internalReplace(orCreate: Boolean): Unit = {
-    runCommand(ReplaceTableAsSelect(
+    runCommand(replaceCommand(orCreate))
+  }
+
+  private[sql] def replaceCommand(orCreate: Boolean): LogicalPlan = {
+    ReplaceTableAsSelect(
       UnresolvedIdentifier(tableName),
       partitioning.getOrElse(Seq.empty) ++ clustering,
       logicalPlan,
       buildTableSpec(),
       writeOptions = options.toMap,
-      orCreate = orCreate))
+      orCreate = orCreate)
   }
 }