fix: reuse the Comet broadcast for a non-AQE DPP subquery whose build side is not a CometNativeExec

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")

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 " +

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",