Skip to content

[Spark] Optimize InferRebalanceAndSortOrders to support skip non-cheap keys and local sort #7551

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ulysses-you wants to merge 2 commits into from
Closed

[Spark] Optimize InferRebalanceAndSortOrders to support skip non-cheap keys and local sort #7551

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
54320bc
Optimize InferRebalanceAndSortOrders to support skip non-cheap keys a…
ulysses-you Jul 9, 2026
0b6003b
comment
ulysses-you Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/extensions/engines/spark/rules.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,8 @@ Kyuubi provides some configs to make these feature easy to use.
| spark.sql.optimizer.zorderUsingOriginalOrdering.enabled | false | When true and `spark.sql.optimizer.rebalanceBeforeZorder.enabled` is true, we do sort by the original ordering i.e. lexicographical order. | 1.6.0 |
| spark.sql.optimizer.inferRebalanceAndSortOrders.enabled | false | When ture, infer columns for rebalance and sort orders from original query, e.g. the join keys from join. It can avoid compression ratio regression. | 1.7.0 |
| spark.sql.optimizer.inferRebalanceAndSortOrdersMaxColumns | 3 | The max columns of inferred columns. | 1.7.0 |
| spark.sql.optimizer.inferRebalanceAndSortOrders.cheapColumnsOnly | true | When true and `spark.sql.optimizer.inferRebalanceAndSortOrders.enabled` is true, only infer the rebalance and sort columns when all inferred columns are cheap expressions, i.e. attributes, foldable values, or field extractions over cheap expressions. This avoids evaluating expensive expressions during the inferred shuffle and sort. | 1.12.0 |
| spark.sql.optimizer.inferRebalanceAndSortOrders.skipSort | false | When true and `spark.sql.optimizer.inferRebalanceAndSortOrders.enabled` is true, only infer the rebalance partition columns and skip inferring the sort orders. Skipping the sort avoids a local sort before writing when only the file layout from rebalance is wanted. | 1.12.0 |
| spark.sql.optimizer.insertRepartitionBeforeWriteIfNoShuffle.enabled | false | When true, add repartition even if the original plan does not have shuffle. | 1.7.0 |
| spark.sql.optimizer.finalStageConfigIsolationWriteOnly.enabled | true | When true, only enable final stage isolation for writing. | 1.7.0 |
| spark.sql.adaptive.rebalancePartitionsAdvisoryPartitionSizeInBytes | none | The advisory partition size for the `RebalancePartitions` operator injected by `insertRepartitionBeforeWrite` and related rules. When set, it is passed to `optAdvisoryPartitionSize` of the `RebalancePartitions` node. It takes precedence over `spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes`. | 1.12.0 |
Expand Down
24 changes: 21 additions & 3 deletions ...xtension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ package org.apache.kyuubi.sql

import scala.annotation.tailrec

import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, Expression, NamedExpression, UnaryExpression}
import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, BoundReference, Expression, ExtractValue, NamedExpression, OuterReference, UnaryExpression}
import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys
import org.apache.spark.sql.catalyst.plans.{FullOuter, Inner, LeftAnti, LeftOuter, LeftSemi, RightOuter}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Generate, LogicalPlan, Project, Sort, SubqueryAlias, View, Window}
Expand Down Expand Up @@ -53,7 +53,16 @@ object InferRebalanceAndSortOrders {
}.toMap
}

