Skip to content

Commit 6853c22

Browse files
brijrajkbrijrajk
authored
[GLUTEN-10992][VL] Fix MatchError for KeyGroupedPartitioning in native shuffle (#12335)
What changes were proposed in this pull request? When Spark 4.0's V2 bucketing shuffle (spark.sql.v2.bucketing.shuffle.enabled=true) is used in a join where only one side reports partitioning, Spark generates a ShuffleExchangeExec with KeyGroupedPartitioning as its output partitioning. The default case _ => in VeloxSparkPlanExecApi.genColumnarShuffleExchange created a ColumnarShuffleExchangeExec for this node without validation. When the query executed, ExecUtil.genShuffleDependency crashed with a scala.MatchError because KeyGroupedPartitioning was missing from its exhaustive match. Changes: VeloxSparkPlanExecApi.genColumnarShuffleExchange: add an explicit case _: KeyGroupedPartitioning => before the default that adds a fallback tag and returns the vanilla ShuffleExchangeExec. This prevents a ColumnarShuffleExchangeExec from being created for an unsupported partitioning type. ExecUtil.genShuffleDependency: add an explicit wildcard case other => that throws GlutenNotSupportException instead of the cryptic scala.MatchError, as a defensive guard for any future unknown partitioning types. How was this patch tested? The existing testGluten("SPARK-41471: shuffle one side: only one side reports partitioning") tests in GlutenKeyGroupedPartitioningSuite (both spark40 and spark41) reproduce the crash exactly — they set V2_BUCKETING_SHUFFLE_ENABLED=true with only one bucketed side, which triggers a ShuffleExchangeExec with KeyGroupedPartitioning output and then call checkAnswer. After this fix these tests pass without MatchError. Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (https://claude.ai/code) Related issue: #10992
1 parent 999a949 commit 6853c22

4 files changed

Lines changed: 102 additions & 0 deletions

File tree

‎backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -471,6 +471,12 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
471471
}
472472
}
473473
}
474+
case _: KeyGroupedPartitioning =>
475+
FallbackTags.add(
476+
shuffle,
477+
ValidationResult.failed(
478+
"KeyGroupedPartitioning is not supported by Gluten native shuffle"))
479+
shuffle.withNewChildren(child :: Nil)
474480
case _ =>
475481
ColumnarShuffleExchangeExec(shuffle, child, null)
476482
}

‎backends-velox/src/main/scala/org/apache/spark/sql/execution/utils/ExecUtil.scala‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ package org.apache.spark.sql.execution.utils
1919
import org.apache.gluten.backendsapi.BackendsApiManager
2020
import org.apache.gluten.columnarbatch.{ColumnarBatches, VeloxColumnarBatches}
2121
import org.apache.gluten.config.ShuffleWriterType
22+
import org.apache.gluten.exception.GlutenNotSupportException
2223
import org.apache.gluten.iterator.Iterators
2324
import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators
2425
import org.apache.gluten.runtime.Runtimes
@@ -172,6 +173,9 @@ object ExecUtil {
172173
// range partitioning fall back to row-based partition id computation
173174
case RangePartitioning(orders, n) =>
174175
new NativePartitioning(GlutenShuffleUtils.RangePartitioningShortName, n)
176+
case other =>
177+
throw new GlutenNotSupportException(
178+
s"Partitioning $other is not supported by native shuffle")
175179
}
176180

177181
val isRoundRobin = newPartitioning.isInstanceOf[RoundRobinPartitioning] &&

‎gluten-ut/spark40/src/test/scala/org/apache/spark/sql/connector/GlutenKeyGroupedPartitioningSuite.scala‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import org.apache.gluten.execution.SortMergeJoinExecTransformer
2121

2222
import org.apache.spark.SparkConf
2323
import org.apache.spark.sql.{DataFrame, GlutenSQLTestsBaseTrait, Row}
24+
import org.apache.spark.sql.catalyst.plans.physical.KeyGroupedPartitioning
2425
import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryTableCatalog}
2526
import org.apache.spark.sql.connector.distributions.Distributions
2627
import org.apache.spark.sql.connector.expressions.Expressions.{bucket, days, identity, years}
@@ -1072,6 +1073,51 @@ class GlutenKeyGroupedPartitioningSuite
10721073
}
10731074
}
10741075

