From 1df1eaf2d60a43ce8dcb8f3bfe7cfe8b6faee662 Mon Sep 17 00:00:00 2001 From: tlfua <0050211@gmail.com> Date: Sat, 7 Nov 2026 01:24:57 -0800 Subject: [PATCH 1/3] fix: reuse the Comet broadcast for a non-AQE DPP subquery whose build side is not a CometNativeExec Without AQE, CometExecRule.rewriteInSubqueryPlan only converted a DPP SubqueryBroadcastExec to CometSubqueryBroadcastExec when the broadcast child was a CometNativeExec. The join's own broadcast converts over CometUnionExec, CometCoalesceExec, CometTakeOrderedAndProjectExec and CometShuffleExchangeExec too, so ReuseExchangeAndSubquery couldn't match the two exchanges and the dimension was broadcast twice. Accept any columnar CometPlan, the same check the AQE path uses in CometPlanAdaptiveDynamicPruningFilters.isNativeBuildSide. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 95bb3fbb-4ffd-407c-a7bb-e912bfe7b965 --- .../apache/comet/rules/CometExecRule.scala | 9 +- ...metPlanAdaptiveDynamicPruningFilters.scala | 2 +- .../apache/comet/exec/CometExecSuite.scala | 93 ++++++++++++++++++- 3 files changed, 99 insertions(+), 5 deletions(-) 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..456df1c86dc 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,14 @@ 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. + if (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..14227ff0549 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,93 @@ class CometExecSuite extends CometTestBase { } } + // 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], + "SELECT value FROM dpp_shape_dim1 WHERE w < 2 " + + "UNION ALL SELECT value FROM dpp_shape_dim2 WHERE w > 8"), + ( + "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"), + ( + "ORDER BY LIMIT", + classOf[CometTakeOrderedAndProjectExec], + "SELECT value FROM dpp_shape_dim1 WHERE w < 4 ORDER BY value LIMIT 2")).foreach { + case (shape, buildClass, dimSql) => + test(s"non-AQE DPP: BHJ reuses broadcast exchange when the build side is $shape") { + 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( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") { + 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") + + val df = sql(s"""SELECT /*+ BROADCAST(d) */ f.v, f.fk + |FROM dpp_shape_fact f JOIN ($dimSql) d + |ON f.fk = d.value""".stripMargin) + val (_, cometPlan) = checkSparkAnswer(df) + + 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: non-atomic type (struct/array) join key") { withTempDir { dir => val path = s"${dir.getAbsolutePath}/data" @@ -850,9 +937,9 @@ 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. + // Query should still produce correct results whether or not the DPP subquery is + // converted. withSQLConf( SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", From 60e29bdac34d9c194af090f20616c47b7fb2c8bc Mon Sep 17 00:00:00 2001 From: tlfua <0050211@gmail.com> Date: Sat, 7 Nov 2026 03:15:49 -0800 Subject: [PATCH 2/3] test: drop the ORDER BY LIMIT DPP reuse case The join's inferred isnotnull filter on the key can't be pushed below the limit, so a native CometFilter is the top operator of that build side and CometTakeOrderedAndProjectExec never sits directly under the broadcast. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 95bb3fbb-4ffd-407c-a7bb-e912bfe7b965 --- .../test/scala/org/apache/comet/exec/CometExecSuite.scala | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) 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 14227ff0549..e078629d0ae 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -479,11 +479,7 @@ class CometExecSuite extends CometTestBase { ( "REPARTITION", classOf[CometShuffleExchangeExec], - "SELECT /*+ REPARTITION(2, value) */ value FROM dpp_shape_dim1 WHERE w < 2"), - ( - "ORDER BY LIMIT", - classOf[CometTakeOrderedAndProjectExec], - "SELECT value FROM dpp_shape_dim1 WHERE w < 4 ORDER BY value LIMIT 2")).foreach { + "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") { withTempDir { dir => From 263c3d3547a7de6b0dd551a9207f2b0251e2141a Mon Sep 17 00:00:00 2001 From: tlfua <0050211@gmail.com> Date: Sat, 7 Nov 2026 09:20:27 -0800 Subject: [PATCH 3/3] fix: keep the Spark broadcast for a non-AQE DPP subquery when Comet broadcast exchange is disabled With spark.comet.exec.broadcastExchange.enabled=false the join's broadcast stays a Spark BroadcastExchangeExec, so converting the DPP subquery to a CometBroadcastExchangeExec breaks Spark's exchange reuse and broadcasts the dimension twice. Only convert the DPP subquery when the config is on, which matches the AQE path requiring a Comet join. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 95bb3fbb-4ffd-407c-a7bb-e912bfe7b965 --- .../apache/comet/rules/CometExecRule.scala | 7 +- .../apache/comet/exec/CometExecSuite.scala | 173 +++++++++++------- 2 files changed, 115 insertions(+), 65 deletions(-) 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 456df1c86dc..a612704f71a 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -974,7 +974,12 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) // 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. - if (cometChild.isInstanceOf[CometPlan] && cometChild.supportsColumnar) { + // + // 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/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index e078629d0ae..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,15 +463,59 @@ 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], - "SELECT value FROM dpp_shape_dim1 WHERE w < 2 " + - "UNION ALL SELECT value FROM dpp_shape_dim2 WHERE w > 8"), + ("UNION ALL", classOf[CometUnionExec], dppUnionDimSql), ( "COALESCE", classOf[CometCoalesceExec], @@ -482,70 +526,72 @@ class CometExecSuite extends CometTestBase { "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") { - 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( - SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") { - 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") - - val df = sql(s"""SELECT /*+ BROADCAST(d) */ f.v, f.fk - |FROM dpp_shape_fact f JOIN ($dimSql) d - |ON f.fk = d.value""".stripMargin) - val (_, cometPlan) = checkSparkAnswer(df) - - 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) + withDppShapeTables() { + val (_, cometPlan) = checkSparkAnswer(dppShapeQuery(dimSql)) - val cometSubqueries = collectWithSubqueries(cometPlan) { - case s: CometSubqueryBroadcastExec => s - } - assert( - cometSubqueries.nonEmpty, - s"Expected CometSubqueryBroadcastExec for DPP:\n${cometPlan.treeString}") + 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 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) + 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" @@ -934,8 +980,7 @@ class CometExecSuite extends CometTestBase { } // Disable Comet broadcast exchange so the join keeps a Spark BroadcastExchangeExec. - // Query should still produce correct results whether or not the DPP subquery is - // converted. + // convertSubqueryBroadcasts then leaves the DPP subquery on Spark too. withSQLConf( SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",