[CARMEL-7613] Add persistWithTempTable to Dataset to replace cache (apache#55)

Author: wangyum

sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -17,7 +17,8 @@
 
 package org.apache.spark.sql
 
-import java.io.{ByteArrayOutputStream, CharArrayWriter, DataOutputStream}
+import java.io.{ByteArrayOutputStream, CharArrayWriter, DataOutputStream, File}
+import java.util.UUID
 
 import scala.annotation.varargs
 import scala.collection.JavaConverters._
@@ -39,7 +40,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.scheduler.RepeatableIterator
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow, QueryPlanningTracker, ScalaReflection, TableIdentifier}
 import org.apache.spark.sql.catalyst.analysis._
-import org.apache.spark.sql.catalyst.catalog.HiveTableRelation
+import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable, CatalogTableType, HiveTableRelation}
 import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.json.{JacksonGenerator, JSONOptions}
@@ -3798,6 +3799,30 @@ class Dataset[T] private[sql](
     this
   }
 
+  def persistWithTempTable(): DataFrame = {
+    persistWithTempTable(UUID.randomUUID().toString.replace('-', '_'))
+  }
+
+  def persistWithTempTable(tableName: String): DataFrame = sparkSession.withActive {
+    val sessionPath = sparkSession.sessionState.catalog.sessionPath
+    val table =
+      CatalogTable(
+        TableIdentifier(tableName),
+        CatalogTableType.TEMPORARY,
+        CatalogStorageFormat.empty.copy(
+          locationUri = Some(sessionPath.suffix(File.separator + tableName).toUri)),
+        StructType(Nil),
+        provider = Some(sparkSession.sessionState.conf.defaultDataSourceName))
+
+    CreateDataSourceTableAsSelectCommand(
+      table,
+      SaveMode.ErrorIfExists,
+      logicalPlan,
+      logicalPlan.output.map(_.name)).run(sparkSession)
+
+    sparkSession.table(tableName)
+  }
+
   /**
    * Persist this Dataset with the default storage level (`MEMORY_AND_DISK`).
    *

sql/core/src/test/scala/org/apache/spark/sql/TemporaryTableSuite.scala

@@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException
 import org.apache.spark.sql.catalyst.catalog.CatalogTableType
 import org.apache.spark.sql.catalyst.parser.ParseException
+import org.apache.spark.sql.execution.FileSourceScanExec
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.test.SharedSparkSession
 
@@ -347,4 +348,11 @@ class TemporaryTableSuite extends QueryTest
     testUnsupportedCommand("ALTER TABLE t1 CHANGE COLUMN id TYPE bigint")
     testUnsupportedCommand("ALTER TABLE t1 SET SERDE 'whatever'")
   }
+
+  test("HADP-51648: persistToTempTable API") {
+    val df = spark.range(10).persistWithTempTable()
+    val location = df.queryExecution.sparkPlan.asInstanceOf[FileSourceScanExec].relation.location
+    assert(location.rootPaths.head.toString.contains(sparkScratch))
+    checkAnswer(df, spark.range(10).toDF())
+  }
 }