1076+
testGluten(
1077+
"GLUTEN-10992: KeyGroupedPartitioning shuffle falls back to vanilla Spark") {
1078+
val items_partitions = Array(identity("id"))
1079+
createTable(items, itemsColumns, items_partitions)
1080+
1081+
sql(
1082+
s"INSERT INTO testcat.ns.$items VALUES " +
1083+
"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
1084+
"(3, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
1085+
"(4, 'cc', 15.5, cast('2020-02-01' as timestamp))")
1086+
1087+
createTable(purchases, purchasesColumns, Array.empty)
1088+
sql(
1089+
s"INSERT INTO testcat.ns.$purchases VALUES " +
1090+
"(1, 42.0, cast('2020-01-01' as timestamp)), " +
1091+
"(3, 19.5, cast('2020-02-01' as timestamp))")
1092+
1093+
// With V2 bucketing shuffle enabled and only one side reporting partitioning, Spark
1094+
// shuffles the other side with a ShuffleExchangeExec whose output partitioning is
1095+
// KeyGroupedPartitioning. Gluten native shuffle does not support it, so the exchange
1096+
// must fall back to vanilla Spark. Offloading it to ColumnarShuffleExchangeExec would
1097+
// crash with a scala.MatchError in ExecUtil.genShuffleDependency (GLUTEN-10992).
1098+
withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
1099+
val df = createJoinTestDF(Seq("id" -> "item_id"))
1100+
val plan = df.queryExecution.executedPlan
1101+
1102+
val keyGroupedShuffles = collect(plan) {
1103+
case s: ShuffleExchangeExec
1104+
if s.outputPartitioning.isInstanceOf[KeyGroupedPartitioning] =>
1105+
s
1106+
}
1107+
assert(
1108+
keyGroupedShuffles.nonEmpty,
1109+
"KeyGroupedPartitioning shuffle should fall back to a vanilla ShuffleExchangeExec")
1110+
1111+
val columnarKeyGroupedShuffles = collectAllShuffles(plan)
1112+
.filter(_.outputPartitioning.isInstanceOf[KeyGroupedPartitioning])
1113+
assert(
1114+
columnarKeyGroupedShuffles.isEmpty,
1115+
"KeyGroupedPartitioning must not be offloaded to ColumnarShuffleExchangeExec")
1116+
1117+
checkAnswer(df, Seq(Row(1, "aa", 40.0, 42.0), Row(3, "bb", 10.0, 19.5)))
1118+
}
1119+
}
1120+
10751121
testGluten("SPARK-41471: shuffle one side: only one side reports partitioning") {
10761122
val items_partitions = Array(identity("id"))
10771123
createTable(items, itemsColumns, items_partitions)

‎gluten-ut/spark41/src/test/scala/org/apache/spark/sql/connector/GlutenKeyGroupedPartitioningSuite.scala‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import org.apache.gluten.execution.SortMergeJoinExecTransformer
2121

2222
import org.apache.spark.SparkConf
2323
import org.apache.spark.sql.{DataFrame, GlutenSQLTestsBaseTrait, Row}
24+
import org.apache.spark.sql.catalyst.plans.physical.KeyGroupedPartitioning
2425
import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryTableCatalog}
2526
import org.apache.spark.sql.connector.distributions.Distributions
2627
import org.apache.spark.sql.connector.expressions.Expressions.{bucket, days, identity, years}
@@ -1072,6 +1073,51 @@ class GlutenKeyGroupedPartitioningSuite
10721073
}
10731074
}
10741075

1076+
testGluten(
1077+
"GLUTEN-10992: KeyGroupedPartitioning shuffle falls back to vanilla Spark") {
1078+
val items_partitions = Array(identity("id"))
1079+
createTable(items, itemsColumns, items_partitions)
1080+
1081+
sql(
1082+
s"INSERT INTO testcat.ns.$items VALUES " +
1083+
"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
1084+
"(3, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
1085+
"(4, 'cc', 15.5, cast('2020-02-01' as timestamp))")
1086+
1087+
createTable(purchases, purchasesColumns, Array.empty)
1088+
sql(
1089+
s"INSERT INTO testcat.ns.$purchases VALUES " +
1090+
"(1, 42.0, cast('2020-01-01' as timestamp)), " +
1091+
"(3, 19.5, cast('2020-02-01' as timestamp))")
1092+
1093+
// With V2 bucketing shuffle enabled and only one side reporting partitioning, Spark
1094+
// shuffles the other side with a ShuffleExchangeExec whose output partitioning is
1095+
// KeyGroupedPartitioning. Gluten native shuffle does not support it, so the exchange
1096+
// must fall back to vanilla Spark. Offloading it to ColumnarShuffleExchangeExec would
1097+
// crash with a scala.MatchError in ExecUtil.genShuffleDependency (GLUTEN-10992).
1098+
withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
1099+
val df = createJoinTestDF(Seq("id" -> "item_id"))
1100+
val plan = df.queryExecution.executedPlan
1101+
1102+
val keyGroupedShuffles = collect(plan) {
1103+
case s: ShuffleExchangeExec
1104+
if s.outputPartitioning.isInstanceOf[KeyGroupedPartitioning] =>
1105+
s
1106+
}
1107+
assert(
1108+
keyGroupedShuffles.nonEmpty,
1109+
"KeyGroupedPartitioning shuffle should fall back to a vanilla ShuffleExchangeExec")
1110+
1111+
val columnarKeyGroupedShuffles = collectAllShuffles(plan)
1112+
.filter(_.outputPartitioning.isInstanceOf[KeyGroupedPartitioning])
1113+
assert(
1114+
columnarKeyGroupedShuffles.isEmpty,
1115+
"KeyGroupedPartitioning must not be offloaded to ColumnarShuffleExchangeExec")
1116+
1117+
checkAnswer(df, Seq(Row(1, "aa", 40.0, 42.0), Row(3, "bb", 10.0, 19.5)))
1118+
}
1119+
}
1120+
10751121
testGluten("SPARK-41471: shuffle one side: only one side reports partitioning") {
10761122
val items_partitions = Array(identity("id"))
10771123
createTable(items, itemsColumns, items_partitions)

0 commit comments

Comments
 (0)