fix: [branch-1.1] keep unconverted DPP filters in a scan's canonical form (#6270)

spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala

@@ -287,18 +287,6 @@ case class CometNativeScanExec(
   }
 
   override def doCanonicalize(): CometNativeScanExec = {
-    // Canonicalize originalPlan but strip its DPP partition filters.
-    // originalPlan carries column selection and schema info needed for
-    // equals/hashCode. But its partitionFilters may have stale DPP
-    // expressions (e.g., SABs not yet converted to TrueLiteral) that
-    // would prevent exchange reuse between otherwise-identical scans.
-    val canonOriginal = if (originalPlan != null) {
-      val stripped = originalPlan.copy(partitionFilters =
-        CometScanUtils.filterUnusedDynamicPruningExpressions(originalPlan.partitionFilters))
-      stripped.doCanonicalize()
-    } else {
-      null
-    }
     CometNativeScanExec(
       nativeOp,
       relation,
@@ -312,7 +300,7 @@ case class CometNativeScanExec(
       QueryPlan.normalizePredicates(dataFilters, output),
       None,
       disableBucketedScan,
-      canonOriginal,
+      CometScanUtils.canonicalizeOriginalPlan(originalPlan),
       SerializedPlan(None),
       null, // Transient scan not needed for canonicalization
       ""

spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala

@@ -20,7 +20,7 @@
 package org.apache.spark.sql.comet
 
 import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, Literal}
-import org.apache.spark.sql.execution.{InSubqueryExec, SubqueryAdaptiveBroadcastExec}
+import org.apache.spark.sql.execution.FileSourceScanExec
 
 object CometScanUtils {
 
@@ -29,20 +29,27 @@ object CometScanUtils {
    * DynamicPruningExpression(Literal.TrueLiteral) during Physical Planning
    */
   def filterUnusedDynamicPruningExpressions(predicates: Seq[Expression]): Seq[Expression] = {
-    // Strip DPP expressions for canonicalization. Matches Spark's
-    // FileSourceScanExec.filterUnusedDynamicPruningExpressions (TrueLiteral).
-    // Also strips unconverted SAB wrappers because AQE stageCache canonicalizes
-    // before our queryStageOptimizerRule converts them, so they would prevent
-    // exchange reuse between otherwise-identical scans.
-    predicates.filterNot {
-      case DynamicPruningExpression(Literal.TrueLiteral) => true
-      case DynamicPruningExpression(
-            InSubqueryExec(_, _: CometSubqueryAdaptiveBroadcastExec, _, _, _, _)) =>
-        true
-      case DynamicPruningExpression(
-            InSubqueryExec(_, _: SubqueryAdaptiveBroadcastExec, _, _, _, _)) =>
-        true
-      case _ => false
+    // Matches Spark's FileSourceScanExec.filterUnusedDynamicPruningExpressions. A DPP filter that
+    // still holds the adaptive broadcast placeholder must stay: AQE keys a query stage by its
+    // exchange as it was before the stage optimizer rules converted that placeholder, so dropping
+    // it makes scans with different DPP filters equal and lets AQE reuse one stage for both.
+    predicates.filterNot(_ == DynamicPruningExpression(Literal.TrueLiteral))
+  }
+
+  /**
+   * Canonicalizes a Comet scan's `originalPlan` (which may be null) without any of its DPP
+   * partition filters. That copy carries the column selection and schema needed for equals and
+   * hashCode, but its DPP filters are stale because no DPP rule rewrites them. The scan's own DPP
+   * filters, passed through [[filterUnusedDynamicPruningExpressions]], carry its DPP identity.
+   */
+  def canonicalizeOriginalPlan(originalPlan: FileSourceScanExec): FileSourceScanExec = {
+    if (originalPlan == null) {
+      null
+    } else {
+      originalPlan
+        .copy(partitionFilters =
+          originalPlan.partitionFilters.filterNot(_.isInstanceOf[DynamicPruningExpression]))
+        .doCanonicalize()
     }
   }
 }

spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala

@@ -5558,6 +5558,87 @@ class CometIcebergNativeSuite
     }
   }
 
+  // TPC-DS q64 shape (#6264): a broadcast of (fact SMJ other) per UNION ALL branch, with the
+  // dimension join above it. The two broadcasts differ only in the DPP filter of the fact scan
+  // inside a child shuffle stage, so they must not be reused for each other.
+  test("AQE DPP - broadcast over a join of DPP scans is not reused across branches (#6264)") {
+    assume(icebergAvailable, "Iceberg not available")
+    withTempIcebergDir { warehouseDir =>
+      val dimDir = new File(warehouseDir, "dim_parquet")
+      val otherDir = new File(warehouseDir, "other_parquet")
+      withSQLConf(
+        "spark.sql.catalog.aqe_cat" -> "org.apache.iceberg.spark.SparkCatalog",
+        "spark.sql.catalog.aqe_cat.type" -> "hadoop",
+        "spark.sql.catalog.aqe_cat.warehouse" -> warehouseDir.getAbsolutePath,
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+        SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+        // A coalesced shuffle read carries per-partition sizes that keep the two broadcasts apart.
+        SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false",
+        CometConf.COMET_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_ENABLED.key -> "true",
+        CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") {
+
+        spark.sql("""
+          CREATE TABLE aqe_cat.db.branch_fact (
+            store_id INT, item INT, amount INT
+          ) USING iceberg PARTITIONED BY (store_id)
+        """)
+        spark
+          .range(1000)
+          .selectExpr(
+            "cast(id % 10 as int) as store_id",
+            "cast(id % 13 as int) as item",
+            "cast(id as int) as amount")
+          .write
+          .format("iceberg")
+          .mode("append")
+          .saveAsTable("aqe_cat.db.branch_fact")
+        // Group 'a' keeps store 1 and group 'b' keeps stores 2 to 4.
+        spark
+          .range(10)
+          .selectExpr(
+            "cast(id as int) as store_id",
+            "case when id = 1 then 'a' when id in (2, 3, 4) then 'b' else 'c' end as grp")
+          .write
+          .parquet(dimDir.getAbsolutePath)
+        spark.read.parquet(dimDir.getAbsolutePath).createOrReplaceTempView("branch_dim")
+        spark
+          .range(13)
+          .selectExpr("cast(id as int) as item", "cast(id * 10 as int) as w")
+          .write
+          .parquet(otherDir.getAbsolutePath)
+        spark.read.parquet(otherDir.getAbsolutePath).createOrReplaceTempView("branch_other")
+
+        val query =
+          """WITH x AS (
+            |  SELECT /*+ MERGE(f, o) */ f.store_id, f.item, f.amount, o.w
+            |  FROM aqe_cat.db.branch_fact f JOIN branch_other o ON f.item = o.item)
+            |SELECT /*+ BROADCAST(x) */ o2.item, x.store_id, x.amount, x.w
+            |FROM branch_other o2 JOIN x ON o2.item = x.item
+            |JOIN branch_dim d ON x.store_id = d.store_id WHERE d.grp = 'a'
+            |UNION ALL
+            |SELECT /*+ BROADCAST(x) */ o2.item, x.store_id, x.amount, x.w
+            |FROM branch_other o2 JOIN x ON o2.item = x.item
+            |JOIN branch_dim d ON x.store_id = d.store_id WHERE d.grp = 'b'""".stripMargin
+        val (_, cometPlan) = checkSparkAnswer(query)
+        assertIcebergNativeScanPresent(cometPlan)
+        assertNoLeftoverCSAB(cometPlan)
+
+        // Also checked on 3.4: the Iceberg scan stays in Comet there and Spark's own DPP rule
+        // plans its filter.
+        val dppScan: PartialFunction[SparkPlan, CometIcebergNativeScanExec] = {
+          case s: CometIcebergNativeScanExec if s.runtimeFilters.nonEmpty => s
+        }
+        assert(
+          collect(cometPlan)(dppScan).size == 2,
+          s"Expected one DPP scan per branch:\n${cometPlan.treeString}")
+        assertNoReusedExchangeOverDppScan(cometPlan)(dppScan)
+
+        spark.sql("DROP TABLE aqe_cat.db.branch_fact")
+      }
+    }
+  }
+
   // Multi-key BHJ where SAB build keys must keep their position so the index lookup
   // selects the right DPP value. A bug in convertSAB key handling would either pick
   // the wrong column or produce SubqueryExec instead of broadcast reuse.

spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala

@@ -1569,6 +1569,129 @@ class CometExecSuite extends CometTestBase {
     }
   }
 
+  // Tables for the #6264 tests. The fact table has 10 store partitions. Dimension group 'a' keeps
+  // store 1 and group 'b' keeps stores 2 to 4, so the two DPP scans prune to 1 and 3 partitions.
+  private def withDppBranchTables(f: => Unit): Unit = {
+    withTempDir { dir =>
+      val path = dir.getAbsolutePath
+      spark
+        .range(1000)
+        .selectExpr(
+          "cast(id % 10 as int) as store_id",
+          "cast(id % 13 as int) as item",
+          "cast(id as int) as amount")
+        .write
+        .partitionBy("store_id")
+        .parquet(s"$path/fact")
+      spark
+        .range(10)
+        .selectExpr(
+          "cast(id as int) as store_id",
+          "case when id = 1 then 'a' when id in (2, 3, 4) then 'b' else 'c' end as grp")
+        .write
+        .parquet(s"$path/dim")
+      spark
+        .range(13)
+        .selectExpr("cast(id as int) as item", "cast(id * 10 as int) as w")
+        .write
+        .parquet(s"$path/other")
+      withTempView("dpp_branch_fact", "dpp_branch_dim", "dpp_branch_other") {
+        spark.read.parquet(s"$path/fact").createOrReplaceTempView("dpp_branch_fact")
+        spark.read.parquet(s"$path/dim").createOrReplaceTempView("dpp_branch_dim")
+        spark.read.parquet(s"$path/other").createOrReplaceTempView("dpp_branch_other")
+        f
+      }
+    }
+  }
+
+  // Coalescing is off: a coalesced shuffle read carries per-partition sizes that would keep the two
+  // branches' parent exchanges apart and hide a wrong reuse.
+  private def withDppBranchConf(f: => Unit): Unit = {
+    withSQLConf(
+      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+      SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+      SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false")(f)
+  }
+
+  private val dppScan: PartialFunction[SparkPlan, CometNativeScanExec] = {
+    case s: CometNativeScanExec
+        if s.partitionFilters.exists(_.isInstanceOf[DynamicPruningExpression]) =>
+      s
+  }
+
+  // A UNION ALL of one CTE joined to dimension groups grp1 and grp2, with a shuffle stage over each
+  // DPP scan.
+  private def dppShuffleQuery(grp1: String, grp2: String): String =
+    s"""WITH x AS (
+       |  SELECT store_id, item, sum(amount) AS s FROM dpp_branch_fact GROUP BY store_id, item),
+       |y AS (SELECT store_id, s % 7 AS b, count(*) AS c FROM x GROUP BY store_id, s % 7)
+       |SELECT y.store_id, y.b, y.c FROM y JOIN dpp_branch_dim d ON y.store_id = d.store_id
+       |WHERE d.grp = '$grp1'
+       |UNION ALL
+       |SELECT y.store_id, y.b, y.c FROM y JOIN dpp_branch_dim d ON y.store_id = d.store_id
+       |WHERE d.grp = '$grp2'""".stripMargin
+
+  // Runs a UNION ALL of one CTE joined to dimension group 'a' and to group 'b'.
+  private def checkDppBranchesKeepTheirRows(query: String): Unit = {
+    withDppBranchConf {
+      if (isSpark35Plus) {
+        val (_, cometPlan) = checkSparkAnswerAndOperator(sql(query))
+        assert(
+          collect(cometPlan)(dppScan).map(_.metrics("numPartitions").value).sorted ==
+            Seq(1L, 3L),
+          "Expected one DPP scan per branch, pruned to 1 and 3 partitions:\n" +
+            cometPlan.treeString)
+        assertNoReusedExchangeOverDppScan(cometPlan)(dppScan)
+      } else {
+        // On 3.4, AQE DPP falls back to Spark.
+        checkSparkAnswer(sql(query))
+      }
+    }
+  }
+
+  // TPC-DS q64 shape: a broadcast of (fact SMJ other), with the dimension join above it. The two
+  // broadcasts differ only in the DPP filter of the fact scan inside a child shuffle stage.
+  test("AQE DPP: broadcast over a join of DPP scans is not reused across branches (#6264)") {
+    withDppBranchTables {
+      checkDppBranchesKeepTheirRows("""WITH x AS (
+          |  SELECT /*+ MERGE(f, o) */ f.store_id, f.item, f.amount, o.w
+          |  FROM dpp_branch_fact f JOIN dpp_branch_other o ON f.item = o.item)
+          |SELECT /*+ BROADCAST(x) */ o2.item, x.store_id, x.amount, x.w
+          |FROM dpp_branch_other o2 JOIN x ON o2.item = x.item
+          |JOIN dpp_branch_dim d ON x.store_id = d.store_id WHERE d.grp = 'a'
+          |UNION ALL
+          |SELECT /*+ BROADCAST(x) */ o2.item, x.store_id, x.amount, x.w
+          |FROM dpp_branch_other o2 JOIN x ON o2.item = x.item
+          |JOIN dpp_branch_dim d ON x.store_id = d.store_id WHERE d.grp = 'b'""".stripMargin)
+    }
+  }
+
+  test("AQE DPP: shuffle over a DPP scan stage is not reused across branches (#6264)") {
+    withDppBranchTables {
+      checkDppBranchesKeepTheirRows(dppShuffleQuery("a", "b"))
+    }
+  }
+
+  // With the same dimension filter in both branches, AQE's stage cache lookup after the DPP
+  // placeholder is converted finds the first branch's stage, so the shuffles are reused as in
+  // Spark. This fails if the scan's canonical form keeps too much, such as originalPlan's DPP
+  // filters or unnormalized partition filters.
+  test("AQE DPP: shuffle over scans with the same DPP filter is reused across branches (#6264)") {
+    withDppBranchTables {
+      withDppBranchConf {
+        val df = sql(dppShuffleQuery("a", "a"))
+        val (sparkPlan, cometPlan) =
+          if (isSpark35Plus) checkSparkAnswerAndOperator(df) else checkSparkAnswer(df)
+        assertSameExchangeReuseAsSpark(sparkPlan, cometPlan)
+        if (isSpark35Plus) {
+          assertExchangeReuseOver(cometPlan, "Expected the shuffle over a DPP scan to be reused")(
+            dppScan)
+        }
+      }
+    }
+  }
+
   // SPARK-34637: previously IgnoreComet(#4045). DPP side broadcast query stage
   // should be created before the main join's broadcast stage.
   test("AQE DPP: broadcast query stage creation order (SPARK-34637)") {

spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala

@@ -68,12 +68,40 @@ abstract class CometTestBase
   protected val shuffleManager: String =
     "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager"
 
+  /** The reused exchanges in `plan`. */
+  protected def reusedExchanges(plan: SparkPlan): Seq[ReusedExchangeExec] =
+    collect(plan) { case r: ReusedExchangeExec => r }
+
+  /** The reused exchanges in `plan` whose child contains a node that `pf` matches. */
+  protected def reusedExchangesOver[T](plan: SparkPlan)(
+      pf: PartialFunction[SparkPlan, T]): Seq[ReusedExchangeExec] =
+    reusedExchanges(plan).filter(exchange => collect(exchange.child)(pf).nonEmpty)
+
+  /** Asserts that some reused exchange in `plan` sits over a node that `pf` matches. */
   protected def assertExchangeReuseOver[T](plan: SparkPlan, clue: String)(
-      pf: PartialFunction[SparkPlan, T]): Unit = {
-    val reused = collect(plan) {
-      case exchange: ReusedExchangeExec if collect(exchange.child)(pf).nonEmpty => exchange
-    }
-    assert(reused.nonEmpty, s"$clue:\n$plan")
+      pf: PartialFunction[SparkPlan, T]): Unit =
+    assert(reusedExchangesOver(plan)(pf).nonEmpty, s"$clue:\n$plan")
+
+  /** Asserts that no reused exchange in `plan` sits over a scan that `dppScan` matches. */
+  protected def assertNoReusedExchangeOverDppScan[T](plan: SparkPlan)(
+      dppScan: PartialFunction[SparkPlan, T]): Unit =
+    assert(
+      reusedExchangesOver(plan)(dppScan).isEmpty,
+      s"An exchange over a DPP scan was reused by another branch:\n$plan")
+
+  /**
+   * Asserts that Spark reuses at least one exchange in `sparkPlan` and that Comet reuses as many
+   * in `cometPlan`.
+   */
+  protected def assertSameExchangeReuseAsSpark(
+      sparkPlan: SparkPlan,
+      cometPlan: SparkPlan): Unit = {
+    val sparkReused = reusedExchanges(sparkPlan).size
+    val cometReused = reusedExchanges(cometPlan).size
+    assert(sparkReused > 0, s"Spark should reuse an exchange:\n$sparkPlan")
+    assert(
+      cometReused == sparkReused,
+      s"expected $sparkReused reused exchanges, got $cometReused:\n" + cometPlan.treeString)
   }
 
   protected def sparkConf: SparkConf = {