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))