def infer(plan: LogicalPlan): Option[PartitioningAndOrdering] = {
def isCheap(e: Expression): Boolean = e match {

@pan3793 pan3793 Jul 10, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this matches Spark's impl, except for PythonUDF, but looks fine

case _: Attribute | _: OuterReference | _: BoundReference => true
case _ if e.foldable => true
case _: Alias | _: ExtractValue => e.children.forall(isCheap)
case _ => false
}

def infer(
plan: LogicalPlan,
onlyInferWithCheapColumns: Boolean): Option[PartitioningAndOrdering] = {
def candidateKeys(
input: LogicalPlan,
output: AttributeSet = AttributeSet.empty): Option[PartitioningAndOrdering] = {
Expand Down Expand Up @@ -107,10 +116,19 @@ object InferRebalanceAndSortOrders {
}
}

candidateKeys(plan).map { case (partitioning, ordering) =>
val partitioningAndSort = candidateKeys(plan).map { case (partitioning, ordering) =>
(
partitioning.filter(_.references.subsetOf(plan.outputSet)),
ordering.filter(_.references.subsetOf(plan.outputSet)))
}
val allCheap = partitioningAndSort.exists {
case (partitionings, sorts) =>
partitionings.forall(isCheap) && sorts.forall(isCheap)
}
if (!onlyInferWithCheapColumns || allCheap) {
partitioningAndSort
} else {
None
}
}
}
19 changes: 19 additions & 0 deletions ...spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,25 @@ object KyuubiSQLConf {
.booleanConf
.createWithDefault(false)

val SKIP_INFER_SORT_ORDERS =
buildConf("spark.sql.optimizer.inferRebalanceAndSortOrders.skipSort")
.doc(s"When true and `${INFER_REBALANCE_AND_SORT_ORDERS.key}` is true, only infer the " +
s"rebalance partition columns and skip inferring the sort orders. Skipping the sort " +
s"avoids a local sort before writing when only the file layout from rebalance is wanted.")
.version("1.12.0")
.booleanConf
.createWithDefault(false)

val INFER_REBALANCE_AND_SORT_ORDERS_WITH_CHEAP_COLUMNS =
buildConf("spark.sql.optimizer.inferRebalanceAndSortOrders.cheapColumnsOnly")
.doc(s"When true and `${INFER_REBALANCE_AND_SORT_ORDERS.key}` is true, only infer the " +
s"rebalance and sort columns when all inferred columns are cheap expressions, i.e. " +
s"attributes, foldable values, or field extractions over cheap expressions. This avoids " +
s"evaluating expensive expressions during the inferred shuffle and sort.")
.version("1.12.0")
.booleanConf
.createWithDefault(true)

val INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS =
buildConf("spark.sql.optimizer.inferRebalanceAndSortOrdersMaxColumns")
.doc("The max columns of inferred columns.")
Expand Down
7 changes: 5 additions & 2 deletions ...ubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,14 +77,17 @@ trait RebalanceBeforeWritingBase extends Rule[LogicalPlan] {
optAdvisoryPartitionSize = advisoryPartitionSize)
} else {
val maxColumns = conf.getConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS)
val inferred = InferRebalanceAndSortOrders.infer(query)
val onlyInferWithCheapColumns = conf.getConf(
KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_WITH_CHEAP_COLUMNS)
val inferred = InferRebalanceAndSortOrders.infer(query, onlyInferWithCheapColumns)
if (inferred.isDefined) {
val (partitioning, ordering) = inferred.get
val rebalance = RebalancePartitions(
partitioning.take(maxColumns),
query,
optAdvisoryPartitionSize = advisoryPartitionSize)
if (ordering.nonEmpty) {
val skipInferOrder = conf.getConf(KyuubiSQLConf.SKIP_INFER_SORT_ORDERS)
if (!skipInferOrder && ordering.nonEmpty) {
val sortOrders = ordering.take(maxColumns).map(o => SortOrder(o, Ascending))
Sort(sortOrders, false, rebalance)
} else {
Expand Down
87 changes: 86 additions & 1 deletion ...extension-spark-3-5/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationComm
import org.apache.spark.sql.hive.HiveUtils
import org.apache.spark.sql.hive.execution.{InsertIntoHiveDirCommand, InsertIntoHiveTable}

import org.apache.kyuubi.sql.KyuubiSQLConf
import org.apache.kyuubi.sql.{InferRebalanceAndSortOrders, KyuubiSQLConf}

class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {

Expand Down Expand Up @@ -305,6 +305,91 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
}
}

test("Skip inferring sort orders") {
def checkShuffleAndSort(dataWritingCommand: LogicalPlan, sSize: Int, rSize: Int): Unit = {
assert(dataWritingCommand.isInstanceOf[DataWritingCommand])
val plan = dataWritingCommand.asInstanceOf[DataWritingCommand].query
assert(plan.collect { case s: Sort => s }.size == sSize)
assert(plan.collect {
case r: RebalancePartitions if r.partitionExpressions.size == rSize => r
}.nonEmpty || rSize == 0)
}

withTable("t", "input1", "input2") {
sql(s"CREATE TABLE t (c1 int, c2 long) USING PARQUET PARTITIONED BY (p string)")
sql(s"CREATE TABLE input1 USING PARQUET AS SELECT * FROM VALUES(1,2),(1,3)")
sql(s"CREATE TABLE input2 USING PARQUET AS SELECT * FROM VALUES(1,3),(1,3)")

val insert =
s"""
|INSERT INTO TABLE t PARTITION(p='a')
|SELECT /*+ broadcast(input2) */ input1.col1, input2.col1
|FROM input1
|JOIN input2
|ON input1.col1 = input2.col1
|""".stripMargin

// when skip is disabled (default), both rebalance and sort orders are inferred
withSQLConf(
KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS.key -> "true",
KyuubiSQLConf.SKIP_INFER_SORT_ORDERS.key -> "false") {
checkShuffleAndSort(sql(insert).queryExecution.analyzed, 1, 1)
}

// when skip is enabled, only the rebalance is inferred and the sort is skipped
withSQLConf(
KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS.key -> "true",
KyuubiSQLConf.SKIP_INFER_SORT_ORDERS.key -> "true") {
checkShuffleAndSort(sql(insert).queryExecution.analyzed, 0, 1)
}
}
}

test("Infer rebalance and sort orders only with cheap columns") {
withTable("input1", "input2") {
sql(s"CREATE TABLE input1 USING PARQUET AS SELECT * FROM VALUES(1,2),(1,3)")
sql(s"CREATE TABLE input2 USING PARQUET AS SELECT * FROM VALUES(1,3),(1,3)")

// the join keys `input1.col1 + 1` / `input2.col1 + 1` are not cheap expressions
val expensivePlan = sql(
s"""
|SELECT input1.col1, input2.col1
|FROM input1
|JOIN input2
|ON input1.col1 + 1 = input2.col1 + 1
|""".stripMargin).queryExecution.analyzed

// when only cheap columns are allowed, the expensive keys are not inferred
assert(InferRebalanceAndSortOrders.infer(expensivePlan, true).isEmpty)
// when the cheap-column restriction is disabled, the expensive keys are inferred
InferRebalanceAndSortOrders.infer(expensivePlan, false) match {
case Some((partitioning, ordering)) =>
assert(partitioning.nonEmpty)
assert(ordering.nonEmpty)
case None => fail("expected inferred columns when cheap-column restriction is disabled")
}

// the join keys `input1.col1` / `input2.col1` are cheap attributes
val cheapPlan = sql(
s"""
|SELECT input1.col1, input2.col1
|FROM input1
|JOIN input2
|ON input1.col1 = input2.col1
|""".stripMargin).queryExecution.analyzed

// cheap columns are inferred regardless of the restriction
Seq(true, false).foreach { onlyCheap =>
InferRebalanceAndSortOrders.infer(cheapPlan, onlyCheap) match {
case Some((partitioning, ordering)) =>
assert(partitioning.nonEmpty)
assert(ordering.nonEmpty)
case None => fail(s"expected inferred columns when onlyCheap=$onlyCheap")
}
}
}
}

test("Test rebalance in InsertIntoHiveDirCommand") {
withSQLConf(
HiveUtils.CONVERT_METASTORE_PARQUET.key -> "false",
Expand Down
24 changes: 21 additions & 3 deletions ...xtension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ package org.apache.kyuubi.sql

import scala.annotation.tailrec

import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, Expression, NamedExpression, UnaryExpression}
import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, BoundReference, Expression, ExtractValue, NamedExpression, OuterReference, UnaryExpression}
import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys
import org.apache.spark.sql.catalyst.plans.{FullOuter, Inner, LeftAnti, LeftOuter, LeftSemi, RightOuter}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Generate, LogicalPlan, Project, Sort, SubqueryAlias, View, Window}
Expand Down Expand Up @@ -53,7 +53,16 @@ object InferRebalanceAndSortOrders {
}.toMap
}

def infer(plan: LogicalPlan): Option[PartitioningAndOrdering] = {
def isCheap(e: Expression): Boolean = e match {
case _: Attribute | _: OuterReference | _: BoundReference => true
case _ if e.foldable => true
case _: Alias | _: ExtractValue => e.children.forall(isCheap)
case _ => false
}

def infer(
plan: LogicalPlan,
onlyInferWithCheapColumns: Boolean): Option[PartitioningAndOrdering] = {
def candidateKeys(
input: LogicalPlan,
output: AttributeSet = AttributeSet.empty): Option[PartitioningAndOrdering] = {
Expand Down Expand Up @@ -107,10 +116,19 @@ object InferRebalanceAndSortOrders {
}
}

candidateKeys(plan).map { case (partitioning, ordering) =>
val partitioningAndSort = candidateKeys(plan).map { case (partitioning, ordering) =>
(
partitioning.filter(_.references.subsetOf(plan.outputSet)),
ordering.filter(_.references.subsetOf(plan.outputSet)))
}
val allCheap = partitioningAndSort.exists {
case (partitionings, sorts) =>
partitionings.forall(isCheap) && sorts.forall(isCheap)
}
if (!onlyInferWithCheapColumns || allCheap) {
partitioningAndSort
} else {
None
}
}
}
19 changes: 19 additions & 0 deletions ...spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,25 @@ object KyuubiSQLConf {
.booleanConf
.createWithDefault(false)

val SKIP_INFER_SORT_ORDERS =
buildConf("spark.sql.optimizer.inferRebalanceAndSortOrders.skipSort")
.doc(s"When true and `${INFER_REBALANCE_AND_SORT_ORDERS.key}` is true, only infer the " +
s"rebalance partition columns and skip inferring the sort orders. Skipping the sort " +
s"avoids a local sort before writing when only the file layout from rebalance is wanted.")
.version("1.12.0")
.booleanConf
.createWithDefault(false)

val INFER_REBALANCE_AND_SORT_ORDERS_WITH_CHEAP_COLUMNS =
buildConf("spark.sql.optimizer.inferRebalanceAndSortOrders.cheapColumnsOnly")
.doc(s"When true and `${INFER_REBALANCE_AND_SORT_ORDERS.key}` is true, only infer the " +
s"rebalance and sort columns when all inferred columns are cheap expressions, i.e. " +
s"attributes, foldable values, or field extractions over cheap expressions. This avoids " +
s"evaluating expensive expressions during the inferred shuffle and sort.")
.version("1.12.0")
.booleanConf
.createWithDefault(true)

val INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS =
buildConf("spark.sql.optimizer.inferRebalanceAndSortOrdersMaxColumns")
.doc("The max columns of inferred columns.")
Expand Down
7 changes: 5 additions & 2 deletions ...ubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,14 +76,17 @@ trait RebalanceBeforeWritingBase extends Rule[LogicalPlan] {
optAdvisoryPartitionSize = advisoryPartitionSize)
} else {
val maxColumns = conf.getConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS)
val inferred = InferRebalanceAndSortOrders.infer(query)
val onlyInferWithCheapColumns = conf.getConf(
KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_WITH_CHEAP_COLUMNS)
val inferred = InferRebalanceAndSortOrders.infer(query, onlyInferWithCheapColumns)
if (inferred.isDefined) {
val (partitioning, ordering) = inferred.get
val rebalance = RebalancePartitions(
partitioning.take(maxColumns),
query,
optAdvisoryPartitionSize = advisoryPartitionSize)
if (ordering.nonEmpty) {
val skipInferOrder = conf.getConf(KyuubiSQLConf.SKIP_INFER_SORT_ORDERS)
if (!skipInferOrder && ordering.nonEmpty) {
val sortOrders = ordering.take(maxColumns).map(o => SortOrder(o, Ascending))
Sort(sortOrders, false, rebalance)
} else {
Expand Down
Loading
Loading