This is an automated email from the ASF dual-hosted git repository.
dtenedor 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 3eebe4202014 [SPARK-57328][SQL] Extract coalesce hint-building helpers
into CoalesceHintUtils
3eebe4202014 is described below
commit 3eebe4202014159801ee8ad5b5dd5174f9c44faf
Author: Daniel Tenedorio <[email protected]>
AuthorDate: Wed Jun 24 16:38:56 2026 -0700
[SPARK-57328][SQL] Extract coalesce hint-building helpers into
CoalesceHintUtils
## What changes were proposed in this pull request?
This PR refactors the COALESCE / REPARTITION / REPARTITION_BY_RANGE /
REBALANCE
hint-building helpers out of the `ResolveHints.ResolveCoalesceHints` rule
and into a
reusable `CoalesceHintUtils` object, so the logic can be shared without
depending on
fixed-point rule state.
- Added `org.apache.spark.sql.catalyst.analysis.CoalesceHintUtils`,
containing the pure,
stateless hint-building helpers lifted out of `ResolveCoalesceHints`:
- `getNumOfPartitions`
- `validateParameters`
- `createRepartition`
- `createRepartitionByRange`
- `transformStringToAttribute`
- These helpers contain no tree-traversal logic and depend only on OSS
classes
(`UnresolvedHint`, `Repartition`, `RepartitionByExpression`, the literal
extractors, and
`QueryCompilationErrors`), so they are safe to call outside the rule.
- `ResolveCoalesceHints` now imports `CoalesceHintUtils._` and delegates to
it. The
`createRebalance` helper stays in the rule (it builds
`RebalancePartitions` inline) and
delegates to the shared `getNumOfPartitions` and `validateParameters`.
The rule retains
its only state dependency, `conf.adaptiveExecutionEnabled`, used to gate
`REBALANCE`.
- Trimmed the now-unused expression imports in `ResolveHints.scala` down to
`StringLiteral`.
No behavior change: the extracted helpers are behavior-identical to the
originals
(`getNumOfPartitions` still re-reads `hint.parameters.tail`, and the nested
`createRepartitionByExpression` closures are preserved as-is).
## How was this patch tested?
Existing hint-resolution test coverage (e.g. join/coalesce hint suites)
pins the behavior;
this is a pure extraction with no functional change.
## Was this patch authored or co-authored using generative AI tooling?
Generated-by: `claude-opus-4-8-thinking-high`
Closes #56477 from dtenedor/refactor-hint-building-helpers.
Authored-by: Daniel Tenedorio <[email protected]>
Signed-off-by: Daniel Tenedorio <[email protected]>
(cherry picked from commit 88a6945e827df67918450a513a9af5265efb9bd7)
Signed-off-by: Daniel Tenedorio <[email protected]>
---
.../sql/catalyst/analysis/CoalesceHintUtils.scala | 123 +++++++++++++++++++++
.../spark/sql/catalyst/analysis/ResolveHints.scala | 101 +----------------
2 files changed, 125 insertions(+), 99 deletions(-)
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
new file mode 100644
index 000000000000..1210ea8351bb
--- /dev/null
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CoalesceHintUtils.scala
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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.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.
+ */
+object CoalesceHintUtils {
+
+ def getNumOfPartitions(hint: UnresolvedHint): (Option[Int], Seq[Expression])
= {
+ hint.parameters match {
+ case Seq(ByteLiteral(numPartitions), _*) =>
+ (Some(numPartitions.toInt), hint.parameters.tail)
+ case Seq(ShortLiteral(numPartitions), _*) =>
+ (Some(numPartitions.toInt), hint.parameters.tail)
+ case Seq(IntegerLiteral(numPartitions), _*) => (Some(numPartitions),
hint.parameters.tail)
+ case _ => (None, hint.parameters)
+ }
+ }
+
+ def validateParameters(hint: String, parms: Seq[Expression]): Unit = {
+ val invalidParams = parms.filter(!_.isInstanceOf[UnresolvedAttribute])
+ if (invalidParams.nonEmpty) {
+ val hintName = hint.toUpperCase(Locale.ROOT)
+ throw QueryCompilationErrors.invalidHintParameterError(hintName,
invalidParams)
+ }
+ }
+
+ /**
+ * This function handles hints for "COALESCE" and "REPARTITION".
+ * The "COALESCE" hint only has a partition number as a parameter. The
"REPARTITION" hint
+ * has a partition number, columns, or both of them as parameters.
+ */
+ def createRepartition(shuffle: Boolean, hint: UnresolvedHint): LogicalPlan =
{
+
+ def createRepartitionByExpression(
+ numPartitions: Option[Int], partitionExprs: Seq[Expression]):
RepartitionByExpression = {
+ val sortOrders = partitionExprs.filter(_.isInstanceOf[SortOrder])
+ if (sortOrders.nonEmpty) {
+ throw
QueryCompilationErrors.invalidRepartitionExpressionsError(sortOrders)
+ }
+ validateParameters(hint.name, partitionExprs)
+ RepartitionByExpression(partitionExprs, hint.child, numPartitions)
+ }
+
+ getNumOfPartitions(hint) match {
+ case (Some(numPartitions), partitionExprs) if partitionExprs.isEmpty =>
+ Repartition(numPartitions, shuffle, hint.child)
+ // The "COALESCE" hint (shuffle = false) must have a partition number
only
+ case _ if !shuffle =>
+ throw QueryCompilationErrors.invalidCoalesceHintParameterError(
+ hint.name.toUpperCase(Locale.ROOT))
+ case (Some(numPartitions), partitionExprs) =>
+ createRepartitionByExpression(Some(numPartitions), partitionExprs)
+ case (None, partitionExprs) =>
+ createRepartitionByExpression(None, partitionExprs)
+ }
+ }
+
+ /**
+ * This function handles hints for "REPARTITION_BY_RANGE".
+ * The "REPARTITION_BY_RANGE" hint must have column names and a partition
number is optional.
+ */
+ def createRepartitionByRange(hint: UnresolvedHint): RepartitionByExpression
= {
+ def createRepartitionByExpression(
+ numPartitions: Option[Int], partitionExprs: Seq[Expression]):
RepartitionByExpression = {
+ validateParameters(hint.name, partitionExprs)
+ val sortOrder = partitionExprs.map {
+ case expr: SortOrder => expr
+ case expr: Expression => SortOrder(expr, Ascending)
+ }
+ RepartitionByExpression(sortOrder, hint.child, numPartitions)
+ }
+
+ getNumOfPartitions(hint) match {
+ case (Some(numPartitions), partitionExprs) =>
+ createRepartitionByExpression(Some(numPartitions), partitionExprs)
+ case (None, partitionExprs) =>
+ createRepartitionByExpression(None, partitionExprs)
+ }
+ }
+
+ /**
+ * This function handles hints for "REBALANCE".
+ */
+ def createRebalance(hint: UnresolvedHint): LogicalPlan = {
+ val (numPartitionsOption, partitionExprs) = getNumOfPartitions(hint)
+ validateParameters(hint.name, partitionExprs)
+ 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
+ }
+ 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 5ff8c4cfc85d..91731e614b84 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
@@ -21,7 +21,7 @@ import java.util.Locale
import scala.collection.mutable
-import org.apache.spark.sql.catalyst.expressions.{Ascending, ByteLiteral,
Expression, IntegerLiteral, ShortLiteral, SortOrder, StringLiteral}
+import org.apache.spark.sql.catalyst.expressions.StringLiteral
import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.CurrentOrigin
@@ -185,104 +185,7 @@ object ResolveHints {
* COALESCE Hint accepts names "COALESCE", "REPARTITION",
"REPARTITION_BY_RANGE" and "REBALANCE".
*/
object ResolveCoalesceHints extends Rule[LogicalPlan] {
- private def getNumOfPartitions(hint: UnresolvedHint): (Option[Int],
Seq[Expression]) = {
- hint.parameters match {
- case Seq(ByteLiteral(numPartitions), _*) =>
- (Some(numPartitions.toInt), hint.parameters.tail)
- case Seq(ShortLiteral(numPartitions), _*) =>
- (Some(numPartitions.toInt), hint.parameters.tail)
- case Seq(IntegerLiteral(numPartitions), _*) => (Some(numPartitions),
hint.parameters.tail)
- case _ => (None, hint.parameters)
- }
- }
-
- private def validateParameters(hint: String, parms: Seq[Expression]): Unit
= {
- val invalidParams = parms.filter(!_.isInstanceOf[UnresolvedAttribute])
- if (invalidParams.nonEmpty) {
- val hintName = hint.toUpperCase(Locale.ROOT)
- throw QueryCompilationErrors.invalidHintParameterError(hintName,
invalidParams)
- }
- }
-
- /**
- * This function handles hints for "COALESCE" and "REPARTITION".
- * The "COALESCE" hint only has a partition number as a parameter. The
"REPARTITION" hint
- * has a partition number, columns, or both of them as parameters.
- */
- private def createRepartition(shuffle: Boolean, hint: UnresolvedHint):
LogicalPlan = {
-
- def createRepartitionByExpression(
- numPartitions: Option[Int], partitionExprs: Seq[Expression]):
RepartitionByExpression = {
- val sortOrders = partitionExprs.filter(_.isInstanceOf[SortOrder])
- if (sortOrders.nonEmpty) {
- throw
QueryCompilationErrors.invalidRepartitionExpressionsError(sortOrders)
- }
- validateParameters(hint.name, partitionExprs)
- RepartitionByExpression(partitionExprs, hint.child, numPartitions)
- }
-
- getNumOfPartitions(hint) match {
- case (Some(numPartitions), partitionExprs) if partitionExprs.isEmpty =>
- Repartition(numPartitions, shuffle, hint.child)
- // The "COALESCE" hint (shuffle = false) must have a partition number
only
- case _ if !shuffle =>
- throw QueryCompilationErrors.invalidCoalesceHintParameterError(
- hint.name.toUpperCase(Locale.ROOT))
- case (Some(numPartitions), partitionExprs) =>
- createRepartitionByExpression(Some(numPartitions), partitionExprs)
- case (None, partitionExprs) =>
- createRepartitionByExpression(None, partitionExprs)
- }
- }
-
- /**
- * This function handles hints for "REPARTITION_BY_RANGE".
- * The "REPARTITION_BY_RANGE" hint must have column names and a partition
number is optional.
- */
- private def createRepartitionByRange(hint: UnresolvedHint):
RepartitionByExpression = {
- def createRepartitionByExpression(
- numPartitions: Option[Int], partitionExprs: Seq[Expression]):
RepartitionByExpression = {
- validateParameters(hint.name, partitionExprs)
- val sortOrder = partitionExprs.map {
- case expr: SortOrder => expr
- case expr: Expression => SortOrder(expr, Ascending)
- }
- RepartitionByExpression(sortOrder, hint.child, numPartitions)
- }
-
- getNumOfPartitions(hint) match {
- case (Some(numPartitions), partitionExprs) =>
- createRepartitionByExpression(Some(numPartitions), partitionExprs)
- case (None, partitionExprs) =>
- createRepartitionByExpression(None, partitionExprs)
- }
- }
-
- private def createRebalance(hint: UnresolvedHint): LogicalPlan = {
- def createRebalancePartitions(
- partitionExprs: Seq[Expression],
- initialNumPartitions: Option[Int]): RebalancePartitions = {
- validateParameters(hint.name, partitionExprs)
- RebalancePartitions(partitionExprs, hint.child, initialNumPartitions)
- }
-
- getNumOfPartitions(hint) match {
- case (Some(numPartitions), partitionExprs) =>
- createRebalancePartitions(partitionExprs, Some(numPartitions))
- case (None, partitionExprs) =>
- createRebalancePartitions(partitionExprs, None)
- }
- }
-
- private 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
- }
- hint.copy(parameters = parameters)
- }
+ import CoalesceHintUtils._
def apply(plan: LogicalPlan): LogicalPlan =
plan.resolveOperatorsWithPruning(
_.containsPattern(UNRESOLVED_HINT), ruleId) {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
