From 6160b9de30335ed74c60da0fe7bf3196acbcef57 Mon Sep 17 00:00:00 2001 From: ulysses-you Date: Tue, 12 Oct 2021 10:50:03 +0800 Subject: [PATCH] [KYUUBI #1214] Refine sql extension and avoid setup data for unused test ### _Why are the changes needed?_ - Refine rule order - Speed up test ### _How was this patch tested?_ - [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible - [ ] Add screenshots for manual tests if appropriate - [ ] [Run test](https://kyuubi.readthedocs.io/en/latest/develop_tools/testing.html#running-tests) locally before make a pull request Closes #1214 from ulysses-you/tests. Closes #1214 261271f0 [ulysses-you] test Authored-by: ulysses-you Signed-off-by: ulysses-you --- .../org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala | 5 +++-- .../apache/spark/sql/FinalStageConfigIsolationSuite.scala | 5 +++++ .../apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala | 5 +++++ .../org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala | 3 +-- .../src/test/scala/org/apache/spark/sql/WatchDogSuite.scala | 4 ++++ 5 files changed, 18 insertions(+), 4 deletions(-) diff --git a/dev/kyuubi-extension-spark-3-1/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala b/dev/kyuubi-extension-spark-3-1/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala index 7e4c78053..e20052a39 100644 --- a/dev/kyuubi-extension-spark-3-1/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala +++ b/dev/kyuubi-extension-spark-3-1/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala @@ -38,13 +38,14 @@ class KyuubiSparkSQLExtension extends (SparkSessionExtensions => Unit) { extensions.injectParser{ case (_, parser) => new ZorderSparkSqlExtensionsParser(parser) } extensions.injectResolutionRule(ResolveZorder) + // a help rule for ForcedMaxOutputRowsRule + extensions.injectResolutionRule(MarkAggregateOrderRule) + // Note that: // InsertZorderBeforeWritingDatasource and InsertZorderBeforeWritingHive // should be applied before // RepartitionBeforeWrite and RepartitionBeforeWriteHive // because we can only apply one of them (i.e. Global Sort or Repartition) - extensions.injectResolutionRule(MarkAggregateOrderRule) - extensions.injectPostHocResolutionRule(InsertZorderBeforeWritingDatasource) extensions.injectPostHocResolutionRule(InsertZorderBeforeWritingHive) extensions.injectPostHocResolutionRule(KyuubiSqlClassification) diff --git a/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala b/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala index 992105a73..8592f4dff 100644 --- a/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala +++ b/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/FinalStageConfigIsolationSuite.scala @@ -23,6 +23,11 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.kyuubi.sql.{FinalStageConfigIsolation, KyuubiSQLConf} class FinalStageConfigIsolationSuite extends KyuubiSparkSQLExtensionTest { + protected override def beforeAll(): Unit = { + super.beforeAll() + setupData() + } + test("final stage config set reset check") { withSQLConf(KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true", "spark.sql.finalStage.adaptive.coalescePartitions.minPartitionNum" -> "1", diff --git a/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala b/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala index 939fea865..7abeee469 100644 --- a/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala +++ b/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/InsertShuffleNodeBeforeJoinSuite.scala @@ -23,6 +23,11 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.kyuubi.sql.KyuubiSQLConf class InsertShuffleNodeBeforeJoinSuite extends KyuubiSparkSQLExtensionTest { + protected override def beforeAll(): Unit = { + super.beforeAll() + setupData() + } + test("force shuffle before join") { def checkShuffleNodeNum(sqlString: String, num: Int): Unit = { var expectedResult: Seq[Row] = Seq.empty diff --git a/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala b/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala index 8ee01764d..34259dbc6 100644 --- a/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala +++ b/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala @@ -42,7 +42,6 @@ trait KyuubiSparkSQLExtensionTest extends QueryTest .enableHiveSupport() .getOrCreate()) } - setupData() super.beforeAll() } @@ -52,7 +51,7 @@ trait KyuubiSparkSQLExtensionTest extends QueryTest _spark.foreach(_.stop) } - private def setupData(): Unit = { + protected def setupData(): Unit = { val self = spark import self.implicits._ spark.sparkContext.parallelize( diff --git a/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/WatchDogSuite.scala b/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/WatchDogSuite.scala index 656e5bcbb..29ddea696 100644 --- a/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/WatchDogSuite.scala +++ b/dev/kyuubi-extension-spark-3-1/src/test/scala/org/apache/spark/sql/WatchDogSuite.scala @@ -23,6 +23,10 @@ import org.apache.kyuubi.sql.KyuubiSQLConf import org.apache.kyuubi.sql.watchdog.MaxHivePartitionExceedException class WatchDogSuite extends KyuubiSparkSQLExtensionTest { + protected override def beforeAll(): Unit = { + super.beforeAll() + setupData() + } case class LimitAndExpected(limit: Int, expected: Int) val limitAndExpecteds = List(LimitAndExpected(1, 1), LimitAndExpected(11, 10))