[KYUUBI #6551] Allow insert zorder when global sort is false and the …

…plan is Repartition or RepartitionByExpression.

# 🔍 Description
## Issue References 🔗

This pull request fixes #6551

## Describe Your Solution 🔧

Update `canInsertZorder` to allow insert zorder when global sort is `false` and the plan is `Repartition` or `RepartitionByExpression`.

## Types of changes 🔖

- [ ] Bugfix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)

## Test Plan 🧪

#### Behavior Without This Pull Request ⚰️

#### Behavior With This Pull Request 🎉

#### Related Unit Tests
/kyuubi-extension-spark-common/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala

---

# Checklist 📝

- [x] This patch was not authored or co-authored using [Generative Tooling](https://www.apache.org/legal/generative-tooling.html)

**Be nice. Be informative.**

Closes #6552 from huangxiaopingRD/6551.

Closes #6551

b597443 [huangxiaoping] Fix code style
6185946 [huangxiaoping] [KYUUBI #6551] Allow insert zorder when when the plan is Repartition or RepartitionByExpression

Authored-by: huangxiaoping <1754789345@qq.com>
Signed-off-by: ulyssesyou <ulyssesyou@apache.org>

extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting33.scala

@@ -46,6 +46,8 @@ trait InsertZorderHelper33 extends Rule[LogicalPlan] with ZorderBuilder {
 
   def canInsertZorder(query: LogicalPlan): Boolean = query match {
     case Project(_, child) => canInsertZorder(child)
+    case _: RepartitionByExpression | _: Repartition
+        if !conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) => true
     // TODO: actually, we can force zorder even if existed some shuffle
     case _: Sort => false
     case _: RepartitionByExpression => false

extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala

@@ -639,6 +639,31 @@ trait ZorderSuiteBase extends KyuubiSparkSQLExtensionTest with ExpressionEvalHel
     }
   }
 
+  test("Allow insert zorder after repartition if zorder using local sort") {
+    withTable("t") {
+      sql(
+        """
+          |CREATE TABLE t (c1 int, c2 string) TBLPROPERTIES (
+          |'kyuubi.zorder.enabled'= 'true',
+          |'kyuubi.zorder.cols'= 'c1,c2')
+          |""".stripMargin)
+      withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false") {
+        val p1 = sql("INSERT INTO TABLE t SELECT /*+ REPARTITION(1) */* FROM VALUES(1,'a')")
+          .queryExecution.analyzed
+        assert(p1.collect {
+          case sort: Sort if !sort.global => sort
+        }.size == 1)
+      }
+      withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "true") {
+        val p2 = sql("INSERT INTO TABLE t SELECT /*+ REPARTITION(1) */* FROM VALUES(1,'a')")
+          .queryExecution.analyzed
+        assert(p2.collect {
+          case sort: Sort if !sort.global => sort
+        }.size == 0)
+      }
+    }
+  }
+
   test("fast approach test") {
     Seq[Seq[Any]](
       Seq(1L, 2L),

extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala

@@ -45,6 +45,8 @@ trait InsertZorderHelper33 extends Rule[LogicalPlan] with ZorderBuilder {
 
   def canInsertZorder(query: LogicalPlan): Boolean = query match {
     case Project(_, child) => canInsertZorder(child)
+    case _: RepartitionByExpression | _: Repartition
+        if !conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) => true
     // TODO: actually, we can force zorder even if existed some shuffle
     case _: Sort => false
     case _: RepartitionByExpression => false

extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala

@@ -640,6 +640,31 @@ trait ZorderSuiteBase extends KyuubiSparkSQLExtensionTest with ExpressionEvalHel
     }
   }
 
+  test("Allow insert zorder after repartition if zorder using local sort") {
+    withTable("t") {
+      sql(
+        """
+          |CREATE TABLE t (c1 int, c2 string) TBLPROPERTIES (
+          |'kyuubi.zorder.enabled'= 'true',
+          |'kyuubi.zorder.cols'= 'c1,c2')
+          |""".stripMargin)
+      withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false") {
+        val p1 = sql("INSERT INTO TABLE t SELECT /*+ REPARTITION(1) */* FROM VALUES(1,'a')")
+          .queryExecution.analyzed
+        assert(p1.collect {
+          case sort: Sort if !sort.global => sort
+        }.size == 1)
+      }
+      withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "true") {
+        val p2 = sql("INSERT INTO TABLE t SELECT /*+ REPARTITION(1) */* FROM VALUES(1,'a')")
+          .queryExecution.analyzed
+        assert(p2.collect {
+          case sort: Sort if !sort.global => sort
+        }.size == 0)
+      }
+    }
+  }
+
   test("fast approach test") {
     Seq[Seq[Any]](
       Seq(1L, 2L),

extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala

@@ -45,6 +45,8 @@ trait InsertZorderHelper33 extends Rule[LogicalPlan] with ZorderBuilder {
 
   def canInsertZorder(query: LogicalPlan): Boolean = query match {
     case Project(_, child) => canInsertZorder(child)
+    case _: RepartitionByExpression | _: Repartition
+        if !conf.getConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED) => true
     // TODO: actually, we can force zorder even if existed some shuffle
     case _: Sort => false
     case _: RepartitionByExpression => false

extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala

@@ -640,6 +640,31 @@ trait ZorderSuiteBase extends KyuubiSparkSQLExtensionTest with ExpressionEvalHel
     }
   }
 
+  test("Allow insert zorder after repartition if zorder using local sort") {
+    withTable("t") {
+      sql(
+        """
+          |CREATE TABLE t (c1 int, c2 string) TBLPROPERTIES (
+          |'kyuubi.zorder.enabled'= 'true',
+          |'kyuubi.zorder.cols'= 'c1,c2')
+          |""".stripMargin)
+      withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "false") {
+        val p1 = sql("INSERT INTO TABLE t SELECT /*+ REPARTITION(1) */* FROM VALUES(1,'a')")
+          .queryExecution.analyzed
+        assert(p1.collect {
+          case sort: Sort if !sort.global => sort
+        }.size == 1)
+      }
+      withSQLConf(KyuubiSQLConf.ZORDER_GLOBAL_SORT_ENABLED.key -> "true") {
+        val p2 = sql("INSERT INTO TABLE t SELECT /*+ REPARTITION(1) */* FROM VALUES(1,'a')")
+          .queryExecution.analyzed
+        assert(p2.collect {
+          case sort: Sort if !sort.global => sort
+        }.size == 0)
+      }
+    }
+  }
+
   test("fast approach test") {
     Seq[Seq[Any]](
       Seq(1L, 2L),