[CARMEL-5912] Support PARALLEL Hint against Delta Table (#935)

Author: xingchaozh

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/QueryPlanningContext.scala

@@ -29,6 +29,12 @@ class QueryPlanningContext {
     new java.util.concurrent.ConcurrentHashMap[String, TableParallelInfo]
 
   def clear(): Unit = tableParallelHintMap.clear()
+
+  def merge(context: Option[QueryPlanningContext]): Unit = {
+    context.foreach(p => {
+      tableParallelHintMap.putAll(p.tableParallelHintMap)
+    })
+  }
 }
 
 object QueryPlanningContext {
@@ -38,10 +44,13 @@ object QueryPlanningContext {
       .recordStats()
       .build()
 
-  def putIfAbsent(executionId: String, context: QueryPlanningContext): Unit = {
-    if (StringUtils.isNotBlank(executionId) &&
-      planningContextCache.getIfPresent(executionId) == null) {
-      planningContextCache.put(executionId, context)
+  def putIfAbsentOrMerge(executionId: String, context: QueryPlanningContext): Unit = {
+    if (StringUtils.isNotBlank(executionId)) {
+      if (getIfPresent(executionId).isEmpty) {
+        planningContextCache.put(executionId, context)
+      } else {
+        merge(executionId, Some(context))
+      }
     }
   }
 
@@ -52,9 +61,18 @@ object QueryPlanningContext {
   }
 
   def invalidate(executionId: String): Unit = {
-    getIfPresent(executionId).foreach(_.clear())
     if (StringUtils.isNotBlank(executionId)) {
+      getIfPresent(executionId).foreach(_.clear())
       planningContextCache.invalidate(executionId)
     }
   }
+
+  def merge(executionId: String, context: Option[QueryPlanningContext]): Unit = {
+    if (StringUtils.isNotBlank(executionId)) {
+      if (getIfPresent(executionId).isEmpty) {
+        planningContextCache.put(executionId, new QueryPlanningContext)
+      }
+      getIfPresent(executionId).foreach(_.merge(context))
+    }
+  }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/QueryPlanningTracker.scala

@@ -83,7 +83,7 @@ object QueryPlanningTracker {
   }
 
   /** Returns the current tracker in scope, based on the thread local variable. */
-  def get: Option[QueryPlanningTracker] = Option(localTracker.get())
+  def getCurrent: Option[QueryPlanningTracker] = Option(localTracker.get())
 
   /** Sets the current tracker for the execution of function f. We assume f is single-threaded. */
   def withTracker[T](tracker: QueryPlanningTracker)(f: => T): T = {

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleExecutor.scala

@@ -128,7 +128,7 @@ abstract class RuleExecutor[TreeType <: TreeNode[_]] extends Logging {
     var curPlan = plan
     val queryExecutionMetrics = RuleExecutor.queryExecutionMeter
     val planChangeLogger = new PlanChangeLogger()
-    val tracker: Option[QueryPlanningTracker] = QueryPlanningTracker.get
+    val tracker: Option[QueryPlanningTracker] = QueryPlanningTracker.getCurrent
     val beforeMetrics = RuleExecutor.getCurrentMetrics()
 
     // Run the structural integrity checker against the initial input

sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala

@@ -24,6 +24,7 @@ import org.apache.spark.SparkContext
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.internal.config.Tests.IS_TESTING
 import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.QueryPlanningContext
 import org.apache.spark.sql.execution.ui.{SparkListenerSQLExecutionEnd, SparkListenerSQLExecutionStart}
 import org.apache.spark.sql.internal.StaticSQLConf.SQL_EVENT_TRUNCATE_LENGTH
 import org.apache.spark.util.Utils
@@ -62,15 +63,28 @@ object SQLExecution {
    * Wrap an action that will execute "queryExecution" to track all Spark jobs in the body so that
    * we can connect them with an execution.
    */
+  def withNewExecutionId[T](
+                             queryExecution: QueryExecution,
+                             name: Option[String] = None)(body: => T): T = {
+    withNewExecutionId(queryExecution, name, false)(body)
+  }
+
   def withNewExecutionId[T](
       queryExecution: QueryExecution,
-      name: Option[String] = None)(body: => T): T = queryExecution.sparkSession.withActive {
+      name: Option[String],
+      inheritParentQueryPlanningContext: Boolean)(body: => T): T =
+    queryExecution.sparkSession.withActive {
     val sparkSession = queryExecution.sparkSession
     val sc = sparkSession.sparkContext
     val oldExecutionId = sc.getLocalProperty(EXECUTION_ID_KEY)
     val executionId = SQLExecution.nextExecutionId
     sc.setLocalProperty(EXECUTION_ID_KEY, executionId.toString)
     executionIdToQueryExecution.put(executionId, queryExecution)
+
+    if (inheritParentQueryPlanningContext) {
+      val parentQueryPlanningContext = QueryPlanningContext.getIfPresent(oldExecutionId)
+      QueryPlanningContext.merge(executionId.toString, parentQueryPlanningContext)
+    }
     try {
       // sparkContext.getCallSite() would first try to pick up any call site that was previously
       // set, then fall back to Utils.getCallSite(); call Utils.getCallSite() directly on

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/CacheParallelHint.scala

@@ -51,20 +51,24 @@ private[sql] object CacheParallelHint
           if (f.partitionSize.nonEmpty || f.partitionNumber.nonEmpty) => f
       }.foreach(r => {
         logicalRelation.catalogTable.foreach(catalogTable => {
-          val tracker: Option[QueryPlanningTracker] = QueryPlanningTracker.get
-          tracker.foreach(f => {
-            SparkContext.getActive.foreach(sc => {
-              val tableParallelInfo =
-                TableParallelInfo(catalogTable.identifier.toString, r.partitionSize,
+          SparkContext.getActive.foreach(sc => {
+            val tableParallelInfo =
+              TableParallelInfo(catalogTable.identifier.toString, r.partitionSize,
                 r.partitionNumber)
 
-              f.queryPlanningContext.tableParallelHintMap.
-                put(catalogTable.identifier.toString, tableParallelInfo)
+            val queryPlanningContext = if (QueryPlanningTracker.getCurrent.isEmpty) {
+              new QueryPlanningContext
+            } else {
+              QueryPlanningTracker.getCurrent.get.queryPlanningContext
+            }
 
-              val executionId = sc.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
-              QueryPlanningContext.putIfAbsent(executionId, f.queryPlanningContext)
-            })
+            queryPlanningContext.tableParallelHintMap.
+              put(catalogTable.identifier.toString, tableParallelInfo)
+
+            val executionId = sc.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
+            QueryPlanningContext.putIfAbsentOrMerge(executionId, queryPlanningContext)
           })
+
         })
       })
 

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala

@@ -227,7 +227,7 @@ object FileSourceStrategy extends Strategy with SQLConfHelper with Logging {
           })
 
           if (tableParallelInfo.isEmpty) {
-            QueryPlanningTracker.get.foreach(f => {
+            QueryPlanningTracker.getCurrent.foreach(f => {
               tableParallelInfo = Option(
                 f.queryPlanningContext.tableParallelHintMap.get(tableIdentifier)
               )

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/delta/DeltaQuerySuite.scala

@@ -3164,6 +3164,101 @@ class DeltaQuerySuite extends QueryTest
     }
   }
 
+  test("Bucket scan in delta CTAS disabled if PARALLEL HINT detected.") {
+    checkBucketingResultInDeltaCTASIfParallelDetected("/*+ PARALLEL(bucket_table, 100) */",
+      targetPartitionNumber = 100, parallelHintEnabled = true,
+      viewName = "bucket_table_v_with_hint")
+  }
+
+  test("Bucket scan in delta CTAS enabled if PARALLEL HINT not detected.") {
+    checkBucketingResultInDeltaCTASIfParallelDetected("",
+      targetPartitionNumber = 100, parallelHintEnabled = true,
+      viewName = "bucket_table_v_without_hint")
+  }
+
+  def checkBucketingResultInDeltaCTASIfParallelDetected(hint: String,
+                                                   targetPartitionNumber: Long = 1024,
+                                                   parallelHintEnabled: Boolean = true,
+                                                   viewName: String): Unit = {
+    withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+      SQLConf.RANGE_JOIN_ENABLED.key -> "false", // Disable range join for testing
+      SQLConf.PARALLEL_HINT_ENABLED.key -> parallelHintEnabled.toString,
+      SQLConf.SQL_UI_PLAN_WITH_METRICS.key -> "false") {
+      withTable("bucket_table", "a1") {
+
+        withView(viewName) {
+          sql(
+            s"""
+               |CREATE TABLE bucket_table(id string, ip_add_int long, dt string) USING PARQUET
+               |PARTITIONED BY (dt)
+               |CLUSTERED BY (ip_add_int)
+               |INTO 10 BUCKETS
+               |
+               |""".stripMargin)
+
+          sql(
+            s"""
+               |insert into bucket_table PARTITION (dt = "2021-12-20")
+               |select 123 as id, 456 as ip_add_int
+       """.stripMargin
+          )
+
+          sql(
+            s"""
+               |create view ${viewName}(id, ip_add_int, dt) as
+               |select id, ip_add_int, dt from bucket_table
+       """.stripMargin
+          )
+
+          val df = sql(
+            s"""
+               |create temp table a1 using delta AS
+               |SELECT ${hint} ip_add_int, count(*) as cnt FROM (
+               |SELECT ip_add_int, count(*) as c FROM (
+               |SELECT ip_add_int, count(*) FROM ${viewName} ev3
+               |WHERE ev3.id !=0 AND ev3.dt>="2021-12-20" group by 1
+               |) ev2 group by 1
+               |) ev group by 1
+       """.stripMargin
+          )
+
+          sql(
+            s"""
+               |select * from a1
+         """.stripMargin
+          ).show()
+
+          // scalastyle:off println
+          println(s"logical plan: ${df.queryExecution.logical}")
+          println(s"optimized plan: ${df.queryExecution.optimizedPlan}")
+          println(df.collect().length)
+          // scalastyle:on println
+
+          // checkRangeJoinResultInner(df, buildSide, false, 0)
+          val planAfter = df.queryExecution.executedPlan
+          // scalastyle:off println
+          println(s"planAfter: ${planAfter}")
+          // scalastyle:on println
+          val scan = collectFirst(planAfter) {
+            case f: FileSourceScanExec => f
+          }
+          assert(scan.isDefined)
+          assert(scan.get.dataFilters.nonEmpty)
+
+          if (hint.contains("PARALLEL") && parallelHintEnabled) {
+            assert(scan.get.tableParallelInfo.isDefined)
+            assert(scan.get.tableParallelInfo.get.partitionNumber.get == targetPartitionNumber)
+            assert(scan.get.bucketedScan == false)
+          } else {
+            assert(scan.get.tableParallelInfo.isEmpty)
+            assert(scan.get.bucketedScan == true)
+          }
+        }
+      }
+    }
+  }
+
+
   test("test analyze delta table") {
     def getCatalogStatistics(tableName: String): CatalogStatistics = {
       getCatalogTable(TableIdentifier(tableName)).stats.get