[SPARK-47270][SQL] Dataset.isEmpty projects CommandResults locally

apache · Nov 1, 2024 · da4df40 · da4df40
1 parent 7b67d5d
commit da4df40
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/...uubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/KyuubiSparkUtil.scala b/...uubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/KyuubiSparkUtil.scala
@@ -24,6 +24,7 @@ import scala.annotation.meta.getter
 
 import org.apache.spark.{SPARK_VERSION, SparkContext}
 import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.kyuubi.SparkDatasetHelper
 import org.apache.spark.util.kvstore.KVIndex
 
 import org.apache.kyuubi.Logging
@@ -48,7 +49,7 @@ object KyuubiSparkUtil extends Logging {
         interruptOnCancel = true)
       debug(s"Execute initialization sql: $sql")
       try {
-        spark.sql(sql).take(1).isEmpty
+        SparkDatasetHelper.commandResultOptimized(spark.sql(sql)).isEmpty
       } finally {
         spark.sparkContext.clearJobGroup()
       }

diff --git a/...uubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala b/...uubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala
@@ -24,7 +24,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.network.util.{ByteUnit, JavaUtils}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{DataFrame, Dataset, Row}
-import org.apache.spark.sql.catalyst.plans.logical.GlobalLimit
+import org.apache.spark.sql.catalyst.plans.logical.{CommandResult, GlobalLimit, LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils
 import org.apache.spark.sql.execution.{CollectLimitExec, CommandResultExec, HiveResult, LocalTableScanExec, QueryExecution, SparkPlan, SQLExecution}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
@@ -290,4 +290,13 @@ object SparkDatasetHelper extends Logging {
     nodeName == "org.apache.spark.sql.execution.command.ExecutedCommandExec" ||
     nodeName == "org.apache.spark.sql.execution.CommandResultExec"
   }
+
+  /** SPARK-47270: Returns a optimized plan for CommandResult, convert to `LocalRelation`. */
+  def commandResultOptimized[T](dataset: Dataset[T]): Dataset[T] = {
+    dataset.logicalPlan match {
+      case c: CommandResult =>
+        Dataset(dataset.sparkSession, LocalRelation(c.output, c.rows))(dataset.encoder)
+      case _ => dataset
+    }
+  }
 }