diff --git a/docs/extensions/engines/spark/rules.md b/docs/extensions/engines/spark/rules.md index 94e0a797f2a..3999336c9d7 100644 --- a/docs/extensions/engines/spark/rules.md +++ b/docs/extensions/engines/spark/rules.md @@ -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 | diff --git a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala index 3b840f2a014..ba80eda0782 100644 --- a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala +++ b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala @@ -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} @@ -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] = { @@ -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 + } } } diff --git a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala index d68c5fb7f92..481ffdf116f 100644 --- a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala +++ b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala @@ -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.") diff --git a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala index b74b33cc6fd..79aaa5c48ef 100644 --- a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala +++ b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala @@ -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 { diff --git a/extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala index 11563292fb8..d543633a7d8 100644 --- a/extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala +++ b/extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala @@ -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 { @@ -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", diff --git a/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala b/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala index 3b840f2a014..ba80eda0782 100644 --- a/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala +++ b/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala @@ -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} @@ -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] = { @@ -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 + } } } diff --git a/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala b/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala index d68c5fb7f92..481ffdf116f 100644 --- a/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala +++ b/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala @@ -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.") diff --git a/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala b/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala index a55766f51f7..5a8c43d969b 100644 --- a/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala +++ b/extensions/spark/kyuubi-extension-spark-4-0/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala @@ -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 { diff --git a/extensions/spark/kyuubi-extension-spark-4-0/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala b/extensions/spark/kyuubi-extension-spark-4-0/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala index de8e463fbd5..8a413b298f3 100644 --- a/extensions/spark/kyuubi-extension-spark-4-0/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala +++ b/extensions/spark/kyuubi-extension-spark-4-0/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala @@ -24,7 +24,7 @@ import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationComm import org.apache.spark.sql.hive.HiveUtils import org.apache.spark.sql.hive.execution.InsertIntoHiveTable -import org.apache.kyuubi.sql.KyuubiSQLConf +import org.apache.kyuubi.sql.{InferRebalanceAndSortOrders, KyuubiSQLConf} class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest { @@ -304,6 +304,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("advisory partition size - finalStage config takes precedence") { withSQLConf( KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true", diff --git a/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala b/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala index 3b840f2a014..ba80eda0782 100644 --- a/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala +++ b/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala @@ -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} @@ -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] = { @@ -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 + } } } diff --git a/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala b/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala index d68c5fb7f92..481ffdf116f 100644 --- a/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala +++ b/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala @@ -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.") diff --git a/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala b/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala index a55766f51f7..5a8c43d969b 100644 --- a/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala +++ b/extensions/spark/kyuubi-extension-spark-4-1/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala @@ -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 { diff --git a/extensions/spark/kyuubi-extension-spark-4-1/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala b/extensions/spark/kyuubi-extension-spark-4-1/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala index de8e463fbd5..8a413b298f3 100644 --- a/extensions/spark/kyuubi-extension-spark-4-1/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala +++ b/extensions/spark/kyuubi-extension-spark-4-1/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala @@ -24,7 +24,7 @@ import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationComm import org.apache.spark.sql.hive.HiveUtils import org.apache.spark.sql.hive.execution.InsertIntoHiveTable -import org.apache.kyuubi.sql.KyuubiSQLConf +import org.apache.kyuubi.sql.{InferRebalanceAndSortOrders, KyuubiSQLConf} class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest { @@ -304,6 +304,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("advisory partition size - finalStage config takes precedence") { withSQLConf( KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true",