Merge branch 'master' of https://github.com/NetEase/kyuubi into spark-2.4

Author: ulysses-you

.scalafmt.conf

@@ -0,0 +1,28 @@
+version = 3.1.1
+runner.dialect=scala212
+project.git=true
+
+align.preset = none
+align.openParenDefnSite = false
+align.openParenCallSite = false
+align.stripMargin = true
+align.tokens = []
+assumeStandardLibraryStripMargin = true
+danglingParentheses.preset = false
+docstrings.style = Asterisk
+docstrings.wrap = no
+importSelectors = singleLine
+indent.extendSite = 2
+literals.hexDigits = Upper
+maxColumn = 100
+newlines.source = keep
+newlines.topLevelStatementBlankLines = []
+optIn.configStyleArguments = false
+rewrite.imports.groups = [
+  ["javax?\\..*"],
+  ["scala\\..*"],
+  ["(?!org\\.apache\\.kyuubi\\.).*"],
+  ["org\\.apache\\.kyuubi\\..*"]
+]
+rewrite.imports.sort = scalastyle
+rewrite.rules = [Imports, SortModifiers]

dev/kyuubi-extension-spark-3-1/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWriting.scala

@@ -26,7 +26,8 @@ import org.apache.spark.sql.types.IntegerType
 
 trait RepartitionBuilderWithRepartitionByExpression extends RepartitionBuilder {
   override def buildRepartition(
-      dynamicPartitionColumns: Seq[Attribute], query: LogicalPlan): LogicalPlan = {
+      dynamicPartitionColumns: Seq[Attribute],
+      query: LogicalPlan): LogicalPlan = {
     if (dynamicPartitionColumns.isEmpty) {
       RepartitionByExpression(
         Seq.empty,
@@ -46,10 +47,11 @@ trait RepartitionBuilderWithRepartitionByExpression extends RepartitionBuilder {
     // Dynamic partition insertion will add repartition by partition column, but it could cause
     // data skew (one partition value has large data). So we add extra partition column for the
     // same dynamic partition to avoid skew.
-    Cast(Multiply(
-      new Rand(Literal(new Random().nextLong())),
-      Literal(partitionNumber.toDouble)
-    ), IntegerType) :: Nil
+    Cast(
+      Multiply(
+        new Rand(Literal(new Random().nextLong())),
+        Literal(partitionNumber.toDouble)),
+      IntegerType) :: Nil
   }
 }
 
@@ -61,8 +63,7 @@ trait RepartitionBuilderWithRepartitionByExpression extends RepartitionBuilder {
  */
 case class RepartitionBeforeWritingDatasource(session: SparkSession)
   extends RepartitionBeforeWritingDatasourceBase
-    with RepartitionBuilderWithRepartitionByExpression {
-}
+  with RepartitionBuilderWithRepartitionByExpression {}
 
 /**
  * For Hive table, there two commands can write data to table
@@ -72,5 +73,4 @@ case class RepartitionBeforeWritingDatasource(session: SparkSession)
  */
 case class RepartitionBeforeWritingHive(session: SparkSession)
   extends RepartitionBeforeWritingHiveBase
-    with RepartitionBuilderWithRepartitionByExpression {
-}
+  with RepartitionBuilderWithRepartitionByExpression {}

dev/kyuubi-extension-spark-3-1/src/main/scala/org/apache/kyuubi/sql/sqlclassification/KyuubiGetSqlClassification.scala

@@ -67,7 +67,6 @@ object KyuubiGetSqlClassification extends Logging {
       }
     }.getOrElse(
       throw new IllegalArgumentException(
-        s"You should restart engine with: ${SQL_CLASSIFICATION_ENABLED.key} true")
-    )
+        s"You should restart engine with: ${SQL_CLASSIFICATION_ENABLED.key} true"))
   }
 }

dev/kyuubi-extension-spark-3-1/src/main/scala/org/apache/kyuubi/sql/watchdog/ForcedMaxOutputRowsRule.scala

