Fix withErrorHandling and remove writing to sink for already executed…

… commands
miland-db · Nov 25, 2024 · 85c9cf8 · 85c9cf8
1 parent c9162f0
commit 85c9cf8
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
@@ -420,7 +420,7 @@ class SparkSession private(
     var result: Option[Seq[Row]] = None
 
     while (sse.hasNext) {
-      sse.withErrorHandling() {
+      sse.withErrorHandling {
         val df = sse.next()
         if (sse.hasNext) {
           df.write.format("noop").mode("overwrite").save()

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecution.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.scripting
 
 import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.apache.spark.sql.catalyst.expressions.Expression
-import org.apache.spark.sql.catalyst.plans.logical.{CommandResult, CompoundBody}
+import org.apache.spark.sql.catalyst.plans.logical.{CommandResult, CompoundBody, MultiResult}
 
 /**
  * SQL scripting executor - executes script and returns result statements.
@@ -53,12 +53,10 @@ class SqlScriptingExecution(
     while (currentStatement.isDefined) {
       currentStatement match {
         case Some(stmt: SingleStatementExec) if !stmt.isExecuted =>
-          withErrorHandling() {
+          withErrorHandling {
             val df = stmt.buildDataFrame(session)
-            if (df.logicalPlan.isInstanceOf[CommandResult]) {
-              // If the statement is not a result, we need to write it to a noop sink to execute it
-              df.write.format("noop").mode("overwrite").save()
-            } else {
+            if (!df.logicalPlan.isInstanceOf[CommandResult]
+              && !df.logicalPlan.isInstanceOf[MultiResult]) {
               // If the statement is a result, we need to return it to the caller
               return Some(df)
             }
@@ -76,7 +74,7 @@ class SqlScriptingExecution(
     throw e
   }
 
-  def withErrorHandling()(f: => Unit): Unit = {
+  def withErrorHandling(f: => Unit): Unit = {
     try {
       f
     } catch {