[SPARK-58485] [SQL] Broadcast hash joins with scalar subquery as join key

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala

@@ -279,11 +279,8 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
 
   @transient private val prepareLock = new Object()
 
-  /**
-   * Finds scalar subquery expressions in this plan node and starts evaluating them.
-   */
-  protected def prepareSubqueries(): Unit = {
-    expressions.foreach {
+   protected def registerSubqueries(exprs: Seq[Expression]): Unit = {
+    exprs.foreach {
       _.collect {
         case e: ExecSubqueryExpression =>
           e.plan.prepare()
@@ -292,6 +289,13 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
     }
   }
 
+   /**
+   * Finds scalar subquery expressions in this plan node and starts evaluating them.
+   */
+  protected def prepareSubqueries(): Unit = {
+    registerSubqueries(expressions)
+  }
+
   /**
    * Blocks the thread until all subqueries finish evaluation and update the results.
    */

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/BroadcastExchangeExec.scala

@@ -193,7 +193,14 @@ case class BroadcastExchangeExec(
             val beforeBuild = System.nanoTime()
             longMetric("collectTime") += NANOSECONDS.toMillis(beforeBuild - beforeCollect)
 
-            // Construct the relation.
+            mode match {
+                case h: HashedRelationBroadcastMode =>
+                registerSubqueries(h.key)
+                case _ =>
+            }
+
+            waitForSubqueries()
+
             val relation = mode.transform(input, Some(numRows))
 
             val dataSize = relation match {

sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala

@@ -27,8 +27,8 @@ import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, Log
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecution}
 import org.apache.spark.sql.execution.datasources.FileScanRDD
-import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
-import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashedRelationBroadcastMode}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 
@@ -2678,4 +2678,45 @@ class SubquerySuite extends SharedSparkSession
 
     assert(exposedAttribute.exprId == outerReferenceAttribute.exprId)
   }
+
+  test("SPARK-58485: BroadcastExchangeExec prepares scalar subqueries in broadcast mode") {
+    withTable("t1", "t2", "t3") {
+      sql("CREATE TABLE t1(a INT) USING PARQUET")
+      sql("INSERT INTO t1 VALUES (1), (2)")
+
+      sql("CREATE TABLE t2(b INT) USING PARQUET")
+      sql("INSERT INTO t2 VALUES (1), (NULL)")
+
+      sql("CREATE TABLE t3(c INT) USING PARQUET")
+      sql("INSERT INTO t3 VALUES (1)")
+
+      withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
+        val df = spark.sql(
+          """
+            |SELECT *
+            |FROM t1
+            |JOIN t2
+            |ON t1.a = COALESCE(
+            |      t2.b,
+            |      (SELECT MIN(c) FROM t3)
+            |   )
+            |""".stripMargin)
+
+        df.collect();
+
+        val exchanges = collect(df.queryExecution.executedPlan) {
+          case e: BroadcastExchangeExec => e
+        }
+
+        assert(exchanges.length == 1)
+
+        exchanges.head.mode match {
+          case h: HashedRelationBroadcastMode =>
+            assert(h.key.exists(_.find(_.isInstanceOf[ScalarSubquery]).isDefined))
+          case other =>
+            fail(s"Expected HashedRelationBroadcastMode, got $other")
+        }
+      }
+    }
+  }
 }