diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 70d108f2506..a612704f71a 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -967,7 +967,19 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } case other => other } - if (cometChild.isInstanceOf[CometNativeExec]) { + // Accept any Comet plan with columnar output, not only CometNativeExec. The join's + // own broadcast converts to CometBroadcastExchangeExec over CometUnionExec, + // CometCoalesceExec, CometTakeOrderedAndProjectExec or CometShuffleExchangeExec + // (each is wrapped in a native placeholder while the join is converted), so + // requiring CometNativeExec here would leave the DPP subquery on a Spark broadcast + // that ReuseExchangeAndSubquery can't match. Same check as the AQE path in + // CometPlanAdaptiveDynamicPruningFilters.isNativeBuildSide. + // + // With Comet broadcast exchange disabled, the join's broadcast stays on Spark, so the + // DPP subquery has to stay on Spark too for the two to match. The AQE path gets the + // same result by requiring a Comet join (isCometJoin). + if (CometConf.COMET_EXEC_BROADCAST_EXCHANGE_ENABLED.get(conf) && + cometChild.isInstanceOf[CometPlan] && cometChild.supportsColumnar) { logInfo( "Converting SubqueryBroadcastExec to " + "CometSubqueryBroadcastExec for DPP exchange reuse") diff --git a/spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala b/spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala index b92fd418366..f045b56bad5 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala @@ -267,7 +267,7 @@ case object CometPlanAdaptiveDynamicPruningFilters // plan makes it fail at execution with "Comet execution only takes Arrow Arrays". So only // build a Comet exchange when the join is Comet AND buildSidePlan is actually native; // otherwise fall back to a Spark broadcast + SubqueryBroadcastExec (correct, loses only DPP - // broadcast reuse). Mirrors the non-AQE guard in CometExecRule.rewriteInSubqueryPlan. + // broadcast reuse). Same check as the non-AQE path in CometExecRule.rewriteInSubqueryPlan. val isComet = isCometJoin && isNativeBuildSide(buildSidePlan) logDebug( s"Matched DPP subquery '${sab.name}' to " + diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 1ba17ea2dae..0fa2298ba64 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -463,6 +463,135 @@ class CometExecSuite extends CometTestBase { } } + // Writes a fact table partitioned on `fk` and two small dimension tables, registers them as + // dpp_shape_fact, dpp_shape_dim1 and dpp_shape_dim2, and runs `f` with AQE off and DPP limited + // to reusing a broadcast. + private def withDppShapeTables(extraConfs: (String, String)*)(f: => Unit): Unit = { + withTempDir { dir => + val path = dir.getAbsolutePath + withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") { + spark + .range(100) + .selectExpr("cast(id % 10 as int) as fk", "id as v") + .write + .partitionBy("fk") + .parquet(s"$path/fact") + spark + .range(0, 5) + .selectExpr("cast(id as int) as value", "id as w") + .write + .parquet(s"$path/dim1") + spark + .range(5, 10) + .selectExpr("cast(id as int) as value", "id as w") + .write + .parquet(s"$path/dim2") + } + + withSQLConf( + Seq( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") ++ + extraConfs: _*) { + spark.read.parquet(s"$path/fact").createOrReplaceTempView("dpp_shape_fact") + spark.read.parquet(s"$path/dim1").createOrReplaceTempView("dpp_shape_dim1") + spark.read.parquet(s"$path/dim2").createOrReplaceTempView("dpp_shape_dim2") + f + } + } + } + + private val dppUnionDimSql = + "SELECT value FROM dpp_shape_dim1 WHERE w < 2 " + + "UNION ALL SELECT value FROM dpp_shape_dim2 WHERE w > 8" + + private def dppShapeQuery(dimSql: String): DataFrame = + sql(s"""SELECT /*+ BROADCAST(d) */ f.v, f.fk + |FROM dpp_shape_fact f JOIN ($dimSql) d + |ON f.fk = d.value""".stripMargin) + + // Build sides whose top operator is a columnar Comet plan but not a CometNativeExec. The join + // converts its broadcast to a CometBroadcastExchangeExec over each of them, so the DPP subquery + // has to as well, or ReuseExchangeAndSubquery can't match the two exchanges (#6815). + Seq( + ("UNION ALL", classOf[CometUnionExec], dppUnionDimSql), + ( + "COALESCE", + classOf[CometCoalesceExec], + "SELECT /*+ COALESCE(1) */ value FROM dpp_shape_dim1 WHERE w < 2"), + ( + "REPARTITION", + classOf[CometShuffleExchangeExec], + "SELECT /*+ REPARTITION(2, value) */ value FROM dpp_shape_dim1 WHERE w < 2")).foreach { + case (shape, buildClass, dimSql) => + test(s"non-AQE DPP: BHJ reuses broadcast exchange when the build side is $shape") { + withDppShapeTables() { + val (_, cometPlan) = checkSparkAnswer(dppShapeQuery(dimSql)) + + val broadcasts = collectWithSubqueries(cometPlan) { + case e: CometBroadcastExchangeExec => e + } + assert( + broadcasts.size == 1, + "Expected exactly 1 CometBroadcastExchangeExec (other reused):\n" + + cometPlan.treeString) + // Without this the test could pass on a build side that was already native. + assert( + buildClass.isInstance(broadcasts.head.child), + s"Expected ${buildClass.getSimpleName} directly under the broadcast:\n" + + cometPlan.treeString) + + val cometSubqueries = collectWithSubqueries(cometPlan) { + case s: CometSubqueryBroadcastExec => s + } + assert( + cometSubqueries.nonEmpty, + s"Expected CometSubqueryBroadcastExec for DPP:\n${cometPlan.treeString}") + + val reused = collectWithSubqueries(cometPlan) { case e: ReusedExchangeExec => e } + assert( + reused.exists(_.child.canonicalized == broadcasts.head.canonicalized), + "Expected the DPP subquery to reuse the join's CometBroadcastExchangeExec:\n" + + cometPlan.treeString) + } + } + } + + test( + "non-AQE DPP: DPP subquery keeps the Spark broadcast when Comet broadcast exchange " + + "is disabled") { + // The join's broadcast stays on Spark, so the DPP subquery must too, or Spark can't reuse + // the join's broadcast and the dimension is broadcast twice. + withDppShapeTables(CometConf.COMET_EXEC_BROADCAST_EXCHANGE_ENABLED.key -> "false") { + val (_, cometPlan) = checkSparkAnswer(dppShapeQuery(dppUnionDimSql)) + + val cometBroadcasts = collectWithSubqueries(cometPlan) { + case e: CometBroadcastExchangeExec => e + } + assert( + cometBroadcasts.isEmpty, + s"Expected no CometBroadcastExchangeExec:\n${cometPlan.treeString}") + + val sparkSubqueries = collectWithSubqueries(cometPlan) { case s: SubqueryBroadcastExec => + s + } + assert( + sparkSubqueries.nonEmpty, + s"Expected a Spark SubqueryBroadcastExec for DPP:\n${cometPlan.treeString}") + + val broadcasts = collectWithSubqueries(cometPlan) { case e: BroadcastExchangeExec => e } + assert( + broadcasts.size == 1, + s"Expected exactly 1 BroadcastExchangeExec (other reused):\n${cometPlan.treeString}") + val reused = collectWithSubqueries(cometPlan) { case e: ReusedExchangeExec => e } + assert( + reused.exists(_.child.canonicalized == broadcasts.head.canonicalized), + "Expected the DPP subquery to reuse the join's BroadcastExchangeExec:\n" + + cometPlan.treeString) + } + } + test("non-AQE DPP: non-atomic type (struct/array) join key") { withTempDir { dir => val path = s"${dir.getAbsolutePath}/data" @@ -850,9 +979,8 @@ class CometExecSuite extends CometTestBase { .parquet(s"$path/dim") } - // Disable Comet broadcast exchange so SubqueryBroadcastExec wraps a Spark - // BroadcastExchangeExec. convertSubqueryBroadcasts should skip it (child isn't - // CometNativeExec). Query should still produce correct results via Spark's standard path. + // Disable Comet broadcast exchange so the join keeps a Spark BroadcastExchangeExec. + // convertSubqueryBroadcasts then leaves the DPP subquery on Spark too. withSQLConf( SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",