You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kyuubi.apache.org by ul...@apache.org on 2026/07/28 03:28:22 UTC
(kyuubi) branch master updated: [KYUUBI #7587] [SPARK] Infer cast-wrapped join keys in InferRebalanceAndSortOrders
This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kyuubi.git
The following commit(s) were added to refs/heads/master by this push:
new ac66880b61 [KYUUBI #7587] [SPARK] Infer cast-wrapped join keys in InferRebalanceAndSortOrders
ac66880b61 is described below
commit ac66880b6184b6cf288b648a259468125d52fdc7
Author: Xiduo You <ul...@apache.org>
AuthorDate: Tue Jul 28 11:28:10 2026 +0800
[KYUUBI #7587] [SPARK] Infer cast-wrapped join keys in InferRebalanceAndSortOrders
### Why are the changes needed?
Two fixes to `InferRebalanceAndSortOrders` (applied across the Spark 3.5, 4.0, 4.1 and 4.2 extension modules):
1. **Support cast-wrapped join keys.** When join key columns have mismatched types (e.g. via a `UNION ALL` branch that widens `int` to `bigint`), the analyzer inserts a cast on the join key such as `cast(genre_tag_id as bigint)`. The outer projection's alias map is keyed by the unary-peeled attribute, so the full cast expression never matched and the key was silently dropped by the `outputSet` filter — only the un-cast keys survived, weakening the inferred rebalance/sort. This adds `m [...]
2. **Gate the cheap-column check by `maxColumns`.** `INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS` is now threaded into `infer` and the truncation is applied *before* the cheap-column check, so only the first `maxColumns` inferred columns need to be cheap (previously an expensive key beyond the limit could discard the whole inference).
### How was this patch tested?
Added unit tests in `RebalanceBeforeWritingSuite` for each affected module:
- cast-wrapped join key resolves to its aliased output attribute,
- a `Cast` over a cheap column is treated as cheap,
- `maxColumns` gates the cheap-column check.
Verified locally with:
- Spark 3.5: `build/mvn -Pspark-3.5 test -pl extensions/spark/kyuubi-extension-spark-3-5 -am -DwildcardSuites=org.apache.spark.sql.RebalanceBeforeWritingSuite`
- Spark 4.0 / 4.1 / 4.2 (JDK 17, `-Pscala-2.13`): same suite per module.
All suites pass; `dev/reformat` applied.
### Was this patch authored or co-authored using generative AI tooling?
Assisted-by: Claude:claude-opus-4-8
Closes #7587 from ulysses-you/ne-infer-cast-wrapped-join-keys.
Closes #7587
3999f3a68 [Xiduo You] [SPARK] Infer cast-wrapped join keys in InferRebalanceAndSortOrders
Authored-by: Xiduo You <ul...@apache.org>
Signed-off-by: Xiduo You <ul...@apache.org>
---
.../kyuubi/sql/InferRebalanceAndSortOrders.scala | 52 +++++++++-----
.../apache/kyuubi/sql/RebalanceBeforeWriting.scala | 9 ++-
.../spark/sql/RebalanceBeforeWritingSuite.scala | 79 +++++++++++++++++++++-
.../kyuubi/sql/InferRebalanceAndSortOrders.scala | 52 +++++++++-----
.../apache/kyuubi/sql/RebalanceBeforeWriting.scala | 9 ++-
.../spark/sql/RebalanceBeforeWritingSuite.scala | 79 +++++++++++++++++++++-
.../kyuubi/sql/InferRebalanceAndSortOrders.scala | 52 +++++++++-----
.../apache/kyuubi/sql/RebalanceBeforeWriting.scala | 9 ++-
.../spark/sql/RebalanceBeforeWritingSuite.scala | 79 +++++++++++++++++++++-
.../kyuubi/sql/InferRebalanceAndSortOrders.scala | 52 +++++++++-----
.../apache/kyuubi/sql/RebalanceBeforeWriting.scala | 9 ++-
.../spark/sql/RebalanceBeforeWritingSuite.scala | 79 +++++++++++++++++++++-
12 files changed, 464 insertions(+), 96 deletions(-)
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 ba80eda078..66bf225df3 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, BoundReference, Expression, ExtractValue, NamedExpression, OuterReference, UnaryExpression}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, BoundReference, Cast, 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}
@@ -38,14 +38,14 @@ object InferRebalanceAndSortOrders {
type PartitioningAndOrdering = (Seq[Expression], Seq[Expression])
- private def getAliasMap(named: Seq[NamedExpression]): Map[Expression, Attribute] = {
- @tailrec
- def throughUnary(e: Expression): Expression = e match {
- case u: UnaryExpression if u.deterministic =>
- throughUnary(u.child)
- case _ => e
- }
+ @tailrec
+ private def throughUnary(e: Expression): Expression = e match {
+ case u: UnaryExpression if u.deterministic =>
+ throughUnary(u.child)
+ case _ => e
+ }
+ private def getAliasMap(named: Seq[NamedExpression]): Map[Expression, Attribute] = {
named.flatMap {
case a @ Alias(child, _) =>
Some((throughUnary(child).canonicalized, a.toAttribute))
@@ -53,16 +53,32 @@ object InferRebalanceAndSortOrders {
}.toMap
}
+ /**
+ * Map an expression to its output attribute via the alias map. Tries the full
+ * expression first, then retries with the unary-peeled key (mirroring how
+ * `getAliasMap` builds its keys), so that a cast-wrapped key such as
+ * `cast(genre_tag_id as bigint)` still maps to its aliased output attribute. Falls
+ * back to the original expression when neither matches.
+ */
+ private def mapThroughAlias(
+ e: Expression,
+ aliasMap: Map[Expression, Attribute]): Expression = {
+ aliasMap.getOrElse(
+ e.canonicalized,
+ aliasMap.getOrElse(throughUnary(e).canonicalized, e))
+ }
+
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 _: Alias | _: ExtractValue | _: Cast => e.children.forall(isCheap)
case _ => false
}
def infer(
plan: LogicalPlan,
- onlyInferWithCheapColumns: Boolean): Option[PartitioningAndOrdering] = {
+ onlyInferWithCheapColumns: Boolean,
+ maxColumns: Int): Option[PartitioningAndOrdering] = {
def candidateKeys(
input: LogicalPlan,
output: AttributeSet = AttributeSet.empty): Option[PartitioningAndOrdering] = {
@@ -92,15 +108,15 @@ object InferRebalanceAndSortOrders {
case agg: Aggregate =>
val aliasMap = getAliasMap(agg.aggregateExpressions)
Some((
- agg.groupingExpressions.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- agg.groupingExpressions.map(o => aliasMap.getOrElse(o.canonicalized, o))))
+ agg.groupingExpressions.map(p => mapThroughAlias(p, aliasMap)),
+ agg.groupingExpressions.map(o => mapThroughAlias(o, aliasMap))))
case s: Sort => Some((s.order.map(_.child), s.order.map(_.child)))
case p: Project =>
val aliasMap = getAliasMap(p.projectList)
candidateKeys(p.child, p.references).map { case (partitioning, ordering) =>
(
- partitioning.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- ordering.map(o => aliasMap.getOrElse(o.canonicalized, o)))
+ partitioning.map(p => mapThroughAlias(p, aliasMap)),
+ ordering.map(o => mapThroughAlias(o, aliasMap)))
}
case f: Filter => candidateKeys(f.child, output)
case s: SubqueryAlias => candidateKeys(s.child, output)
@@ -109,8 +125,8 @@ object InferRebalanceAndSortOrders {
case w: Window =>
val aliasMap = getAliasMap(w.windowExpressions)
Some((
- w.partitionSpec.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- w.orderSpec.map(_.child).map(o => aliasMap.getOrElse(o.canonicalized, o))))
+ w.partitionSpec.map(p => mapThroughAlias(p, aliasMap)),
+ w.orderSpec.map(_.child).map(o => mapThroughAlias(o, aliasMap))))
case _ => None
}
@@ -118,8 +134,8 @@ object InferRebalanceAndSortOrders {
val partitioningAndSort = candidateKeys(plan).map { case (partitioning, ordering) =>
(
- partitioning.filter(_.references.subsetOf(plan.outputSet)),
- ordering.filter(_.references.subsetOf(plan.outputSet)))
+ partitioning.filter(_.references.subsetOf(plan.outputSet)).take(maxColumns),
+ ordering.filter(_.references.subsetOf(plan.outputSet)).take(maxColumns))
}
val allCheap = partitioningAndSort.exists {
case (partitionings, sorts) =>
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 79aaa5c48e..66da4fe06c 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
@@ -79,16 +79,19 @@ trait RebalanceBeforeWritingBase extends Rule[LogicalPlan] {
val maxColumns = conf.getConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS)
val onlyInferWithCheapColumns = conf.getConf(
KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_WITH_CHEAP_COLUMNS)
- val inferred = InferRebalanceAndSortOrders.infer(query, onlyInferWithCheapColumns)
+ val inferred = InferRebalanceAndSortOrders.infer(
+ query,
+ onlyInferWithCheapColumns,
+ maxColumns)
if (inferred.isDefined) {
val (partitioning, ordering) = inferred.get
val rebalance = RebalancePartitions(
- partitioning.take(maxColumns),
+ partitioning,
query,
optAdvisoryPartitionSize = advisoryPartitionSize)
val skipInferOrder = conf.getConf(KyuubiSQLConf.SKIP_INFER_SORT_ORDERS)
if (!skipInferOrder && ordering.nonEmpty) {
- val sortOrders = ordering.take(maxColumns).map(o => SortOrder(o, Ascending))
+ val sortOrders = ordering.map(o => SortOrder(o, Ascending))
Sort(sortOrders, false, rebalance)
} else {
rebalance
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 e1ed7df936..88f535d8f0 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
@@ -360,9 +360,9 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
|""".stripMargin).queryExecution.analyzed
// when only cheap columns are allowed, the expensive keys are not inferred
- assert(InferRebalanceAndSortOrders.infer(expensivePlan, true).isEmpty)
+ assert(InferRebalanceAndSortOrders.infer(expensivePlan, true, 10).isEmpty)
// when the cheap-column restriction is disabled, the expensive keys are inferred
- InferRebalanceAndSortOrders.infer(expensivePlan, false) match {
+ InferRebalanceAndSortOrders.infer(expensivePlan, false, 10) match {
case Some((partitioning, ordering)) =>
assert(partitioning.nonEmpty)
assert(ordering.nonEmpty)
@@ -380,7 +380,7 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
// cheap columns are inferred regardless of the restriction
Seq(true, false).foreach { onlyCheap =>
- InferRebalanceAndSortOrders.infer(cheapPlan, onlyCheap) match {
+ InferRebalanceAndSortOrders.infer(cheapPlan, onlyCheap, 10) match {
case Some((partitioning, ordering)) =>
assert(partitioning.nonEmpty)
assert(ordering.nonEmpty)
@@ -390,6 +390,79 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
}
}
+ test("Infer cast-wrapped join keys (aliased through a cast)") {
+ withTable("t1", "t2") {
+ // genre_tag_id types differ: int vs bigint. The analyzer inserts a cast on the
+ // left join key, and the outer projection aliases it to `medal_type`.
+ sql("CREATE TABLE t1 (user_id int, genre_tag_id int) USING PARQUET")
+ sql("CREATE TABLE t2 (user_id int, genre_tag_id bigint) USING PARQUET")
+ val plan = sql(
+ """
+ |SELECT a.user_id, a.genre_tag_id as medal_type
+ |FROM (SELECT user_id, genre_tag_id FROM t1) a
+ |LEFT JOIN (SELECT user_id, genre_tag_id FROM t2) b
+ |ON a.user_id = b.user_id and a.genre_tag_id = b.genre_tag_id
+ |""".stripMargin).queryExecution.analyzed
+
+ // The cast-wrapped key resolves to its aliased output attribute, so both join
+ // keys are inferred (before the fix only `user_id` survived).
+ InferRebalanceAndSortOrders.infer(plan, true, 10) match {
+ case Some((partitioning, ordering)) =>
+ assert(partitioning.map(_.asInstanceOf[Attribute].name) === Seq("user_id", "medal_type"))
+ assert(ordering.map(_.asInstanceOf[Attribute].name) === Seq("user_id", "medal_type"))
+ case None => fail("expected the cast-wrapped join key to be inferred")
+ }
+ }
+ }
+
+ test("isCheap treats a cast over a cheap column as cheap") {
+ withTable("t1", "t2") {
+ // int vs bigint forces a cast on the left join key; the column is selected
+ // directly (no alias), so the cast survives to the cheap check.
+ sql("CREATE TABLE t1 (col1 int) USING PARQUET")
+ sql("CREATE TABLE t2 (col1 bigint) USING PARQUET")
+ val plan = sql(
+ """
+ |SELECT t1.col1
+ |FROM t1
+ |LEFT JOIN t2
+ |ON t1.col1 = t2.col1
+ |""".stripMargin).queryExecution.analyzed
+
+ InferRebalanceAndSortOrders.infer(plan, true, 10) match {
+ case Some((partitioning, _)) => assert(partitioning.nonEmpty)
+ case None => fail("expected a cast over a cheap column to be inferred")
+ }
+ }
+ }
+
+ test("maxColumns gates the cheap-column check") {
+ withTable("t1", "t2") {
+ sql("CREATE TABLE t1 (a int, b int, c int) USING PARQUET")
+ sql("CREATE TABLE t2 (a int, b int, c int) USING PARQUET")
+ // Three join keys; the third (`c + 1`) is an expensive expression.
+ val plan = sql(
+ """
+ |SELECT t1.a, t1.b, t1.c
+ |FROM t1
+ |LEFT JOIN t2
+ |ON t1.a = t2.a and t1.b = t2.b and t1.c + 1 = t2.c + 1
+ |""".stripMargin).queryExecution.analyzed
+
+ // With maxColumns = 2, only the first two (cheap) keys are considered, so
+ // inference succeeds.
+ InferRebalanceAndSortOrders.infer(plan, true, 2) match {
+ case Some((partitioning, _)) =>
+ assert(partitioning.map(_.asInstanceOf[Attribute].name) === Seq("a", "b"))
+ case None => fail("expected inference to succeed when the expensive key is truncated")
+ }
+
+ // With maxColumns = 3, the expensive third key is included and the cheap-only
+ // restriction discards everything.
+ assert(InferRebalanceAndSortOrders.infer(plan, true, 3).isEmpty)
+ }
+ }
+
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 ba80eda078..66bf225df3 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, BoundReference, Expression, ExtractValue, NamedExpression, OuterReference, UnaryExpression}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, BoundReference, Cast, 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}
@@ -38,14 +38,14 @@ object InferRebalanceAndSortOrders {
type PartitioningAndOrdering = (Seq[Expression], Seq[Expression])
- private def getAliasMap(named: Seq[NamedExpression]): Map[Expression, Attribute] = {
- @tailrec
- def throughUnary(e: Expression): Expression = e match {
- case u: UnaryExpression if u.deterministic =>
- throughUnary(u.child)
- case _ => e
- }
+ @tailrec
+ private def throughUnary(e: Expression): Expression = e match {
+ case u: UnaryExpression if u.deterministic =>
+ throughUnary(u.child)
+ case _ => e
+ }
+ private def getAliasMap(named: Seq[NamedExpression]): Map[Expression, Attribute] = {
named.flatMap {
case a @ Alias(child, _) =>
Some((throughUnary(child).canonicalized, a.toAttribute))
@@ -53,16 +53,32 @@ object InferRebalanceAndSortOrders {
}.toMap
}
+ /**
+ * Map an expression to its output attribute via the alias map. Tries the full
+ * expression first, then retries with the unary-peeled key (mirroring how
+ * `getAliasMap` builds its keys), so that a cast-wrapped key such as
+ * `cast(genre_tag_id as bigint)` still maps to its aliased output attribute. Falls
+ * back to the original expression when neither matches.
+ */
+ private def mapThroughAlias(
+ e: Expression,
+ aliasMap: Map[Expression, Attribute]): Expression = {
+ aliasMap.getOrElse(
+ e.canonicalized,
+ aliasMap.getOrElse(throughUnary(e).canonicalized, e))
+ }
+
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 _: Alias | _: ExtractValue | _: Cast => e.children.forall(isCheap)
case _ => false
}
def infer(
plan: LogicalPlan,
- onlyInferWithCheapColumns: Boolean): Option[PartitioningAndOrdering] = {
+ onlyInferWithCheapColumns: Boolean,
+ maxColumns: Int): Option[PartitioningAndOrdering] = {
def candidateKeys(
input: LogicalPlan,
output: AttributeSet = AttributeSet.empty): Option[PartitioningAndOrdering] = {
@@ -92,15 +108,15 @@ object InferRebalanceAndSortOrders {
case agg: Aggregate =>
val aliasMap = getAliasMap(agg.aggregateExpressions)
Some((
- agg.groupingExpressions.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- agg.groupingExpressions.map(o => aliasMap.getOrElse(o.canonicalized, o))))
+ agg.groupingExpressions.map(p => mapThroughAlias(p, aliasMap)),
+ agg.groupingExpressions.map(o => mapThroughAlias(o, aliasMap))))
case s: Sort => Some((s.order.map(_.child), s.order.map(_.child)))
case p: Project =>
val aliasMap = getAliasMap(p.projectList)
candidateKeys(p.child, p.references).map { case (partitioning, ordering) =>
(
- partitioning.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- ordering.map(o => aliasMap.getOrElse(o.canonicalized, o)))
+ partitioning.map(p => mapThroughAlias(p, aliasMap)),
+ ordering.map(o => mapThroughAlias(o, aliasMap)))
}
case f: Filter => candidateKeys(f.child, output)
case s: SubqueryAlias => candidateKeys(s.child, output)
@@ -109,8 +125,8 @@ object InferRebalanceAndSortOrders {
case w: Window =>
val aliasMap = getAliasMap(w.windowExpressions)
Some((
- w.partitionSpec.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- w.orderSpec.map(_.child).map(o => aliasMap.getOrElse(o.canonicalized, o))))
+ w.partitionSpec.map(p => mapThroughAlias(p, aliasMap)),
+ w.orderSpec.map(_.child).map(o => mapThroughAlias(o, aliasMap))))
case _ => None
}
@@ -118,8 +134,8 @@ object InferRebalanceAndSortOrders {
val partitioningAndSort = candidateKeys(plan).map { case (partitioning, ordering) =>
(
- partitioning.filter(_.references.subsetOf(plan.outputSet)),
- ordering.filter(_.references.subsetOf(plan.outputSet)))
+ partitioning.filter(_.references.subsetOf(plan.outputSet)).take(maxColumns),
+ ordering.filter(_.references.subsetOf(plan.outputSet)).take(maxColumns))
}
val allCheap = partitioningAndSort.exists {
case (partitionings, sorts) =>
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 5a8c43d969..8d71213129 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
@@ -78,16 +78,19 @@ trait RebalanceBeforeWritingBase extends Rule[LogicalPlan] {
val maxColumns = conf.getConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS)
val onlyInferWithCheapColumns = conf.getConf(
KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_WITH_CHEAP_COLUMNS)
- val inferred = InferRebalanceAndSortOrders.infer(query, onlyInferWithCheapColumns)
+ val inferred = InferRebalanceAndSortOrders.infer(
+ query,
+ onlyInferWithCheapColumns,
+ maxColumns)
if (inferred.isDefined) {
val (partitioning, ordering) = inferred.get
val rebalance = RebalancePartitions(
- partitioning.take(maxColumns),
+ partitioning,
query,
optAdvisoryPartitionSize = advisoryPartitionSize)
val skipInferOrder = conf.getConf(KyuubiSQLConf.SKIP_INFER_SORT_ORDERS)
if (!skipInferOrder && ordering.nonEmpty) {
- val sortOrders = ordering.take(maxColumns).map(o => SortOrder(o, Ascending))
+ val sortOrders = ordering.map(o => SortOrder(o, Ascending))
Sort(sortOrders, false, rebalance)
} else {
rebalance
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 ff516964be..e3d2ac143f 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
@@ -359,9 +359,9 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
|""".stripMargin).queryExecution.analyzed
// when only cheap columns are allowed, the expensive keys are not inferred
- assert(InferRebalanceAndSortOrders.infer(expensivePlan, true).isEmpty)
+ assert(InferRebalanceAndSortOrders.infer(expensivePlan, true, 10).isEmpty)
// when the cheap-column restriction is disabled, the expensive keys are inferred
- InferRebalanceAndSortOrders.infer(expensivePlan, false) match {
+ InferRebalanceAndSortOrders.infer(expensivePlan, false, 10) match {
case Some((partitioning, ordering)) =>
assert(partitioning.nonEmpty)
assert(ordering.nonEmpty)
@@ -379,7 +379,7 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
// cheap columns are inferred regardless of the restriction
Seq(true, false).foreach { onlyCheap =>
- InferRebalanceAndSortOrders.infer(cheapPlan, onlyCheap) match {
+ InferRebalanceAndSortOrders.infer(cheapPlan, onlyCheap, 10) match {
case Some((partitioning, ordering)) =>
assert(partitioning.nonEmpty)
assert(ordering.nonEmpty)
@@ -389,6 +389,79 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
}
}
+ test("Infer cast-wrapped join keys (aliased through a cast)") {
+ withTable("t1", "t2") {
+ // genre_tag_id types differ: int vs bigint. The analyzer inserts a cast on the
+ // left join key, and the outer projection aliases it to `medal_type`.
+ sql("CREATE TABLE t1 (user_id int, genre_tag_id int) USING PARQUET")
+ sql("CREATE TABLE t2 (user_id int, genre_tag_id bigint) USING PARQUET")
+ val plan = sql(
+ """
+ |SELECT a.user_id, a.genre_tag_id as medal_type
+ |FROM (SELECT user_id, genre_tag_id FROM t1) a
+ |LEFT JOIN (SELECT user_id, genre_tag_id FROM t2) b
+ |ON a.user_id = b.user_id and a.genre_tag_id = b.genre_tag_id
+ |""".stripMargin).queryExecution.analyzed
+
+ // The cast-wrapped key resolves to its aliased output attribute, so both join
+ // keys are inferred (before the fix only `user_id` survived).
+ InferRebalanceAndSortOrders.infer(plan, true, 10) match {
+ case Some((partitioning, ordering)) =>
+ assert(partitioning.map(_.asInstanceOf[Attribute].name) === Seq("user_id", "medal_type"))
+ assert(ordering.map(_.asInstanceOf[Attribute].name) === Seq("user_id", "medal_type"))
+ case None => fail("expected the cast-wrapped join key to be inferred")
+ }
+ }
+ }
+
+ test("isCheap treats a cast over a cheap column as cheap") {
+ withTable("t1", "t2") {
+ // int vs bigint forces a cast on the left join key; the column is selected
+ // directly (no alias), so the cast survives to the cheap check.
+ sql("CREATE TABLE t1 (col1 int) USING PARQUET")
+ sql("CREATE TABLE t2 (col1 bigint) USING PARQUET")
+ val plan = sql(
+ """
+ |SELECT t1.col1
+ |FROM t1
+ |LEFT JOIN t2
+ |ON t1.col1 = t2.col1
+ |""".stripMargin).queryExecution.analyzed
+
+ InferRebalanceAndSortOrders.infer(plan, true, 10) match {
+ case Some((partitioning, _)) => assert(partitioning.nonEmpty)
+ case None => fail("expected a cast over a cheap column to be inferred")
+ }
+ }
+ }
+
+ test("maxColumns gates the cheap-column check") {
+ withTable("t1", "t2") {
+ sql("CREATE TABLE t1 (a int, b int, c int) USING PARQUET")
+ sql("CREATE TABLE t2 (a int, b int, c int) USING PARQUET")
+ // Three join keys; the third (`c + 1`) is an expensive expression.
+ val plan = sql(
+ """
+ |SELECT t1.a, t1.b, t1.c
+ |FROM t1
+ |LEFT JOIN t2
+ |ON t1.a = t2.a and t1.b = t2.b and t1.c + 1 = t2.c + 1
+ |""".stripMargin).queryExecution.analyzed
+
+ // With maxColumns = 2, only the first two (cheap) keys are considered, so
+ // inference succeeds.
+ InferRebalanceAndSortOrders.infer(plan, true, 2) match {
+ case Some((partitioning, _)) =>
+ assert(partitioning.map(_.asInstanceOf[Attribute].name) === Seq("a", "b"))
+ case None => fail("expected inference to succeed when the expensive key is truncated")
+ }
+
+ // With maxColumns = 3, the expensive third key is included and the cheap-only
+ // restriction discards everything.
+ assert(InferRebalanceAndSortOrders.infer(plan, true, 3).isEmpty)
+ }
+ }
+
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 ba80eda078..66bf225df3 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, BoundReference, Expression, ExtractValue, NamedExpression, OuterReference, UnaryExpression}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, BoundReference, Cast, 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}
@@ -38,14 +38,14 @@ object InferRebalanceAndSortOrders {
type PartitioningAndOrdering = (Seq[Expression], Seq[Expression])
- private def getAliasMap(named: Seq[NamedExpression]): Map[Expression, Attribute] = {
- @tailrec
- def throughUnary(e: Expression): Expression = e match {
- case u: UnaryExpression if u.deterministic =>
- throughUnary(u.child)
- case _ => e
- }
+ @tailrec
+ private def throughUnary(e: Expression): Expression = e match {
+ case u: UnaryExpression if u.deterministic =>
+ throughUnary(u.child)
+ case _ => e
+ }
+ private def getAliasMap(named: Seq[NamedExpression]): Map[Expression, Attribute] = {
named.flatMap {
case a @ Alias(child, _) =>
Some((throughUnary(child).canonicalized, a.toAttribute))
@@ -53,16 +53,32 @@ object InferRebalanceAndSortOrders {
}.toMap
}
+ /**
+ * Map an expression to its output attribute via the alias map. Tries the full
+ * expression first, then retries with the unary-peeled key (mirroring how
+ * `getAliasMap` builds its keys), so that a cast-wrapped key such as
+ * `cast(genre_tag_id as bigint)` still maps to its aliased output attribute. Falls
+ * back to the original expression when neither matches.
+ */
+ private def mapThroughAlias(
+ e: Expression,
+ aliasMap: Map[Expression, Attribute]): Expression = {
+ aliasMap.getOrElse(
+ e.canonicalized,
+ aliasMap.getOrElse(throughUnary(e).canonicalized, e))
+ }
+
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 _: Alias | _: ExtractValue | _: Cast => e.children.forall(isCheap)
case _ => false
}
def infer(
plan: LogicalPlan,
- onlyInferWithCheapColumns: Boolean): Option[PartitioningAndOrdering] = {
+ onlyInferWithCheapColumns: Boolean,
+ maxColumns: Int): Option[PartitioningAndOrdering] = {
def candidateKeys(
input: LogicalPlan,
output: AttributeSet = AttributeSet.empty): Option[PartitioningAndOrdering] = {
@@ -92,15 +108,15 @@ object InferRebalanceAndSortOrders {
case agg: Aggregate =>
val aliasMap = getAliasMap(agg.aggregateExpressions)
Some((
- agg.groupingExpressions.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- agg.groupingExpressions.map(o => aliasMap.getOrElse(o.canonicalized, o))))
+ agg.groupingExpressions.map(p => mapThroughAlias(p, aliasMap)),
+ agg.groupingExpressions.map(o => mapThroughAlias(o, aliasMap))))
case s: Sort => Some((s.order.map(_.child), s.order.map(_.child)))
case p: Project =>
val aliasMap = getAliasMap(p.projectList)
candidateKeys(p.child, p.references).map { case (partitioning, ordering) =>
(
- partitioning.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- ordering.map(o => aliasMap.getOrElse(o.canonicalized, o)))
+ partitioning.map(p => mapThroughAlias(p, aliasMap)),
+ ordering.map(o => mapThroughAlias(o, aliasMap)))
}
case f: Filter => candidateKeys(f.child, output)
case s: SubqueryAlias => candidateKeys(s.child, output)
@@ -109,8 +125,8 @@ object InferRebalanceAndSortOrders {
case w: Window =>
val aliasMap = getAliasMap(w.windowExpressions)
Some((
- w.partitionSpec.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- w.orderSpec.map(_.child).map(o => aliasMap.getOrElse(o.canonicalized, o))))
+ w.partitionSpec.map(p => mapThroughAlias(p, aliasMap)),
+ w.orderSpec.map(_.child).map(o => mapThroughAlias(o, aliasMap))))
case _ => None
}
@@ -118,8 +134,8 @@ object InferRebalanceAndSortOrders {
val partitioningAndSort = candidateKeys(plan).map { case (partitioning, ordering) =>
(
- partitioning.filter(_.references.subsetOf(plan.outputSet)),
- ordering.filter(_.references.subsetOf(plan.outputSet)))
+ partitioning.filter(_.references.subsetOf(plan.outputSet)).take(maxColumns),
+ ordering.filter(_.references.subsetOf(plan.outputSet)).take(maxColumns))
}
val allCheap = partitioningAndSort.exists {
case (partitionings, sorts) =>
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 5a8c43d969..8d71213129 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
@@ -78,16 +78,19 @@ trait RebalanceBeforeWritingBase extends Rule[LogicalPlan] {
val maxColumns = conf.getConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS)
val onlyInferWithCheapColumns = conf.getConf(
KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_WITH_CHEAP_COLUMNS)
- val inferred = InferRebalanceAndSortOrders.infer(query, onlyInferWithCheapColumns)
+ val inferred = InferRebalanceAndSortOrders.infer(
+ query,
+ onlyInferWithCheapColumns,
+ maxColumns)
if (inferred.isDefined) {
val (partitioning, ordering) = inferred.get
val rebalance = RebalancePartitions(
- partitioning.take(maxColumns),
+ partitioning,
query,
optAdvisoryPartitionSize = advisoryPartitionSize)
val skipInferOrder = conf.getConf(KyuubiSQLConf.SKIP_INFER_SORT_ORDERS)
if (!skipInferOrder && ordering.nonEmpty) {
- val sortOrders = ordering.take(maxColumns).map(o => SortOrder(o, Ascending))
+ val sortOrders = ordering.map(o => SortOrder(o, Ascending))
Sort(sortOrders, false, rebalance)
} else {
rebalance
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 ff516964be..e3d2ac143f 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
@@ -359,9 +359,9 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
|""".stripMargin).queryExecution.analyzed
// when only cheap columns are allowed, the expensive keys are not inferred
- assert(InferRebalanceAndSortOrders.infer(expensivePlan, true).isEmpty)
+ assert(InferRebalanceAndSortOrders.infer(expensivePlan, true, 10).isEmpty)
// when the cheap-column restriction is disabled, the expensive keys are inferred
- InferRebalanceAndSortOrders.infer(expensivePlan, false) match {
+ InferRebalanceAndSortOrders.infer(expensivePlan, false, 10) match {
case Some((partitioning, ordering)) =>
assert(partitioning.nonEmpty)
assert(ordering.nonEmpty)
@@ -379,7 +379,7 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
// cheap columns are inferred regardless of the restriction
Seq(true, false).foreach { onlyCheap =>
- InferRebalanceAndSortOrders.infer(cheapPlan, onlyCheap) match {
+ InferRebalanceAndSortOrders.infer(cheapPlan, onlyCheap, 10) match {
case Some((partitioning, ordering)) =>
assert(partitioning.nonEmpty)
assert(ordering.nonEmpty)
@@ -389,6 +389,79 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
}
}
+ test("Infer cast-wrapped join keys (aliased through a cast)") {
+ withTable("t1", "t2") {
+ // genre_tag_id types differ: int vs bigint. The analyzer inserts a cast on the
+ // left join key, and the outer projection aliases it to `medal_type`.
+ sql("CREATE TABLE t1 (user_id int, genre_tag_id int) USING PARQUET")
+ sql("CREATE TABLE t2 (user_id int, genre_tag_id bigint) USING PARQUET")
+ val plan = sql(
+ """
+ |SELECT a.user_id, a.genre_tag_id as medal_type
+ |FROM (SELECT user_id, genre_tag_id FROM t1) a
+ |LEFT JOIN (SELECT user_id, genre_tag_id FROM t2) b
+ |ON a.user_id = b.user_id and a.genre_tag_id = b.genre_tag_id
+ |""".stripMargin).queryExecution.analyzed
+
+ // The cast-wrapped key resolves to its aliased output attribute, so both join
+ // keys are inferred (before the fix only `user_id` survived).
+ InferRebalanceAndSortOrders.infer(plan, true, 10) match {
+ case Some((partitioning, ordering)) =>
+ assert(partitioning.map(_.asInstanceOf[Attribute].name) === Seq("user_id", "medal_type"))
+ assert(ordering.map(_.asInstanceOf[Attribute].name) === Seq("user_id", "medal_type"))
+ case None => fail("expected the cast-wrapped join key to be inferred")
+ }
+ }
+ }
+
+ test("isCheap treats a cast over a cheap column as cheap") {
+ withTable("t1", "t2") {
+ // int vs bigint forces a cast on the left join key; the column is selected
+ // directly (no alias), so the cast survives to the cheap check.
+ sql("CREATE TABLE t1 (col1 int) USING PARQUET")
+ sql("CREATE TABLE t2 (col1 bigint) USING PARQUET")
+ val plan = sql(
+ """
+ |SELECT t1.col1
+ |FROM t1
+ |LEFT JOIN t2
+ |ON t1.col1 = t2.col1
+ |""".stripMargin).queryExecution.analyzed
+
+ InferRebalanceAndSortOrders.infer(plan, true, 10) match {
+ case Some((partitioning, _)) => assert(partitioning.nonEmpty)
+ case None => fail("expected a cast over a cheap column to be inferred")
+ }
+ }
+ }
+
+ test("maxColumns gates the cheap-column check") {
+ withTable("t1", "t2") {
+ sql("CREATE TABLE t1 (a int, b int, c int) USING PARQUET")
+ sql("CREATE TABLE t2 (a int, b int, c int) USING PARQUET")
+ // Three join keys; the third (`c + 1`) is an expensive expression.
+ val plan = sql(
+ """
+ |SELECT t1.a, t1.b, t1.c
+ |FROM t1
+ |LEFT JOIN t2
+ |ON t1.a = t2.a and t1.b = t2.b and t1.c + 1 = t2.c + 1
+ |""".stripMargin).queryExecution.analyzed
+
+ // With maxColumns = 2, only the first two (cheap) keys are considered, so
+ // inference succeeds.
+ InferRebalanceAndSortOrders.infer(plan, true, 2) match {
+ case Some((partitioning, _)) =>
+ assert(partitioning.map(_.asInstanceOf[Attribute].name) === Seq("a", "b"))
+ case None => fail("expected inference to succeed when the expensive key is truncated")
+ }
+
+ // With maxColumns = 3, the expensive third key is included and the cheap-only
+ // restriction discards everything.
+ assert(InferRebalanceAndSortOrders.infer(plan, true, 3).isEmpty)
+ }
+ }
+
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-2/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala b/extensions/spark/kyuubi-extension-spark-4-2/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala
index ba80eda078..66bf225df3 100644
--- a/extensions/spark/kyuubi-extension-spark-4-2/src/main/scala/org/apache/kyuubi/sql/InferRebalanceAndSortOrders.scala
+++ b/extensions/spark/kyuubi-extension-spark-4-2/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, BoundReference, Expression, ExtractValue, NamedExpression, OuterReference, UnaryExpression}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, BoundReference, Cast, 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}
@@ -38,14 +38,14 @@ object InferRebalanceAndSortOrders {
type PartitioningAndOrdering = (Seq[Expression], Seq[Expression])
- private def getAliasMap(named: Seq[NamedExpression]): Map[Expression, Attribute] = {
- @tailrec
- def throughUnary(e: Expression): Expression = e match {
- case u: UnaryExpression if u.deterministic =>
- throughUnary(u.child)
- case _ => e
- }
+ @tailrec
+ private def throughUnary(e: Expression): Expression = e match {
+ case u: UnaryExpression if u.deterministic =>
+ throughUnary(u.child)
+ case _ => e
+ }
+ private def getAliasMap(named: Seq[NamedExpression]): Map[Expression, Attribute] = {
named.flatMap {
case a @ Alias(child, _) =>
Some((throughUnary(child).canonicalized, a.toAttribute))
@@ -53,16 +53,32 @@ object InferRebalanceAndSortOrders {
}.toMap
}
+ /**
+ * Map an expression to its output attribute via the alias map. Tries the full
+ * expression first, then retries with the unary-peeled key (mirroring how
+ * `getAliasMap` builds its keys), so that a cast-wrapped key such as
+ * `cast(genre_tag_id as bigint)` still maps to its aliased output attribute. Falls
+ * back to the original expression when neither matches.
+ */
+ private def mapThroughAlias(
+ e: Expression,
+ aliasMap: Map[Expression, Attribute]): Expression = {
+ aliasMap.getOrElse(
+ e.canonicalized,
+ aliasMap.getOrElse(throughUnary(e).canonicalized, e))
+ }
+
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 _: Alias | _: ExtractValue | _: Cast => e.children.forall(isCheap)
case _ => false
}
def infer(
plan: LogicalPlan,
- onlyInferWithCheapColumns: Boolean): Option[PartitioningAndOrdering] = {
+ onlyInferWithCheapColumns: Boolean,
+ maxColumns: Int): Option[PartitioningAndOrdering] = {
def candidateKeys(
input: LogicalPlan,
output: AttributeSet = AttributeSet.empty): Option[PartitioningAndOrdering] = {
@@ -92,15 +108,15 @@ object InferRebalanceAndSortOrders {
case agg: Aggregate =>
val aliasMap = getAliasMap(agg.aggregateExpressions)
Some((
- agg.groupingExpressions.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- agg.groupingExpressions.map(o => aliasMap.getOrElse(o.canonicalized, o))))
+ agg.groupingExpressions.map(p => mapThroughAlias(p, aliasMap)),
+ agg.groupingExpressions.map(o => mapThroughAlias(o, aliasMap))))
case s: Sort => Some((s.order.map(_.child), s.order.map(_.child)))
case p: Project =>
val aliasMap = getAliasMap(p.projectList)
candidateKeys(p.child, p.references).map { case (partitioning, ordering) =>
(
- partitioning.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- ordering.map(o => aliasMap.getOrElse(o.canonicalized, o)))
+ partitioning.map(p => mapThroughAlias(p, aliasMap)),
+ ordering.map(o => mapThroughAlias(o, aliasMap)))
}
case f: Filter => candidateKeys(f.child, output)
case s: SubqueryAlias => candidateKeys(s.child, output)
@@ -109,8 +125,8 @@ object InferRebalanceAndSortOrders {
case w: Window =>
val aliasMap = getAliasMap(w.windowExpressions)
Some((
- w.partitionSpec.map(p => aliasMap.getOrElse(p.canonicalized, p)),
- w.orderSpec.map(_.child).map(o => aliasMap.getOrElse(o.canonicalized, o))))
+ w.partitionSpec.map(p => mapThroughAlias(p, aliasMap)),
+ w.orderSpec.map(_.child).map(o => mapThroughAlias(o, aliasMap))))
case _ => None
}
@@ -118,8 +134,8 @@ object InferRebalanceAndSortOrders {
val partitioningAndSort = candidateKeys(plan).map { case (partitioning, ordering) =>
(
- partitioning.filter(_.references.subsetOf(plan.outputSet)),
- ordering.filter(_.references.subsetOf(plan.outputSet)))
+ partitioning.filter(_.references.subsetOf(plan.outputSet)).take(maxColumns),
+ ordering.filter(_.references.subsetOf(plan.outputSet)).take(maxColumns))
}
val allCheap = partitioningAndSort.exists {
case (partitionings, sorts) =>
diff --git a/extensions/spark/kyuubi-extension-spark-4-2/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala b/extensions/spark/kyuubi-extension-spark-4-2/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala
index 5a8c43d969..8d71213129 100644
--- a/extensions/spark/kyuubi-extension-spark-4-2/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala
+++ b/extensions/spark/kyuubi-extension-spark-4-2/src/main/scala/org/apache/kyuubi/sql/RebalanceBeforeWriting.scala
@@ -78,16 +78,19 @@ trait RebalanceBeforeWritingBase extends Rule[LogicalPlan] {
val maxColumns = conf.getConf(KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS)
val onlyInferWithCheapColumns = conf.getConf(
KyuubiSQLConf.INFER_REBALANCE_AND_SORT_ORDERS_WITH_CHEAP_COLUMNS)
- val inferred = InferRebalanceAndSortOrders.infer(query, onlyInferWithCheapColumns)
+ val inferred = InferRebalanceAndSortOrders.infer(
+ query,
+ onlyInferWithCheapColumns,
+ maxColumns)
if (inferred.isDefined) {
val (partitioning, ordering) = inferred.get
val rebalance = RebalancePartitions(
- partitioning.take(maxColumns),
+ partitioning,
query,
optAdvisoryPartitionSize = advisoryPartitionSize)
val skipInferOrder = conf.getConf(KyuubiSQLConf.SKIP_INFER_SORT_ORDERS)
if (!skipInferOrder && ordering.nonEmpty) {
- val sortOrders = ordering.take(maxColumns).map(o => SortOrder(o, Ascending))
+ val sortOrders = ordering.map(o => SortOrder(o, Ascending))
Sort(sortOrders, false, rebalance)
} else {
rebalance
diff --git a/extensions/spark/kyuubi-extension-spark-4-2/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala b/extensions/spark/kyuubi-extension-spark-4-2/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala
index ff516964be..e3d2ac143f 100644
--- a/extensions/spark/kyuubi-extension-spark-4-2/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala
+++ b/extensions/spark/kyuubi-extension-spark-4-2/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala
@@ -359,9 +359,9 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
|""".stripMargin).queryExecution.analyzed
// when only cheap columns are allowed, the expensive keys are not inferred
- assert(InferRebalanceAndSortOrders.infer(expensivePlan, true).isEmpty)
+ assert(InferRebalanceAndSortOrders.infer(expensivePlan, true, 10).isEmpty)
// when the cheap-column restriction is disabled, the expensive keys are inferred
- InferRebalanceAndSortOrders.infer(expensivePlan, false) match {
+ InferRebalanceAndSortOrders.infer(expensivePlan, false, 10) match {
case Some((partitioning, ordering)) =>
assert(partitioning.nonEmpty)
assert(ordering.nonEmpty)
@@ -379,7 +379,7 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
// cheap columns are inferred regardless of the restriction
Seq(true, false).foreach { onlyCheap =>
- InferRebalanceAndSortOrders.infer(cheapPlan, onlyCheap) match {
+ InferRebalanceAndSortOrders.infer(cheapPlan, onlyCheap, 10) match {
case Some((partitioning, ordering)) =>
assert(partitioning.nonEmpty)
assert(ordering.nonEmpty)
@@ -389,6 +389,79 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
}
}
+ test("Infer cast-wrapped join keys (aliased through a cast)") {
+ withTable("t1", "t2") {
+ // genre_tag_id types differ: int vs bigint. The analyzer inserts a cast on the
+ // left join key, and the outer projection aliases it to `medal_type`.
+ sql("CREATE TABLE t1 (user_id int, genre_tag_id int) USING PARQUET")
+ sql("CREATE TABLE t2 (user_id int, genre_tag_id bigint) USING PARQUET")
+ val plan = sql(
+ """
+ |SELECT a.user_id, a.genre_tag_id as medal_type
+ |FROM (SELECT user_id, genre_tag_id FROM t1) a
+ |LEFT JOIN (SELECT user_id, genre_tag_id FROM t2) b
+ |ON a.user_id = b.user_id and a.genre_tag_id = b.genre_tag_id
+ |""".stripMargin).queryExecution.analyzed
+
+ // The cast-wrapped key resolves to its aliased output attribute, so both join
+ // keys are inferred (before the fix only `user_id` survived).
+ InferRebalanceAndSortOrders.infer(plan, true, 10) match {
+ case Some((partitioning, ordering)) =>
+ assert(partitioning.map(_.asInstanceOf[Attribute].name) === Seq("user_id", "medal_type"))
+ assert(ordering.map(_.asInstanceOf[Attribute].name) === Seq("user_id", "medal_type"))
+ case None => fail("expected the cast-wrapped join key to be inferred")
+ }
+ }
+ }
+
+ test("isCheap treats a cast over a cheap column as cheap") {
+ withTable("t1", "t2") {
+ // int vs bigint forces a cast on the left join key; the column is selected
+ // directly (no alias), so the cast survives to the cheap check.
+ sql("CREATE TABLE t1 (col1 int) USING PARQUET")
+ sql("CREATE TABLE t2 (col1 bigint) USING PARQUET")
+ val plan = sql(
+ """
+ |SELECT t1.col1
+ |FROM t1
+ |LEFT JOIN t2
+ |ON t1.col1 = t2.col1
+ |""".stripMargin).queryExecution.analyzed
+
+ InferRebalanceAndSortOrders.infer(plan, true, 10) match {
+ case Some((partitioning, _)) => assert(partitioning.nonEmpty)
+ case None => fail("expected a cast over a cheap column to be inferred")
+ }
+ }
+ }
+
+ test("maxColumns gates the cheap-column check") {
+ withTable("t1", "t2") {
+ sql("CREATE TABLE t1 (a int, b int, c int) USING PARQUET")
+ sql("CREATE TABLE t2 (a int, b int, c int) USING PARQUET")
+ // Three join keys; the third (`c + 1`) is an expensive expression.
+ val plan = sql(
+ """
+ |SELECT t1.a, t1.b, t1.c
+ |FROM t1
+ |LEFT JOIN t2
+ |ON t1.a = t2.a and t1.b = t2.b and t1.c + 1 = t2.c + 1
+ |""".stripMargin).queryExecution.analyzed
+
+ // With maxColumns = 2, only the first two (cheap) keys are considered, so
+ // inference succeeds.
+ InferRebalanceAndSortOrders.infer(plan, true, 2) match {
+ case Some((partitioning, _)) =>
+ assert(partitioning.map(_.asInstanceOf[Attribute].name) === Seq("a", "b"))
+ case None => fail("expected inference to succeed when the expensive key is truncated")
+ }
+
+ // With maxColumns = 3, the expensive third key is included and the cheap-only
+ // restriction discards everything.
+ assert(InferRebalanceAndSortOrders.infer(plan, true, 3).isEmpty)
+ }
+ }
+
test("advisory partition size - finalStage config takes precedence") {
withSQLConf(
KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true",