@@ -34,29 +34,30 @@ object ForcedMaxOutputRowsConstraint {
 }
 
 /*
-* Add ForcedMaxOutputRows rule for output rows limitation
-* to avoid huge output rows of non_limit query unexpectedly
-* mainly applied to cases as below:
-*
-* case 1:
-* {{{
-*   SELECT [c1, c2, ...]
-* }}}
-*
-* case 2:
-* {{{
-*   WITH CTE AS (
-*   ...)
-* SELECT [c1, c2, ...] FROM CTE ...
-* }}}
-*
-* The Logical Rule add a GlobalLimit node before root project
-* */
+ * Add ForcedMaxOutputRows rule for output rows limitation
+ * to avoid huge output rows of non_limit query unexpectedly
+ * mainly applied to cases as below:
+ *
+ * case 1:
+ * {{{
+ *   SELECT [c1, c2, ...]
+ * }}}
+ *
+ * case 2:
+ * {{{
+ *   WITH CTE AS (
+ *   ...)
+ * SELECT [c1, c2, ...] FROM CTE ...
+ * }}}
+ *
+ * The Logical Rule add a GlobalLimit node before root project
+ * */
 case class ForcedMaxOutputRowsRule(session: SparkSession) extends Rule[LogicalPlan] {
 
   private def isChildAggregate(a: Aggregate): Boolean = a
-    .aggregateExpressions.exists(p => p.getTagValue(ForcedMaxOutputRowsConstraint.CHILD_AGGREGATE)
-    .contains(ForcedMaxOutputRowsConstraint.CHILD_AGGREGATE_FLAG))
+    .aggregateExpressions.exists(p =>
+      p.getTagValue(ForcedMaxOutputRowsConstraint.CHILD_AGGREGATE)
+        .contains(ForcedMaxOutputRowsConstraint.CHILD_AGGREGATE_FLAG))
 
   private def isView: Boolean = {
     val nestedViewDepth = AnalysisContext.get.nestedViewDepth
@@ -65,7 +66,7 @@ case class ForcedMaxOutputRowsRule(session: SparkSession) extends Rule[LogicalPl
 
   private def canInsertLimitInner(p: LogicalPlan): Boolean = p match {
 
-    case Aggregate(_, Alias(_, "havingCondition")::Nil, _) => false
+    case Aggregate(_, Alias(_, "havingCondition") :: Nil, _) => false
     case agg: Aggregate => !isChildAggregate(agg)
     case _: RepartitionByExpression => true
     case _: Distinct => true
@@ -87,17 +88,19 @@ case class ForcedMaxOutputRowsRule(session: SparkSession) extends Rule[LogicalPl
 
     maxOutputRowsOpt match {
       case Some(forcedMaxOutputRows) => canInsertLimitInner(p) &&
-        !p.maxRows.exists(_ <= forcedMaxOutputRows) &&
-        !isView
+          !p.maxRows.exists(_ <= forcedMaxOutputRows) &&
+          !isView
       case None => false
     }
   }
 
   override def apply(plan: LogicalPlan): LogicalPlan = {
     val maxOutputRowsOpt = conf.getConf(KyuubiSQLConf.WATCHDOG_FORCED_MAXOUTPUTROWS)
     plan match {
-      case p if p.resolved && canInsertLimit(p, maxOutputRowsOpt) => Limit(
-        maxOutputRowsOpt.get, plan)
+      case p if p.resolved && canInsertLimit(p, maxOutputRowsOpt) =>
+        Limit(
+          maxOutputRowsOpt.get,
+          plan)
       case _ => plan
     }
   }
@@ -110,33 +113,33 @@ case class MarkAggregateOrderRule(session: SparkSession) extends Rule[LogicalPla
     // mark child aggregate
     a.aggregateExpressions.filter(_.resolved).foreach(_.setTagValue(
       ForcedMaxOutputRowsConstraint.CHILD_AGGREGATE,
-      ForcedMaxOutputRowsConstraint.CHILD_AGGREGATE_FLAG)
-    )
+      ForcedMaxOutputRowsConstraint.CHILD_AGGREGATE_FLAG))
   }
 
   private def findAndMarkChildAggregate(plan: LogicalPlan): LogicalPlan = plan match {
     /*
-    * The case mainly process order not aggregate column but grouping column as below
-    * SELECT c1, COUNT(*) as cnt
-    * FROM t1
-    * GROUP BY c1
-    * ORDER BY c1
-    * */
-    case a: Aggregate if a.aggregateExpressions
-      .exists(x => x.resolved && x.name.equals("aggOrder")) => markChildAggregate(a)
+     * The case mainly process order not aggregate column but grouping column as below
+     * SELECT c1, COUNT(*) as cnt
+     * FROM t1
+     * GROUP BY c1
+     * ORDER BY c1
+     * */
+    case a: Aggregate
+        if a.aggregateExpressions
+          .exists(x => x.resolved && x.name.equals("aggOrder")) =>
+      markChildAggregate(a)
       plan
 
-    case _ => plan.children.foreach(_.foreach {
+    case _ =>
+      plan.children.foreach(_.foreach {
         case agg: Aggregate => markChildAggregate(agg)
         case _ => Unit
-      }
-    )
+      })
       plan
   }
 
   override def apply(plan: LogicalPlan): LogicalPlan = conf.getConf(
-    KyuubiSQLConf.WATCHDOG_FORCED_MAXOUTPUTROWS
-  ) match {
+    KyuubiSQLConf.WATCHDOG_FORCED_MAXOUTPUTROWS) match {
     case Some(_) => findAndMarkChildAggregate(plan)
     case _ => plan
   }

dev/kyuubi-extension-spark-3-1/src/main/scala/org/apache/kyuubi/sql/watchdog/MaxHivePartitionStrategy.scala

@@ -37,25 +37,28 @@ case class MaxHivePartitionStrategy(session: SparkSession)
   override def apply(plan: LogicalPlan): Seq[SparkPlan] = {
     conf.getConf(KyuubiSQLConf.WATCHDOG_MAX_HIVEPARTITION) match {
       case Some(maxHivePartition) => plan match {
-        case ScanOperation(_, _, relation: HiveTableRelation) if relation.isPartitioned =>
+          case ScanOperation(_, _, relation: HiveTableRelation) if relation.isPartitioned =>
             relation.prunedPartitions match {
-              case Some(prunedPartitions) => if (prunedPartitions.size > maxHivePartition) {
-                throw new MaxHivePartitionExceedException(
-                  s"""
-                    |SQL job scan hive partition: ${prunedPartitions.size}
-                    |exceed restrict of hive scan maxPartition $maxHivePartition
-                    |You should optimize your SQL logical according partition structure
-                    |or shorten query scope such as p_date, detail as below:
-                    |Table: ${relation.tableMeta.qualifiedName}
-                    |Owner: ${relation.tableMeta.owner}
-                    |Partition Structure: ${relation.partitionCols.map(_.name).mkString(" -> ")}
-                    |""".stripMargin)
-              } else {
-                Nil
-              }
-              case _ => val totalPartitions = session
-                .sessionState.catalog.externalCatalog.listPartitionNames(
-                relation.tableMeta.database, relation.tableMeta.identifier.table)
+              case Some(prunedPartitions) =>
+                if (prunedPartitions.size > maxHivePartition) {
+                  throw new MaxHivePartitionExceedException(
+                    s"""
+                       |SQL job scan hive partition: ${prunedPartitions.size}
+                       |exceed restrict of hive scan maxPartition $maxHivePartition
+                       |You should optimize your SQL logical according partition structure
+                       |or shorten query scope such as p_date, detail as below:
+                       |Table: ${relation.tableMeta.qualifiedName}
+                       |Owner: ${relation.tableMeta.owner}
+                       |Partition Structure: ${relation.partitionCols.map(_.name).mkString(" -> ")}
+                       |""".stripMargin)
+                } else {
+                  Nil
+                }
+              case _ =>
+                val totalPartitions = session
+                  .sessionState.catalog.externalCatalog.listPartitionNames(
+                    relation.tableMeta.database,
+                    relation.tableMeta.identifier.table)
                 if (totalPartitions.size > maxHivePartition) {
                   throw new MaxHivePartitionExceedException(
                     s"""
@@ -66,12 +69,12 @@ case class MaxHivePartitionStrategy(session: SparkSession)
                        |Owner: ${relation.tableMeta.owner}
                        |Partition Structure: ${relation.partitionCols.map(_.name).mkString(" -> ")}
                        |""".stripMargin)
-                   } else {
+                } else {
                   Nil
                 }
             }
-        case _ => Nil
-      }
+          case _ => Nil
+        }
       case _ => Nil
     }
   }

dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala

@@ -23,13 +23,14 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.kyuubi.sql.{FinalStageConfigIsolation, KyuubiSQLConf}
 
 class FinalStageConfigIsolationSuite extends KyuubiSparkSQLExtensionTest {
-  protected override def beforeAll(): Unit = {
+  override protected def beforeAll(): Unit = {
     super.beforeAll()
     setupData()
   }
 
   test("final stage config set reset check") {
-    withSQLConf(KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true",
+    withSQLConf(
+      KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true",
       "spark.sql.finalStage.adaptive.coalescePartitions.minPartitionNum" -> "1",
       "spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes" -> "100") {
       // use loop to double check final stage config doesn't affect the sql query each other
@@ -79,8 +80,10 @@ class FinalStageConfigIsolationSuite extends KyuubiSparkSQLExtensionTest {
   }
 
   test("final stage config isolation") {
-    def checkPartitionNum(sqlString: String, previousPartitionNum: Int,
-                          finalPartitionNum: Int): Unit = {
+    def checkPartitionNum(
+        sqlString: String,
+        previousPartitionNum: Int,
+        finalPartitionNum: Int): Unit = {
       val df = sql(sqlString)
       df.collect()
       val shuffleReaders = collect(df.queryExecution.executedPlan) {
@@ -99,7 +102,8 @@ class FinalStageConfigIsolationSuite extends KyuubiSparkSQLExtensionTest {
       assert(df.rdd.partitions.length === finalPartitionNum)
     }
 
-    withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+    withSQLConf(
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
       SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
       SQLConf.SHUFFLE_PARTITIONS.key -> "3",
       KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true",
@@ -156,8 +160,7 @@ class FinalStageConfigIsolationSuite extends KyuubiSparkSQLExtensionTest {
             |) t1 ON t0.c2 = t1.c2
             |""".stripMargin,
           3,
-          1
-        )
+          1)
 
         // one shuffle reader
         checkPartitionNum(
@@ -169,8 +172,7 @@ class FinalStageConfigIsolationSuite extends KyuubiSparkSQLExtensionTest {
             |) t1 ON t0.c1 = t1.c1
             |""".stripMargin,
           1,
-          1
-        )
+          1)
       }
     }
   }

dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/RepartitionBeforeWritingSuite.scala

@@ -33,8 +33,7 @@ class RepartitionBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
             assert(r.optNumPartitions ===
               spark.sessionState.conf.getConf(KyuubiSQLConf.INSERT_REPARTITION_NUM))
             r
-        }.size == 1
-      )
+        }.size == 1)
     }
 
     // It's better to set config explicitly in case of we change the default value.
@@ -75,8 +74,7 @@ class RepartitionBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
       assert(
         df.queryExecution.analyzed.collect {
           case r: RepartitionByExpression => r
-        }.isEmpty
-      )
+        }.isEmpty)
     }
 
     withSQLConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true") {
@@ -132,7 +130,8 @@ class RepartitionBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
       }.size == 1)
     }
 
-    withSQLConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true",
+    withSQLConf(
+      KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true",
       KyuubiSQLConf.DYNAMIC_PARTITION_INSERTION_REPARTITION_NUM.key -> "2") {
       Seq("USING PARQUET", "").foreach { storage =>
         withTable("tmp1") {
@@ -149,7 +148,8 @@ class RepartitionBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
   }
 
   test("OptimizedCreateHiveTableAsSelectCommand") {
-    withSQLConf(HiveUtils.CONVERT_METASTORE_PARQUET.key -> "true",
+    withSQLConf(
+      HiveUtils.CONVERT_METASTORE_PARQUET.key -> "true",
       HiveUtils.CONVERT_METASTORE_CTAS.key -> "true") {
       withTable("t") {
         val df = sql(s"CREATE TABLE t STORED AS parquet AS SELECT 1 as a")