[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

b597443c3 [huangxiaoping] Fix code style
618594667 [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>
This commit is contained in:
huangxiaoping 2024-07-23 09:36:21 +08:00 committed by ulyssesyou
parent 063a192c7a
commit ec232c18b5
No known key found for this signature in database
GPG Key ID: 4C500BC62D576766
6 changed files with 81 additions and 0 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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