This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new d1186d18fb9d [SPARK-57993][SQL] Support advisory partition size in
rebalance hint
d1186d18fb9d is described below
commit d1186d18fb9dbad07258469b1ccd5fda856521f3
Author: Zhen Wang <[email protected]>
AuthorDate: Wed Jul 15 11:46:05 2026 +0800
[SPARK-57993][SQL] Support advisory partition size in rebalance hint
### What changes were proposed in this pull request?
This PR adds support for a new `REBALANCE_BY_SIZE` SQL hint.
The hint works like `REBALANCE`, but requires a positive advisory partition
size as its first parameter, optionally followed by partition columns. For
example:
```sql
SELECT /*+ REBALANCE_BY_SIZE(134217728) */ * FROM t;
SELECT /*+ REBALANCE_BY_SIZE(134217728, c) */ * FROM t;
SELECT /*+ REBALANCE_BY_SIZE('128m') */ * FROM t;
SELECT /*+ REBALANCE_BY_SIZE('128m', c) */ * FROM t;
```
The implementation resolves `REBALANCE_BY_SIZE` to `RebalancePartitions`
with `optAdvisoryPartitionSize`, so AQE can use the hint-specific advisory
partition size for the rebalance shuffle.
### Why are the changes needed?
A common use case of `REBALANCE` is small file control before writing query
results. Users may want to insert a rebalance before the write and set a larger
advisory partition size for the final output stage, so each output
partition/file has a reasonable target size.
Today this requires changing the session-level
`spark.sql.adaptive.advisoryPartitionSizeInBytes` configuration. That is too
coarse-grained because it may also affect preceding AQE stages, including
shuffle coalescing and read/computation parallelism before the final write
stage.
`REBALANCE_BY_SIZE` provides a query-level way to apply a target partition
size only to the rebalance shuffle, which is useful for output file size
control without changing the advisory size used by earlier stages.
### Does this PR introduce _any_ user-facing change?
Yes, a new `REBALANCE_BY_SIZE` hint.
### How was this patch tested?
Added unit tests
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57073 from wForget/SPARK-57993.
Authored-by: Zhen Wang <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit d1d6c78bd151a52ea9c8de0b6fedd2f69de212a0)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../src/main/resources/error/error-conditions.json | 6 +++
docs/sql-performance-tuning.md | 9 +++-
docs/sql-ref-syntax-qry-select-hints.md | 17 +++++-
.../sql/catalyst/analysis/CoalesceHintUtils.scala | 62 +++++++++++++++++++---
.../spark/sql/catalyst/analysis/ResolveHints.scala | 5 +-
.../spark/sql/errors/QueryCompilationErrors.scala | 9 ++++
.../sql/catalyst/analysis/ResolveHintsSuite.scala | 47 ++++++++++++++++
.../adaptive/AdaptiveQueryExecSuite.scala | 26 +++++++++
8 files changed, 168 insertions(+), 13 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json
b/common/utils/src/main/resources/error/error-conditions.json
index 3e6345d490c1..ca6360d9fc92 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -4685,6 +4685,12 @@
],
"sqlState" : "42613"
},
+ "INVALID_REBALANCE_BY_SIZE_HINT_PARAMETER" : {
+ "message" : [
+ "The first parameter for <hintName> must be a positive byte-size
literal, but got <advisoryPartitionSize>."
+ ],
+ "sqlState" : "22023"
+ },
"INVALID_RECURSIVE_CTE" : {
"message" : [
"Invalid recursive definition found. Recursive queries must contain an
UNION or an UNION ALL statement with 2 children. The first child needs to be
the anchor term without any recursive references. Any top level inner CTE must
not contain self references."
diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md
index bbb449360287..0294f05641d9 100644
--- a/docs/sql-performance-tuning.md
+++ b/docs/sql-performance-tuning.md
@@ -141,7 +141,8 @@ Coalesce hints allow Spark SQL users to control the number
of output files just
tuning and reducing the number of output files. The "COALESCE" hint only has a
partition number as a
parameter. The "REPARTITION" hint has a partition number, columns, or
both/neither of them as parameters.
The "REPARTITION_BY_RANGE" hint must have column names and a partition number
is optional. The "REBALANCE"
-hint has an initial partition number, columns, or both/neither of them as
parameters.
+hint has an initial partition number, columns, or both/neither of them as
parameters. The "REBALANCE_BY_SIZE"
+hint requires an advisory partition size, optionally followed by columns.
```sql
SELECT /*+ COALESCE(3) */ * FROM t;
@@ -155,6 +156,10 @@ SELECT /*+ REBALANCE */ * FROM t;
SELECT /*+ REBALANCE(3) */ * FROM t;
SELECT /*+ REBALANCE(c) */ * FROM t;
SELECT /*+ REBALANCE(3, c) */ * FROM t;
+SELECT /*+ REBALANCE_BY_SIZE(134217728) */ * FROM t;
+SELECT /*+ REBALANCE_BY_SIZE(134217728, c) */ * FROM t;
+SELECT /*+ REBALANCE_BY_SIZE('128m') */ * FROM t;
+SELECT /*+ REBALANCE_BY_SIZE('128m', c) */ * FROM t;
```
For more details please refer to the documentation of [Partitioning
Hints](sql-ref-syntax-qry-select-hints.html#partitioning-hints).
@@ -556,4 +561,4 @@ SET
'spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled' 'tru
+- * Filter (7)
+- * ColumnarToRow (6)
+- BatchScan (5)
-```
\ No newline at end of file
+```
diff --git a/docs/sql-ref-syntax-qry-select-hints.md
b/docs/sql-ref-syntax-qry-select-hints.md
index 861e7fbae668..c9c591d6728e 100644
--- a/docs/sql-ref-syntax-qry-select-hints.md
+++ b/docs/sql-ref-syntax-qry-select-hints.md
@@ -33,8 +33,9 @@ Hints give users a way to suggest how Spark SQL to use
specific approaches to ge
Partitioning hints allow users to suggest a partitioning strategy that Spark
should follow. `COALESCE`, `REPARTITION`,
and `REPARTITION_BY_RANGE` hints are supported and are equivalent to
`coalesce`, `repartition`, and
-`repartitionByRange` [Dataset
APIs](api/scala/org/apache/spark/sql/Dataset.html), respectively. The
`REBALANCE` can only
-be used as a hint .These hints give users a way to tune performance and
control the number of output files in Spark SQL.
+`repartitionByRange` [Dataset
APIs](api/scala/org/apache/spark/sql/Dataset.html), respectively. `REBALANCE`
and
+`REBALANCE_BY_SIZE` can only be used as hints. These hints give users a way to
tune performance and
+control the number of output files in Spark SQL.
When multiple partitioning hints are specified, multiple nodes are inserted
into the logical plan, but the leftmost hint
is picked by the optimizer.
@@ -56,6 +57,10 @@ is picked by the optimizer.
The `REBALANCE` hint can be used to rebalance the query result output
partitions, so that every partition is of a reasonable size (not too small and
not too big). It can take column names as parameters, and try its best to
partition the query result by these columns. This is a best-effort: if there
are skews, Spark will split the skewed partitions, to make these partitions not
too big. This hint is useful when you need to write the result of this query to
a table, to avoid too small/bi [...]
+* **REBALANCE_BY_SIZE**
+
+ The `REBALANCE_BY_SIZE` hint works like `REBALANCE`, but requires an
advisory partition size as its first parameter. This hint is ignored if AQE is
not enabled.
+
#### Examples
```sql
@@ -79,6 +84,14 @@ SELECT /*+ REBALANCE(c) */ * FROM t;
SELECT /*+ REBALANCE(3, c) */ * FROM t;
+SELECT /*+ REBALANCE_BY_SIZE(134217728) */ * FROM t;
+
+SELECT /*+ REBALANCE_BY_SIZE(134217728, c) */ * FROM t;
+
+SELECT /*+ REBALANCE_BY_SIZE('128m') */ * FROM t;
+
+SELECT /*+ REBALANCE_BY_SIZE('128m', c) */ * FROM t;
+
-- multiple partitioning hints
EXPLAIN EXTENDED SELECT /*+ REPARTITION(100), COALESCE(500),
REPARTITION_BY_RANGE(3, c) */ * FROM t;
== Parsed Logical Plan ==
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CoalesceHintUtils.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CoalesceHintUtils.scala
index 1210ea8351bb..c846d42565a9 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CoalesceHintUtils.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CoalesceHintUtils.scala
@@ -19,13 +19,14 @@ package org.apache.spark.sql.catalyst.analysis
import java.util.Locale
-import org.apache.spark.sql.catalyst.expressions.{Ascending, ByteLiteral,
Expression, IntegerLiteral, ShortLiteral, SortOrder, StringLiteral}
+import org.apache.spark.network.util.JavaUtils
+import org.apache.spark.sql.catalyst.expressions.{Ascending, ByteLiteral,
Expression, IntegerLiteral, LongLiteral, ShortLiteral, SortOrder, StringLiteral}
import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan,
RebalancePartitions, Repartition, RepartitionByExpression, UnresolvedHint}
import org.apache.spark.sql.errors.QueryCompilationErrors
/**
* Helper functions used to build the logical plans for the "COALESCE",
"REPARTITION",
- * "REPARTITION_BY_RANGE" and "REBALANCE" hints.
+ * "REPARTITION_BY_RANGE", "REBALANCE" and "REBALANCE_BY_SIZE" hints.
*/
object CoalesceHintUtils {
@@ -40,6 +41,39 @@ object CoalesceHintUtils {
}
}
+ def getAdvisorySizeOfPartitions(hint: UnresolvedHint): (Long,
Seq[Expression]) = {
+ val (advisoryPartitionSize, partitionExprs) = hint.parameters match {
+ case Seq(ByteLiteral(advisoryPartitionSize), _*) =>
+ (advisoryPartitionSize.toLong, hint.parameters.tail)
+ case Seq(ShortLiteral(advisoryPartitionSize), _*) =>
+ (advisoryPartitionSize.toLong, hint.parameters.tail)
+ case Seq(IntegerLiteral(advisoryPartitionSize), _*) =>
+ (advisoryPartitionSize.toLong, hint.parameters.tail)
+ case Seq(LongLiteral(advisoryPartitionSize), _*) =>
+ (advisoryPartitionSize, hint.parameters.tail)
+ case Seq(StringLiteral(advisoryPartitionSize), _*) =>
+ val sizeInBytes = try {
+ JavaUtils.byteStringAsBytes(advisoryPartitionSize)
+ } catch {
+ case _: IllegalArgumentException =>
+ throw
QueryCompilationErrors.invalidRebalanceBySizeHintParameterError(
+ hint.name.toUpperCase(Locale.ROOT),
+ advisoryPartitionSize)
+ }
+ (sizeInBytes, hint.parameters.tail)
+ case _ =>
+ throw QueryCompilationErrors.invalidRebalanceBySizeHintParameterError(
+ hint.name.toUpperCase(Locale.ROOT),
+ hint.parameters.headOption.map(_.sql).getOrElse("empty"))
+ }
+ if (advisoryPartitionSize <= 0) {
+ throw QueryCompilationErrors.invalidRebalanceBySizeHintParameterError(
+ hint.name.toUpperCase(Locale.ROOT),
+ advisoryPartitionSize.toString)
+ }
+ (advisoryPartitionSize, partitionExprs)
+ }
+
def validateParameters(hint: String, parms: Seq[Expression]): Unit = {
val invalidParams = parms.filter(!_.isInstanceOf[UnresolvedAttribute])
if (invalidParams.nonEmpty) {
@@ -111,12 +145,24 @@ object CoalesceHintUtils {
RebalancePartitions(partitionExprs, hint.child, numPartitionsOption)
}
- def transformStringToAttribute(hint: UnresolvedHint): UnresolvedHint = {
- // for all the coalesce hints, it's safe to transform the string literal
to an attribute as
- // all the parameters should be column names.
- val parameters = hint.parameters.map {
- case StringLiteral(name) => UnresolvedAttribute(name)
- case e => e
+ /**
+ * This function handles hints for "REBALANCE_BY_SIZE".
+ */
+ def createRebalanceBySize(hint: UnresolvedHint): LogicalPlan = {
+ val (advisoryPartitionSize, partitionExprs) =
getAdvisorySizeOfPartitions(hint)
+ validateParameters(hint.name, partitionExprs)
+ RebalancePartitions(partitionExprs, hint.child, None,
Some(advisoryPartitionSize))
+ }
+
+ def transformStringToAttribute(
+ hint: UnresolvedHint,
+ skipFirstParameter: Boolean = false): UnresolvedHint = {
+ // For all the coalesce hints, string literal parameters should be
transformed to attributes,
+ // except for the first parameter of REBALANCE_BY_SIZE, which is the
advisory partition size.
+ val parameters = hint.parameters.zipWithIndex.map {
+ case (StringLiteral(name), index) if !skipFirstParameter || index > 0 =>
+ UnresolvedAttribute(name)
+ case (e, _) => e
}
hint.copy(parameters = parameters)
}
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveHints.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveHints.scala
index 91731e614b84..2998dd7e4f68 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveHints.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveHints.scala
@@ -182,7 +182,8 @@ object ResolveHints {
}
/**
- * COALESCE Hint accepts names "COALESCE", "REPARTITION",
"REPARTITION_BY_RANGE" and "REBALANCE".
+ * COALESCE Hint accepts names "COALESCE", "REPARTITION",
"REPARTITION_BY_RANGE", "REBALANCE"
+ * and "REBALANCE_BY_SIZE".
*/
object ResolveCoalesceHints extends Rule[LogicalPlan] {
import CoalesceHintUtils._
@@ -198,6 +199,8 @@ object ResolveHints {
createRepartitionByRange(transformStringToAttribute(hint))
case "REBALANCE" if conf.adaptiveExecutionEnabled =>
createRebalance(transformStringToAttribute(hint))
+ case "REBALANCE_BY_SIZE" if conf.adaptiveExecutionEnabled =>
+ createRebalanceBySize(transformStringToAttribute(hint,
skipFirstParameter = true))
case _ => hint
}
}
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
index 230065f5848a..62bc7b511d8d 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
@@ -1151,6 +1151,15 @@ private[sql] object QueryCompilationErrors extends
QueryErrorsBase with Compilat
messageParameters = Map("hintName" -> hintName))
}
+ def invalidRebalanceBySizeHintParameterError(
+ hintName: String, advisoryPartitionSize: String): Throwable = {
+ new AnalysisException(
+ errorClass = "INVALID_REBALANCE_BY_SIZE_HINT_PARAMETER",
+ messageParameters = Map(
+ "hintName" -> hintName,
+ "advisoryPartitionSize" -> advisoryPartitionSize))
+ }
+
def starExpandDataTypeNotSupportedError(attributes: Seq[String]): Throwable
= {
new AnalysisException(
errorClass = "_LEGACY_ERROR_TEMP_1050",
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveHintsSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveHintsSuite.scala
index 1c36728663f8..54b0827717c9 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveHintsSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveHintsSuite.scala
@@ -373,4 +373,51 @@ class ResolveHintsSuite extends AnalysisTest {
}
}
}
+
+ test("SPARK-57993: Support specify advisory partition size for rebalance") {
+ withSQLConf(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "1024") {
+ Seq(
+ Seq(Literal(2048)) -> Some(2048L),
+ Seq(Literal(4096L)) -> Some(4096L),
+ Seq(Literal("2k")) -> Some(2048L),
+ Seq(Literal("4k")) -> Some(4096L),
+ Seq(Literal("2k"), Literal("a")) -> Some(2048L)).foreach {
+ case (param, advisoryPartitionSize) =>
+ assert(
+ UnresolvedHint("REBALANCE_BY_SIZE", param, testRelation).analyze
+ .asInstanceOf[RebalancePartitions]
+ .optAdvisoryPartitionSize == advisoryPartitionSize)
+ }
+
+ // invalid parameters for REBALANCE_BY_SIZE hint
+ Seq(
+ Seq(Literal(-1)) -> "-1",
+ Seq(Literal(0)) -> "0",
+ Nil -> "empty",
+ Seq(Literal("a")) -> "a",
+ Seq(UnresolvedAttribute("a")) -> "a").foreach { case (params,
advisoryPartitionSize) =>
+ checkError(
+ exception = intercept[AnalysisException] {
+ UnresolvedHint("REBALANCE_BY_SIZE", params, testRelation).analyze
+ },
+ condition = "INVALID_REBALANCE_BY_SIZE_HINT_PARAMETER",
+ parameters = Map(
+ "hintName" -> "REBALANCE_BY_SIZE",
+ "advisoryPartitionSize" -> advisoryPartitionSize))
+ }
+ }
+
+ // REBALANCE_BY_SIZE hint should be ignored when AQE is disabled
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ Seq(
+ Seq(Literal(2048)),
+ Seq(Literal(4096L)),
+ Seq(Literal("4k")),
+ Seq(Literal("2k"), Literal("a"))).foreach { params =>
+ checkAnalysisWithoutViewWrapper(
+ UnresolvedHint("REBALANCE_BY_SIZE", params, table("TaBlE")),
+ testRelation)
+ }
+ }
+ }
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
index 8cf6fbf921da..c3b8de4f37c1 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -2537,6 +2537,32 @@ class AdaptiveQueryExecSuite
}
}
+ test("SPARK-57993: Use specified advisory partition size in
REBALANCE_BY_SIZE") {
+ withTempView("v") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key ->
"true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "5",
+ SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "10000") {
+
+ spark.sparkContext.parallelize(
+ (1 to 10).map(i => TestData(if (i > 4) 5 else i, i.toString)), 3)
+ .toDF("c1", "c2").createOrReplaceTempView("v")
+
+ val (_, adaptive) =
+ runAdaptiveAndVerifyResult("SELECT /*+ REBALANCE_BY_SIZE('150b', c1)
*/ * FROM v")
+ val read = collect(adaptive) {
+ case read: AQEShuffleReadExec => read
+ }
+ assert(read.size == 1)
+
assert(read.head.partitionSpecs.count(_.isInstanceOf[PartialReducerPartitionSpec])
== 2)
+ assert(read.head.partitionSpecs.size == 4)
+ }
+ }
+ }
+
test("SPARK-35888: join with a 0-partition table") {
withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
