@@ -25,7 +25,7 @@ import org.scalatest.Tag
2525import org .apache .spark .sql .CometTestBase
2626import org .apache .spark .sql .catalyst .TableIdentifier
2727import 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 }
2929import org .apache .spark .sql .execution .adaptive .AQEShuffleReadExec
3030import 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