diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala index 4f8f66eb59690..ffcb39153f59e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala +++ b/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. */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/BroadcastExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/BroadcastExchangeExec.scala index 8c695f4f3958d..fadb5dbdb9e08 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/BroadcastExchangeExec.scala +++ b/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 { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala index d5cc3e6c0c00d..c1a58e79d997c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala +++ b/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") + } + } + } + } }