[SPARK-24940][SQL] Coalesce and Repartition Hint for SQL Queries

## What changes were proposed in this pull request?

Many Spark SQL users in my company have asked for a way to control the number of output files in Spark SQL. The users prefer not to use function repartition(n) or coalesce(n, shuffle) that require them to write and deploy Scala/Java/Python code. We propose adding the following Hive-style Coalesce and Repartition Hint to Spark SQL:
```
... SELECT /*+ COALESCE(numPartitions) */ ...
... SELECT /*+ REPARTITION(numPartitions) */ ...
```
Multiple such hints are allowed. Multiple nodes are inserted into the logical plan, and the optimizer will pick the leftmost hint.
```
INSERT INTO s SELECT /*+ REPARTITION(100), COALESCE(500), COALESCE(10) */ * FROM t

== Logical Plan ==
'InsertIntoTable 'UnresolvedRelation `s`, false, false
+- 'UnresolvedHint REPARTITION, [100]
   +- 'UnresolvedHint COALESCE, [500]
      +- 'UnresolvedHint COALESCE, [10]
         +- 'Project [*]
            +- 'UnresolvedRelation `t`

== Optimized Logical Plan ==
InsertIntoHadoopFsRelationCommand ...
+- Repartition 100, true
   +- HiveTableRelation ...
```

## How was this patch tested?

All unit tests. Manual tests using explain.

Author: John Zhuge <[email protected]>

Closes apache#21911 from jzhuge/SPARK-24940.

Author: jzhuge

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -145,6 +145,7 @@ class Analyzer(
   lazy val batches: Seq[Batch] = Seq(
     Batch("Hints", fixedPoint,
       new ResolveHints.ResolveBroadcastHints(conf),
+      ResolveHints.ResolveCoalesceHints,
       ResolveHints.RemoveAllHints),
     Batch("Simple Sanity Check", Once,
       LookupFunctions),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveHints.scala

@@ -20,10 +20,12 @@ package org.apache.spark.sql.catalyst.analysis
 import java.util.Locale
 
 import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.catalyst.expressions.Literal
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.CurrentOrigin
 import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.IntegerType
 
 
 /**
@@ -102,6 +104,32 @@ object ResolveHints {
     }
   }
 
+  /**
+   * COALESCE Hint accepts name "COALESCE" and "REPARTITION".
+   * Its parameter includes a partition number.
+   */
+  object ResolveCoalesceHints extends Rule[LogicalPlan] {
+    private val COALESCE_HINT_NAMES = Set("COALESCE", "REPARTITION")
+
+    def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperators {
+      case h: UnresolvedHint if COALESCE_HINT_NAMES.contains(h.name.toUpperCase(Locale.ROOT)) =>
+        val hintName = h.name.toUpperCase(Locale.ROOT)
+        val shuffle = hintName match {
+          case "REPARTITION" => true
+          case "COALESCE" => false
+        }
+        val numPartitions = h.parameters match {
+          case Seq(Literal(numPartitions: Int, IntegerType)) =>
+            numPartitions
+          case Seq(numPartitions: Int) =>
+            numPartitions
+          case _ =>
+            throw new AnalysisException(s"$hintName Hint expects a partition number as parameter")
+        }
+        Repartition(numPartitions, shuffle, h.child)
+    }
+  }
+
   /**
    * Removes all the hints, used to remove invalid hints provided by the user.
    * This must be executed after all the other hint rules are executed.

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveHintsSuite.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.analysis
 
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions.Literal
 import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
 import org.apache.spark.sql.catalyst.plans.Inner
 import org.apache.spark.sql.catalyst.plans.logical._
@@ -120,4 +121,38 @@ class ResolveHintsSuite extends AnalysisTest {
       testRelation.where('a > 1).select('a).select('a).analyze,
       caseSensitive = false)
   }
+
+  test("coalesce and repartition hint") {
+    checkAnalysis(
+      UnresolvedHint("COALESCE", Seq(Literal(10)), table("TaBlE")),
+      Repartition(numPartitions = 10, shuffle = false, child = testRelation))
+    checkAnalysis(
+      UnresolvedHint("coalesce", Seq(Literal(20)), table("TaBlE")),
+      Repartition(numPartitions = 20, shuffle = false, child = testRelation))
+    checkAnalysis(
+      UnresolvedHint("REPARTITION", Seq(Literal(100)), table("TaBlE")),
+      Repartition(numPartitions = 100, shuffle = true, child = testRelation))
+    checkAnalysis(
+      UnresolvedHint("RePARTITion", Seq(Literal(200)), table("TaBlE")),
+      Repartition(numPartitions = 200, shuffle = true, child = testRelation))
+
+    val errMsgCoal = "COALESCE Hint expects a partition number as parameter"
+    assertAnalysisError(
+      UnresolvedHint("COALESCE", Seq.empty, table("TaBlE")),
+      Seq(errMsgCoal))
+    assertAnalysisError(
+      UnresolvedHint("COALESCE", Seq(Literal(10), Literal(false)), table("TaBlE")),
+      Seq(errMsgCoal))
+    assertAnalysisError(
+      UnresolvedHint("COALESCE", Seq(Literal(1.0)), table("TaBlE")),
+      Seq(errMsgCoal))
+
+    val errMsgRepa = "REPARTITION Hint expects a partition number as parameter"
+    assertAnalysisError(
+      UnresolvedHint("REPARTITION", Seq(UnresolvedAttribute("a")), table("TaBlE")),
+      Seq(errMsgRepa))
+    assertAnalysisError(
+      UnresolvedHint("REPARTITION", Seq(Literal(true)), table("TaBlE")),
+      Seq(errMsgRepa))
+  }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala

@@ -593,6 +593,33 @@ class PlanParserSuite extends AnalysisTest {
       parsePlan("SELECT /*+ MAPJOIN(t) */ a from t where true group by a order by a"),
       UnresolvedHint("MAPJOIN", Seq($"t"),
         table("t").where(Literal(true)).groupBy('a)('a)).orderBy('a.asc))
+
+    comparePlans(
+      parsePlan("SELECT /*+ COALESCE(10) */ * FROM t"),
+      UnresolvedHint("COALESCE", Seq(Literal(10)),
+        table("t").select(star())))
+
+    comparePlans(
+      parsePlan("SELECT /*+ REPARTITION(100) */ * FROM t"),
+      UnresolvedHint("REPARTITION", Seq(Literal(100)),
+        table("t").select(star())))
+
+    comparePlans(
+      parsePlan(
+        "INSERT INTO s SELECT /*+ REPARTITION(100), COALESCE(500), COALESCE(10) */ * FROM t"),
+      InsertIntoTable(table("s"), Map.empty,
+        UnresolvedHint("REPARTITION", Seq(Literal(100)),
+          UnresolvedHint("COALESCE", Seq(Literal(500)),
+            UnresolvedHint("COALESCE", Seq(Literal(10)),
+              table("t").select(star())))), overwrite = false, ifPartitionNotExists = false))
+
+    comparePlans(
+      parsePlan("SELECT /*+ BROADCASTJOIN(u), REPARTITION(100) */ * FROM t"),
+      UnresolvedHint("BROADCASTJOIN", Seq($"u"),
+        UnresolvedHint("REPARTITION", Seq(Literal(100)),
+          table("t").select(star()))))
+
+    intercept("SELECT /*+ COALESCE(30 + 50) */ * FROM t", "mismatched input")
   }
 
   test("SPARK-20854: select hint syntax with expressions") {

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

@@ -59,4 +59,14 @@ class DataFrameHintSuite extends AnalysisTest with SharedSQLContext {
       )
     )
   }
+
+  test("coalesce and repartition hint") {
+    check(
+      df.hint("COALESCE", 10),
+      UnresolvedHint("COALESCE", Seq(10), df.logicalPlan))
+
+    check(
+      df.hint("REPARTITION", 100),
+      UnresolvedHint("REPARTITION", Seq(100), df.logicalPlan))
+  }
 }

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

