Skip to content

Commit 7ef3bbd

Browse files
committed
fix: skip broadcast coalescing for schema mismatches
1 parent a0c8f89 commit 7ef3bbd

2 files changed

Lines changed: 32 additions & 1 deletion

File tree

‎spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -396,6 +396,13 @@ object Utils extends CometTypeShim with Logging {
396396
}
397397
while (reader.loadNextBatch()) {
398398
val sourceRoot = reader.getVectorSchemaRoot
399+
if (targetRoot != null && targetRoot.getSchema != sourceRoot.getSchema) {
400+
logWarning(
401+
"Arrow schemas differ during BroadcastExchange coalescing; skipping coalesce")
402+
targetRoot.close()
403+
targetRoot = null
404+
return (buffers, 0L, 0L)
405+
}
399406
if (targetRoot == null) {
400407
targetRoot = VectorSchemaRoot.create(sourceRoot.getSchema, allocator)
401408
targetRoot.allocateNew()

‎spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ import org.scalatest.Tag
2525
import org.apache.spark.sql.CometTestBase
2626
import org.apache.spark.sql.catalyst.TableIdentifier
2727
import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
28-
import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometSortMergeJoinExec}
28+
import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometSortMergeJoinExec, CometUnionExec}
2929
import org.apache.spark.sql.execution.adaptive.AQEShuffleReadExec
3030
import org.apache.spark.sql.internal.SQLConf
3131

@@ -657,6 +657,30 @@ class CometJoinSuite extends CometTestBase {
657657
}
658658
}
659659

660+
test("Broadcast coalescing falls back when union children have different nullability") {
661+
withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
662+
withParquetTable(Seq((1, 10), (2, 20), (3, 30)), "t") {
663+
val (_, cometPlan) = checkSparkAnswerAndOperator(
664+
sql("""
665+
|SELECT /*+ BROADCAST(b) */ p._1, b.v
666+
|FROM t p JOIN (
667+
| SELECT _1 AS k, 99 AS v FROM t
668+
| UNION ALL
669+
| SELECT _1 AS k, _2 + 1 AS v FROM t
670+
|) b ON p._1 = b.k
671+
|""".stripMargin),
672+
Seq(
673+
classOf[CometBroadcastExchangeExec],
674+
classOf[CometBroadcastHashJoinExec],
675+
classOf[CometUnionExec]))
676+
677+
val broadcast = collect(cometPlan) { case b: CometBroadcastExchangeExec => b }.head
678+
assert(broadcast.metrics("numCoalescedBatches").value == 0L)
679+
assert(broadcast.metrics("numCoalescedRows").value == 0L)
680+
}
681+
}
682+
}
683+
660684
// Reproducer for SPARK-43113: full outer SMJ with a join filter that references
661685
// a nullable column should not match when the filter evaluates to NULL.
662686
test("SPARK-43113: Full outer SMJ with NULL in join filter") {

0 commit comments

Comments
 (0)