diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting33.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting33.scala index 083c80afd00..4ae2057dce6 100644 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting33.scala +++ b/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 diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala index 476402fba1c..82ca8d4d2e8 100644 --- a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala +++ b/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), diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala index b3f98ec6d7f..73ed5e253bb 100644 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala +++ b/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 diff --git a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala index 2d3eec95722..26c8ebd1de9 100644 --- a/extensions/spark/kyuubi-extension-spark-3-4/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala +++ b/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), diff --git a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala index b3f98ec6d7f..73ed5e253bb 100644 --- a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/zorder/InsertZorderBeforeWriting.scala +++ b/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 diff --git a/extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala index 2d3eec95722..26c8ebd1de9 100644 --- a/extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/ZorderSuiteBase.scala +++ b/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),