@@ -27,6 +27,7 @@ import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart}
 import org.apache.spark.sql.catalyst.util.StringUtils
 import org.apache.spark.sql.execution.aggregate
 import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec}
+import org.apache.spark.sql.execution.datasources.FilePartition
 import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, CartesianProductExec, SortMergeJoinExec}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
@@ -2797,4 +2798,36 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
       checkAnswer(df, Seq(Row(3, 99, 1)))
     }
   }
+
+
+  test("SPARK-24940: coalesce and repartition hint") {
+    withTempView("nums1") {
+      val numPartitionsSrc = 10
+      spark.range(0, 100, 1, numPartitionsSrc).createOrReplaceTempView("nums1")
+      assert(spark.table("nums1").rdd.getNumPartitions == numPartitionsSrc)
+
+      withTable("nums") {
+        sql("CREATE TABLE nums (id INT) USING parquet")
+
+        Seq(5, 20, 2).foreach { numPartitions =>
+          sql(
+            s"""
+               |INSERT OVERWRITE TABLE nums
+               |SELECT /*+ REPARTITION($numPartitions) */ *
+               |FROM nums1
+             """.stripMargin)
+          assert(spark.table("nums").inputFiles.length == numPartitions)
+
+          sql(
+            s"""
+               |INSERT OVERWRITE TABLE nums
+               |SELECT /*+ COALESCE($numPartitions) */ *
+               |FROM nums1
+             """.stripMargin)
+          // Coalesce can not increase the number of partitions
+          assert(spark.table("nums").inputFiles.length == Seq(numPartitions, numPartitionsSrc).min)
+        }
+      }
+    }
+  }
 }